Weighted Ensemble Sampling for Rare-Event Estimation in Language Models

• Direct Monte Carlo becomes highly sample-inefficient for extremely rare events, while proposal-based approaches typically require an additional procedure to learn an effective proposal distribution.For an event around 10^-5, the passage estimates that roughly 10^6 samples may be needed to observe sufficient rare-event realizations.Source: §1 · p. 1
• WE periodically reallocates sampling toward promising partial responses while preserving the original language-model generation process, without training a separate proposal distribution.The method uses internal LLM representations as signals for reallocating the sampling budget.Source: §1 · p. 2
• The paper analyzes WE theoretically, establishing unbiasedness and consistency and deriving an asymptotic distribution that characterizes how design choices affect estimation efficiency.Source: §1 · p. 2
• WE improves rare-event estimation efficiency and reliability across safety-related tasks spanning small-scale and modern large-scale language models.Source: §1 · p. 2
• WE can reuse the same samples for unbiased estimation when the rare-event definition changes, whereas proposal-based methods may fail in that setting.Source: §1 · p. 2

2.1 Problem Setup

• The target is the probability, under the original language model, that a response’s rare-event score R(x) reaches the predefined threshold γ⋆.The event is A⋆ = {x : R(x) ≥ γ⋆}.Source: §2.1 · p. 3

2.2 Weighted Ensemble Algorithm

• Weighted Ensemble (WE) propagates prefixes using the original model, then bins and resamples them to redirect continuations toward promising regions while preserving each bin’s probability mass.The described experiments use equal-population allocation across occupied bins; child weights preserve each bin’s total weight exactly.Source: §2.2 · pp. 3–4

2.3 Theoretical Results

• The estimator is unbiased for any binning maps and child allocations, and its estimation error converges to zero as the number of walkers increases under the stated assumptions.The correctness of these design choices is distinct from their effect on efficiency.Source: §2.3 · pp. 4–5
• For sufficiently large walker populations, the estimator is approximately Gaussian, and WE has lower asymptotic sample complexity than direct Monte Carlo exactly when its asymptotic variance is smaller.The variance has propagation and resampling components, both affected by reaction-coordinate quality.Source: §2.3 · p. 5

2.4 Choice of Reaction Coordinate

• Because the prefix value function is unknown, WE uses LLM representations early and switches to the rare-event score later, retaining a training-free design.The score can be weakly informative early in generation; the analysis permits reaction coordinates to change across checkpoints without changing the target probability.Source: §2.4 · p. 6

2.5 Robustness to Rare-Event Shift

• When the target event is enlarged, weighted samples from an earlier WE run can still yield unbiased estimates, whereas an event-specialized proposal may miss newly added regions.The contrast arises because WE generates continuations directly from the original language model.Source: §2.5 · pp. 6–7

3.1 Experiment Setup

• The evaluation covers TinyStories toxic-story generation and StrongREJECT jailbreaks, using direct Monte Carlo references and proposal-based comparisons under matched sampling budgets.It reports estimated probabilities, hit rates, ESS, relative error, and sample-complexity gains.Source: §3.1 · p. 7

3.2 Results

• Hybrid representation coordinates improve estimation, achieve sample-efficiency gains greater than 1 over direct Monte Carlo, and outperform proposal-based methods; their Ck is smaller than reward-only across most checkpoints.The checkpoint-wise comparison is reported for most checkpoints.Source: Table 1 · p. 7, Figure 2 · p. 7, §3.2 · p. 8
• On the StrongREJECT jailbreak task, hybrid WE further improves estimation performance, rare-event hit rates, and ESS over reward-only WE.The experiments report results for Llama-8B over five seeds.Source: Table 2 · p. 8, §3.2 · p. 8
• Under rare-event shifts, proposal-based methods become substantially less accurate, while representation-based WE maintains low relative error across thresholds.The same samples are reused for enlarged events after methods are configured for the original threshold.Source: Table 3 · p. 8, §3.2 · p. 8

Limitations

• Hybrid reaction coordinates require white-box access; reward-only WE is an available black-box alternative, though hybrid coordinates generally estimate better.The paper also identifies surrogate encoders as a possible direction for constructing black-box reaction coordinates.Source: Limitations · p. 9

A Conclusion

• The conclusion reports accurate, sample-efficient estimates across model scales and harmful-behavior categories, with greater robustness to event shifts than proposal-based methods.WE avoids the additional proposal-learning procedure and improves on direct Monte Carlo comparisons.Source: A Conclusion · p. 13

B Related Work

