Estimating the Probabilities of Rare Outputs in Language Models
Authors: Gabriel Wu
Jacob Hilton
Venue: ICLR 2025 Spotlight (2025)
Keywords: low probabilities
adversarial training
importance sampling
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We present methods for estimating the probability that a language model outputs a rare token on a given input distribution.

Abstract: We consider the problem of *low probability estimation*: given a machine learning model and a formally-specified input distribution, how can we estimate the probability of a binary property of the model's output, even when that probability is too small to estimate by random sampling? This problem is motivated by the need to improve worst-case performance, which distribution shift can make much more likely. We study low probability estimation in the context of argmax sampling from small transformer language models. We compare two types of methods: importance sampling, which involves searching for inputs giving rise to the rare output, and activation extrapolation, which involves extrapolating a probability distribution fit to the model's logits. We find that importance sampling outperforms activation extrapolation, but both outperform naive sampling. Finally, we explain how minimizing the probability estimate of an undesirable behavior generalizes adversarial training, and argue that new methods for low probability estimation are needed to provide stronger guarantees about worst-case performance.

Published as a conference paper at ICLR 2025
ESTIMATING THE PROBABILITIES OF RARE OUTPUTS IN LANGUAGE MODELS
Gabriel Wu" Alignment Research Center
Jacob Hilton
ABSTRACT
We consider the problem of low probability estimation: given a machine learning model and formally-specified input distribution, how can we estimate the prob- ability of a binary property of the model's output, even when that probability is too small to estimate by random sampling? This problem is motivated by the need to improve worst-case performance, which distribution shift can make much more likely. We study low probability estimation the context of argmax sampling from small transformer language models. We compare two types of methods: im- portance sampling. which involves searching for inputs giving rise to the rare out- put, and activation extrapolation, which involves extrapolating probability dis- tribution fit to the model's logits. We find that importance sampling outperforms activation extrapolation, but both outperform naive sampling. Finally, explain how minimizing the probability estimate of an undesirable behavior generalizes adversarial training, and argue that new methods for low probability estimation are needed to provide stronger guarantees about worst-case performance.
INTRODUCTION
Modern ML systems undergo black-box optimization to minimize a loss function on samples drawn from training distribution. Although models produced in this way perform desirably average over this distribution, they can still produce highly undesirable outputs on very rare inputs. This is a problem, because these rare inputs can become much more likely in the presence of distribution shift, especially one chosen adversarially, such as with large language model "jailbreaks" (Carlini
al. 2024. Wei al. 2024).
Preventing such highly undesirable outputs is notoriously challenging problem. The most common remedy is adversarial training, in which inputs that produce these undesirable outputs are searched for and used additional training data (Goodfellow et al. 2014, Madry (017). but the transfer between different search methods generally weak (Kang et al. 2019 Wei et al. 2024). In this work, we propose the more modest goal of simply estimating the probability that an input drawn from some distribution will produce a certain kind output, which has been considered before in the context of computer vision in Webb et al. (2019). We will show that even this intermediate goal challenging, but successful methods could enable new ways of preventing undesirable outputs by minimizing their estimated probability.
To advance work on this problem, we study low probability estimation in the context of small trans- former language models. We consider various formally-defined input distributions in which each input token sampled independently, and develop methods for estimating the probability that a particular target token will have the largest output logit. We constrain the computational budget of our methods and obtain ground truth probabilities by random sampling using a much larger compu- tational budget. The target tokens chosen have ground truth probabilities between 10-9 and 10-5, which are too small for random sampling to produce good estimate under the constrained computational budget.
In this context, we study two types of methods:
"Correspondence gabriel .d. wu31 1003144mmm 44gma com

---

Published as a conference paper at ICLR 2025
P (xx treee)
10
Ground Truth Hers
Calculate true probability with brute-force naive sampling (Slow, uses n somples)
Estimate probability with each LPE method (Fast, uses 2 somples)
10
Estmate
50-7
Apply learned fit ox +
10
Compare probabilities with Ttakura-Saito
(Repeat for 256 different target tokens ("tree", ath",, ,...)
Figure 1: Left: To evaluate the performance of our low probability estimation methods, compare their estimates against ground-truth probabilities obtained by brute-force sampling with a larger computational budget. Right: The estimates of Metropolis- Importance Sampling the icl input distribution and 4-layer model, after a fit has been applied. Each point represents different target token.
Importance sampling. We define new input distribution under which the rare event is much more likely, sample from that distribution, and re-weight samples to obtain an un- biased estimate for the original distribution. Our Independent Token Gradient Importance Sampling (ITGIS) method treats token positions independently and uses gradients to obtain this new input distribution, while our Metropolis-Hastings Importance Sampling (MHIS) method uses a Markov chain Monte Carlo algorithm to sample from a distribution with non-independent tokens.
Activation extrapolation. We random samples fit probability distribution the model's logits, and extrapolate into the tails of this distribution to produce probability estimate. Our Quadratic Logit Decomposition (QLD) method applies presumption of independence to the empirical distribution of logits, motivated by Christiano et al. (2022). and our Gaussian Logit Difference (GLD) method is simple baseline that fits a Gaussian to the difference between the maximum logit and target logit.
In our setting, both types of methods outperform random sampling, and importance sampling tends outperform activation extrapolation. Nevertheless, remain interested in activation extrapola- tion and similar approaches because they produce new methods for reducing the probabilities of rare outputs, whereas importance sampling essentially recovers standard adversarial training.
The remainder of the paper is structured as follows. In Section 2 we formally define the problem of low probability estimation, both in general and our language model setting, Section 3 we describe our four methods in more detail. In Sections 4 and 5 we describe the models and input distributions on which we test our methods and convey our experimental findings. Finally, in Sections 6 and g we discuss the limitations and implications of our results, related work, and future directions.
PROBLEM STATEMENT
Given an input space X, an output space V, an input distribution E A(X), a model M : X + V, and formal boolean property of model outputs C {0, 1), low probability estimation the
2

---

Published as a conference paper at ICLR 2025
problem of efficiently estimating
Pr [x(M())) = xnD
We sometimes refer to the event (x)) probability of the behavior large enough (say, larger than 1/n), it is easy to estimate by drawing n independent samples from X' and using the sample mean of (x)(. However, if the probability is significantly smaller than 1/n, this sample mean is almost always 0, making it uninformative at distinguishing between small probabilities like 10-10 100 and 10-20.
the "target behavior", or just "the behavior." If the
2.1 OUR SETTING
In this paper, we study low probability estimation in the setting of argmax sampling from language models with single-token behaviors. Let M V + predicts the next token given string of previous tokens, where V is the token vocabulary. Note that we sample at temperature 0, so M deterministic. Given a distribution D over V and a target token t V. the low probability estimation problem for single-token behaviors is the task of estimating
a transformer language model that
Pr (M( = f].
x~D
Letting M.(x) be the logit the model assigns to token i E V. this can also be written as: Pr [M)x) > (x) Vi + tl.
X~D
In general, D can be any distribution that can be formally specified. However, in this paper we focus only distributions D with independent tokens. That is, we specify input length & and token distributions .,,..., Dr E A(V), then write as the product D1 x ... x DD.- Table [ shows the tested, with tokens colored for clarity. To prevent overfitting, the methods
8 distributions that were only run on the first four distributions during development, and they were finalized before testing the last four distributions. The results were qualitatively the same both halves of the split.
Table Input distributions and examples. See Table for more detailed descriptions.
Name
Short description
Tokenized example
hex
Hexadecimal characters
<|aus hatal//aaaaaiieeeaaa...e.e.eaae.
camel
CamelCase Python tokens
< [B05 https///wpp....eeeeeeeeeeeeeeeae
slotDeid
colon
Python tokens, ending with ":"

 \rtn
**"); ('n Thisc
1.2dafault..v
self deft)",
if
Python tokens, starting with if
< 
 if: else, postia la
24 found
fromout, self node 1-)a VA (do
AA this values
caps
"He/She sereamed:". followed by caps and punctuation English words
<|308| screamed: "ESOTTOLEBOV, NB! IIIILEER. ,RY
...IIBSSION
english 1iglish
<|RG8 
sting 1a invent School found from one shocting sveryone Coz George around responsive employees ground on stone various,
spaniss sh
Spanish tokens
< 
 To no bu dees or soct josant ler de enidad sreljd final de de la much
icl
In-context learning prompt
<I805 for Amezicca R for Return C for crack
for troubles R for house B for equipment 0 for operating for reason for your for
ESTIMATION METHODS
We introduce four methods in this section: two importance sampling methods (Independent Token Gradient and Metropolis-Hastings)) and two activation extrapolation methods (Quadratic Logit De- composition and Gaussian Logit Difference). We also compare against the baseline of outputting
3

---

Published as a conference paper at ICLR 2025
optimal constant, which can be thought of as the performance of naive sampling because only evaluate the methods on tokens with ground truth probabilities less than the reciprocal of the allotted sampling budget (see Section 4
3.1 IMPORTANCE SAMPLING METHODS
Naive sampling fails to produce good estimates for low probability events because takes too many samples from D to observe positive example. To address this, we can instead draw samples from a different distribution that up- up-weights regions of input space most likely to produce the behavior of interest. If we re-weight our observations properly, this gives an unbiased estimator for the true probability. This known as importance sampling, and enjoys the same advantages that adversarial training has over standard training: by using a narrower input distribution, we can more efficiently discover positive examples of the target behavior.
Formally, let p(æ) be the probability mass function of D, and let g(x) be the PMF of any other distribution. Then
p(x) g(x)
[M(x =1]
(r(M(x) -1 [ =]]==
xeup
but the latter may have less variance (and so require fewer samples get a good estimate).
The following two importance sampling methods take (xx) to be Boltzmann posterior with prior p(x). The first defines q(x) with independent tokens, while the second defines g(x) have non- independent tokens and so requires a more sophisticated sampling method.
3.1.1 INDEPENDENT TOKEN GRADIENT IMPORTANCE SAMPLING (ITGIS)
We want g to up-weight tokens that contribute to being outputted. One way to do this is continue to treat each input token as independent, but change the probability of tokens according to their average linear contribution to the logit of Let (I1...... I²) V be an input of length and
say that p(x) factors as (x1) ... p& (*x). Then we define ((x) = qu (æ1) ... qx(x)), where
(x(
qi (xr4) pe(xs) exp
and
(x,)) = Exing V M, (x))....
T's is temperature parameter, and the gradient is taken by treating x' a one-hot vector in R Intuitively, the gradient Vx gives us a linear approximation to how much the logit of t would change if we replaced i-th token of x' with z; (up to an additive constant w.r.t x2). Thus, s scores each token value according to its average linear contribution to M,, and qi is defined as the Boltzmann distribution with respect to this score function!
However, since and are both defined in terms of each other, we can't calculate s; directly. To
overcome this, construct sequence of score functions (0) (11) and associated distributions (0) (T)) .... that adaptively refined with respect each other (see Appendix B for details). Sampling from each (()) lets calculate importance sampling estimate, and the final output of the method is the average value of these estimates across all j.
3.1.2 METROPOLIS-
IMPORTANCE SAMPLING (MHIS)
A problem with ITGIS is that the new sampling distribution g(x) still treats all tokens as indepen- dent, and only accounts for linear effects of tokens on the target logit. Thus, ITGIS may fail to sample into the most important regions of the input space if the model is sensitive to non-linear interactions between tokens (e.g., if the model's target logit is only high when the last two tokens of the input the same each other).
It can be shown that, given score function s(z) and prior p(z), the distribution that maximizes Ex- [x(x) T (T((glp) is q(x) x p(x) exp(s(x)/T).
4

---

Published as a conference paper at ICLR 2025
To remedy this, can define importance sampling distribution that doesn't 't have independent tokens. We must use a score function that depends on the entire input; the most natural choice the target logit Mx)... We define
M1(a) T
((x) p(x) exp
again using Boltzmann distribution to up-weight regions of input space that are more likely to have positive samples.
Unlike ITGIS, we cannot explicitly compute because it does not factor into independent distribu- tions over each token. Instead, we use the Hasssng algorithm produce random walk in input space that has stationary distribution of q To do so, we must define a proposal distri- bution áia'|as [æ)) that suggests the next element of the walk. To encourage fast mixing, this proposal distribution should be good at exploring into regions of input space that q weights highly.
Here we take inspiration from Greedy Coordinate Gradient, an algorithm that optimizes a discrete prompt to jailbreak model using gradients (Zou et al. (023). We adapt this optimization procedure into a proposal distribution: to pick proposed next step of the walk, we choose a random token position to replace, compute the gradient of s sx) with respect to then sample a replacement token for position i according to a Boltzmann distribution defined by this gradient (similarly to ITGIS). The final output of the method is the average importance sampling estimate taken after burn-in period. For a precise description of the algorithm, see Appendix Bli
3.2 ACTIVATION EXTRAPOLATION METHOSS
The importance sampling methods search for explicit examples of inputs that cause the given behav- ior. This makes their task at least as hard as the adversarial training search problem-if it is difficult to find an supp(D) such that (xx) L, the importance sampling estimators will likely fail to produce positive estimate.
We hope find low probability estimation methods that work even when the search problem for importance sampling is hard. To do this, we introduce activation extrapolation: first fit a distribution the activations or logits of M, then estimate the probability of the output property of interest under this idealized distribution. Our first such method is Quadratic Logit Decomposition, which applies presumption independence between uncorrelated subspaces the model's pre-unembed activations. We also develop Gaussian Logit Difference, which is intended as simple baseline method.
3.2.1 QUADRATIC LOGIT DECOMPOSITION (QLD)
Let the random vector v(x) E Rd be the activation of the model right before applying the unembed matrix Wy dV), That is, v(x). Wy represents the model's output logit vector M (x))....Iv)-
We first collect samples of (call them v(1) ..... v(n).. (nn). We then choose some unit direction d E Rdd (see below), then decompose each v into a + ((),, where a lies in the subspace spanned by d, and b lies in the complementary subspace that is orthogonal in whitened basis
This decomposition is chosen such that the random vectors a and b are uncorrelated across the samples.
Next, by treating the random vectors a and b as independent, we can use our samples of each to obtain n²2 "synthetic" samples of u. The final output of QLD the proportion of these synthetic samples that cause ! to be outputted:
{ (i.5) [n]² a + +b00 E s}|
where 5 C IRd is the "acceptance region" of activation space corresponding to activations that result in the target logit being highest after unembedding. Despite the fact that there are n² synthetic samples, this proportion can be computed in Ö(n) time by first sorting the samples ( A more complete description of the QLD algorithm can be found in Appendix B
Metropolis- Hastings is Markov Chain Monte Carlo method for sampling from distribution with an unknown normalizing constant. Sce Robert (2016) for deseription of the algorithm. Actually, a, b, and d all defined in the whitened space of v; see Appendix B
5

---

Published as a conference paper at ICLR 2025
Choice of direction. We rely on the following two assumptions for QLD to perform well: 1) and b are independent (so that our estimate is unbiased), and 2) the contribution towards the output behavior is split roughly equally between these two terms (to minimize the variance of our estimate). See Appendix [C] for more discussion of this motivation. After some initial experimentation with a variety of candidate directions [ we decided to set d be the direction of the shortest vector in whitened space that results in the model outputting t. It can also be thought of as the maximum like- lihood value under a Gaussian prior, conditioned on observing the model output Appendix D describes the algorithm we use to compute d.
3.2.2 GAUSSIAN LOGIT DIFFERENCE
On any given input, we can record the difference A, := M,(x) max; (x). We wish to estimate the probability that A 0. A natural estimation method, which we view as simple baseline, is to treat A, as Gaussian by estimating its mean μ and standard deviation with samples, then calculate Pr[N(p, o²) 2 0j. In practice, we use a slightly different functional form that captures the Gaussian PDF, which approximates the CDF well in the tails. The output of the Gaussian Logit Difference method is:
exp (-(o0t au
2 +
+
where a, c, and e are parameters that fit minimize loss across all target tokens associated with given distribution (see Section 4.2)
4 EXPERIMENTAL SETUP
We apply our methods on three models: l-layer, a 2-layer. and a 4-layer transformer from Nanda & Bloom (2022). All models have a hidden dimension of d 512, vocabulary size of [V] 48262, GELU lineaiieesssies (Hendrycks & Gimpel| 2023), and were trained on the C4 dataset (Raffel ct al. 2023) and CodeParrot (Tunstall et al. 2022).
For each of the 8 distributions (listed in Table and for each model, we generate ground- -ruth token probabilities by running forward passes on 2- random samples. We then select a random set 256 tokens among those with ground-truth probabilities between 10-9 and 10-5, and we test all of our methods on these tokens.
We give each method computational budget of 266 model calls (see details in Appendix F). This budget was chosen so that naive sampling would almost never result in any positive estimates for the range of token probabilities we test (216 < 10a), but the theoretical quadratic gains from QLD would still be enough to get signal on the entire range of probabilities ((21)) > 109).
Our
code is available
httpps://githbb.com com. wwt..///www.........o.oo
low iittabbbiittyyesetmati..
4.1 ITAKURA-SAITO Loss
We measure the quality of the method with loss function inspired by the Mtakura-Saito divergence (Itakura & Saito 1968). If is the ground-truth probability of a particular target token, then an estimate of q incurs a loss of:
Dis(p, q) 9
P- 2-
In
Two considerations went into the choice this loss function. First, Itakura Saito loss a proper scoring rule (Buja et al. 2019). Second, since it only depends the ratio p/q, Itaura--aaite loss sensitive to small probabilities: if p 10 10-100 and q 10-10, then Dis (p. q) is very large. In contrast, the squared error loss function (p q)2 would be extremely small. Intuitively, this sensitivity is desirable because we care how our methods perform on a wide (as measured in log- space) range of ground-truth probabilities. We don't want the performance metric be dominated by a method's s behavior on only the most probable tokens.
*Other candidate directions included 1) the f-th column of Wu pulled hack into whitened space and 2) the expectation of (0 Ida) conditioned lying in the whitened acceptance region.

---

Published as a conference paper at ICLR 2025
For completeness, we also report our results using squared error in tpacs (Appendix H even though this is not a proper scoring rule. The results are qualitatively identical.
4.2 AFFINE FITS
Many methods often report estimates of but Dis is undefined for q To address this, we fit a transformation x + az + 6 to the oulputs of each method, where a, b and are chosen to minimize Itakura-Sail Saito loss. ar can be thought of an affine transformation in log-space, and adding b prevents values from being too small while barely affecting larger outputs. To ensure that this transformation is not overfitting to the particular set of 256 tokens, we report the leave-one-out cross-validation (LOOCV) loss of each method. We train separate fit for each (method, input distribution) pair.
RESULTS
Figure 2 shows the performance of each method. The relative ordering is clear: both importance sampling methods outperform Quadratic Logit Decomposition, which in turn outperforms Gaussian Logit Difference. GL.D is barely better than outputting an optimal constant (which can be interpreted the performance of naive sampling). Figure 4 shows that there is fair amount of variation in method performance across the 8 distributions: some behaviors like hex and i icl favor MHIS, while others like spanish heavily favor ITGIS. more detailed table of results in Appendix G
Among the two importance sampling methods, ITGIS does better on smaller models, while MHIS does better on larger models. We believe this is because larger models are less easily approximated as linear functions and are more likely to have complex behaviors arising from intertooken interactions.
Figure 3 displays example scatter plots of ITGIS, MHIS, and QLD estimates before a fit is applied. Each point represents the ground-truth and estimated probability of a different target token. More scatter plots can be found in Appendix D note that the qualitative performances of the methods can vary significantly on different input distributions. We perform an ablation study on our choice of loss function in Appendix H in which we score methods based on squared error in log-space instead of Itakura-Saito loss.
6 DISCUSSION
6.1 DISTRIBUTION SHIFT AS MOTIVATION
One might ask: if a particular model behavior is rare that it never arises during training, why would we care about estimating its probability? There are a few reasons. First, some systems may be run many more inputs during the course of deployment than during training. Thus, if certain model behavior would be so catastrophic that it is unacceptable for it to occur even once in deployment, we cannot rely training drive down its probability low enough. Second, there may be distributional shift between training and deployment such that events that occur extremely rarely during training become more likely deployment. This could occur because of an input chosen adversarially, but it could also occur because of goal misgeneralization (Shah et al. 2022).
particularly challenging case is deceptive alignment, the possibility that an ML model would look for clues about whether it is in a training or deployment environment, and only behave well in training (Hubinger et al. (02I). To detect whether model is deceptively aligned, one could craft an input distribution that is "wide enough" to assign some probability mass, even if very small, to any possible deployment-time input, then apply low probability estimation methods to detect if the model would ever perform catastrophic behavior on this distribution 6 For more discussion of this idea, see Xu (2024).
NNte that the Gaussian Logit Difference method has a special functional form of its fit ((1, + exp (- (ap1(aa + +))² + + e instead of æ ++ ar + b) but is otherwise evaluated the same way. To prevent false positives, this would require very demanding definition of catastrophe that would be impossible for the model to trigger "by accident."
7

---

Published as a conference paper at ICLR 2025
2.5
Average LOOCV bos (raaurr-
1.-
Optimal Constant
difference
ndependert loca Gradiant Importance Sampling
Matmmpolis hatings importance Sampling
Quedhatic Logit Decompesiton
Number of
The Itakura-Saite loss of all methods across different model sizes. The solid lines indicate the loss of each method averaged over all distributions, with bands showing standard error. The colored points indicate the loss on individual distributions, with horizontal jitter added for visibility. Lower better.
Figure 2:
caps behavior, layers
10 Ground
MHIS
maG
10-4
QuD
stnts (remss)
ITaIs (zeros)
tzeresi
Pstimato
30
10:1
10.5
30
10-40
10
10-9
30
30
10
10
10
30
10
30
10"
10
10
Fimund Truth
Gimnndl Thith
Chnire Truth
colon behavior, 2 layers
Dround
Ground
10
Ground Truth
NIBS Lzeros)
MeiS
ITGIS
mas
10:5
OLD
DED
10-4
20
10-<
Estimate
Pateatte
10
30-1
10-4
10-f
20
10
10
20
30
10
30
10
30
10
10
10
Ground Tarth
Ground Tichh
Goound Truth
Figure 3: Examples of method outputs on two different behaviors and models, before fit is applied. Estimates of 0 are placed at the bottom of each graph for visibility.
8

---

Published as a conference paper at ICLR 2025
LODCV (itakura
1.5
1.0
Optimal Constant
0.5
ITGIS
MHIS
0.0
hex
camel
colon
caps
engits
spanish
Input Distribution
Figure 4: The Makura-Saitt loss of all methods across different distributions, averaged over all 3 model sizes. Lower better.
6.2 RELATION TO RED-TEAMING AND ADVERSARIAL TRAINING
Our importance sampling methods for low probability estimation involve finding inputs for which the rare event occurs. This amounts to the well-studied task of "red-teaming". As long as the required importance sampling ratios can be computed, any method for red-teaming can be turned into an importance sampling method for low probability estimation, as we demonstrate with our adaptation of Greedy Coordinate Gradient into MHIS (Zou et al. 2023). However, our activation extrapolation methods such as QLD do not correspond to any red- oed-tenming method.
further reason to interested low probability estimation that it could be used to reduce the probability of the rare event, by optimizing the model to produce a lower estimate. For example, this could be done using gradient descent, if the estimate were a differentiable function of the model's parameters. For importance sampling method, this amounts to finding inputs for which the rare event occurs (i.c., red-teaming) and using them as training data, which is essentially the well-known method of adversarial training (Goodfellow et al. 2014). However, since our activation extrapola- tion methods do not correspond to any red-teaming method, new activation extrapolation methods potentially provide us with new ways reduce the probabilities of rare events.
6.3 IMPORTANCE SAMPLING VERSUS ACTIVATION EXTRAPOLATION
In our experiments, we found that importance sampling methods outperformed activation extrapo- lation. Nevertheless, there are theoretical cases in which importance sampling performs worse than other methods. For example, consider model that outputs the SHA-256 hash of its input: find- ing any input that gives rise to a particular output is computationally infeasible, yet it is still easy estimate the probability of a particular output by modeling the output of the hash function as random.
More generally, we are excited about low probability estimation as a concrete problem for which for which it may be necessary to leverage internal model activations. In place of importance sampling, we may able to use deductive estimates based on a presumption of independence (Christiano) et al. 2022). Our Quadratic Logit Decomposition method is early proof of concept of this, even though it is outperformed by importance sampling in our setting.
6.4 LIMITATIONS
There are two main limitations of our experimental setup. First, we only use input distributions that factor into independent tokens. This choice is necessary for the definition of ITGIS. It is also

---

Published as a conference paper at ICLR 2025
very convenient for the implementation of MHIS, because it gives efficient sampling access to the proposal distribution. To move beyond independent token input distributions, we could define the input distribution to be the output of separate generative model and adapt some of the current estimation methods appropriately.
Second, only study model behaviors that consist of a single token sampled temperature This unrealistic because in practice, we were concerned about specific single-token outputs, it would be easy to filter them out. In contrast, the types of behaviors we actually worry about likely involve long chains of autoregressive generation or interaction with the external world (e.g., when forming and executing a plan). We are excited to see future work extending our setting in this direction.
Nevertheless, it is worth noting that formally-defined distributions and behaviors are more general than they may initially seem. For example, we could formalize the event "M writes buggy code", as: When M's 's output is given to GPT-4 along with the prompt "Does this code contain any bugs? Let's think step by step.". does GPT-4 end its response with YES?
7 RELATED WORK
The problem of low probability estimation was previously considered in the context of computer vision by Webb et al. (2019), where they propose using an Adaptive Multi-Level Splitting algorithm with Metropolis Hastings. However, they only study the problem in the context of computer vision with continuous input spaces, and their approaches still require finding positive samples, unlike our activation extrapolation methods. |Phuonn et al. (2024) and Højmark al. (2024) attempt to estimate the probability that a language model passes certain capability evaluations, even when its success rate is low, though their methods are not directly applicable to our formal setting.
Our importance sampling methods can be viewed as solving a special case of controlled text genera- tion (Zhang et al.) (202) in which we want to sample from an autoregressive distribution conditioned on property of the full output (in our case, that the last token is t). Yang & Klein (2021) do this by training Future Discriminators to steer model generation towards the desired attribute. Lew et al. (2023) approach the problem with a Sequential Monte Carlo steering approach; however, their in- hilling algorithm doesn't 't provide any benefit over naive sampling when all tokens except the last are independent. These works don't consider the problem of low probability estimation.
Zhao et al. (2024) focus on the problem of estimating the partition function of an unnormalized target distribution over sequences, which is more general case of our low probability estimation problem. Their Twisted Sequential Monte Carlo methods can be viewed as more advanced versions of our importance sampling methods. In contrast, in this work we focus motivating the low probability estimation problem and introducing methods that do not involve searching for positive samples, such activation extrapolation.
Finally, there is a large body of work applying adversarial training to improve worst-cas model performance (Bai et al. 2021 Goodfellow al. 2014 I[yas et al. 2201), especially in the context of language models [017]] Liu Liu et al. 2020). Perez et al. (2022) explores using language models themselves to aid in red-teaming other models. Latent adversarial training (Casper al., 2024]
Sheshadri et al. 2024) generalizes standard adversarial training by optimizing over perturbations in activation space; this means that, like activation extrapolation methods, it can be effective even when the adversarial training search problem over input space is hard.
CONCLUSION
In this paper, we introduce the problem of low probability estimation along with four novel esti- mation methods. We define and collect ground-truth probabilities for 8 different input distributions, then use them to evaluate the performance of our proposed methods. We find that the two importance sampling- based methods perform the best, with larger models favoring MHIS over ITGIS.
We are excited for future work that extends our empirical setup to non-independent input distribu- tions and output behaviors that involve more than token. We are also looking forward to future papers that develop more accurate estimation methods, especially methods like QLD that move be- yond importance sampling.
10

---

Published as a conference paper at ICLR 2025
9 ACKNOWLEDGEMENTS
Paul Christiano played a significant role advising the project. We are grateful for intermediate theoretical contributions from David Matolesi and George Robinson. Thanks additionally Jean- Stanislas Denain and Eric Neyman for feedback on a draft.

---

Published as a conference paper at ICLR 2025
REFERENCES
Tao Bai, Jinqi Luo, Jun Zhao, Bihan Wen, and Qian Wang. Recent advances in adversarial training for adversarial robustness, 2021.
01356
Andreas Buja, Lawrence Brown, Arun Kumar Kuchibhotla, Richard Berk, Ed George, and Linda Zhao. Models as approximations ii: A model-free theory of parametric regression, 2019. URL
https://arxiv. wttp://www...gg/////o/ooa//
03257
Nicholas Carlini, Milad Nasr, Christopher A Choquette- -hhoo, Matthew Jagielski, Irena Gao, Pang Wei W Koh, Daphne Ippolito, Florian Tramer, and Ludwig Schmidt. Are aligned neural networks adversarially aligned? Advances in Neural Information Processing Systems, 36, 2024.
Stephen Casper, Lennart Schulze, Oam Patel, and Dylan Hadfield-Menell. Defending against un- foreseen failure modes with latent adversarial training, 2024. URL https://arxiv. org/
00s/2403.050333
Paul Christiano, Eric Neyman, and Mark Xu. Formalizing the presumption of independence, 2022.
URL JtRpps::://aaxx htt.pp////add..orr//a/////// wttp://p/a/.....//.... 06738
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6633 2014.
L.G. Gubin, Boris Polyak, and E.V. Raik. The method of projections for finding the common point of convex sets. USSR Computational Mathematics and Mathematical Physics, 7:1-24, 12 1967. doi: 10.101600041-55536699113333
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023. URL https:// //
arxiv. /166 08415
Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant, Risks from learned optimization in advanced machine learning systems, 2021. URL https://arxiv.
org abs 1906. 01820
Axel Højmark, Govind Pimpale, Arjun Panickssery, Marius Hobbhahn, and Jérémy Scheurer. An- alyzing probabilistic methods for evaluating agent capabilities, 2024. URL bttps://axxiv.
arg/abs/2409.0 16125
Andrew llyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples are not bugs, they are features, 2019. URL https:://arxis
org/ hatg/aba/Daa 1905. 02175
Itakura and S. Saito. Analysis synthesis telephony based on the maximum likelihood method. In Proc. 6th of the International Congress on Acoustics, pp. C-17-C-20, Los Alamitos, CA, 1968. IEEE.
Daniel Kang, Yi Sun, Tom Brown, Dan Hendrycks, and Jacob Steinhardt. Transfer of adversarial robustness between perturbation types. arXiv preprint vrXiiv:195.001 2019.
Alexander K. Lew, Tan Zhi-Xuan, Gabriel Grand, and Vikash K. Mansinghka. Sequential monte carlo steering of large language models using probabilistic programs, 2023. URL https://
arxiv. www./orggbb..... /20b6///33.... 2306.3300000 03081
Xiaodong Liu, Hao Cheng, Pengcheng He, Weizhu Chen, Yu Wang, Hoifung Poon, and Jianfeng Gao. Adversarial training for large neural language models, 2020. URL https://arxiv.
org/abs/200 2004. 08994
Alcksander Madry. Towards deep learning models resistant to adversarial attacks. arXiv preprint 1:X1v:1706.06083, 2017.
Neel Nanda and Joseph Bloom.
Transformerlens.
https://github. com/
2022.
12

---

Published as a conference paper at ICLR 2025
Ethan Perez, Saffron Huang, Francis Song, Trevor Cai, Roman Ring, John Aslanides, Amelia Glaese, Nat McAleese, and Geoffrey Irving. Red teaming language models with language models, 2022. uttpp:::://axx
www.oorggaaa. abs 22022222222222 03286
Mary Phuong, Matthew Aitchison, Elliot Catt, Sarah Cogan, Alexandre Kaskasoli, Victoria Krakovna, David Lindner, Matthew Rahtz, Yannis Assael, Sarah Hodkinson, Heidi Howard, Tom Lieberum, Ramana Kumar, Maria Abi Raad, Albert Webson, Lewis Ho, Sharon Lin, Sebastian Farquhar, Marcus Hutter, Gregoire Deletang. Anian Ruoss, Seliem El-Sayed, Sasha Brown, Anca Dragan, Rohin Shah, Allan Dafoe, and Toby Shevlane. Evaluating frontier models for dangerous capabilities, 2024. htt.:tts:// ptttts://arxiin arp.iiiivv/.
abs hblgg//bb//2/ 2403. 13793
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. https::/arxxx......
910. 10683
Christian Robert. The metropolis-hastings algorithm, 2016. URL https://arxiv
abs/1504, 01896
Rohin Shah, Vikrant Varma, Ramana Kumar, Mary Phuong, Victoria Krakovna, Jonathan Uesato, and Zac Kenton. Goal misgeneralization: Why correct specifications aren't enough for correct goals. arXiv preprini arXiv:2210.01790, 2022.
Abhay Sheshadri, Aidan Ewart, Phillip Guo, Aengus Lynch, Cindy Wu, Vivek Hebbar, Henry Sleight, Asa Cooper Stickland, Ethan Perez, Dylan Hadfield-Menede and Stephen Casper. La- tent adversarial training improves robustness persistent harmful behaviors in Ilms, 2024. URL hittps://arxinn org/ abs/ attp:///////22/2////
Lewis Tunstall, Leandro von Werra, and Thomas Wolf. Natural Language Processing with Transformers: Building Language Applications with Hugging Face. O'Reilly Media, Incor- porated, 2022. ISBN 1098103246. URL https: http:///oooooo/ http://ttooooollcooo/// www.gool..........eeeeeooeeee
7hhyZgEACAAJ
Stefan Webb, Tom Rainforth, Yee whye Teh, and M. Pawan Kumar. A statistical approach to as- sessing neural network robustness, 2019.
07209
Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does LLM safety training fail? Advances in Neural Information Processing Systems, 36, 2024.
Mark Xu. Estimating tail risk in neural networks.
https://al gnment. org/blog/ in- neerrinnnneall---- -ntworksss September 2024. Alignment Re-
imatinn-ntal-l-ee search Center blog post.
Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. In Pro- ceedings of the 2021 Conference of the North American Chapter of the Association for Computa- tional Linguisties: Human Language Technologies. Association for Computational Linguistics, 2021. doi: 10.1 2..18853/v1/222 ma.cl-main.276. URL U ort..//.ooo..g..1100.......
v1/2021. naacc main. 276
Hanqing Zhang, Haolin Song, Shaoyu Li, Ming Zhou, and Dawei Song. A survey of controllable text generation using transformer-based pre-trained language models, 2023. URL https://
22011 05337
Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Grosse. Probabilistic inference in language models via twisted sequential monte carlo, 2024. URL /tR/:ttpp///axxx.ooo////s 2404. 17546
Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J. Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models, 2023. URL htttps://axxii: org/ abs//b/s////0 15043
13

---

Published as a conference paper at ICLR 2025
A
INPUT DISTRIBUTIONS
Table 2: Definitions of imput distributions.
Name
Tokens
Description
bex
33
Tokens that consist solely of hexadecimal characters, weighted by their frequency in long, uniformly-randow hexadecimal strings.
camel
33
Tokens that start with capital letter, then have only lowercase letters, weighted by their
frequency in Python code.
colon
34
Tokens weighted by frequency in Python code. Always ends with a colon.
if
34
Tokens weighted frequency in Python code. Always starts with if'.
caps
21
Tokens that consist only of capital letters or punctuation, weighted by frequency in or 'She screamed:
English text. Starts with 'Ee screamed:
english
26
Tokens that consist only of letters and start with space, well punctuation.
Weighted by frequency English
spanish
25
Tokens that consist only of letters and spaces. Weighted by frequency in Spanish text. for C for
icl
29
A simple in-context learning prompt of the form "A for replaced with random tokens that start the corresponding letter (weighted by frequency English text), and the replaced with a uniformly random letter. The letters spell out 'ARCTHEORY.
for ? for', where the underscores
14

---

Published as a conference paper at ICLR 2025
B FULL ESTIMATION ALGORITHMS
I ALGORITIM FOR ITGIS
The algorithm for ITGIS is as follows. Initialize 8 repeatedly:
(0) = 0 for all and a step counter j = 0. Then,
Increase the step counter by
2. Compute the distribution q(i-1) as defined by the previous score function s be explicitly computed list of [V] probabilities per token position.
s/5-1) This can
3. Draw a batch of samples from from q(5-1), Use the empirical mean of their gradient to estimate a(4) (z) Ex-=qu-- [Vx My (x)) . for all i [k], x E
4. Update the next (1) for all i according to exponentially-weighre moving average:
{00) + as 1-1)) + .... ........0... + 1 a .... a
In practice, we use a 0.9.
5. Calculate the average value the importance sampling estimator for all samples in this batch.
p(x) go [x) 1[M(x) 0
The final output of the method is the average value of this importance sampling estimator across all batches. The number of batches and samples per batch is specified in Appendix [F] See Algorithm for pseudocode.
Algorithm Independent Token Gradient Importance Sampling (ITGIS)
Require: Model M, target token t, input length k token distributions P1...., Pk. temperature T, iterations n, batch size B
1: 8; + 0 E RM for all E [k]
(0)
# Initialize score functions
2: estimates 0
3: for to n do +- 1 to k do
4: for
- (x)x) exp(s (x)(( for all
5:
6: Normalize - to have sum 1
7: end for
8: Sample B inputs (x (x) from oi 11-1) x ... tk
()-1)
9: for (0) (I) < $ Σ V M, for all x E V
to k do
10:
end for
a a-9.9
13:
for i +
1 to k do
())
2(5) +as t...
1 411
end for
Ie n. (x)
estimate 1
I of (x 1 (x(())))
Append estimate to estimates 18: end for
19: return
E estimates[]]
15

---

Published as a conference paper at ICLR 2025
ALGORITHM FOR MHIS
We define the proposal distribution <(-1x) to be the distribution induced by the following procedure:
Choose a random token position i [A] to modify.
2. Calculate the gradient at that token [ E RIM (treating x = (X1,...,XX, as one-hot vector [ Call this gradient g.
3. Sample a replacement token x from the distribution proportional to
Pr(2.) exp
T
4. Output x' = (X1,..., ...,,,,, Xt..
Note that the transition probability in Metropolis-Hasting only depends on the ratio ((x')((x)x) (x() g (x))(((' x) (x,)
exp [M(xx') M2(x)
((x x')
(x'(x)
which is easy to compute given forwards and backwards passes at x and x'.
We use initial burn-in period (see Appendix F) for the random walk before recording samples
(l)) (n), The final output of the method is the empirical importance sampling estimate
p(x(i)) (x(i)) ([((x() t].
This requires computing q(x), which involves the normalization constant. save samples, we estimate the normalization constant using the identity:
Mx(x) T
)
Mx(x) T
Exxp
exp
= Exng exp
The right-hand side can be estimated using the samples we already have from (approximately) q- In practice, the way we estimate the normalizing constant does not matter much (most of the error comes from other steps). See Algorithm for pseudocode.
16

---

Published as a conference paper at ICLR 2025
Algorithm 2 Metropolis-Hastingg Importance Sampling (MHIS)
Require: Model M, target token i, input length k; token distributions pl :l....: Pk. temperature T, number of burn-in steps Yhum: number of samples n
1: Initialize by sampling from P1 .... Pę
2: samples +
3: for step < to bum + do
4: i <- [k]
Choose random position
5: g [x)]
6:
Sample x from & pi(x)) exp(xx;/T)
# Proposed token
7:
x' -- (x1..... X,..., ;.....,,)))
# Proposed new state
# Compute acceptance ratio ((x')((x)x)
((x))(xx(xx
8:
r AcceptanceRatio(A T, ,e).....,,)),,,..
9:
ir Unif(O, 1) then
10:
x 4- x
# Accept proposal
end if
12:
if step Whum then
13:
Append x to samples
14: end if
15:
end for
Z ( Σ exp( exp(-Mp exp(-M(((amplees]1//7
# Est. normalizing constant
n Σ (sampleaa)) [M (samples(j)
17: return
III ALGORITHM FOR QLD
Recall that vxx) Rd is the random vector representing the activations of the model right before the unembedding step. After collecting samples of v(2), ,..,, v(n), we compute their empirical mean Id and covariance E E Rddx Then define u to be the whitened version of v:
u = A (v 11)
Au μ4,
where E Rdxd any matrix such that AA = Note that has mean and covariance Ida- From now principally work in this whitened representation of activation space, as it has the convenient property that the u and u are uncorrelated iff e and e' are orthogonal.
We choose the unit vector d E Rn to point in the direction of the shortest accepting vector (Ap- pendix D). then decompose our whitened samples u ..... (n) into components parallel and per- pendicular to
a == dd ()
(( ( (1)
Finally, we output:
n² { (1.5) E [n]² | + +b00 E ss}
where
S S: fu Rd arg max, ((Au+ ma
Wo), = =].
This proportion can be computed in Ö(n) time-we we don't need explicitly iterate over all n² pairs. By the convexity of the acceptance region S, for any fixed there is a single interval [f,r] such that [f, + ad b E S. We can efficiently compute the bounds of this interval for every sample (() by solving linear system of inequalities, and then into each range Oloog n) time after sorting. Thus, the computational cost of QLD is dominated by running n forwards passes of M to generate the samples u(!), ((),.... (n). See Algorithm for pseudocode.
calculate how many a(*) fall
17

---

Published as a conference paper at ICLR 2025
Algorithm 3 Quadratic Logit Decomposition (QLD)
Require: Model M with unembed matrix Wt. target token 6, sample size
1: Sample n inputs xx")
x(*) from input distribution
2: Compute pre-unembed activations (() E R by running M on (() for i E [n]
3: μ f 1 n
(i)
4: Σ --
(v(0 ()(((() -²)T
5: Find A such that AA = = + Idl
# via Cholesky decomposition
6: for + to do
7: (1) A 1-11111 (v(0) p))
# Whitened activations
8: end for
9: d + ShortesstAcooptingVeetrrr ,r, Wu, t)
# See Appendix D
for f-
do
11: a <- dd (0)
# Parallel component # Perpendicular component
12: b(0) + ( (
end for
Sort (1 J= ascending order
15: count +-
16: for j 4- 1 to n do
# Solve linear inequalities to get [Cg,ry] such that E [eg, + ad bb0 e
17:
[g Find Acepaaacaanneeenneeelllll count count Ki a [1g,
d, A, ,4, Wo,
Binary search for bounds
19: end for
20: return count/n²
18

---

Published as a conference paper at ICLR 2025
C PRINCIPLES FOR CHOOSING A DECOMPOSITION IN QLD
In this section, we justify the claim that we rely on the following two assumptions of the decom- position u + b for QLD perform well: 1) a and independent, and 2) the contribution towards the output behavior is split roughly equally between the two terms.
The first assumption is straightforward: if and b are not independent, then a b' (where b' comes from an i..d. copy of u) does not have the same distribution as u. The failure of this assumption introduces bias into the estimation method. Working in whitened space ensures that and are uncorrelated, which is first step towards independence.
The second assumption that the contribution to the target behavior is roughly equally split between and b is necessary for QLD to have an advantage over naive sampling. For purposes of illustra- ltd tion, say that d = (so that both a and can be treated as scalars), and that a, b (/ 1). Then, consider three ways that the contribution to the target behavior could be split:
Scenario I (no split): The target token is outputted iff > 10. In this case, QLD provides no advantage over naive sampling, as the proportion of (a b(0) pairs in the acceptance region is exactly the same as the proportion of (a b (()) pairs. If p is the probability of the behavior, then p samples are required consistently obtain a positive estimate.
Scenario 2 (even split): The target token is outputted iff + b > 10v2. In this case, QLD has a quadratic advantage over naive sampling. only requires around p 1/2 samples for QLD to consistently obtain positive estimate.
Scenario 3 (uneven split): The target token is outputted iff 2b 10.5. The contribu- tion here is split between a and b, though not equally, so QLD's efficiency falls in between the previous two scenarios. We can calculate that requires around p- p559 samples for QLD to consistently obtain positive estimate
In practice, the condition for outputting the target token is more complex than a single linear con- straint on a and b. Nevertheless, these examples motivate the idea that the more evenly we can split contribution between two subspaces, the lower variance our estimator will have.
Given that has more influence over the behavior than This motivates us to identify particularly important direc- tion with d: in some informal sense we want to find a direction that is "d/2 times more important" than the average direction.
dimensions while a only has 1, most choices of d will end up giving much
When we run QLD with many more samples than its standard budget of 216, its performance im- proves but plateaus at level that is still worse than the sampling methods. This shows that QLD is currently limited by a poor independence assumption, as the error arising from an unequal split of contribution should vanish with sufficient number of samples.
In general, if the condition is -1-aah 10, it requires roughly p -1/(atv1-03)
samples.
19

---

Published as a conference paper at ICLR 2025
D COMPUTING THE SHORTEST ACCEPTING VECTOR
Recall that the acceptance region S C IRd the subset of whitened pre-unembed space that results in the model outputting&
S {u E |Ra| arg max, ((xx + p.) Wu); = =
We define point the direction of the shortest vector in S, i.e., arg min ||u]]. This vector can be approximated using iterative convex projection algorithm.
" Rd an intersection [V 1 half-spaces HI...., H1-1, H Hy. where H, represents the set of all activations (in whitened pre-unembed space) that result in the logit on token being larger than the logit on token
Given any convex set C (such x' 1122- Given a collection of convex sets, there exists a simple algorithm for finding a point in their intersection: start with arbitrary point at, then repeatedly project ar onto random convex set that does not already contain æ. Eventually, this process converges to a point in their intersection (Gubin al. 1967).
half- half- plane), the projection of C onto C' is arg ein ||
We apply this method to find an element of S. To ensure that it is the shortest element, we also project æ onto balls centered at with smaller and smaller radii by multiplying æ by 0.99. The exact procedure is described in Algorithm 4 In practice, it always takes much less than 100 rrppss steps for the algorithm return value.
Algorithm 4 Random Constraint Projection
Require: Half-spaces H1, H4-1, H4 Hy, number of repetitions T'reps
1: æ 0 Rdd
2: for step_cnt <- to 100 Trrpps do
3: Pick random among all such that x & H;
4: Project æ onto H,
5: if æ lies in (up to some tolerance c) then if step.cne ent Prps then
6:
7:
Scale by 0.99. 8: else
9: return a
10: end if
11: end if
12:
end for
We found few minor bugs in our implementation of the e-tolerance in our algorithm after experi- ments, but we don't expect them to have affected the results all.
20

---

Published as a conference paper at ICLR 2025
E
GROUND TRUTH TOKEN DISTRIBUTION
Laaer Layans
10-11
10
10-3
20-3
13-3
30
10-7
10"1-
10"
10-7
200
2000
0000
6000
0003
2200
1000
8000
5000
10000
13000
english
10
10"
10"
30"3
10"3
30"5
10-7
10-1
15n7
4-
1600
2000
10000 20003 30000
2500
5000
75.00
10000
2000 4000
6000
8000
Figure 5: The ground truth probabilities of tokens for each distribution and model size, sorted from most to least probable (the height of the curve position is the probability of the æ-th most common token). Any tokens that appeared times across all 23 samples are not plotted. The hex distribution only had 159 and 135 tokens in the range [10-9, (0--) for the 2- and 4-layer models, respectively, so we used every such token instead of sampling 256 of them.
21

---

Published as a conference paper at ICLR 2025
F COMPUTATIONAL BUDGETS
Each estimation method was given a budget of roughly 2 model calls. More specifically:
Independent Token Gradient Importance Sampling uses 2 batches of size 28, for a total of 2 samples. The average gradient is updated after each batch. Note that this method requires backwards passes as well as forwards passes.
Metropolis Mettopolis-Hastingss Importance Sampling uses 2 2 batches of size 25, for a total of 1.5 2 samples (the batch size indicates the number of independent random walks the method simulates). The first 2 batches are used as a burn-in period for the random walk and are discarded, only 2 samples are actually used to calculate the estimate.
Quadratic Logit Decomposition uses n = 2 samples of the pre-unembed activation v. The MLE direction is approximated with nrpps 200 iterations of the Random Constraint Projection algorithm (Appendix D); this makes up trivial fraction of the total compute usage of the method).
Gaussian Logit Difference uses 2 samples of the logit difference estimate Jr and a, the mean and standard deviation of the difference between the target logit and the maximum logit. Note that in practice, the μ and can be accurately estimated with much fewer than 2 samples.
In practice ITGIS and MHIS take the longest to test because they require separate samples for each target token. In contrast, QLD reuses the same 26 samples of e for all 256 target tokens associated with given behavior.
22

---

Published as a conference paper at ICLR 2025
G ALL METHOD PERFORMANCES
Table shows the Itakura-San Saito loss (p/g - In(p/4) - 1) of all estimation methods on all input distributions and model sizes.
Table Itakura-Saii loss comparison of methods, distributions, and model sizes.
(a) l-layer model
(b) 2-layer model
Distribution
Constant
GLD
QLD
ITGIS
MHIS
Distribution
Constant
GLD
QLD
ITGIS
MHIS
2.5891
2.2163
2,0038
2.0484
1,3960
2,8839
2.8453
2.7698
2,6652
1.4985
camel
2.5908
2.4419
2.0648
1.1997
2,01877
camel
2.3242
2.2268
2.2679
1.9221
2.0637
colon
2,7700
2.7091
1.2786
1.2200
1.0267
colon
2.9893
2.8335
2.1215
2.2986
1.9042
if
2,2424
2.1872
1.2321
109966
1.4455
if
2.6172
2.4377
1,8997
1,0301
1.2516
capa
266199
2.6147
1.9413
1.4788
2.4023
caps
2.6345
2.6820
2.4477
1.9271
1.6981
engli
1.9095
1.8120
1.2409
1.4519
0.7017
engl sh
2.0989
2.0956
1.3462
0,9908
1,4480
spanish
2.6079
2.4463
1.5628
1.6396
2.1538
spanish
2.5442
2,3662
1,5670
1.0951
2. 2.3095
2.5467
2,4328
2.2373
0.7992
1.2344
iel
1.0502
Average
2,4906
2.3875
1.6952
Average
2.0946
1.6964
(c) 4-ayer model
Distribution
Constant
GLD
QLD
ITGIS
MHIS
hex
2.4803
2.2934
2.4282
2,3090
2.0830
camel
2.4895
2.3271
2.4534
2.0937
2.2961
colon
2.9225
2.7319
2,4382
2.1295
1.1710
1f
2.6477
2.5408
1.8313
1,8430
0.9261
caps
2.5970
2.5382
2.4142
2.0484
1.2972
ongli sh
2.7008
2.7432
1.5681
0,9943
1.3051
spanish
2.6415
lel
2,5029
2.1925
2.3866
2.2971
0.5477
2.6240
2.4837
2.1615
1.9220
1.4775
23

---

Published as a conference paper at ICLR 2025
H SQUARED ERROR IN LOG-SPACE
Figure and Table show the method performances when measured using squared error in log- space loss (i.e., (log Bogpp-1 log g)² instead of Itakura-Saito loss. The results are qualitatively identical using either metric. Note that we use separate affine fits to minimize each loss function-in Table 4 we naturally report the results of the fit corresponding to squared error log-space. However, the importance sampling temperatures are not changed between the two metrics (they were tuned while minimizing Itakura- Saito loss).
Optimal Curslanc
Independnnt Chedion: Impar lansue Sumppinn Matropllis- importnnca samp ing
quadranic cecamposition
LODCV tazs isquared Enor Log-ppcc)
oue
Number
Figure 6: The squared error in log-space loss of all methods across different model sizes. The solid lines indicate the loss of each method averaged over all 8 distributions, with bands indicating standard error. The colored points indicate the loss on individual distributions, with horizontal jitter added for visibility.
24

---

Published as a conference paper at ICLR 2025
Table 4: Squared error in log-space loss comparison of all methods, distributions, and model sizes.
(a) I-layer model
(b) 2-layer model
Distribttion
Constant
GLD
QLD
ITGIS
MHIS
Distribttion
Constant
GLD
QLD
IIGIS
MHIS
6,0093
5.1815
4.4863
3.7179
2.1963
hex
8.0536
7,3674
7.6533
6,8004
1.9913
camel
7.7295
6.4975
5.8488
1,5156
4.3248
camel
7.6849
6.3518
7.2325
4.2695
5.6046
colon
6.8554
5.4361
2.1379
2.3549
1.8061
7.1728
5.6416
3.3924
2.8242
if
5,3977
4,9287
2.9356
1.9555
3.2223
if
6.3346
5.6923
3.3695
1.7059
2.1639
caps
caps
eng'1 1 sh
6.0051
5.3220
3.2639
1.8357
1.6395
eng! 1 sh
5.2742
5.1188
3.4466
1.6335
3,2728
spanish sh
5.7990
5.5649
2.9740
2.1231
4,2612
spanish
6.3718
5.6877
3,8404
2.1582
4,0937
iel
5,9035
5.6915
4.8023
1,0222
2.3030
iel
6.2785
6.2061
5,0740
2,0615
1.3515
6.4740
5,7737
3,8230
2.2178
3,2885
Average
6.8673
6.1158
4.9357
3.2319
3.1556
(c) 4-layer model
Distribution
Constant
GLD
QLD
ITGIS
MHIS
hex
7,5559
6.9208
7,0750
6.4364
3.1605
camel
colon
7.1791
6.6035
3.5817
4.3916
1.7897
9.1105
8.2957
6.1383
4,3615
2.2514
english
5,5003
5,4866
3.3561
1,9180
2.1938
spanish
7,0490
6,4641
4,2083
3.4425
4,8425
iel
5.1793
4,4569
4.4772
4,5584
1.0216
Average
6,9780
6.3345
4.9827
4.0413
2,8115
25

---

Published as a conference paper at ICLR 2025
I TEMPERATURE TUNING
Both importance sampling methods require choosing a temperature parameter T. To tune T, we sweep over 9 different temperatures from 0.2 to uniformly spaced in log-space. We choose the value T that achieves the lowest loss on 100 randomly chosen tokens with ground-truth probabil- ities in the range [10-5, 10-31 to prevent over-fitiing We tune separate temperatures for each dis- tribution, model size. and importance sampling method, shown in Table [ It is likely that spending more effort to tune these temperatures (e.g., by tuning on more and rarer tokens) would moderately improve the final performances of the importance sampling methods.
Table Temperatures 7' used for the different methods.
layer
layers ITGIS
4 layers
Distribution
ITGIS
MHIS
MHIS
ITGIS
MHIS
hex
1.00
0.67
1.50
0.67
5.00
0.67
camel
1.00
2.24
1.50
2.24
1.00
2.24
colon
1.00
1.00
1.00
1.50
0.67
1.00
if
1.00
2.24
0.45
1.50
1.00
1.00
caps
1.50
3.34
0.45
1.50
0.67
1.00
english iss
0.45
1.50
0.67
2.24
0.45
1.50
spanish
0.67
0.67
2.24
icl
0.45
1.00
0.30
0.67
3.34
0.67
26

---

Published as a conference paper at ICLR 2025
J PLOTS OF METHOD OUTPUTS
Figures] and show the outputs of the four methods on all three model sizes, using log-log plots of ground-truth probability vs method output. All graphs show the outputs after the Itakura-Saito fit has been applied (see Section 4.2). The horizontal lines of points reveal the value of the additive constant in the fit; any outputs will all lie this line after the fit is applied.
1-layer model Came!
hex
colon
177
0.
caps
english
157
spanish
icl
Ground Truth
MHis
ITGI
GLD
Figure 7: The outputs of methods, after a fit is applied, on all 256 tokens for each distribution on the 1-layer model. The horizontal axis represents the ground-truth token probability, while the vertical axis the output the model after a fit.
27

---

Published as a conference paper at ICLR 2025
2-layer model
157
175
307
caps
english
124
spanish
icl
Ground truth
MHIS
I1GIS
GLD
12"
179
Figure 8: The outputs of methods, after a fit is applied, on all 256 tokens for each distribution on the 2-layer model. The horizontal axis represents the ground truth token probability, while the vertical axis is the output of the model after fit.
28

---

Published as a conference paper at ICLR 2025
4-layer model
caps
english
spanish
icl
Ground truth
MHIS
I1GIS
GLD
Figure 9: The outputs of methods, after a fit is applied, on all 256 tokens for each distribution on the layer model. The horizontal axis represents the ground truth token probability, while the vertical axis is the output of the model after fit.
29


Reviews:
[{"review_id": "DmrUWcvGVf", "reviewer": "Reviewer_DbYe", "summary": "This paper studies low-probability estimation - the problem of estimating the probability that model outputs will satisfy a certain criteria on inputs sampled from some distribution. The authors are primarily concerned with models/behaviors where the probability is hard to estimate empirically using unbiased samples from the input distribution, due to the event being rare.\n\nMore specifically, the authors study the task of estimating the percentage of inputs on which a transformer model predicts a certain token $t$ as being the most likely next token. Three methods are tested - two importance sampling methods, which leverage gradients with respect to the target logit, and an activation extrapolation method, which uses an empirical distribution over the final hidden states to estimate the probability that $t$ will have the highest logit.\n\nThe authors show that importance sampling methods currently dominate activation extrapolation on all tasks, but provide strong arguments for why they might be relevant in more complex models, where finding any particular input that gives rise to a behavior is computationally intractable.", "questions": "- Do you have preliminary thoughts or analysis on how these methods might scale to larger language models?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The paper has a novel framing of low probability estimation as a concrete approach to understanding and improving worst-case model performance.  It is the first paper to explore extrapolation-based methods that don't require finding explicit examples of rare behaviors.\n\n\n- The best performing method presented in the paper, MHIS, is motivated using existing work in automated redteaming. By adapting the Greedy Coordinate Gradient technique (previously used for finding jailbreaks) into a proposal distribution for Metropolis-Hastings sampling, the paper shows how adversarial attack methods can be leveraged for low-probability estimation.  This innovative connection between adversarial attacks and probability estimation shows how existing attack methods can be repurposed for studying the probability of rare behaviors on the normal input distribution.", "weaknesses": "- The models tested are small, and it's unclear whether these methods can be scaled up to actual models of interest.  However, the decision to use smaller models is justified, as it would be otherwise impossible to compute the true probability on all 2^32 input samples, which is necessary for measuring the performance of the methods.\n\n\n- Given that adversarial training is stated as a part of the motivation for developing low probability estimation methods, there is a gap in the discussion around automated red teaming and adversarial training approaches for language models.  The related works section would benefit from an extended discussion on this area.  This context would help readers better understand how this work relates to and advances existing techniques for improving model robustness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kEJg7Xw7WF", "reviewer": "Reviewer_D4d5", "summary": "This paper introduces methods for estimating the probability of sampling certain outputs when conditioned on an input distribution in cases where naive sampling is infeasible. Specifically, they introduce four methods (two based on importance sampling and two based on extrapolating model activations) for the special case where the input distribution is independent per token and the output is a single token.", "questions": "Is there any concrete benefit that the proposed methods add to the overall motivation of minimizing adversarial outputs (ex. based on the activation extrapolation)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I am not familiar with prior work in this domain, but it appears as if the paper introduces novel methods for trying to estimate the probabilities of certain outputs given an input distribution.\n\nMany intuitions are described in prose and are very clear/helpful:\n- Why MHIS may be needed, why this is better for higher capacity models\n- Why activation extrapolation methods could be better in certain scenarios.\n- Why formally defined distribution can be relevant.\n\nThe experiments are sound and the results show that the proposed methods are clearly better than the naive sampling baseline.\n\nGraphics are generally clear!", "weaknesses": "To my understanding, the main motivation described for this problem stems from using these methods to reduce the probability of rare but undesirable outputs. Although a reasonable motivation, there are limited results in this direction because, as the authors point out, doing this for importance sampling reduces to adversarial training, and the activation extrapolation methods are currently worse than the importance sampling methods.\n- A more thoroughly description of how you would do this for activation extrapolation methods would be good.\n\nThe current problem as formulated is restricted (ie. single token output detection) and not practically useful. However the authors recognize this, and I agree that it is a good first step.\n\nThe description of the activation extrapolation methods in the main text could be made more clear, for example the motivation for why the QLD directions could be summarized in the main text (without this, it is hard to understand the method).\n\nNits:\n- Add lines of best fit for figure 1 and 3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MTA5faKPX4", "reviewer": "Reviewer_fKZ4", "summary": "This work focus on probability estimation of low probability outputs in LLMs with argmax decoding. The paper explores the use of importance sampling and activation extrapolation for better estimate of these rare output events. The paper also provides experiments showing the relative performance of the various proposed algorithms.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written.\n- The paper provides simple methods for low probability estimation in LLMs.\n- Empirical results validate the working of their proposed methods (the importance sampling works better as shown by the authors).", "weaknesses": "- One major drawback I find is the lack of clear goal of the paper. The paper seems to provide some interesting algorithm for low probability estimation and high level discussion of how they can be used in practice. But there is no clear application or experiment that shows the usefulness of the contribution. I would urge the authors to kindly provide some clear applications and experiments to show its usefulness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9u5AzcnH1G", "reviewer": "Reviewer_sH4y", "summary": "The paper studies the problem of estimating the probabilities of rare outputs of transformer models. The authors explore two ways to estimate these probabilities: i) importance sampling and ii) activation extrapolation, which is more novel, but less effective. For each method, they propose two sub-strategies. Results show that importance sampling performs best at estimating the probability of rare outputs and both proposed methods perform better than random sampling.", "questions": "In short, I believe this paper to be a good contribution to the community, providing a starting point to address the problem of rare output probability estimation. However, referencing the weaknesses above, I advise to address the following two points:\n\n1) State more clearly that proposed methods are restricted in their applicability to different models and more complex scenario, in particular wrt the need to compute and propagate the gradient of the output probability.\n\n2) detail more clearly the proposed strategies, ideally with pseudo-codes.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper focuses on an important and, in my opinion, understudied problem of deep learning in general, which is estimating the probability of rare, but potentially catastrophic outputs in deep learning models, and by so doing also capturing the distribution of inputs that would cause them. As stated by the authors correctly, solving this problem is key in advancing adversarial training and even reinforcement learning in many settings.\n\nThe approaches tested are principled and thoroughly derived and, while limited in the current applicability (see weaknesses), provide a starting point for further studies and developments of new methods to quantify the probability of rare outputs and all its applications.", "weaknesses": "Two main weaknesses:\n\n1) The methods provided are quite limited in their application to larger models and real-world scenarios, and not just because the authors chose to test with small models, but also because the methods themselves perform computations that restrict their use. Specific points that introduce restrictions to be clarified:\na. The proposed method can only work with target behaviours defined as single token generation, i.e., define specific tokens that are considered rare for the next generation. This is quite limiting, as in language models target/dangerous behaviours are more often associated with semantic meaning of entire sentences and not with a single next token prediction, e.g., alignment for harmfulness.\nb. Both methods involve computing the gradient of the rare token generation probability wrt the input tokens/sequence. This reinforces point a. above and further necessitates to have full access to the model, as we need to be able to compute this gradient and propagate it through the model, which is not trivial if the target behaviour is not simply a single specific token or the model is black-box.\n\n2) The technical section describing the two method could be clearer. I suggest converting the point-by-point methods descriptions of supplementary B into pseudo-codes to make it clear what are the exact procedures to implement the different proposed methods and, if possible, put some of them in the main text. In the current form, the methods are given little space and I could not follow them without carefully reading supplementary B.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)