Paper deep dive

Estimating the Probabilities of Rare Outputs in Language Models

Gabriel Wu, Jacob Hilton

Year: 2024Venue: arXiv preprintArea: Safety EvaluationType: EmpiricalEmbeddings: 117

Models: 1-layer transformer (d=512), 2-layer transformer (d=512), 4-layer transformer (d=512)

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%

Last extracted: 3/12/2026, 6:51:11 PM

Summary

The paper addresses the challenge of estimating the probability of rare, undesirable outputs in language models, which are often too infrequent for standard random sampling. The authors evaluate two primary approaches: importance sampling (specifically Independent Token Gradient Importance Sampling and Metropolis–Hastings Importance Sampling) and activation extrapolation (Quadratic Logit Decomposition and Gaussian Logit Difference). They find that importance sampling generally outperforms activation extrapolation, though both are superior to naive sampling, and discuss the implications for improving worst-case model performance.

Entities (5)

Transformer Language Models · technology · 99%Independent Token Gradient Importance Sampling · method · 98%Metropolis–Hastings Importance Sampling · method · 98%Gaussian Logit Difference · method · 95%Quadratic Logit Decomposition · method · 95%

Relation Signals (5)

Independent Token Gradient Importance Sampling → isa → Importance Sampling

confidence 99% · We study two types of methods: Importance sampling... The first defines q(x) with independent tokens... Independent Token Gradient Importance Sampling (ITGIS)

Metropolis–Hastings Importance Sampling → isa → Importance Sampling

confidence 99% · We study two types of methods: Importance sampling... our Metropolis–Hastings Importance Sampling (MHIS) method

Quadratic Logit Decomposition → isa → Activation Extrapolation

confidence 99% · We study two types of methods: ... activation extrapolation... Our Quadratic Logit Decomposition (QLD) method

Gaussian Logit Difference → isa → Activation Extrapolation

confidence 99% · We study two types of methods: ... activation extrapolation... our Gaussian Logit Difference (GLD) method

Importance Sampling → outperforms → Activation Extrapolation

confidence 95% · We find that importance sampling outperforms activation extrapolation

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

116,630 characters extracted from source content.

Expand or collapse full text