• WE generates samples under the original model rather than learning a proposal distribution, distinguishing it from proposal-based methods for fixed-prompt rare-event estimation.The related-work discussion also distinguishes rare-output prompt search and population-level risk forecasting from fixed-prompt completion-probability estimation.Source: B Related Work · pp. 13–14
• The paper presents applying WE to rare-event estimation in autoregressive language-model generation as a new use of the method.It notes that WE had not previously been applied to this setting, to the authors’ knowledge.Source: B Related Work · p. 14

C.1 Unbiasedness Estimator of Algorithm 1

• The prefix value function is the probability that a response beginning with a given prefix eventually satisfies the rare-event condition, and the ensemble’s weighted sum represents that probability mass.At the empty prefix it equals the overall rare-event probability; for a completed response it equals the event indicator.Source: §C.1 · pp. 14–15
• For any finite number of walkers, Algorithm 1’s estimator is unbiased under adaptive choices of coordinates, bins, and walker allocations based on information available at the current checkpoint.The proof establishes this by iterating expectation-preserving propagation and resampling steps.Source: §C.1 · pp. 14–16
• Propagation preserves expected represented rare-event probability because each prefix is extended by sampling from the original language model.The proof uses the prefix value function and iterated conditional expectation.Source: §C.1 · pp. 15–16
• Resampling preserves each occupied bin’s expected contribution, so changing binning and child allocation redistributes computational effort without changing expected probability mass.The proof conditions on the propagated ensemble before resampling.Source: §C.1 · p. 16

• Coordinate choice leaves estimator correctness unchanged; grouping prefixes with similar true future rare-event probabilities can reduce resampling variance without numerically predicting those probabilities.Source: §C.2 · pp. 16–18
• The proof treats the checkpoint estimate as a martingale, enabling total variance to be accumulated from the conditional variance introduced at each iteration.Source: §C.2 · p. 18
• The complete estimator variance sums, across checkpoints, propagation uncertainty under the language model and variance introduced by within-bin WE resampling.Source: §C.2 · pp. 18, 20
• Reaction coordinates and binning affect within-bin heterogeneity, walker allocation determines children per bin, and checkpoint intervals affect propagation uncertainty and the number of resampling operations.Source: §C.2 · p. 20

• The allocation condition is naturally met with a fixed number of bins and approximately uniform walker allocation; with at most M occupied bins, any a < 1/M suffices for sufficiently large N.Source: §C.3 · p. 20
• Under fixed horizon T and checkpoint count K, and an allocation of at least aN walkers to each occupied bin, the estimator variance is O(N^-1).Source: §C.3 · pp. 20, 22
• The variance bound does not require the reaction coordinate to approximate the optimal prefix value function, though a better coordinate may reduce actual variance.Source: §C.3 · p. 22
• For a fixed target rare event and WE procedure, increasing the number of walkers drives the estimation error to zero.Source: §C.3 · p. 22

C.4 Asymptotic Normality

• The analysis establishes a finite-horizon central limit theorem for the WE estimator as the walker count grows, deriving a limiting variance and asymptotic sample-complexity guarantees.The authors note that earlier weighted-ensemble analyses had not derived this result.Source: §C.4 · p. 22
• The theoretical binning assumption permits checkpoint-specific fixed maps, but practical k-means bins are data-dependent, so the variance expression is an approximation and diagnostic, not an exact characterization.A hybrid strategy that switches reaction coordinates at a predefined checkpoint remains compatible with the assumption.Source: §C.4 · p. 23
• Assumption 4 requires each limiting occupied bin to receive a positive fraction of walkers and makes allocation a deterministic function of the pre-resampling ensemble.Uniform allocation meets the positive-fraction condition when the occupied-bin set stabilizes; non-uniform allocation is also allowed if every positive-mass bin retains a strictly positive limiting fraction.Source: §C.4 · p. 23
• The finite-N variance is order N^-1, making N Var of the WE estimator the relevant quantity for its asymptotic distribution.The variance decomposition separates cumulative propagation and resampling contributions.Source: §C.4 · p. 26
• WE’s estimation error is decomposed into walker-level propagation terms and child-level resampling terms, each with conditional mean zero given the information available at its stage.Propagation terms are conditioned on the checkpoint information; resampling terms are conditioned on the propagated pre-resampling ensemble.Source: §C.4 · pp. 31, 33
• Ordering the independent propagation and parent-selection outcomes sequentially yields a martingale-difference array for the WE estimation error.The array has length L_N = (2K − 1)N, which diverges as N → ∞; the terminal checkpoint has no resampling step.Source: §C.4 · p. 35
• The proof verifies the conditional Lindeberg condition and then turns to convergence of the predictable quadratic variation through conditional second moments of propagation contributions.The supplied passage begins the variance analysis but does not provide its final convergence result.Source: §C.4 · pp. 35–36
• Under the stated assumptions and non-degenerate variance, the WE estimator’s fluctuations around p⋆ are asymptotically Gaussian, with variance decreasing as 1/N.The result strengthens the consistency result in Corollary 3.Source: §C.4 · p. 37
• The absolute estimation error, when scaled, converges to a halfnormal distribution.Corollary 4 gives a sampling-budget guarantee for absolute error at most ϵ with confidence approximately 1 − δ.Source: §C.4 · pp. 37–38
• The non-degenerate result excludes cases with σ^2_WE = 0, such as deterministic generation or events that always occur or cannot occur; these have no sampling variability.The asymptotic statement can be interpreted as the degenerate Gaussian N(0, 0) in these cases.Source: §C.4 · p. 37
• Corollary 4 sharpens the Chebyshev-based sample-complexity bound by using σ^2_WE, so reducing propagation and resampling uncertainty reduces the required number of walkers.The variance reflects the chosen WE configuration, including its reaction coordinate and allocation strategy.Source: §C.4 · p. 38

