Research map / Learning theory and algorithms / Sample complexity
Estimating the Probabilities of Rare Outputs in Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies low probability estimation for argmax outputs of small transformer LMs under independent-token input distributions, comparing importance sampling (ITGIS, MHIS) with activation extrapolation (QLD, GLD) against naive sampling.
Problem
Estimating the probability that a language model outputs a target token under a formally specified input distribution when that probability (1e-9 to 1e-5) is too small for random sampling within a fixed budget.
Approach
Two families of estimators: importance sampling that up-weights inputs likely to produce the rare token (gradient-based independent-token proposal ITGIS, and Metropolis-Hastings MCMC over a Boltzmann posterior MHIS), and activation extrapolation that fits distributions to pre-unembed activations (Quadratic Logit Decomposition, which uses a presumed independence between whitened subspaces, and a Ga
Key results
- Averaged over distributions, MHIS (1.48 IS loss on 4-layer) and ITGIS (1.92) outperform QLD (2.16) and GLD (2.48), with constant baseline at 2.62 on the 4-layer model.
- ITGIS does better on the 1-layer model (avg 1.37 vs MHIS 1.55) while MHIS does better on the 2- and 4-layer models, which the authors attribute to larger models having more inter-token interactions.
- Both importance sampling and QLD use a budget of 2^16 model calls, and the authors argue naive sampling cannot give positive estimates for probabilities below 1e-5.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use importance sampling (ITGIS or MHIS) instead of activation extrapolation (QLD or GLD) | Low-probability estimation · Language-model rare-token behaviors | Itakura–Saito loss | Helped |
| Use QLD instead of GLD | Activation extrapolation · Language-model rare-token behaviors | Itakura–Saito loss | Helped |
| Use an optimal constant estimate instead of GLD | Low-probability estimation · Language-model rare-token behaviors | Itakura–Saito loss | No clear effect |
| Use MHIS instead of ITGIS | Low-probability estimation · Smaller models | Itakura–Saito loss | Hurt |
| Use MHIS instead of ITGIS | Low-probability estimation · Larger models | Itakura–Saito loss | Helped |
| Use ITGIS rather than MHIS | Low-probability estimation · spanish distribution | Itakura–Saito loss | Mixed |
| Use MHIS rather than ITGIS | Low-probability estimation · hex and icl distributions | Itakura–Saito loss | Mixed |
| Increase QLD's sample budget beyond 2^16 | QLD · Language-model rare-token behaviors | Itakura–Saito loss | Helped |
| Score methods using squared error in log-space instead of Itakura–Saito loss (ablation) | Compared estimation methods · All tested distributions and model sizes | Squared error in log-space | No clear effect |
Surprising findings
- Activation extrapolation (QLD) is outperformed by importance sampling in this setting, even though importance sampling essentially reduces to red-teaming and adversarial training.
- QLD plateaus in performance with more samples than its standard budget, which the authors attribute to its independence assumption rather than sample size.
Limitations
- Only independent-token input distributions are studied, which limits generality.
- Only single-token behaviors at temperature 0 are evaluated, which is unrealistic for real-world harms.
- Experiments use small transformers (1-, 2-, and 4-layer) rather than production-scale models.
Experimental setup
Datasets: C4 CodeParrot
Benchmarks: 8 synthetic input distributions (hex, camel, colon, if, caps, english, spanish, icl) over small transformer models
Metrics: Itakura-Saito loss squared error in log-space leave-one-out cross-validation loss
Baselines: Naive sampling (optimal constant) Gaussian Logit Difference Quadratic Logit Decomposition
Builds on: Greedy Coordinate Gradient adversarial training Adaptive Multi-Level Splitting
Closest papers by meaning
- A Probabilistic Perspective on Unlearning and Alignment for Large Language Models (ICLR 2025 · similarity 0.87)
- Large language models suffer from a curse of ambiguity (NeurIPS 2026 · similarity 0.86)
- Learning Parametric Distributions from Samples and Preferences (ICML 2025 · similarity 0.86)
- Statistical Estimation of Adversarial Risk in Large Language Models under Best-of-N Sampling (NeurIPS 2026 · similarity 0.86)
- Reducing the Probability of Undesirable Outputs in Language Models Using Probabilistic Inference (NeurIPS 2025 · similarity 0.86)
- Efficient Prediction of Pass@k Scaling in Large Language Models (NeurIPS 2026 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.