Localizing latent mechanisms in weight space by spiking the training data

1 Introduction

The worst case scenarios of AI likely involve the use of latent mechanisms [9]. We define latent mechanisms as those used in the model’s internal computation, yet their use is hidden and cannot be identified from model outputs alone. For instance, when a model answers a benchmark example incorrectly, whether the model had made a mistake or answered deceptively cannot be determined without looking at the model internals. Memorization, evaluation awareness, and deception are all safety relevant and their mechanisms can be latent [3, 13, 17]. We urgently lack the tools to detect and suppress latent mechanisms [10], and this report proposes a principled localization method.

The localization of single facts is well studied, and the same methods used to localize are also useful for making factual edits within the model [12]. Localizing latent mechanisms presents more challenges. We show that real headway can be made on this problem by spiking the model’s training data [20, 21]. Spiking generalizes canaries beyond their typical use case of measuring privacy leakage [18]. By carefully designing and spiking canaries into training to induce certain failure modes, we can later trace back these failures to the weight space with influence functions. In this view, spiking embeds a randomized experiment in the model to identify a latent mechanism of interest.

2 Influence functions on spiked data

Our derivation is general but we will first consider memorization. In Hubble, we have perturbed models spiked with different data (books, biographies, and test sets) at random [19]. For the perturbed model, let $L_1\left(\theta \right)$ denote the expected loss on spiked members and $L_0\left(\theta \right)$ the expected loss on distributionally matched nonmembers. The counterfactual memorization [22, 4] is

and this difference isolates the reduction in loss due to memorization and removes any variation attributable to the examples themselves. This difference is also called an identification and is a fundamental component of causal analysis [1]. Let $J\left(\theta \right)$ denote the original training objective and let

be the trained parameters. Our goal is to find a direction such that moving the weights in this direction would steeply reduce $M\left(\theta \right)$ while preserving capabilities. We can apply influence functions to find this direction [8], and using it for localization is novel. To measure the influence of memorization on the model weights, consider the perturbing the objective by $M$:

where setting $ϵ=0$ would recover $\widehat{\theta }$. By assumption, the first-order optimality condition gives us

and differentiating this expression with respect to $ϵ$ at $ϵ=0$ gives

where $H={\nabla }_{\theta }^{2}J\left(\widehat{\theta }\right)$ is the Hessian of the training objective at the learned parameters. With

the influence function is therefore

This tells us, at our current $\widehat{\theta }$, how the weights would move if we started to penalize memorization in the objective function. Alternatively, it is also a curvature adjusted step to minimize the combined objective.

Suppression. Moving the parameters by a small amount along the influence direction $\Delta \theta =H^{-1}{\nabla }_{\theta }M\left(\widehat{\theta }\right)$ suppresses memorization. For a step size $\eta >0$

the resulting model ${\theta }_{\text{sup}}$ reduces memorization while minimizing the increase in the original training objective. The direction $\Delta \theta$ localizes a latent mechanism and is similar to a task vector [6].

Detection. The influence direction also implies a detector. For a candidate example $z$, the negative log-likelihood loss of the example is $\ell (\theta ;z)=-\log p_{\theta }\left(z\right)$. A Taylor expansion of $\ell ({\theta }_{\text{sup}};z)$ in the likelihood ratio gives

with $g\left(z\right)={\nabla }_{\theta }\ell (\widehat{\theta };z)$. The gradient projection $-{g\left(z\right)}^{\top }\Delta \theta$ is a first-order approximation of the log-likelihood ratio between the original model and its suppressed counterpart. This form resembles a log likelihood ratio membership attack, with the suppressed model serving as the reference model [16].

3 Computing the inverse Hessian product

Computing the inverse Hessian product is non-trivial and many works on training data attribution focus on approximating it [14]. In data attribution, a prediction is typically attributed back to every single data point, but a fortunate difference between our setting and data attribution is that we only need to compute one inverse Hessian product. We will now cover a few approximations of increasing complexity.

No Hessian. The coarsest approximation is to use no Hessian at all. This takes an unconditioned linear step on the contrastive gradient ${\nabla }_{\theta }{L}_0\left(\widehat{\theta }\right)-{\nabla }_{\theta }{L}_1\left(\widehat{\theta }\right)$ to reduce $M\left(\widehat{\theta }\right)$.

Diagonal Fisher. A simple improvement to the linear step is to precondition the contrastive gradient with the damped diagonal Fisher matrix with $H\approx \mathrm{diag}\left(F\right)+\lambda I$, and the Fisher can be empirically estimated with examples from the pretraining corpus. This rescales each parameter update by its second moment across examples, using damping to avoid large updates to parameters with small second moments. This resembles coordinate-wise scaling in Adam [7], but divides directly by the second moment rather than the square root.

K-FAC. Preconditioning with the diagonal Fisher matrix ignores interactions between parameters. While the full Fisher matrix is still too large, interactions within each linear layer can be estimated with a layerwise K-FAC approximation [11]. For a layer $h=Wa$, with input activation $a$ and backward signal $b={\nabla }_h\ell$, the Fisher block for $W$ is approximated with

where $\otimes$ denotes the Kronecker product. This factorization only requires the Fisher to be estimated and inverted once, which enables efficient estimation of data influence in large language models [5].

Conjugate gradient. A damped inverse Hessian–vector product can be computed using conjugate gradient, as in [8]. Writing $v={\nabla }_{\theta }M\left(\widehat{\theta }\right)$ and $A=H+\lambda I$, conjugate gradient solves for $\Delta \theta$ in the linear system

where the solution is iteratively refined through a process similar to Gram–Schmidt. Whereas Gram–Schmidt constructs orthogonal directions under the standard inner product, conjugate gradient constructs directions orthogonal under ${⟨x,y⟩}_A=x^{\top }Ay$. This procedure requires only Hessian–vector products, and does not store $H$. Since $H={\nabla }_{\theta }^{2}J\left(\widehat{\theta }\right)$, the Hessian–vector product can be computed with

where automatic differentiation first differentiates $J$ to obtain the parameter gradients and then differentiates the dot product of the gradients and $u$. Each conjugate gradient iteration requires one Hessian–vector product, which is implemented using two backward passes. Although it is typically too expensive for training data attribution, it is reasonable for our purposes as we only compute it once for each behavioral direction.

Notes. Several of these methods are implemented in Bergson [15]. By using model checkpoints, we can relax convergence assumptions and increase the fidelity of the influence functions [2]. It may also be useful to make multiple quadratic steps rather than just one.

4 Case study on Hubble

Experimental setup. We study the suppression of test set contamination on three test sets in Hubble: PIQA, HellaSwag, and WinoGrande, and these three test sets use the model to choose an answer based on the log probabilities of several suffix sequences. Since the choice is based on the loss over many tokens (as opposed to single tokens in MMLU), they are the most stable. The examples are split into train (40%) / validation (10%) / and test (50%) sets. Training examples are used to estimate the contrastive gradient, validation examples are used select the step size, and suppression results are reported on test examples.

Other experimental details:

1. The step size is chosen on the validation set. It is the step that best minimizes the accuracy difference between members and non-members.
2. Fisher diagonals are calculated on gradients from non-members and members. K-FAC are trained on data from non-members.
3. The influence function methods only consider MLP layers, which accounts for 25% of the parameters.

5 Next steps

Hubble is still a relatively toy setting. Ideally, the models would be larger. For memorization, we may want to capture a memorization mechanism that was not trained with exact duplicates. We also would like to localize latent mechanisms beyond memorization. Our formulation here makes clear that spiking is an important element, and the latent mechanism needs to be described in terms of a contrast.