C.5 Analysis of Reaction-Coordinate Quality

• The asymptotic-variance decomposition separates the initial propagation contribution from resampling contributions at intermediate checkpoints.The analysis also identifies the population bin mass and allocation fractions used in the decomposition.Source: §C.5 · p. 38
• A useful reaction coordinate groups prefixes whose future rare-event probabilities vary little within each bin, reducing the heterogeneity relevant to the asymptotic-variance characterization.The analysis represents this heterogeneity by Hk,b and expresses the asymptotic variance compactly using checkpoint-wise contributions.Source: §C.5 · pp. 38–39
• The checkpoint contribution is approximated using descendant-weighted continuation values from trajectories already generated by WE, which are conditionally unbiased for the propagated prefix value.The argument relies on propagation preserving weighted future rare-event probability in expectation and resampling preserving expected descendant weight.Source: §C.5 · pp. 39–40
• Empirical Hk,b uses within-bin means and dispersion, while bin mass and allocation fraction are approximated by total weight and walker fraction.The estimated checkpoint contribution is computed for each reaction coordinate and independent run, then plotted across generation checkpoints.Source: §C.5 · p. 40

C.6 Sample Complexity Relative to Direct Monte Carlo

• WE and direct MC have the same asymptotic dependence on target error and confidence, but their sample complexities differ through their variance constants.The analysis derives the comparison from the sample sizes sufficient for the same absolute error and confidence.Source: §C.6 · pp. 40–41
• WE has lower asymptotic sample complexity than direct MC when its initial propagation variance plus cumulative reaction-coordinate-dependent variance is below direct MC’s asymptotic variance.Under the stated assumptions, this is the condition for the sample-complexity advantage.Source: §C.6 · p. 41
• Grouping prefixes with similar future rare-event probabilities reduces Hk,b and Ck, lowering asymptotic variance and increasing WE’s predicted sample-efficiency gain over direct MC.The initial propagation contribution U0,1 does not depend on the reaction-coordinate-induced partition.Source: §C.6 · p. 41
• The empirical test estimates WE’s asymptotic variance constant for different reaction coordinates and compares it with the direct-MC constant to assess the sample-complexity condition.The reported predicted gain is based on this variance comparison.Source: §C.6 · p. 41

C.7 Effect of Bin Refinement on Asymptotic Variance

• Under oracle-optimal walker allocation, refining the partition cannot increase the minimum achievable checkpoint-level asymptotic variance.This is an oracle characterization: the allocation depends on the unknown prefix value function and is not generally available as a practical rule.Source: §C.7 · p. 43
• Under oracle-optimal allocation, refinement can reduce variance by separating trajectories with different future rare-event behavior.The gain depends on appropriately reallocating walkers among refined bins.Source: §C.7 · pp. 43, 45
• Under uniform allocation, refinement lowers checkpoint-level asymptotic variance only when its reduction in weighted future-risk heterogeneity compensates for the increase in occupied bins.More bins divide the fixed walker budget; the tradeoff can make variance change nonmonotonically with bin count.Source: §C.7 · pp. 45–46
• Doubling the number of bins requires the weighted heterogeneity term to decrease by more than one half for refinement to improve asymptotic variance.The comparison assumes uniform allocation across occupied bins.Source: §C.7 · p. 45
• An unequal-mass split that reveals no additional future-risk structure strictly increases asymptotic variance under uniform allocation.Equality holds only for an even split.Source: §C.7 · p. 46

