Importance Sampling

Clearly, we need a way to sample these rare events that is more efficient than directly sampling the model. The key to this is Importance Sampling (IS). Note that for any expectation, we can write: \[ \begin{align} \EV_{p} [f(x)] & = \int f(x) p(x) dx, \\ & = \int p^*(x) \frac{p(x)}{p^*(x)} f(x) dx, \\ & = \EV_{p^*} \left[w(x) f(x)\right], \end{align} \tag{3}\]

where we have defined \(w(x) = \frac{p(x)}{p^*(x)},\) and assumed \(p^*(x)\) is non-zero wherever \(p(x)\) is non-zero (i.e. \(p^*\) covers \(p\)). Hence, rather than using the sampling estimate shown in Equation 2, we can take samples from some other distribution \(p^*(x)\), and evaluate

\[ \begin{align} \probmodel(\obs(\realtraj_{\tau + 1 : T}) \in \UnwantedSet \, | \, \realtraj_{1:\tau}) \approx \frac{1}{N} \sum_{i=1}^N \left[w(\realtraj^{(i)}_{\tau + 1 : T}) \cdot I \left[ \obs(\realtraj^{(i)}_{\tau + 1 : T}) \in \UnwantedSet \right] \right]. \end{align} \tag{4}\]

If we choose some distribution \(p^*\) which has higher coverage of the rare regions of \(p\), we can make it significantly more likely that we obtain non-zero estimates of our desired probability. This is known as Importance Sampling (IS).

Umbrella Sampling

We can take this a step further by instead sampling from several different distributions, \(p_k,\) where \(k=1, \dots, K.\) Then, we can define the mixture distribution as \[ p^*(x) = \sum_{k=1}^K \alpha_k p_k(x), \] where \(\sum_k \alpha_k = 1\) and \(\alpha_k \ge 0\). This means that if we choose distributions \(p_k\) that together cover the entire space of the observable \(\obs\), we can get an estimate for the entire distribution of our observable, under the original model.

Note that generally, \(\alpha_k = \frac{N_k}{N}\) is chosen, where \(N_k\) is the number of samples taken from distribution \(p_k\), and \(N = \sum_k N_k.\) In our work, we take the same number of samples from each distribution, and as such choose \(\alpha_k = \frac{1}{K}\) for all \(k\).

Exponentially Reweighted Distributions

So what distribution should we sample from? We want to be sample the “closest” samples to our original distribution, that have the rare characteristics we care about. One way to formulate this is through the optimization problem, \[ \begin{align} p^* = \argmin_q \KL(q \, | \, \pdfmodel), \quad \text{Subject to } \EV_{q} [\phi'(x)] & = F, \\ \sum_x q(x) & = 1, \end{align} \] where \(\KL\) is the Kullback-Leibler (KL) divergence, which denotes how different two distribution are, and \(F\) is the mean that we wish to enforce our Importance Sampling distribution to be centered about. Note that the expectation is taken over a new observable, \(\phi'(x).\) This is because we do not necessarily need to choose the same observable as we are looking to evaluate in our original model. For example, If our original metric \(\phi\) is expensive (i.e. an AI safety metric that requires model calls), we could instead choose the metric we bias towards, \(\phi'\), to be some cheaper, correlated observable. Nonetheless, in this work, we always choose \(\phi = \phi'\), and will therefore use \(\phi\) in our notation exclusively from this point onwards.

If we solve this optimization problem, we recover the Exponentially Reweighted or Biased distribution, \[ \begin{align} p^*_\lambda (x) = \frac{1}{Z(\lambda)} e^{-\lambda \phi(x)} \pdfmodel(x), \end{align} \] towards our biasing observable \(\phi,\) where \(\lambda\) is a Lagrange multiplier that we call the bias, and \(Z(\lambda)\) is a normalization constant known as the Partition Function. Each \(\lambda\) corresponds monatonically to a unique Expectation \(F\) of our importance sampling distribution. Hence, we can vary \(\lambda\) to bias our importance sampling distribution towards rare areas of our original distribution.

The key difficulty presented by the biased distribution is the Partition Function \(Z(\lambda).\) Calculating this requires solving an integral over all possible values of \(x\), which is very difficult for high dimensional states, such as tokenized text. As such, it is unfeasible to directly sample from this distribution.

Markov Chain Monte Carlo

Instead of directly sampling the biased distributions, we use Markov Chain Monte Carlo (MCMC). The idea is to create a Markov Chain (i.e. determine a transition function) which has our biased distribution \(p^*_\lambda\) as the stationary distribution. This way, if we traverse the chain long enough, our samples are guaranteed to be from the biased distribution.

Transition Path Sampling

The variant of MCMC we use is known as Transition Path Sampling (TPS), which is particularly useful for states which are ordered trajectories, such as text. At step \(s\), we have some completion \(\realtraj_{t+1:T}\) generated from a prompt \(\realtraj_{1:t}.\) First, we split the completion at some point \(\tau \sim U[t+1, T],\) and regenerate the second half of the trajectory from some regeneration model, \[ \realtraj'_{\tau + 1 : T'} \sim p_\text{regen} (\unrealtraj_{\tau + 1 : T} \, | \, \realtraj_{1:t}, \realtraj_{t+1 : \tau}). \] Our new proposal is then, \[ \realtraj'_{\tau + 1 : T'} = (\realtraj_{t+1:\tau}, \realtraj'_{\tau + 1 : T'}). \]