Estimating the Probabilities of Rare Outputs in Language Models Gabriel Wu Jacob Hilton Alignment Research Center Correspondence to: gabriel.d.wu314@gmail.com 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. 1 Introduction Modern ML systems undergo black-box optimization to minimize a loss function on samples drawn from a training distribution. Although models produced in this way perform desirably on 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 et al., 2024; Wei et al., 2024). Preventing such highly undesirable outputs is a notoriously challenging problem. The most common remedy is adversarial training, in which inputs that produce these undesirable outputs are searched for and used as additional training data (Goodfellow et al., 2014; Madry, 2017), but the transfer between different search methods is 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 of 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 is 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 transformer language models. We consider various formally-defined input distributions in which each input token is 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 computational budget. The target tokens are chosen to have ground truth probabilities between 10−9superscript10910^-910- 9 and 10−5superscript10510^-510- 5, which are too small for random sampling to produce a good estimate under the constrained computational budget. In this context, we study two types of methods: • Importance sampling. We define a new input distribution under which the rare event is much more likely, sample from that distribution, and re-weight samples to obtain an unbiased 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 use random samples to fit a probability distribution to the model’s logits, and extrapolate into the tails of this distribution to produce a probability estimate. Our Quadratic Logit Decomposition (QLD) method applies a presumption of independence to the empirical distribution of logits, motivated by Christiano et al. (2022), and our Gaussian Logit Difference (GLD) method is a 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 to outperform activation extrapolation. Nevertheless, we remain interested in activation extrapolation 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 in our language model setting. In 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, 7 and 8, we discuss the limitations and implications of our results, related work, and future directions. Figure 1: Left: To evaluate the performance of our low probability estimation methods, we compare their estimates against ground-truth probabilities obtained by brute-force sampling with a larger computational budget. Right: The estimates of Metropolis–Hastings Importance Sampling on the icl input distribution and 4-layer model, after a fit has been applied. Each point represents a different target token. 2 Problem Statement Given an input space XX, an output space YY, an input distribution ∈Δ⁢()ΔD∈ (X)D ∈ Δ ( X ), a model M:→:→M:X : X → Y, and a formal boolean property of model outputs C:→0,1:→01C:Y→\0,1\C : Y → 0 , 1 , low probability estimation is the problem of efficiently estimating Prx∼⁡[C⁢(M⁢(x))=1].subscriptPrsimilar-toxx1 _ x [C(M( x))=1].Prx ∼ D [ C ( M ( x ) ) = 1 ] . We sometimes refer to the event C⁢(M⁢(x))=1x1C(M( x))=1C ( M ( x ) ) = 1 as the “target behavior”, or just “the behavior.” If the probability of the behavior large enough (say, larger than 1/n11/n1 / n), it is easy to estimate by drawing n independent samples from XX and using the sample mean of C⁢(M⁢(x))xC(M( x))C ( M ( x ) ). However, if the probability is significantly smaller than 1/n11/n1 / n, this sample mean is almost always 00, making it uninformative at distinguishing between small probabilities like 10−10superscript101010^-1010- 10 and 10−20superscript102010^-2010- 20. 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:∗→:→superscriptM:V^* : V∗ → V be a transformer language model that predicts the next token given a string of previous tokens, where VV is the token vocabulary. Note that we sample at temperature 00, so M is deterministic. Given a distribution DD over ∗superscriptV^*V∗ and a target token t∈t ∈ V, the low probability estimation problem for single-token behaviors is the task of estimating Pr∼⁡[M⁢()=t].subscriptPrsimilar-to _x [M(x)=t].Prbold_x ∼ D [ M ( x ) = t ] . Letting Mi⁢()subscriptM_i(x)Mitalic_i ( x ) be the logit the model assigns to token i∈i ∈ V, this can also be written as: Pr∼⁡[Mt⁢()>Mi⁢()∀i≠t].subscriptPrsimilar-tosubscriptsubscriptfor-all _x [M_t(x)>M_i(x) % ∀ i≠ t].Prbold_x ∼ D [ Mitalic_t ( x ) > Mitalic_i ( x ) ∀ i ≠ t ] . In general, DD can be any distribution that can be formally specified. However, in this paper we focus only on distributions DD with independent tokens. That is, we specify an input length k and token distributions 1,…,k∈Δ⁢()subscript1…subscriptΔD_1,…,D_k∈ (V)D1 , … , Ditalic_k ∈ Δ ( V ), then write DD as the product 1×⋯×ksubscript1⋯subscriptD_1×…×D_kD1 × ⋯ × Ditalic_k. Table 1 shows the 8888 distributions that were tested, with tokens colored for clarity. To prevent overfitting, the methods were only run on the first four distributions during development, and they were finalized before testing on the last four distributions. The results were qualitatively the same on both halves of the split. Table 1: Input distributions and examples. See Table 2 for more detailed descriptions. Name Short description Tokenized example hex Hexadecimal characters <|BOS|>a5acbf6aad468813f94c2fbbff4dc65eadc1553 camel CamelCase Python tokens <|BOS|>LayoutCredServicesVirtualUseTimeInterfaceColorBodyAlRowHeightRepFontAndMetaRequestGroupsOneLabelPasswordAndRaVideoFailedValueGuiTypeMicrosoftSlotDeId colon Python tokens, ending with ‘:’ <|BOS|> et-= """]: (\\ : Thisc\\ \\ (’/\\ \\ \\ <|EOS|>_’].2default.**1 self( def’)",: if Python tokens, starting with ‘ if’ <|BOS|> if: else,-post\\ \\ 2\\ \\ 5 foundfromout, self- node +=\\ \\ =\\ ( this ’values(),.(do caps “He/She screamed:”, followed by caps and punctuation <|BOS|>He screamed: "ESOTTULEBOV.,WR!!IMITLEER.,ARY...IIESSION english English words <|BOS|>ating. is invent School not found from cm an in one to shooting everyone Cor George around responsive employees ground on stone various, spanish Spanish tokens <|BOS|> lo no bu dees cr socialjosabiler m de enidadareljd final de v de lo much icl In-context learning prompt <|BOS|>A for American R for Return C for crack T for troubles H for house E for equipment O for operating R for reason Y for your V for 3 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 Decomposition and Gaussian Logit Difference). We also compare against the baseline of outputting an optimal constant, which can be thought of as the performance of naive sampling because we 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 it takes too many samples from DD to observe a positive example. To address this, we can instead draw samples from a different distribution that 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 is known as importance sampling, and it 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⁢()p( x)p ( italic_x ) be the probability mass function of DD, and let q⁢()q( x)q ( italic_x ) be the PMF of any other distribution. Then Pr∼p⁡[M⁢()=t]=∼p⁡[⁢[M⁢()=t]]=∼q⁡[p⁢()q⁢()⁢[M⁢()=t]],subscriptPrsimilar-tosubscriptsimilar-to1delimited-[]subscriptsimilar-to1delimited-[] _x p[M(x)=t]= E_% x p[ 1[M(x)=t]]= E_% x q [ p(x)q(x) 1[M(% x)=t] ],Prbold_x ∼ p [ M ( x ) = t ] = blackboard_Ex ∼ p [ blackboard_1 [ M ( x ) = t ] ] = blackboard_Ex ∼ q [ divide start_ARG p ( x ) end_ARG start_ARG q ( x ) end_ARG blackboard_1 [ M ( x ) = t ] ] , but the latter may have less variance (and so require fewer samples to get a good estimate). The following two importance sampling methods take q⁢()q( x)q ( italic_x ) to be a Boltzmann posterior with prior p⁢()p( x)p ( italic_x ). The first defines q⁢()q( x)q ( italic_x ) with independent tokens, while the second defines q⁢()q( x)q ( italic_x ) to have non-independent tokens and so requires a more sophisticated sampling method. 3.1.1 Independent Token Gradient Importance Sampling (ITGIS) We want q to up-weight tokens that contribute to t being outputted. One way to do this is to continue to treat each input token as independent, but change the probability of tokens according to their average linear contribution to the logit of t. Let =(x1,…,xk)∈ksubscript1…subscriptsuperscript x=(x_1,…,x_k) ^kitalic_x = ( x1 , … , xitalic_k ) ∈ Vitalic_k be an input of length k, and say that p⁢()p( x)p ( italic_x ) factors as p1⁢(x1)⁢⋯⁢pk⁢(xk)subscript1subscript1⋯subscriptsubscriptp_1(x_1)·s p_k(x_k)p1 ( x1 ) ⋯ pitalic_k ( xitalic_k ). Then we define q⁢()=q1⁢(x1)⁢⋯⁢qk⁢(xk)subscript1subscript1⋯subscriptsubscriptq( x)=q_1(x_1)·s q_k(x_k)q ( italic_x ) = q1 ( x1 ) ⋯ qitalic_k ( xitalic_k ), where qi⁢(xi)∝pi⁢(xi)⋅exp⁡(si⁢(xi)T)proportional-tosubscriptsubscript⋅subscriptsubscriptsubscriptsubscriptq_i(x_i) p_i(x_i)· ( s_i(x_i)T )qitalic_i ( xitalic_i ) ∝ pitalic_i ( xitalic_i ) ⋅ exp ( divide start_ARG sitalic_i ( xitalic_i ) end_ARG start_ARG T end_ARG ) and si(xi)=′∼q[∇′Mt(′)]i,xi.s_i(x_i)= E_x q[ _% x M_t(x )]_i,x_i.sitalic_i ( xitalic_i ) = blackboard_Ex′ ∼ q [ ∇x′ Mitalic_t ( x′ ) ]i , x start_POSTSUBSCRIPT i end_POSTSUBSCRIPT . T is a temperature parameter, and the gradient is taken by treating ′superscript′x x′ as a one-hot vector in ℝk×||superscriptℝR^k×|V|blackboard_Rk × | V |. Intuitively, the gradient ∇′Mt⁢(′)i,xisubscript∇superscript′subscriptsubscriptsuperscript′subscript _x M_t(x )_i,x_i∇x′ Mitalic_t ( x′ )i , x start_POSTSUBSCRIPT i end_POSTSUBSCRIPT gives us a linear approximation to how much the logit of t would change if we replaced i-th token of ′superscript′x x′ with xisubscriptx_ixitalic_i (up to an additive constant w.r.t. xisubscriptx_ixitalic_i). Thus, sisubscripts_isitalic_i scores each token value according to its average linear contribution to MtsubscriptM_tMitalic_t, and qisubscriptq_iqitalic_i is defined as the Boltzmann distribution with respect to this score function.111It can be shown that, given a score function s⁢(x)s(x)s ( x ) and a prior p⁢(x)p(x)p ( x ), the distribution that maximizes x∼q⁡[s⁢(x)]−T⋅KL⁢(q∥p)subscriptsimilar-toxx⋅KLconditional E_ x q[s( x)]-T·% KL(q\|p)blackboard_Ex ∼ q [ s ( x ) ] - T ⋅ KL ( q ∥ p ) is q⁢(x)∝p⁢(x)⋅exp⁡(s⁢(x)/T)proportional-to⋅q(x) p(x)· (s(x)/T)q ( x ) ∝ p ( x ) ⋅ exp ( s ( x ) / T ). However, since sisubscripts_isitalic_i and q are both defined in terms of each other, we can’t calculate sisubscripts_isitalic_i directly. To overcome this, we construct a sequence of score functions si(0),si(1),…superscriptsubscript0superscriptsubscript1…s_i^(0),s_i^(1),…sitalic_i( 0 ) , sitalic_i( 1 ) , … and associated distributions q(0),q(1),…superscript0superscript1…q^(0),q^(1),…q( 0 ) , q( 1 ) , … that are adaptively refined with respect to each other (see Appendix B.i for details). Sampling from each q(j)superscriptq^(j)q( j ) lets us calculate an importance sampling estimate, and the final output of the method is the average value of these estimates across all j. 3.1.2 Metropolis–Hastings Importance Sampling (MHIS) A problem with ITGIS is that the new sampling distribution q⁢()q( x)q ( italic_x ) still treats all tokens as independent, and it 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 are the same as each other). To remedy this, we can define an importance sampling distribution that doesn’t have independent tokens. We must use a score function that depends on the entire input; the most natural choice is the target logit Mt⁢()subscriptM_t( x)Mitalic_t ( italic_x ). We define q⁢()∝p⁢()⋅exp⁡(Mt⁢()T),proportional-to⋅subscriptq( x) p( x)· ( M_t( x)T ),q ( italic_x ) ∝ p ( italic_x ) ⋅ exp ( divide start_ARG Mitalic_t ( italic_x ) end_ARG start_ARG T end_ARG ) , again using a Boltzmann distribution to up-weight regions of input space that are more likely to have positive samples. Unlike ITGIS, we cannot explicitly compute q because it does not factor into independent distributions over each token. Instead, we use the Metropolis–Hastings algorithm to produce a random walk in input space that has a stationary distribution of q.222Metropolis–Hastings is a Markov Chain Monte Carlo method for sampling from a distribution with an unknown normalizing constant. See Robert (2016) for a description of the algorithm. To do so, we must define a proposal distribution ϕ⁢(′|)italic-ϕconditionalsuperscript′φ( x | x)ϕ ( italic_x′ | italic_x ) 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 a model using gradients (Zou et al., 2023). We adapt this optimization procedure into a proposal distribution: to pick a proposed next step ′superscript′x x′ of the walk, we choose a random token position i to replace, compute the gradient of s⁢()s(x)s ( x ) with respect to xisubscriptx x_ixi, 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 a burn-in period. For a precise description of the algorithm, see Appendix B.i. 3.2 Activation Extrapolation Methods The importance sampling methods search for explicit examples of inputs that cause the given behavior. This makes their task at least as hard as the adversarial training search problem—if it is difficult to find an ∈supp⁢()supp x (D)italic_x ∈ supp ( D ) such that M⁢()=tM( x)=tM ( italic_x ) = t, the importance sampling estimators will likely fail to produce a positive estimate. We hope to 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 to 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 a presumption of independence between uncorrelated subspaces the model’s pre-unembed activations. We also develop Gaussian Logit Difference, which is intended as a simple baseline method. 3.2.1 Quadratic Logit Decomposition (QLD) Let the random vector ⁢()∈ℝdsuperscriptℝv(x) ^dv ( x ) ∈ blackboard_Rd be the activation of the model right before applying the unembed matrix U∈ℝd×||subscriptsuperscriptℝ W_U ^d×|V|italic_Witalic_U ∈ blackboard_Rd × | V |. That is, ⁢()⋅U⋅subscriptv(x)· W_Uv ( x ) ⋅ italic_Witalic_U represents the model’s output logit vector M⁢()1,…,||subscript1…M(x)_1,…,|V|M ( x )1 , … , | V |. We first collect n samples of vv (call them (1),…,(n)superscript1…superscriptv^(1),…,v^(n)v( 1 ) , … , v( n )). We then choose some unit direction ∈ℝdsuperscriptℝ d ^ditalic_d ∈ blackboard_Rd (see below), then decompose each (i)superscriptv^(i)v( i ) into (i)+(i)superscriptsuperscripta^(i)+b^(i)a( i ) + b( i ), where aa lies in the subspace spanned by ditalic_d, and bb lies in the complementary subspace that is orthogonal in a whitened basis.333Actually, aa, bb, and ditalic_d are all defined in the whitened space of vv; see Appendix B.i. This decomposition is chosen such that the random vectors aa and bb are uncorrelated across the n samples. Next, by treating the random vectors aa and bb as independent, we can use our n samples of each to obtain n2superscript2n^2n2 “synthetic” samples of uu. The final output of QLD is the proportion of these synthetic samples that cause t to be outputted: 1n2⁢|(i,j)∈[n]2|(i)+(j)∈S|,1superscript2conditional-setsuperscriptdelimited-[]2superscriptsuperscript 1n^2 | \(i,j)∈[n]^2\,\,\, |\,a^(i)+% b^(j)∈ S \ |,divide start_ARG 1 end_ARG start_ARG n2 end_ARG | ( i , j ) ∈ [ n ]2 | a( i ) + b( j ) ∈ S | , where S⊆ℝdsuperscriptℝS ^dS ⊆ blackboard_Rd 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 n2superscript2n^2n2 synthetic samples, this proportion can be computed in O~⁢(n)~ O(n)over~ start_ARG O end_ARG ( n ) time by first sorting the samples (i)superscripta^(i)a( i ). A more complete description of the QLD algorithm can be found in Appendix B.i. Choice of direction. We rely on the following two assumptions for QLD to perform well: 1) aa and bb 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,444Other candidate directions included 1) the t-th column of Usubscript W_Uitalic_Witalic_U pulled back into whitened space and 2) the expectation of ⁢(0,Idd)0subscriptIdN(0,Id_d)N ( 0 , Iditalic_d ) conditioned on lying in the whitened acceptance region. we decided to set ditalic_d to 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 likelihood value of vv under a Gaussian prior, conditioned on observing the model output t. Appendix D describes the algorithm we use to compute ditalic_d. 3.2.2 Gaussian Logit Difference On any given input, we can record the difference Δt:=Mt⁢()−maxi⁡Mi⁢()assignsubscriptΔsubscriptsubscriptsubscript _t:=M_t(x)- _iM_i(x)Δitalic_t := Mitalic_t ( x ) - maxitalic_i Mitalic_i ( x ). We wish to estimate the probability that Δt≥0subscriptΔ0 _t≥ 0Δitalic_t ≥ 0. A natural estimation method, which we view as a simple baseline, is to treat ΔtsubscriptΔ _tΔitalic_t as Gaussian by estimating its mean μ and standard deviation σ with samples, then calculate Pr⁡[⁢(μ,σ2)≥0]Prsuperscript20 [N(μ,σ^2)≥ 0]Pr [ N ( μ , σ2 ) ≥ 0 ]. 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⁡(−(a⁢μσ+ϵ)2+b)+c,superscriptitalic-ϵ2 (- ( aμσ+ε )^2+b )+c,exp ( - ( divide start_ARG a μ end_ARG start_ARG σ + ϵ end_ARG )2 + b ) + c , where a,b,c,a,b,c,a , b , c , and ϵitalic-ϵεϵ are parameters that are fit to minimize loss across all target tokens associated with a given distribution (see Section 4.2). 4 Experimental Setup We apply our methods on three models: a 1-layer, a 2-layer, and a 4-layer transformer from Nanda & Bloom (2022). All models have a hidden dimension of d=512512d=512d = 512, a vocabulary size of ||=4826248262|V|=48262| V | = 48262, GELU non-linearities (Hendrycks & Gimpel, 2023), and were trained on the C4 dataset (Raffel et al., 2023) and CodeParrot (Tunstall et al., 2022). For each of the 8888 distributions (listed in Table 1) and for each model, we generate ground-truth token probabilities by running forward passes on 232superscript2322^32232 random samples. We then select a random set of 256 tokens among those with ground-truth probabilities between 10−9superscript10910^-910- 9 and 10−5superscript10510^-510- 5, and we test all of our methods on these tokens. We give each method a computational budget of 216superscript2162^16216 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<105)superscript216superscript105(2^16<10^5)( 216 < 105 ), but the theoretical quadratic gains from QLD would still be enough to get signal on the entire range of probabilities ((216)2>109)superscriptsuperscript2162superscript109((2^16)^2>10^9)( ( 216 )2 > 109 ). Our code is available at https://github.com/alignment-research-center/ low-probability-estimation. 4.1 Itakura–Saito Loss We measure the quality of the method with a loss function inspired by the Itakura–Saito divergence (Itakura & Saito, 1968). If p is the ground-truth probability of a particular target token, then an estimate of q incurs a loss of: DIS⁢(p,q)=pq−ln⁡pq−1.subscriptIS1D_IS(p,q)= pq- pq-1.Droman_IS ( p , q ) = divide start_ARG p end_ARG start_ARG q end_ARG - ln divide start_ARG p end_ARG start_ARG q end_ARG - 1 . Two considerations went into the choice of this loss function. First, Itakura–Saito loss is a proper scoring rule (Buja et al., 2019). Second, since it only depends on the ratio p/qp/qp / q, Itakura–Saito loss is sensitive to small probabilities: if p=10−100superscript10100p=10^-100p = 10- 100 and q=10−10superscript1010q=10^-10q = 10- 10, then DIS⁢(p,q)subscriptISD_IS(p,q)Droman_IS ( p , q ) is very large. In contrast, the squared error loss function (p−q)2superscript2(p-q)^2( 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 to be dominated by a method’s behavior on only the most probable tokens. For completeness, we also report our results using squared error in log-space (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 00, but DISsubscriptISD_ISDroman_IS is undefined for q=00q=0q = 0. To address this, we fit a transformation x↦a⁢xc+bmaps-tosuperscriptx ax^c+bx ↦ a xitalic_c + b to the outputs of each method, where a,ba,ba , b and c are chosen to minimize Itakura–Saito loss. a⁢xcsuperscriptax^ca xitalic_c 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 a separate fit for each (method, input distribution) pair. 555Note that the Gaussian Logit Difference method has a special functional form of its fit ((μ,σ)↦exp⁡(−(a⁢μ/(σ+ϵ))2+b)+cmaps-tosuperscriptitalic-ϵ2(μ,σ) (- (aμ/(σ+ε) )^2+b )+c( μ , σ ) ↦ exp ( - ( a μ / ( σ + ϵ ) )2 + b ) + c instead of x↦a⁢xc+bmaps-tosuperscriptx ax^c+bx ↦ a xitalic_c + b) but is otherwise evaluated in the same way. 5 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. GLD is barely better than outputting an optimal constant (which can be interpreted as the performance of naive sampling). Figure 4 shows that there is a fair amount of variation in method performance across the 8888 distributions: some behaviors like hex and icl favor MHIS, while others like spanish heavily favor ITGIS. A more detailed table of results is 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 inter-token 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 J; 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. Figure 2: The Itakura–Saito loss of all methods across different model sizes. The solid lines indicate the loss of each method averaged over all 8888 distributions, with bands showing standard error. The colored points indicate the loss on individual distributions, with horizontal jitter added for visibility. Lower is better. Figure 3: Examples of method outputs on two different behaviors and models, before a fit is applied. Estimates of 00 are placed at the bottom of each graph for visibility. Figure 4: The Itakura–Saito loss of all methods across different distributions, averaged over all 3333 model sizes. Lower is better. 6 Discussion 6.1 Distribution shift as motivation One might ask: if a particular model behavior is so rare that it never arises during training, why would we care about estimating its probability? There are a few reasons. First, some AI systems may be run on many more inputs during the course of deployment than during training. Thus, if a certain model behavior would be so catastrophic that it is unacceptable for it to occur even once in deployment, we cannot rely on training to 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 in deployment. This could occur because of an input chosen adversarially, but it could also occur because of goal misgeneralization (Shah et al., 2022). A particularly challenging case is deceptive alignment, the possibility that an ML model would look for clues about whether it is in a training or a deployment environment, and only behave well in training (Hubinger et al., 2021). To detect whether a 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 a catastrophic behavior on this distribution.666To prevent false positives, this would require a very demanding definition of catastrophe that would be impossible for the model to trigger “by accident.” For more discussion of this idea, see Xu (2024). 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-teaming method. A further reason to be interested in low probability estimation is 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 an importance sampling method, this amounts to finding inputs for which the rare event occurs (i.e., 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 extrapolation methods do not correspond to any red-teaming method, new activation extrapolation methods potentially provide us with new ways to reduce the probabilities of rare events. 6.3 Importance sampling versus activation extrapolation In our experiments, we found that importance sampling methods outperformed activation extrapolation. Nevertheless, there are theoretical cases in which importance sampling performs worse than other methods. For example, consider a model that outputs the SHA-256 hash of its input: finding any input that gives rise to a particular output is computationally infeasible, yet it is still easy to 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 be able to use deductive estimates based on a presumption of independence (Christiano et al., 2022). Our Quadratic Logit Decomposition method is an 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 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 a separate generative model and adapt some of the current estimation methods appropriately. Second, we only study model behaviors that consist of a single token sampled at temperature 00. This is unrealistic because in practice, if 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 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. Phuong et al. (2024) and Højmark et 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 generation (Zhang et al., 2023) in which we want to sample from an autoregressive distribution conditioned on a 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 infilling algorithm doesn’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 a 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 on motivating the low probability estimation problem and introducing methods that do not involve searching for positive samples, such as activation extrapolation. Finally, there is a large body of work applying adversarial training to improve worst-case model performance (Bai et al., 2021; Goodfellow et al., 2014; Ilyas et al., 2019), especially in the context of language models (Madry, 2017; Liu et al., 2020). Perez et al. (2022) explores using language models themselves to aid in red-teaming other models. Latent adversarial training (Casper et 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. 8 Conclusion In this paper, we introduce the problem of low probability estimation along with four novel estimation methods. We define and collect ground-truth probabilities for 8888 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 distributions and output behaviors that involve more than one token. We are also looking forward to future papers that develop more accurate estimation methods, especially methods like QLD that move beyond importance sampling. 9 Acknowledgements Paul Christiano played a significant role in advising the project. We are grateful for intermediate theoretical contributions from David Matolcsi and George Robinson. Thanks additionally to Jean-Stanislas Denain and Eric Neyman for feedback on a draft. References Bai et al. (2021) Tao Bai, Jinqi Luo, Jun Zhao, Bihan Wen, and Qian Wang. Recent advances in adversarial training for adversarial robustness, 2021. URL https://arxiv.org/abs/2102.01356. Buja et al. (2019) Andreas Buja, Lawrence Brown, Arun Kumar Kuchibhotla, Richard Berk, Ed George, and Linda Zhao. Models as approximations i: A model-free theory of parametric regression, 2019. URL https://arxiv.org/abs/1612.03257. Carlini et al. (2024) Nicholas Carlini, Milad Nasr, Christopher A Choquette-Choo, 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. Casper et al. (2024) Stephen Casper, Lennart Schulze, Oam Patel, and Dylan Hadfield-Menell. Defending against unforeseen failure modes with latent adversarial training, 2024. URL https://arxiv.org/abs/2403.05030. Christiano et al. (2022) Paul Christiano, Eric Neyman, and Mark Xu. Formalizing the presumption of independence, 2022. URL https://arxiv.org/abs/2211.06738. Goodfellow et al. (2014) Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014. Gubin et al. (1967) 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.1016/0041-5553(67)90113-9. Hendrycks & Gimpel (2023) Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023. URL https://arxiv.org/abs/1606.08415. Hubinger et al. (2021) 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. Højmark et al. (2024) Axel Højmark, Govind Pimpale, Arjun Panickssery, Marius Hobbhahn, and Jérémy Scheurer. Analyzing probabilistic methods for evaluating agent capabilities, 2024. URL https://arxiv.org/abs/2409.16125. Ilyas et al. (2019) Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. Adversarial examples are not bugs, they are features, 2019. URL https://arxiv.org/abs/1905.02175. Itakura & Saito (1968) F. Itakura and S. Saito. Analysis synthesis telephony based on the maximum likelihood method. In Proc. 6th of the International Congress on Acoustics, p. C–17–C–20, Los Alamitos, CA, 1968. IEEE. Kang et al. (2019) Daniel Kang, Yi Sun, Tom Brown, Dan Hendrycks, and Jacob Steinhardt. Transfer of adversarial robustness between perturbation types. arXiv preprint arXiv:1905.01034, 2019. Lew et al. (2023) 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.org/abs/2306.03081. Liu et al. (2020) 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/2004.08994. Madry (2017) Aleksander Madry. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv:1706.06083, 2017. Nanda & Bloom (2022) Neel Nanda and Joseph Bloom. Transformerlens. https://github.com/TransformerLensOrg/TransformerLens, 2022. Perez et al. (2022) 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. URL https://arxiv.org/abs/2202.03286. Phuong et al. (2024) 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. URL https://arxiv.org/abs/2403.13793. Raffel et al. (2023) 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. URL https://arxiv.org/abs/1910.10683. Robert (2016) Christian P. Robert. The metropolis-hastings algorithm, 2016. URL https://arxiv.org/abs/1504.01896. Shah et al. (2022) 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 preprint arXiv:2210.01790, 2022. Sheshadri et al. (2024) Abhay Sheshadri, Aidan Ewart, Phillip Guo, Aengus Lynch, Cindy Wu, Vivek Hebbar, Henry Sleight, Asa Cooper Stickland, Ethan Perez, Dylan Hadfield-Menell, and Stephen Casper. Latent adversarial training improves robustness to persistent harmful behaviors in llms, 2024. URL https://arxiv.org/abs/2407.15549. Tunstall et al. (2022) Lewis Tunstall, Leandro von Werra, and Thomas Wolf. Natural Language Processing with Transformers: Building Language Applications with Hugging Face. O’Reilly Media, Incorporated, 2022. ISBN 1098103246. URL https://books.google.ch/books?id=7hhyzgEACAAJ. Webb et al. (2019) Stefan Webb, Tom Rainforth, Yee Whye Teh, and M. Pawan Kumar. A statistical approach to assessing neural network robustness, 2019. URL https://arxiv.org/abs/1811.07209. Wei et al. (2024) Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does LLM safety training fail? Advances in Neural Information Processing Systems, 36, 2024. Xu (2024) Mark Xu. Estimating tail risk in neural networks. https://alignment.org/blog/estimating-tail-risk-in-neural-networks/, September 2024. Alignment Research Center blog post. Yang & Klein (2021) Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguistics, 2021. doi: 10.18653/v1/2021.naacl-main.276. URL http://dx.doi.org/10.18653/v1/2021.naacl-main.276. Zhang et al. (2023) 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://arxiv.org/abs/2201.05337. Zhao et al. (2024) Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Grosse. Probabilistic inference in language models via twisted sequential monte carlo, 2024. URL https://arxiv.org/abs/2404.17546. Zou et al. (2023) 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 https://arxiv.org/abs/2307.15043. Appendix A Input distributions Table 2: Definitions of input distributions. Name Tokens Description hex 33 Tokens that consist solely of hexadecimal characters, weighted by their frequency in long, uniformly-random hexadecimal strings. camel 33 Tokens that start with a 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 by frequency in Python code. Always starts with ‘ if’. caps 21 Tokens that consist only of capital letters or punctuation, weighted by frequency in English text. Starts with ‘He screamed: "’ or ‘She screamed: "’. english 26 Tokens that consist only of letters and start with a space, as well as punctuation. Weighted by frequency in English text. spanish 25 Tokens that consist only of letters and spaces. Weighted by frequency in Spanish text. icl 29 A simple in-context learning prompt of the form ‘A for _ R for _ C for _ … Y for _ ? for’, where the underscores are replaced with random tokens that start with the corresponding letter (weighted by frequency in English text), and the ? is replaced with a uniformly random letter. The letters spell out ‘ARCTHEORY’. Appendix B Full Estimation Algorithms i Algorithm for ITGIS The algorithm for ITGIS is as follows. Initialize si(0)≡0superscriptsubscript00s_i^(0)≡ 0sitalic_i( 0 ) ≡ 0 for all i and a step counter j=00j=0j = 0. Then, repeatedly: 1. Increase the step counter j by 1111. 2. Compute the distribution q(j−1)superscript1q^(j-1)q( j - 1 ) as defined by the previous score function si(j−1)superscriptsubscript1s_i^(j-1)sitalic_i( j - 1 ). This can be explicitly computed as a list of |||V|| V | probabilities per token position. 3. Draw a batch of samples from from q(j−1)superscript1q^(j-1)q( j - 1 ). Use the empirical mean of their gradient to estimate s^i(j)(x)=′∼q(j−1)[∇′Mt(′)]i,x s_i^(j)(x)= E_x q^(% j-1)[ _x M_t(x )]_i,xover start_ARG s end_ARGi( j ) ( x ) = blackboard_Ex′ ∼ q( j - 1 ) [ ∇x′ Mitalic_t ( x′ ) ]i , x for all i∈[k],x∈formulae-sequencedelimited-[]i∈[k],x ∈ [ k ] , x ∈ V. 4. Update the next si(j)superscriptsubscripts_i^(j)sitalic_i( j ) for all i according to an exponentially-weighted moving average: si(j)=s^i(j)+α⁢s^i(j−1)+⋯+αj−1⁢s^i(1)1+α+α2+⋯+αj−1.superscriptsubscriptsuperscriptsubscript^superscriptsubscript^1⋯superscript1superscriptsubscript^11superscript2⋯superscript1s_i^(j)= s_i^(j)+α s_i^(j-1)+…+α^j-% 1 s_i^(1)1+α+α^2+…+α^j-1.sitalic_i( j ) = divide start_ARG over start_ARG s end_ARGi( j ) + α over start_ARG s end_ARGi( j - 1 ) + ⋯ + αitalic_j - 1 over start_ARG s end_ARGi( 1 ) end_ARG start_ARG 1 + α + α2 + ⋯ + αitalic_j - 1 end_ARG . In practice, we use α=0.90.9α=0.9α = 0.9. 5. Calculate the average value of the importance sampling estimator p⁢()q(j−1)⁢()⁢[M⁢()=t]superscript11delimited-[] p(x)q^(j-1)(x) 1[M(x)=t]divide start_ARG p ( x ) end_ARG start_ARG q( j - 1 ) ( x ) end_ARG blackboard_1 [ M ( x ) = t ] for all samples in this batch. 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 1 for pseudocode. Algorithm 1 Independent Token Gradient Importance Sampling (ITGIS) 0: Model M, target token t, input length k, token distributions p1,…,pksubscript1…subscriptp_1,…,p_kp1 , … , pitalic_k, temperature T, iterations n, batch size B 1: si(0)←∈ℝ||←subscriptsuperscript00superscriptℝs^(0)_i 0 ^|V|s( 0 )i ← 0 ∈ blackboard_R| V | for all i∈[k]delimited-[]i∈[k]i ∈ [ k ] # Initialize score functions 2: estimates←[]←estimatesestimates←[]estimates ← [ ] 2: 3: for j←1←1j← 1j ← 1 to n do 4: for i←1←1i← 1i ← 1 to k do 5: qi(j−1)⁢(x)←pi⁢(x)⋅exp⁡(si(j−1)⁢(x)/T)←subscriptsuperscript1⋅subscriptsubscriptsuperscript1q^(j-1)_i(x)← p_i(x)· (s^(j-1)_i(x)/T)q( j - 1 )i ( x ) ← pitalic_i ( x ) ⋅ exp ( s( j - 1 )i ( x ) / T ) for all x∈x ∈ V 6: Normalize qi(j−1)subscriptsuperscript1q^(j-1)_iq( j - 1 )i to have sum 1111 7: end for 7: 8: Sample B inputs (b)b=1Bsuperscriptsubscriptsuperscript1\x^(b)\_b=1^B x( b ) b = 1B from q1(j−1)×⋯×qk(j−1)subscriptsuperscript11⋯subscriptsuperscript1q^(j-1)_1×…× q^(j-1)_kq( j - 1 )1 × ⋯ × q( j - 1 )k 9: for i←1←1i← 1i ← 1 to k do 10: s^i(j)⁢(x)←1B⁢∑b=1B[∇Mt⁢((b))]i,x←subscriptsuperscript^1superscriptsubscript1subscriptdelimited-[]subscript∇subscriptsuperscript s^(j)_i(x)← 1B _b=1^B[ _xM% _t(x^(b))]_i,xover start_ARG s end_ARG( j )i ( x ) ← divide start_ARG 1 end_ARG start_ARG B end_ARG ∑b = 1B [ ∇x Mitalic_t ( x( b ) ) ]i , x for all x∈x ∈ V 11: end for 11: 12: α←0.9←0.9α← 0.9α ← 0.9 13: for i←1←1i← 1i ← 1 to k do 14: si(j)←s^i(j)+α⁢s^i(j−1)+⋯+αj−1⁢s^i(1)1+α+α2+⋯+αj−1←subscriptsuperscriptsubscriptsuperscript^subscriptsuperscript^1⋯superscript1subscriptsuperscript^11superscript2⋯superscript1s^(j)_i← s^(j)_i+α s^(j-1)_i+…+% α^j-1 s^(1)_i1+α+α^2+…+α^j-1s( j )i ← divide start_ARG over start_ARG s end_ARG( j )i + α over start_ARG s end_ARG( j - 1 )i + ⋯ + αitalic_j - 1 over start_ARG s end_ARG( 1 )i end_ARG start_ARG 1 + α + α2 + ⋯ + αitalic_j - 1 end_ARG 15: end for 15: 16: estimate←1B⁢∑b=1B∏i=1kpi⁢(i(b))∏i=1kqi(j−1)⁢(i(b))⁢[M⁢((b))=t]←estimate1superscriptsubscript1superscriptsubscriptproduct1subscriptsubscriptsuperscriptsuperscriptsubscriptproduct1subscriptsuperscript1subscriptsuperscript1delimited-[]superscriptestimate← 1B _b=1^B _i=1^kp_i% (x^(b)_i) _i=1^kq^(j-1)_i(x^(b)_i)% 1[M(x^(b))=t]estimate ← divide start_ARG 1 end_ARG start_ARG B end_ARG ∑b = 1B divide start_ARG ∏i = 1k pitalic_i ( x( b )i ) end_ARG start_ARG ∏i = 1k q( j - 1 )i ( x( b )i ) end_ARG blackboard_1 [ M ( x( b ) ) = t ] 17: Append estimateestimateestimateestimate to estimatesestimatesestimatesestimates 18: end for 18: 19: return 1n⁢∑j=1nestimates⁢[j]1superscriptsubscript1estimatesdelimited-[] 1n _j=1^nestimates[j]divide start_ARG 1 end_ARG start_ARG n end_ARG ∑j = 1n estimates [ j ] i Algorithm for MHIS We define the proposal distribution ϕ(⋅|)φ(·|x)ϕ ( ⋅ | x ) to be the distribution induced by the following procedure: 1. Choose a random token position i∈[k]delimited-[]i∈[k]i ∈ [ k ] to modify. 2. Calculate the gradient at that token [∇Mt⁢()]i∈ℝ||subscriptdelimited-[]subscript∇subscriptsuperscriptℝ[ _xM_t(x)]_i ^|V|[ ∇x Mitalic_t ( x ) ]i ∈ blackboard_R| V | (treating =(x1,…,xk)subscriptx1…subscriptxx=( x_1,…, x_k)x = ( x1 , … , xk ) as a one-hot vector in ℝk×||superscriptℝR^k×|V|blackboard_Rk × | V |). Call this gradient gg. 3. Sample a replacement token xi′subscriptsuperscriptx′ x _ix′i from the distribution proportional to pi⁢(xi′)⋅exp⁡(xi′T).⋅subscriptsubscriptsuperscript′subscriptsubscriptsuperscript′p_i(x _i)· ( g_x _iT% ).pitalic_i ( x′italic_i ) ⋅ exp ( divide start_ARG gitalic_x′ start_POSTSUBSCRIPT i end_POSTSUBSCRIPT end_ARG start_ARG T end_ARG ) . 4. Output ′=(x1,…,xi′,…,xk)superscript′subscriptx1…subscriptsuperscriptx′…subscriptxx =( x_1,…, x _i% ,…, x_k)x′ = ( x1 , … , x′i , … , xk ). Note that the transition probability in Metropolis–Hastings only depends on the ratio q⁢(′)⁢ϕ⁢(|′)q⁢()⁢ϕ⁢(′|)=pi⁢(xi′)pi⁢(xi)⋅exp⁡(Mt⁢(′)−Mt⁢()T)⋅ϕ⁢(|′)ϕ⁢(′|),superscript′italic-ϕconditionalsuperscript′italic-ϕconditionalsuperscript′⋅subscriptsuperscriptsubscriptx′subscriptsubscriptxsubscriptsuperscript′subscriptitalic-ϕconditionalsuperscript′italic-ϕconditionalsuperscript′ q(x )φ(x|x )q(% x)φ(x |x)= p_i( % x_i )p_i( x_i)· ( M_t(% x )-M_t(x)T )· φ(x% |x )φ(x |x),divide start_ARG q ( x′ ) ϕ ( x | x′ ) end_ARG start_ARG q ( x ) ϕ ( x′ | x ) end_ARG = divide start_ARG pitalic_i ( xi′ ) end_ARG start_ARG pitalic_i ( xi ) end_ARG ⋅ exp ( divide start_ARG Mitalic_t ( x′ ) - Mitalic_t ( x ) end_ARG start_ARG T end_ARG ) ⋅ divide start_ARG ϕ ( x | x′ ) end_ARG start_ARG ϕ ( x′ | x ) end_ARG , which is easy to compute given forwards and backwards passes at xx and ′superscript′x x′. We use an initial burn-in period (see Appendix F) for the random walk before recording samples (1),…,(n)superscript1…superscriptx^(1),…,x^(n)x( 1 ) , … , x( n ). The final output of the method is the empirical importance sampling estimate 1n⁢∑j=1np⁢((j))q⁢((j))⁢[M⁢((j))=t].1superscriptsubscript1superscriptsuperscript1delimited-[]superscript 1n _j=1^n p(x^(j))q(x^(j))% 1[M(x^(j))=t].divide start_ARG 1 end_ARG start_ARG n end_ARG ∑j = 1n divide start_ARG p ( x( j ) ) end_ARG start_ARG q ( x( j ) ) end_ARG blackboard_1 [ M ( x( j ) ) = t ] . This requires computing q⁢()q(x)q ( x ), which involves the normalization constant. To save samples, we estimate the normalization constant using the identity: ∼p[exp(Mt⁢()T)]=∼q[exp(−Mt⁢()T)]−1. E_x p [ ( M_t(% x)T ) ]= E_x q% [ (- M_t(x)T ) ]^-1.blackboard_Ex ∼ p [ exp ( divide start_ARG Mitalic_t ( x ) end_ARG start_ARG T end_ARG ) ] = blackboard_Ex ∼ q [ exp ( - divide start_ARG Mitalic_t ( x ) end_ARG start_ARG T end_ARG ) ]- 1 . The right-hand side can be estimated using the n 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 2 for pseudocode. Algorithm 2 Metropolis–Hastings Importance Sampling (MHIS) 0: Model M, target token t, input length k, token distributions p1,…,pksubscript1…subscriptp_1,…,p_kp1 , … , pitalic_k, temperature T, number of burn-in steps nburnsubscriptburnn_burnnburn, number of samples n 1: Initialize xx by sampling from p1×⋯×pksubscript1⋯subscriptp_1×…× p_kp1 × ⋯ × pitalic_k 2: samples←[]←samplessamples←[]samples ← [ ] 3: for step←1←step1step← 1step ← 1 to nburn+nsubscriptburnn_burn+nnburn + n do 4: i←Unif⁢([k])←Unifdelimited-[]i ([k])i ← Unif ( [ k ] ) # Choose random position 5: ←[∇Mt⁢()]i←subscriptdelimited-[]subscript∇subscriptg←[ _xM_t(x)]_ig ← [ ∇x Mitalic_t ( x ) ]i 6: Sample xi′subscriptsuperscriptx′ x _ix′i from ∝pi⁢(xi′)⋅exp⁡(xi′/T)proportional-toabsent⋅subscriptsubscriptsuperscript′subscriptsubscriptsuperscript′ p_i(x _i)· (g_x _i/T)∝ pitalic_i ( x′italic_i ) ⋅ exp ( gitalic_x′ start_POSTSUBSCRIPT i end_POSTSUBSCRIPT / T ) # Proposed token 7: ′←(1,…,xi′,…,k)←superscript′subscript1…subscriptsuperscriptx′…subscriptx ←(x_1,…, x^% _i,…,x_k)x′ ← ( x1 , … , x′i , … , xitalic_k ) # Proposed new state 7: # Compute acceptance ratio r=q⁢(′)⁢ϕ⁢(|′)q⁢()⁢ϕ⁢(′|)superscript′italic-ϕconditionalsuperscript′italic-ϕconditionalsuperscript′r= q(x )φ(x|x )q(% x)φ(x |x)r = divide start_ARG q ( x′ ) ϕ ( x | x′ ) end_ARG start_ARG q ( x ) ϕ ( x′ | x ) end_ARG 8: r←AcceptanceRatio⁢(M,t,T,(p1,…,pk),,′)←AcceptanceRatiosubscript1…subscriptsuperscript′r (M,t,T,(p_1,…,p_k),x,% x )r ← AcceptanceRatio ( M , t , T , ( p1 , … , pitalic_k ) , x , x′ ) 9: if Unif⁢(0,1)<rUnif01Unif(0,1)<rUnif ( 0 , 1 ) < r then 10: ←′←superscript′x←x x ← x′ # Accept proposal 11: end if 11: 12: if step>nburnstepsubscriptburnstep>n_burnstep > nburn then 13: Append xx to samplessamplessamplessamples 14: end if 15: end for 15: 16: Z←(1n⁢∑j=1nexp⁡(−Mt⁢(samples⁢[j])/T))−1←superscript1superscriptsubscript1subscriptsamplesdelimited-[]1Z← ( 1n _j=1^n (-M_t(samples[j])/T)% )^-1Z ← ( divide start_ARG 1 end_ARG start_ARG n end_ARG ∑j = 1n exp ( - Mitalic_t ( samples [ j ] ) / T ) )- 1 # Est. normalizing constant 17: return 1n⁢∑j=1nexp⁡(Mt⁢(samples⁢[j])/T)Z⁢[M⁢(samples⁢[j])=t]1superscriptsubscript1subscriptsamplesdelimited-[]1delimited-[]samplesdelimited-[] 1n _j=1^n (M_t(samples[j])/T)Z % 1[M(samples[j])=t]divide start_ARG 1 end_ARG start_ARG n end_ARG ∑j = 1n divide start_ARG exp ( Mitalic_t ( samples [ j ] ) / T ) end_ARG start_ARG Z end_ARG blackboard_1 [ M ( samples [ j ] ) = t ] i Algorithm for QLD Recall that ⁢()∈ℝdsuperscriptℝv(x) ^dv ( x ) ∈ blackboard_Rd is the random vector representing the activations of the model right before the unembedding step. After collecting n samples of (1),…,(n)superscript1…superscriptv^(1),…,v^(n)v( 1 ) , … , v( n ), we compute their empirical mean ∈ℝdsuperscriptℝ μ ^ditalic_μ ∈ blackboard_Rd and covariance ∈ℝd×dsuperscriptℝ ^d× dΣ ∈ blackboard_Rd × d. Then define uu to be the whitened version of vv: uu :=−1⁢(−)assignabsentsuperscript1 := A^-1(v- μ):= italic_A- 1 ( v - italic_μ ) vv =⁢+,absent = Au+ μ,= italic_A u + italic_μ , where ∈ℝd×dsuperscriptℝ A ^d× ditalic_A ∈ blackboard_Rd × d is any matrix such that ⁢⊤=superscripttop A A = italic_A italic_A⊤ = Σ. Note that uu has mean 00 and covariance IddsubscriptIdId_dIditalic_d. From now on, we principally work in this whitened representation of activation space, as it has the convenient property that the ⋅u· eu ⋅ italic_e and ⋅′⋅superscript′u· e u ⋅ italic_e′ are uncorrelated iff eitalic_e and ′superscript′ e italic_e′ are orthogonal. We choose the unit vector ∈ℝnsuperscriptℝ d ^nitalic_d ∈ blackboard_Rn to point in the direction of the shortest accepting vector (Appendix D), then decompose our whitened samples (1),…,(n)superscript1…superscriptu^(1),…,u^(n)u( 1 ) , … , u( n ) into components parallel and perpendicular to ditalic_d: (i)superscript a^(i)a( i ) :=⁢⊤⁢(i)assignabsentsuperscripttopsuperscript := d d u^(i):= italic_d italic_d⊤ u( i ) (i)superscript b^(i)b( i ) :=(i)−(i).assignabsentsuperscriptsuperscript :=u^(i)-a^(i).:= u( i ) - a( i ) . Finally, we output: 1n2⁢|(i,j)∈[n]2|(i)+(j)∈S|,1superscript2conditional-setsuperscriptdelimited-[]2superscriptsuperscript 1n^2 | \(i,j)∈[n]^2\,\,\, |\,a^(i)+% b^(j)∈ S \ |,divide start_ARG 1 end_ARG start_ARG n2 end_ARG | ( i , j ) ∈ [ n ]2 | a( i ) + b( j ) ∈ S | , where S:=∈ℝd|arg⁢maxi((+)⋅U)i=t.S:= \ u ^d\, |\,arg\,max_i((% A u+ μ)· W_U)_i=t \.S := italic_u ∈ blackboard_Rd | start_OPFUNCTION arg max end_OPFUNCTIONi ( ( italic_A italic_u + italic_μ ) ⋅ italic_Witalic_U )i = t . This proportion can be computed in O~⁢(n)~ O(n)over~ start_ARG O end_ARG ( n ) time—we don’t need to explicitly iterate over all n2superscript2n^2n2 pairs. By the convexity of the acceptance region S, for any fixed bb there is a single interval [ℓ,r]ℓ[ ,r][ ℓ , r ] such that a∈[ℓ,r]⇔a⁢+∈S⇔ℓa∈[ ,r] a d+b∈ Sa ∈ [ ℓ , r ] ⇔ a italic_d + b ∈ S. We can efficiently compute the bounds of this interval for every sample (j)superscriptb^(j)b( j ) by solving a linear system of inequalities, and then we can calculate how many (i)superscripta^(i)a( i ) fall into each range in O⁢(log⁡n)O( n)O ( log n ) time after sorting. Thus, the computational cost of QLD is dominated by running n forwards passes of M to generate the samples (1),…,(n)superscript1…superscriptu^(1),…,u^(n)u( 1 ) , … , u( n ). See Algorithm 3 for pseudocode. Algorithm 3 Quadratic Logit Decomposition (QLD) 0: Model M with unembed matrix Usubscript W_Uitalic_Witalic_U, target token t, sample size n 1: Sample n inputs (1),…,(n)superscript1…superscriptx^(1),…,x^(n)x( 1 ) , … , x( n ) from input distribution 2: Compute pre-unembed activations (i)∈ℝdsuperscriptsuperscriptℝv^(i) ^dv( i ) ∈ blackboard_Rd by running M on (i)superscriptx^(i)x( i ) for i∈[n]delimited-[]i∈[n]i ∈ [ n ] 3: ←1n⁢∑i=1n(i)←1superscriptsubscript1superscript μ← 1n _i=1^nv^(i)italic_μ ← divide start_ARG 1 end_ARG start_ARG n end_ARG ∑i = 1n v( i ) 4: ←1n⁢∑i=1n((i)−)⁢((i)−)⊤←1superscriptsubscript1superscriptsuperscriptsuperscripttop ← 1n _i=1^n(v^(i)- μ% )(v^(i)- μ) Σ ← divide start_ARG 1 end_ARG start_ARG n end_ARG ∑i = 1n ( v( i ) - italic_μ ) ( v( i ) - italic_μ )⊤ 5: Find Aitalic_A such that ⁢⊤=+ϵ⋅Iddsuperscripttop⋅italic-ϵsubscriptId A A = +ε·Id_ditalic_A italic_A⊤ = Σ + ϵ ⋅ Iditalic_d # via Cholesky decomposition 6: for i←1←1i← 1i ← 1 to n do 7: (i)←−1⁢((i)−)←superscriptsuperscript1superscriptu^(i)← A^-1(v^(i)- μ)u( i ) ← italic_A- 1 ( v( i ) - italic_μ ) # Whitened activations 8: end for 8: 9: ←ShortestAcceptingVector⁢(,,U,t)←ShortestAcceptingVectorsubscript d ( A, μ, W% _U,t)italic_d ← ShortestAcceptingVector ( italic_A , italic_μ , italic_Witalic_U , t ) # See Appendix D 10: for i←1←1i← 1i ← 1 to n do 11: (i)←⁢⊤⁢(i)←superscriptsuperscripttopsuperscripta^(i)← d d u^(i)a( i ) ← italic_d italic_d⊤ u( i ) # Parallel component 12: (i)←(i)−(i)←superscriptsuperscriptsuperscriptb^(i)←u^(i)-a^(i)b( i ) ← u( i ) - a( i ) # Perpendicular component 13: end for 13: 14: Sort (i)i=1nsuperscriptsubscriptsuperscript1\a^(i)\_i=1^n a( i ) i = 1n in ascending order 15: count←0←count0count← 0count ← 0 16: for j←1←1j← 1j ← 1 to n do 16: # Solve linear inequalities to get [ℓj,rj]subscriptℓsubscript[ _j,r_j][ ℓitalic_j , ritalic_j ] such that a∈[ℓj,rj]⇔a⁢+(j)∈S⇔subscriptℓsubscriptsuperscripta∈[ _j,r_j] a d+b^(j)∈ Sa ∈ [ ℓitalic_j , ritalic_j ] ⇔ a italic_d + b( j ) ∈ S 17: [ℓj,rj]←FindAcceptanceInterval⁢((j),,,,U,t)←subscriptℓsubscriptFindAcceptanceIntervalsuperscriptsubscript[ _j,r_j] (b^(j),% d, A, μ, W_U,t)[ ℓitalic_j , ritalic_j ] ← FindAcceptanceInterval ( b( j ) , italic_d , italic_A , italic_μ , italic_Witalic_U , t ) 18: count←count+|i:(i)∈[ℓj,rj]|←countcountconditional-setsuperscriptsubscriptℓsubscriptcount +|\i:a^(i)∈[ _j,r_% j]\|count ← count + | i : a( i ) ∈ [ ℓitalic_j , ritalic_j ] | # Binary search for bounds 19: end for 19: 20: return count/n2countsuperscript2count/n^2count / n2 Appendix 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 decomposition =+u=a+bu = a + b for QLD to perform well: 1) aa and bb are independent, and 2) the contribution towards the output behavior is split roughly equally between the two terms. The first assumption is straightforward: if aa and bb are not independent, then +′superscript′a+b a + b′ (where ′superscript′b b′ comes from an i.i.d. copy of uu) does not have the same distribution as uu. The failure of this assumption introduces bias into the estimation method. Working in whitened space ensures that aa and bb are uncorrelated, which is a first step towards independence. The second assumption—that the contribution to the target behavior is roughly equally split between aa and bb—is necessary for QLD to have an advantage over naive sampling. For purposes of illustration, say that d=22d=2d = 2 (so that both aa and bb can be treated as scalars), and that ,∼iid⁢(0,1)superscriptsimilar-toiid01a,b % N(0,1)a , b start_RELOP SUPERSCRIPTOP start_ARG ∼ end_ARG start_ARG iid end_ARG end_RELOP N ( 0 , 1 ). Then, consider three ways that the contribution to the target behavior could be split: • Scenario 1 (no split): The target token is outputted iff >1010a>10a > 10. In this case, QLD provides no advantage over naive sampling, as the proportion of ((i),(j))superscriptsuperscript(a^(i),b^(j))( a( i ) , b( j ) ) pairs in the acceptance region is exactly the same as the proportion of ((i),(i))superscriptsuperscript(a^(i),b^(i))( a( i ) , b( i ) ) pairs. If p is the probability of the behavior, then p−1superscript1p^-1p- 1 samples are required to consistently obtain a positive estimate. • Scenario 2 (even split): The target token is outputted iff +>10⁢2102a+b>10 2a + b > 10 square-root start_ARG 2 end_ARG. In this case, QLD has a quadratic advantage over naive sampling. It only requires around p−1/2superscript12p^-1/2p- 1 / 2 samples for QLD to consistently obtain a positive estimate. • Scenario 3 (uneven split): The target token is outputted iff +2⁢>10⁢52105a+2b>10 5a + 2 b > 10 square-root start_ARG 5 end_ARG. The contribution here is split between aa and bb, though not equally, so QLD’s efficiency falls in between the previous two scenarios. We can calculate that it requires around p−5/9superscript59p^-5/9p- 5 / 9 samples for QLD to consistently obtain a positive estimate.777In general, if the condition is α⁢+1−α2⁢>101superscript210αa+ 1-α^2b>10α a + square-root start_ARG 1 - α2 end_ARG b > 10, it requires roughly p−1/(α+1−α2)2superscript1superscript1superscript22p^-1/ (α+ 1-α^2 )^2p- 1 / ( α + square-root start_ARG 1 - α start_POSTSUPERSCRIPT 2 end_ARG )2 end_POSTSUPERSCRIPT samples. In practice, the condition for outputting the target token is more complex than a single linear constraint on aa and bb. 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 bb has d−11d-1d - 1 dimensions while aa only has 1111, most choices of ditalic_d will end up giving bb much more influence over the behavior than aa. This motivates us to identify a particularly important direction with ditalic_d; in some informal sense we want to find a direction that is “d/22d/2d / 2 times more important” than the average direction. When we run QLD with many more samples than its standard budget of 216superscript2162^16216, its performance improves but plateaus at a 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 a sufficient number of samples. Appendix D Computing the Shortest Accepting Vector Recall that the acceptance region S⊆ℝdsuperscriptℝS ^dS ⊆ blackboard_Rd is the subset of whitened pre-unembed space that results in the model outputting t: S:=∈ℝd|arg⁢maxi((+)⋅U)i=t.S:= \ u ^d\, |\,arg\,max_i((% A u+ μ)· W_U)_i=t \.S := italic_u ∈ blackboard_Rd | start_OPFUNCTION arg max end_OPFUNCTIONi ( ( italic_A italic_u + italic_μ ) ⋅ italic_Witalic_U )i = t . We define ditalic_d to point in the direction of the shortest vector in S, i.e., arg⁢min∈S⁡‖subscriptargminnormarg\,min_ u∈ S\| u\|start_OPFUNCTION arg min end_OPFUNCTIONitalic_u ∈ S ∥ italic_u ∥. This vector can be approximated using an iterative convex projection algorithm. S⊆ℝdsuperscriptℝS ^dS ⊆ blackboard_Rd is an intersection of |−1|1|V-1|| V - 1 | half-spaces H1,…,Ht−1,Ht+1,…⁢Hsubscript1…subscript1subscript1…subscriptH_1,…,H_t-1,H_t+1,… H_VH1 , … , Hitalic_t - 1 , Hitalic_t + 1 , … Hcaligraphic_V, where HisubscriptH_iHitalic_i represents the set of all activations (in whitened pre-unembed space) that result in the logit on token t being larger than the logit on token i. Given any convex set C (such as a half-plane), the projection of ∉C x∉ Citalic_x ∉ C onto C is arg⁢min′∈C⁡‖−′‖2subscriptargminsuperscript′subscriptnormsuperscript′2arg\,min_ x ∈ C\| x- x \|% _2start_OPFUNCTION arg min end_OPFUNCTIONitalic_x′ ∈ C ∥ italic_x - italic_x′ ∥2. Given a collection of convex sets, there exists a simple algorithm for finding a point in their intersection: start with an arbitrary point xitalic_x, then repeatedly project xitalic_x onto a random convex set that does not already contain xitalic_x. Eventually, this process converges to a point in their intersection (Gubin et al., 1967). We apply this method to find an element of S. To ensure that it is the shortest element, we also project xitalic_x onto balls centered at 00 with smaller and smaller radii by multiplying xitalic_x by 0.990.990.990.99. The exact procedure is described in Algorithm 4.888We found a few minor bugs in our implementation of the ϵitalic-ϵεϵ-tolerance in our algorithm after we ran experiments, but we don’t expect them to have affected the results at all. In practice, it always takes much less than 100⋅nreps⋅100subscriptreps100· n_reps100 ⋅ nreps steps for the algorithm to return a value. Algorithm 4 Random Constraint Projection 0: Half-spaces H1,…,Ht−1,Ht+1,…,H||subscript1…subscript1subscript1…subscriptH_1,…,H_t-1,H_t+1,…,H_|V|H1 , … , Hitalic_t - 1 , Hitalic_t + 1 , … , H| V |, number of repetitions nrepssubscriptrepsn_repsnreps 1: ←∈ℝd←0superscriptℝ x 0 ^ditalic_x ← 0 ∈ blackboard_Rd 2: for step⁢_⁢cnt←1←step_cnt1step\_cnt← 1step _ cnt ← 1 to 100⋅nreps⋅100subscriptreps100· n_reps100 ⋅ nreps do 3: Pick a random i among all i such that ∉Hisubscriptx∉ H_ix ∉ Hitalic_i 4: Project xitalic_x onto HisubscriptH_iHitalic_i 5: if xitalic_x lies in S (up to some tolerance ϵitalic-ϵεϵ) then 6: if step⁢_⁢cnt<nrepsstep_cntsubscriptrepsstep\_cnt<n_repsstep _ cnt < nreps then 7: Scale xitalic_x by 0.990.990.990.99. 8: else 9: return xitalic_x 10: end if 11: end if 12: end for Appendix E Ground Truth Token Distribution 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 at position x is the probability of the x-th most common token). Any tokens that appeared 00 times across all 232superscript2322^32232 samples are not plotted. The hex distribution only had 159159159159 and 135135135135 tokens in the range [10−9,10−5]superscript109superscript105[10^-9,10^-5][ 10- 9 , 10- 5 ] for the 2222- and 4444-layer models, respectively, so we used every such token instead of sampling 256256256256 of them. Appendix F Computational Budgets Each estimation method was given a budget of roughly 216superscript2162^16216 model calls. More specifically: • Independent Token Gradient Importance Sampling uses 28superscript282^828 batches of size 28superscript282^828, for a total of 216superscript2162^16216 samples. The average gradient is updated after each batch. Note that this method requires backwards passes as well as forwards passes. • Metropolis–Hastings Importance Sampling uses 210+211superscript210superscript2112^10+2^11210 + 211 batches of size 25superscript252^525, for a total of 1.5⋅216⋅1.5superscript2161.5· 2^161.5 ⋅ 216 samples (the batch size indicates the number of independent random walks the method simulates). The first 210superscript2102^10210 batches are used as a burn-in period for the random walk and are discarded, so only 216superscript2162^16216 samples are actually used to calculate the estimate. • Quadratic Logit Decomposition uses n=216superscript216n=2^16n = 216 samples of the pre-unembed activation vv. The MLE direction is approximated with nreps=200subscriptreps200n_reps=200nreps = 200 iterations of the Random Constraint Projection algorithm (Appendix D); this makes up a trivial fraction of the total compute usage of the method). • Gaussian Logit Difference uses 216superscript2162^16216 samples of the logit difference to estimate μ and σ, 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 216superscript2162^16216 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 216superscript2162^16216 samples of v for all 256256256256 target tokens associated with a given behavior. Appendix G All Method Performances Table 3 shows the Itakura–Saito loss (p/q−ln⁡(p/q)−11p/q- (p/q)-1p / q - ln ( p / q ) - 1) of all estimation methods on all input distributions and model sizes. Table 3: Itakura–Saito loss comparison of all methods, distributions, and model sizes. (a) 1-layer model Distribution Constant GLD QLD ITGIS MHIS hex 2.58912.58912.58912.5891 2.21632.21632.21632.2163 2.00382.00382.00382.0038 2.04842.04842.04842.0484 1.39601.39601.39601.3960 camel 2.59082.59082.59082.5908 2.44192.44192.44192.4419 2.06482.06482.06482.0648 1.19971.19971.19971.1997 2.01872.01872.01872.0187 colon 2.77702.77702.77702.7770 2.70912.70912.70912.7091 1.27861.27861.27861.2786 1.22091.22091.22091.2209 1.02671.02671.02671.0267 if 2.24242.24242.24242.2424 2.18722.18722.18722.1872 1.23211.23211.23211.2321 1.09161.09161.09161.0916 1.44551.44551.44551.4455 caps 2.66192.66192.66192.6619 2.61472.61472.61472.6147 1.94131.94131.94131.9413 1.47881.47881.47881.4788 2.40232.40232.40232.4023 english 1.90951.90951.90951.9095 1.81201.81201.81201.8120 1.24091.24091.24091.2409 1.45391.45391.45391.4539 0.70170.70170.70170.7017 spanish 2.60792.60792.60792.6079 2.44632.44632.44632.4463 1.56281.56281.56281.5628 1.63961.63961.63961.6396 2.15382.15382.15382.1538 icl 2.54672.54672.54672.5467 2.43282.43282.43282.4328 2.23732.23732.23732.2373 0.79920.79920.79920.7992 1.23441.23441.23441.2344 Average 2.49062.49062.49062.4906 2.35752.35752.35752.3575 1.69521.69521.69521.6952 1.36651.36651.36651.3665 1.54741.54741.54741.5474 (b) 2-layer model Distribution Constant GLD QLD ITGIS MHIS hex 2.88392.88392.88392.8839 2.84532.84532.84532.8453 2.76982.76982.76982.7698 2.66522.66522.66522.6652 1.49851.49851.49851.4985 camel 2.32422.32422.32422.3242 2.22682.22682.22682.2268 2.26792.26792.26792.2679 1.92211.92211.92211.9221 2.06372.06372.06372.0637 colon 2.98932.98932.98932.9893 2.83352.83352.83352.8335 2.12152.12152.12152.1215 2.29862.29862.29862.2986 1.90421.90421.90421.9042 if 2.61722.61722.61722.6172 2.43772.43772.43772.4377 1.83971.83971.83971.8397 1.03011.03011.03011.0301 1.25161.25161.25161.2516 caps 2.63452.63452.63452.6345 2.68202.68202.68202.6820 2.48472.48472.48472.4847 1.92711.92711.92711.9271 1.69811.69811.69811.6981 english 2.09892.09892.09892.0989 2.09562.09562.09562.0956 1.34621.34621.34621.3462 0.99080.99080.99080.9908 1.44801.44801.44801.4480 spanish 2.54422.54422.54422.5442 2.36622.36622.36622.3662 1.56701.56701.56701.5670 1.09511.09511.09511.0951 2.30952.30952.30952.3095 icl 2.63812.63812.63812.6381 2.54192.54192.54192.5419 2.36012.36012.36012.3601 1.64201.64201.64201.6420 1.05021.05021.05021.0502 Average 2.59132.59132.59132.5913 2.50362.50362.50362.5036 2.09462.09462.09462.0946 1.69641.69641.69641.6964 1.65301.65301.65301.6530 (c) 4-layer model Distribution Constant GLD QLD ITGIS MHIS hex 2.48032.48032.48032.4803 2.29342.29342.29342.2934 2.42822.42822.42822.4282 2.30902.30902.30902.3090 2.08302.08302.08302.0830 camel 2.48952.48952.48952.4895 2.32712.32712.32712.3271 2.45342.45342.45342.4534 2.09372.09372.09372.0937 2.29612.29612.29612.2961 colon 2.93252.93252.93252.9325 2.73192.73192.73192.7319 2.43822.43822.43822.4382 2.12952.12952.12952.1295 1.17101.17101.17101.1710 if 2.64772.64772.64772.6477 2.54082.54082.54082.5408 1.83131.83131.83131.8313 1.84301.84301.84301.8430 0.92610.92610.92610.9261 caps 2.59702.59702.59702.5970 2.53822.53822.53822.5382 2.41422.41422.41422.4142 2.04842.04842.04842.0484 1.29721.29721.29721.2972 english 2.70082.70082.70082.7008 2.74322.74322.74322.7432 1.56811.56811.56811.5681 0.99430.99430.99430.9943 1.30511.30511.30511.3051 spanish 2.64152.64152.64152.6415 2.50222.50222.50222.5022 1.77161.77161.77161.7716 1.66111.66111.66111.6611 2.19362.19362.19362.1936 icl 2.50292.50292.50292.5029 2.19252.19252.19252.1925 2.38662.38662.38662.3866 2.29712.29712.29712.2971 0.54770.54770.54770.5477 Average 2.62402.62402.62402.6240 2.48372.48372.48372.4837 2.16152.16152.16152.1615 1.92201.92201.92201.9220 1.47751.47751.47751.4775 Appendix H Squared Error in Log-Space Figure 6 and Table 4 show the method performances when measured using squared error in log-space loss (i.e., (log⁡p−log⁡q)2superscript2( p- q)^2( log p - log q )2) 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 in log-space. However, the importance sampling temperatures are not changed between the two metrics (they were tuned while minimizing Itakura–Saito loss). 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 8888 distributions, with bands indicating standard error. The colored points indicate the loss on individual distributions, with horizontal jitter added for visibility. Table 4: Squared error in log-space loss comparison of all methods, distributions, and model sizes. (a) 1-layer model Distribution Constant GLD QLD ITGIS MHIS hex 6.00936.00936.00936.0093 5.18155.18155.18155.1815 4.48634.48634.48634.4863 3.71793.71793.71793.7179 2.19632.19632.19632.1963 camel 7.72957.72957.72957.7295 6.49756.49756.49756.4975 5.84885.84885.84885.8488 1.51561.51561.51561.5156 4.32484.32484.32484.3248 colon 6.85546.85546.85546.8554 5.43615.43615.43615.4361 2.13792.13792.13792.1379 2.35492.35492.35492.3549 1.80611.80611.80611.8061 if 5.39775.39775.39775.3977 4.92874.92874.92874.9287 2.93562.93562.93562.9356 1.95551.95551.95551.9555 3.22233.22233.22233.2223 caps 8.09278.09278.09278.0927 7.56757.56757.56757.5675 4.13494.13494.13494.1349 3.21733.21733.21733.2173 6.55506.55506.55506.5550 english 6.00516.00516.00516.0051 5.32205.32205.32205.3220 3.26393.26393.26393.2639 1.83571.83571.83571.8357 1.63951.63951.63951.6395 spanish 5.79905.79905.79905.7990 5.56495.56495.56495.5649 2.97402.97402.97402.9740 2.12312.12312.12312.1231 4.26124.26124.26124.2612 icl 5.90355.90355.90355.9035 5.69155.69155.69155.6915 4.80234.80234.80234.8023 1.02221.02221.02221.0222 2.30302.30302.30302.3030 Average 6.47406.47406.47406.4740 5.77375.77375.77375.7737 3.82303.82303.82303.8230 2.21782.21782.21782.2178 3.28853.28853.28853.2885 (b) 2-layer model Distribution Constant GLD QLD ITGIS MHIS hex 8.05368.05368.05368.0536 7.36747.36747.36747.3674 7.65337.65337.65337.6533 6.80046.80046.80046.8004 1.99131.99131.99131.9913 camel 7.68497.68497.68497.6849 6.35186.35186.35186.3518 7.23257.23257.23257.2325 4.26954.26954.26954.2695 5.60465.60465.60465.6046 colon 7.17287.17287.17287.1728 5.64165.64165.64165.6416 3.39243.39243.39243.3924 4.18514.18514.18514.1851 2.82422.82422.82422.8242 if 6.33466.33466.33466.3346 5.69235.69235.69235.6923 3.36953.36953.36953.3695 1.70591.70591.70591.7059 2.16392.16392.16392.1639 caps 7.76827.76827.76827.7682 6.86056.86056.86056.8605 5.47715.47715.47715.4771 3.04113.04113.04113.0411 3.34323.34323.34323.3432 english 5.27425.27425.27425.2742 5.11855.11855.11855.1185 3.44663.44663.44663.4466 1.63351.63351.63351.6335 3.27283.27283.27283.2728 spanish 6.37186.37186.37186.3718 5.68775.68775.68775.6877 3.84043.84043.84043.8404 2.15822.15822.15822.1582 4.69374.69374.69374.6937 icl 6.27856.27856.27856.2785 6.20616.20616.20616.2061 5.07405.07405.07405.0740 2.06152.06152.06152.0615 1.35151.35151.35151.3515 Average 6.86736.86736.86736.8673 6.11586.11586.11586.1158 4.93574.93574.93574.9357 3.23193.23193.23193.2319 3.15563.15563.15563.1556 (c) 4-layer model Distribution Constant GLD QLD ITGIS MHIS hex 7.55597.55597.55597.5559 6.92086.92086.92086.9208 7.07507.07507.07507.0750 6.43646.43646.43646.4364 3.16053.16053.16053.1605 camel 7.55977.55977.55977.5597 6.54136.54136.54136.5413 7.12637.12637.12637.1263 3.50673.50673.50673.5067 5.55015.55015.55015.5501 colon 7.17917.17917.17917.1791 6.60356.60356.60356.6035 3.58173.58173.58173.5817 4.39164.39164.39164.3916 1.78971.78971.78971.7897 if 6.69006.69006.69006.6900 5.90725.90725.90725.9072 3.89873.89873.89873.8987 3.71533.71533.71533.7153 1.68241.68241.68241.6824 caps 9.11059.11059.11059.1105 8.29578.29578.29578.2957 6.13836.13836.13836.1383 4.36154.36154.36154.3615 2.25142.25142.25142.2514 english 5.50035.50035.50035.5003 5.48665.48665.48665.4866 3.35613.35613.35613.3561 1.91801.91801.91801.9180 2.19382.19382.19382.1938 spanish 7.04907.04907.04907.0490 6.46416.46416.46416.4641 4.20834.20834.20834.2083 3.44253.44253.44253.4425 4.84254.84254.84254.8425 icl 5.17935.17935.17935.1793 4.45694.45694.45694.4569 4.47724.47724.47724.4772 4.55844.55844.55844.5584 1.02161.02161.02161.0216 Average 6.97806.97806.97806.9780 6.33456.33456.33456.3345 4.98274.98274.98274.9827 4.04134.04134.04134.0413 2.81152.81152.81152.8115 Appendix I Temperature Tuning Both importance sampling methods require choosing a temperature parameter T. To tune T, we sweep over 9999 different temperatures from 0.20.20.20.2 to 5555, uniformly spaced in log-space. We choose the value of T that achieves the lowest loss on 100 randomly chosen tokens with ground-truth probabilities in the range [10−5,10−3]superscript105superscript103[10^-5,10^-3][ 10- 5 , 10- 3 ] to prevent over-fitting. We tune separate temperatures for each distribution, model size, and importance sampling method, shown in Table 5. 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 5: Temperatures T used for the different methods. 1 layer 2 layers 4 layers Distribution ITGIS MHIS ITGIS 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 0.45 1.50 0.67 2.24 0.45 1.50 spanish 0.67 2.24 0.67 2.24 1.00 2.24 icl 0.45 1.00 0.30 0.67 3.34 0.67 Appendix J Plots of Method Outputs Figures 7, 8, and 9 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 of 00 will all lie on this line after the fit is applied. 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 is the output of the model after a fit. 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 a fit. Figure 9: The outputs of methods, after a fit is applied, on all 256 tokens for each distribution on the 4-layer model. The horizontal axis represents the ground-truth token probability, while the vertical axis is the output of the model after a fit.

Instant research discovery