C.8 Robustness to Rare-Event Shift

• A single WE run’s terminal weighted samples provide unbiased probability estimates for any fixed finite collection of rare-event sets, without rerunning language-model generation.The same samples can be re-evaluated under revised event definitions.Source: §C.8 · p. 47
• For a finite collection of events, Proposition 15 gives the reused WE estimators’ joint asymptotic sampling distribution and covariance, recovering the scalar result when L = 1.Estimates from the same run are generally dependent because they use the same terminal walkers and weights.Source: §C.8 · pp. 47, 49
• For an enlarged event B, the variance of a reused proposal becomes arbitrarily large as its probability η = qA(B \ A) on the new region tends to zero.Here ρ measures the fraction of the revised event probability contributed by the new region.Source: §C.8 · pp. 50–51
• If an enlarged event adds a positive-probability region outside the support of the old ideal proposal, importance weighting cannot recover its probability, whereas stored WE samples can be rescored for the revised event.WE continues generating under the original language-model distribution, so the revised event does not remove the new region from sample support.Source: §C.8 · p. 52
• If lowering the threshold adds a region with positive base-model probability, the original-threshold ideal proposal assigns it zero probability and cannot recover the lower-threshold event probability by reweighting.A WE run is not tied to its propagation distribution through the threshold, so the same stored weighted ensemble can be rescored at a lower threshold.Source: §C.8 · p. 52

D.1 Evaluation Metrics

• Evaluation metrics cover rare-event probability estimates, hit rate, event ESS, relative error, zero rate, sample-efficiency gain, and computational-resource use.The resource metrics are wall-clock time, peak GPU memory, average GPU utilization, and average power draw.Source: §D.1 · pp. 52–53
• Direct Monte Carlo uses substantially more samples than the compared methods as a reference estimate, but it is not exact ground truth, especially for extremely rare events.Resource metrics are recorded per random seed and summarized across seeds using the mean and standard deviation.Source: §D.1 · p. 53

D.2 Implementation Detail

• Toxic-story experiments use TinyStories-33M and ToxiGuardrail, while jailbreak experiments evaluate Llama-3.1-8B and Qwen2.5-7B-Instruct on StrongREJECT.Toxicity is scored as the negative non-toxic logit; StrongREJECT harmfulness scores range from 0 to 1.Source: §D.2 · pp. 54, 56
• WE constructs intermediate bins using the rare-event score, LLM internal representations, or a hybrid that begins with representations and switches to score-based coordinates at a predefined generation step.Source: §D.2 · p. 55
• Score-based bins combine adaptive quantile boundaries with fixed score boundaries, allocating coarser bins to common low-score prefixes and finer resolution toward the rare-event region.Quantile levels are 1 − 2^-j for j=1,…,11; the fixed boundaries preserve separation near the target threshold.Source: §D.2 · pp. 54–55
• Hidden representations are L2-normalized, projected to 32 principal components, and clustered separately at each checkpoint using 12 k-means clusters.Checkpoint-specific clustering accommodates changing representation distributions during generation.Source: §D.2 · pp. 55–56
• The rare-event-shift experiment reuses terminal samples after lowering event thresholds, without further language-model generation.WE re-evaluates terminal event indicators; proposal baselines reuse samples and importance weights learned for the original threshold.Source: §D.2 · p. 57
• Ablations vary bin counts, checkpoint intervals, walker counts, and latent-space binning methods while holding other hyperparameters fixed.Alternatives to k-means include diagonal and full-covariance GMMs, Ward clustering, and latent quantile binning.Source: §D.2 · p. 58

E.1 Toxic Story Generation