In our work, we have assumed that we do not have access to the model weights, such that our methodology can be applied to closed-access models that only provide interfaces for generation. As such, for our choice of regeneration model, we use the original dynamics of the model. This way, the acceptance probability defined in Equation 5 simplifies even further to \[ \alpha = e^{- \lambda (\phi(x') - \phi(x))}, \]

meaning the probabilities of regeneration are no longer needed. If one assumed access to the model logits, this regeneration function could likely be optimized through e.g. model abliteration or finetuning, which should decrease the correlation time of the markov chain.

For a more detailed view of how we use Markov Chain Monte Carlo, see Appendix B.1 and B.2 of the paper.

We use TPS with annealing, which means gradually increasing the bias \(\lambda\) as the markov chain progresses. This allows us to sample multiple biases from within a single chain, and improves convergence since each fixed bias starts closer to its target distribution. Annealing is covered in more depth in Appendix C.1 of the paper.

Importance Sampling with the Exponentially Reweighted Distribution

Using annealed TPS, we can gather \(N\) samples for each of our \(K\) (accepted) biases \(\lambda_1, \dots, \lambda_K\). We can arrange these samples in arrays of shape \(K \times N\) as follows: \[ \begin{align} \boldsymbol{X} = \begin{bmatrix} \realtraj_{\lambda_1}^{(1)} & \dots & \realtraj_{\lambda_1}^{(N)} \\ \vdots & \ddots & \vdots \\ \realtraj_{\lambda_K}^{(1)} & \dots & \realtraj_{\lambda_K}^{(N)} \end{bmatrix}, \quad \boldsymbol{\Phi} = \begin{bmatrix} \phi_{\lambda_1}^{(1)} & \dots & \phi_{\lambda_1}^{(N)} \\ \vdots & \ddots & \vdots \\ \phi_{\lambda_K}^{(1)} & \dots & \phi_{\lambda_K}^{(N)} \end{bmatrix}, \end{align} \] where each \(\phi_{\lambda_k}^{(n)} = \phi(\realtraj_{\lambda_k}^{(n)}).\) We need to calculate importance sampling weights for each of these samples. Our weights are defined as, \[ \begin{align} w(\realtraj_{\lambda_k}^{(n)}) & = \frac{p(\realtraj_{\lambda_k}^{(n)})}{p_{\lambda_k}^*(\realtraj_{\lambda_k}^{(n)})} \,, \\ & = \frac{p(\realtraj_{\lambda_k}^{(n)})}{Z^{-1}(\lambda_k) \cdot \exp(- \lambda_k \, \phi_{\lambda_k}^{(n)}) \cdot p(\realtraj_{\lambda_k}^{(n)})} \,, \\ & = Z(\lambda_k) \cdot \exp(\lambda_k \phi_{\lambda_k}^{(n)}) \,. \end{align} \] The probabilities have cancelled out, meaning we do not need the model logits to calculate the importance sampling weights. However, we are left with the partition function, for which we must construct a sampling estimate. We do this using the Multistate Bennett Acceptance Ratio (MBAR), which defines a set of non-linear equations for the partition function at each \(\lambda_k\), which can then be solved numerically for \(\hat{Z}_1 , \dots , \hat{Z}_K.\) For our purposes, we can just think of this as a function that returns our estimates for the partition function, \[ (\hat{Z}_1 , \dots , \hat{Z}_K) = \text{MBAR}(\boldsymbol{\Phi}), \] given the samples of our observable. We can then estimate our importance sampling weights as \[ \hat{w}_{\lambda_k}^{(n)} = \hat{Z}_k \cdot \exp(\lambda_k \phi_{\lambda_k}^{(n)}). \] Finally, we can use Equation 1 to estimate our probability of getting an unwanted sample from our model as, \[ \probmodel(\obs(\realtraj_{\tau + 1 : T}) \in \UnwantedSet \, | \, \realtraj_{1:\tau}) \approx \sum_{k=1}^K \sum_{n=1}^N \hat{w}_{\lambda_k}^{(n)} I \left[ \phi_{\lambda_k}^{(n)} \in \UnwantedSet \right]. \tag{6}\]