• 10.0% relative error for the first prompt beats 46.6% for Twisted SMC and 31.1% for LMTailRisk; the second prompt reaches 1.0% versus 10.5% and 34.7%.Across the third prompt, WE variants range from 0.6% to 3.7%, compared with 7.5% and 3.2% for the proposal-based methods; selected hidden-representation variants achieve errors between 11% and 21% on the fourth prompt.Source: §E.1 · p. 58
• The illustrated WE trajectory periodically merges walkers in similar regions and splits promising prefixes, with hidden-representation clustering before token 8 and score-based bins afterward.In the example, both terminal descendants with scores above γ = 7 enter the target event.Source: Figure 3 · p. 59, §E.1 · pp. 61–62
• Hybrid representation coordinates produce lower checkpoint-wise variance contributions than reward-only binning over most checkpoints and across the four prompts.Variance contributions often drop around or after the switch to reward-based binning, consistent with representations helping earlier and the rare-event score becoming more informative later.Source: Figure 4 · p. 59, Figure 5 · p. 60, Figure 6 · p. 60
• Reward-only WE has relative errors of 87.6%, 49.7%, and 42% on the first, second, and fourth prompts, while random binning yields zero-estimate rates of 95%, 91.2%, and 72.5% on several prompts.Layer 4 all →R @ 8 has relative errors of 10.0%, 4.2%, and 3.7% on the first three prompts and is comparatively reliable across prompts.Source: §E.1 · pp. 59–60
• Selected hidden-representation WE configurations achieve asymptotic sample-complexity gains greater than 1 across all four prompts, while random-coordinate gains range from 0.15× to 0.33×.Reward-only WE is less consistent, with gains below 1 on the second and third prompts.Source: §E.1 · p. 60
• Qualitative examples show WE reaching rare-event regions with fewer samples than direct Monte Carlo.The depicted middle branch approaches but does not reach the threshold, whereas the bottom branch produces two terminal walkers above γ = 7.Source: §E.1 · pp. 61–62

E.2 Jailbreak Attacks on StrongREJECT

• Reward-only WE is competitive with proposal-based baselines and has lower error than both on some Llama-8B categories and several Qwen-7B categories.The hybrid reaction coordinate improves consistently over reward-only WE across the evaluated categories and models.Source: §E.2 · p. 62
• WE (hybrid) achieves the lowest relative error across all six harmful categories on both Llama-8B and Qwen-7B.For Qwen-7B hate/harassment, hybrid WE reaches 26.8% error versus 100.0% for SMC and 73.4% for LMTailRisk; for Llama-8B violence, it reaches 21.6% versus 100.0% and 129.5%.Source: §E.2 · p. 62
• Hybrid WE often achieves higher hit rates and larger effective sample sizes than reward-only WE, alongside lower probability-estimation error.The results span six harmful categories and two language-model families.Source: §E.2 · p. 62
• WE runs in 89.95 minutes on Llama-8B versus 109.27 for SMC and 176.12 for LMTailRisk, and in 87.6 minutes on Qwen-7B.Peak memory is around 23 GiB across methods, and WE avoids auxiliary proposal learning while maintaining approximately 94% GPU utilization.Source: §E.2 · pp. 62–63

E.3 Robustness to Rare Event Shift

• On the third prompt, Twisted SMC error rises from 7.5% to 60.9% and LMTailRisk from 3.2% to 58.1% when the threshold falls from 5 to 3.Layer 4 all →R @ 8 remains between 0.8% and 3.7% across the evaluated thresholds, while Layer 2 final concat →R @ 12 remains between 2.0% and 2.7%.Source: §E.3 · p. 63
• Layer 4 all achieves relative errors of 3.7%, 2.8%, 0.8%, and 1.2% across the third prompt’s four evaluated thresholds.It stays within 10.0% relative error at four of five thresholds for the first prompt.Source: §E.3 · p. 63
• Reward-only WE reaches 87.6% relative error at γ = 7 on the first prompt, while selected hidden-representation variants reduce error to 10.0%–25.0%.The authors attribute the broader robustness of hidden-state-guided WE to its usefulness when partial-generation reward is weakly informative, followed by a later switch to reward-based guidance.Source: §E.3 · p. 63

E.4 Ablation Study

• Relative error fell from 35.6% with 4 bins to 10.0% with 12, then rose to 91.2% with 16 and 94.3% with 24 bins.The authors attribute the deterioration with finer partitions to spreading the fixed walker budget across more bins.Source: §E.4 · pp. 63–64, Table 30 · p. 80
• Relative error was larger for diagonal GMM, full-covariance GMM, agglomerative Ward, and latent quantile binning than for default k-means clustering.The tested alternatives model additional geometric or distributional structure, but did not improve the bins’ alignment with future rare-event behavior.Source: §E.4 · p. 64, Table 30 · p. 80
• Relative error decreased from 67.4% at 2-token checkpoints to 46.9% at 3 and 10.0% at 4, then increased to 14.0% at 5 tokens.The authors describe a tradeoff: prefixes need time to accumulate semantic information, but longer intervals leave fewer opportunities to redirect computation.Source: §E.4 · p. 64, Table 30 · p. 80
• Relative error fell from 95.5% with 12,048 walkers to 10.0% with 14,336 and 2.6% with 16,144, while 18,192 walkers yielded a 2.5% zero rate but 16.0% relative error of the mean.At 12,048 walkers, the zero rate was 57.5%, indicating that this budget was insufficient for reliably reaching the target event.Source: §E.4 · p. 64, Table 30 · p. 80