Anthropic recently published the paper Studying Large Language Model Generalization with Influence Functions, which describes a scalable technique for measuring which training examples were most influential for a particular set of weights/outputs of a trained model. This can help us better understand model generalization, offering insights into the emergent properties of AI systems. For instance, influence functions could help us answer questions like "is the model asking not to be shut down because it has generalized that this is a generically good strategy for pursuing some goal, or simply because texts where AIs ask not to be shut down are commonly found in the training corpus?". 

In this post, I aim to summarize the approximations used in the paper to calculate the influence of different training examples and outline how the approximations can be implemented in PyTorch to form the basis of further research on influence functions by the AI safety community.  

(Note: most formulae are copied or adapted from the original paper, with a few additional derivation steps / simplified notation used for clarity in some places.)

Deriving the exact form of the influence function

Before we go into approximations, it is necessary to understand what specifically we are trying to measure. 

Given some element $z_m=(x_m,y_m)$ of a dataset $D=\{z_i{\}}_{i=1}^N$, we define the response function as the optimal solution ${\theta }^{\ast }$ (weights that minimize expected loss $L$) as a function of the weighting $ϵ$ of this example. 

${\theta }^{\ast }\left(ϵ\right)={\text{argmin}}_{\theta \in R^D}\frac{1}{N}\sum L(z_i,\theta )+ϵL(z_m,\theta )$

We define the influence $I_{{\theta }^{\ast }}\left(z_m\right)$ of $z_m$ on ${\theta }^{\ast }$ using the first-order Taylor approximation to the response function at $ϵ=0$.

$\Delta \theta ={\theta }^{\ast }\left(ϵ\right)-{\theta }^{\ast }\left(0\right)\approx \frac{\partial {\theta }^{\ast }\left(ϵ\right)}{\partial ϵ}{|}_{ϵ=0}\cdot ϵ=I_{{\theta }^{\ast }}\left(z_m\right)\cdot ϵ$

We can get $\frac{\partial {\theta }^{\ast }\left(ϵ\right)}{\partial ϵ}$ the following way:

We know ${\theta }^{\ast }$ is a minimum of $\frac{1}{N}\sum L(z_i,\theta )+ϵL(z_m,\theta )$ and so the gradient wrt $\theta$ is zero at that point

${\nabla }_{\theta }\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast })\right)=0$

Differentiating each side wrt $ϵ$:

(The LHS both directly depends on $ϵ$, and indirectly via ${\theta }^{\ast }$, so we use the Implicit Function Theorem $u\left(x\right)=f(x,g(x\left)\right)=0\to$ $\frac{du}{dx}=\frac{\partial f}{\partial g}\frac{\partial g}{\partial x}+\frac{\partial f}{\partial x}$) 

${\nabla }_{\theta }^2\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast })\right)\cdot \frac{\partial {\theta }^{\ast }}{\partial ϵ}+\frac{\partial }{\partial ϵ}{\nabla }_{\theta }\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast })\right)=0$

The second term can be simplified: $\frac{\partial }{\partial ϵ}{\nabla }_{\theta }\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast })\right)={\nabla }_{\theta }L(z_m,{\theta }^{\ast })$

And so we can rearrange to get an expression for $\frac{\partial {\theta }^{\ast }}{\partial ϵ}$:

$\frac{\partial {\theta }^{\ast }}{\partial ϵ}=-{\nabla }_{\theta }^2\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast }){)}^{-1}{\nabla }_{\theta }L\right(z_m,{\theta }^{\ast })$

This tells us how the optimal parameters change with a perturbation $ϵ$ to the weighting of an added data point $z_m$. The change is proportional to the negative product of the inverse Hessian of the loss on all the data and the gradient of the loss on the data point in question with respect to the model parameters (both evaluated at the optimal parameters).

For simplicity, as in the paper, we'll denote ${\nabla }_{\theta }^2\left(\frac{1}{N}\sum L(z_i,{\theta }^{\ast })+ϵL(z_m,{\theta }^{\ast })\right)$ as $H$.

Therefore, $I_{{\theta }^{\ast }}\left(z_m\right)\approx -H^{-1}{\nabla }_{\theta }L(z_m,{\theta }^{\ast })$ (This corresponds to Equation 3 in the paper). 

Influence on some function of the model weights

So far, we have derived an expression for the influence of an added data point on the parameters $\theta$. However, we are more interested in the influence of particular data points on some measurable properties of the model, such as the output logits or validation loss. We can see this as some function $f\left({\theta }^{\ast }\right)$ of the trained parameters. 

 By the chain rule $\frac{\partial f\left({\theta }^{\ast }\right(ϵ\left)\right)}{\partial ϵ}={\nabla }_{\theta }f\left({\theta }^{\ast }\right)\frac{\partial {\theta }^{\ast }\left(ϵ\right)}{\partial ϵ}$ and so

$I_f\left(z_m\right)\approx -{\nabla }_{\theta }f\left({\theta }^{\ast }{)}^T{H}^{-1}{\nabla }_{\theta }L\right(z_m,{\theta }^{\ast })$ (This corresponds to Equation 5 in the paper).

Problems with this expression

• Hessian could have zeros and be not invertible (optimal parameters could be underspecified by loss function in case of overparameterized models)
• We often don't train to convergence, so the first derivative of the loss wrt the parameters is then not necessarily zero, as previously assumed

The paper mentions that because of these problems, "past works have found influence functions to be inaccurate for modern neural networks."

How do we fix this?

One approach is to define a new objective that:

• Has a single defined optimum in parameter space 
• Is fully optimized when the model stops training

This is what the proximal Bregman objective (PBO) attempts to define. 

${\theta }^s\left(ϵ\right)={\text{argmin}}_{\theta \in R^D}\frac{1}{N}\sum \left(L\right(z_i,\theta )-L(z_i,{\theta }^s)-{\nabla }_{y_i}L(z_i,{\theta }^s{)}^T(y_i-y_i^s))$

$+ϵL(z_m,\theta )+\frac{\lambda }{2}||\theta -{\theta }^s|{|}^2$

($y_i$ here is the output of the model at the parameters $\theta$ on input $x_i$, and $y_i^s$ is the output of the model at parameters ${\theta }^s$ on input $x_i$)

The PBO basically introduces a penalty for diverging too far from the initialized parameters, so there is some defined optimum that balances moving too far from the parameters at initialization and achieving good loss. 

So we can redefine the gradients used in $I_f$ in terms of this new loss function that considers both the loss given a new training data point and the divergence from current parameters.

From Bae et al.'s 2022 paper If Influence Functions are the Answer, Then What is the Question?:

...while influence functions for neural networks are often a poor match to LOO [Leave One Out] retraining, they are a much better match to what we term the proximal Bregman response function (PBRF). Intuitively, the PBRF approximates the effect of removing a data point while trying to keep the predictions consistent with those of the (partially) trained model.

...

In addition, although the PBRF may not necessarily align with LOO retraining due to the warm-start^[1], proximity, and non-convergence gaps, the motivating use cases for influence functions typically do not rely on exact LOO retraining. This means that the PBRF can be used in place of LOO retraining for many tasks such as identifying influential or mislabelled examples 

Applying the Implicit Function Theorem to the PBO, we can obtain an influence function with respect to the PBO objective (Equation 9 in the paper):

$I_{{\theta }^s}\left(z_m\right)=\frac{d{\theta }^s}{dϵ}=-(G+\lambda I{)}^{-1}{\nabla }_{\theta }L(z_m,{\theta }^s)$ 

$I_{f_{{\theta }^s}}\left(z_m\right)=-{\nabla }_{\theta }f\left({\theta }^s{)}^T\right(G+\lambda I{)}^{-1}{\nabla }_{\theta }L(z_m,{\theta }^s)$

Where $G$ is the Gauss-Newton Hessian $G=E\left[J^T{H}_{y}J\right]$. $J$ is the Jacobian - the first derivative of the network's outputs with respect to the parameters, and $H_y$ is the Hessian of the loss with respect to the network's outputs. 

Efficient calculation

So, we want to get $-{\nabla }_{\theta }f\left({\theta }^s{)}^T\right(G+\lambda I{)}^{-1}{\nabla }_{\theta }L(z_m,{\theta }^s)$. 

Let's assume we have the following:

1. A trained network with parameters ${\theta }^s$
2. An observable property of the network, $f\left({\theta }^s\right)$, for instance, its output logits $y_{\text{target}}$ for some chosen input $x_{\text{target}}$
3. The training dataset $D$ 
4. The loss function $L$ the model was trained on 

The key ingredients needed to calculate $I_f$ are:

1. The gradient of the property $f$ of interest with respect to the parameters $\theta$, evaluated at ${\theta }^s$
2. A way of getting the inverse damped Gauss-Newton Hessian vector product
3. The gradient of the loss $L$ on the training data points (which we want to calculate the influence of) with respect to the parameters $\theta$, evaluated at ${\theta }^s$

Notice that only key ingredient 1 depends on the property of interest. We can pre-compute ingredients 2 and 3 and then use this to test a bunch of different properties (for example, find the most influential training examples for a bunch of different model input-output pairs). 

We can also calculate the influence as a batched operation over many training data points (batch ${\nabla }_{\theta }L(z_m,{\theta }^s)$ over multiple $z_m$'s) to increase efficiency via vectorization. 

Which leaves the final key question: how do we get $(G+\lambda I{)}^{-1}v$?

Kronecker-Factored Approximate Curvature (KFAC)

Originally introduced in the 2015 paper Optimizing Neural Networks with Kronecker-factored Approximate Curvature by Martens and Grosse, KFAC is an approximation to the Fischer information matrix (FIM) $=E_{x\sim p\left(x\right),y\sim P(y|x;\theta )}\left[{\nabla }_{\theta }\log p\right(y|x;\theta \left){\nabla }_{\theta }\log p\right(y|x;\theta {)}^{\top }]$ that can be inverted very efficiently. In the case of many models where the loss is given by the negative log probability associated with a simple predictive distribution, the FIM is equal to the Gauss-Newton Hessian.

KFAC for MLP models involves the following:

Given a fully connected model with $L$ layers, let's assume each layer $l$'s output is:

$a_l={\varphi }_l\left(W_l{a}_{l-1}\right)$

where $a_{l-1}\in R^M$, $W_l\in R^{P\times M}$ and ${\varphi }_l$ is a nonlinear activation function. ^[2]

When we backpropagate $\log p(y|x;\theta )$ to get ${\nabla }_{\theta }\log p(y|x;\theta )$, we need to calculate the derivative of the $\log p(y|x;\theta )$ with respect to intermediate stages of the computation at each layer.  So as we go backward through the computational graph, once we get to the output of $W_l{a}_{l-1}$ , we'll have computed ${\nabla }_{W_l{a}_{l-1}}\log p(y|x;\theta )$. 

By the chain rule, and using the fact that ${\nabla }_{W_l}{W}_l{a}_{l-1}=a_{l-1}$, ${\nabla }_{W_l}\log p(y|x;\theta )={\nabla }_{W_l{a}_{l-1}}\log p(y|x;\theta )\cdot a_{l-1}^T$

This means we can decompose the gradient^[3] of the log-likelihood loss on some data point $(x,y)$ with respect to the weight matrix $W$ into the intermediate gradients of the loss with respect to the output of applying the weight matrix and the activations prior to that layer.

Working with gradients of weight matrices is inconvenient though, as we end up with 3D tensors for the Jacobian. We can instead consider ${\theta }_l$, the unrolled weight matrix for layer $l$. 

Then, defining $Dv={\nabla }_v\log p(y|x;\theta )$, $s_l=W_l{a}_{l-1}$ , and $\otimes$ as the Kronecker product:

$D{\theta }_l=a_{l-1}\otimes D{s}_l$

So far, so exact... But now, time for approximations. KFAC makes things simpler by assuming:

• Gradients $D{\theta }_l$ are uncorrelated between different layers
• Activations $a_l$ are independent of pre-activation gradients $D{\theta }_l$

This allows us to write down a simple block-diagonal approximation for $\text{G}$:

$G_l=E\left[D{\theta }_{l}D{\theta }_l^T\right]=E[a_{l-1}{a}_{l-1}^T\otimes D{s}_{l}D{s}_l^T]\approx E\left[a_{l-1}{a}_{l-1}\right]\otimes E\left[D{s}_{l}D{s}_l^T\right]=A_{l-1}\otimes S_l$

Where $A_{l-1}$ and $S_l$ are uncentered covariance matrices for the layer's input activations and pre-nonlinearity gradients, respectively. 

This structure enables us to efficiently get the inverse (approximate) Gauss-Newton Hessian vector product:

Let $V_l$ denote the entries of $v$ for layer $l$, reshaped to match $W_l$, and let $v_l=\text{vec}\left(V_l\right)$

Using various Kronecker product identities, we can compute the inverse (approximate) Gauss-Newton Hessian vector product as: 

${\hat{\text{G}}}_l^{-1}{v}_l=\text{vec}\left({\text{S}}_l^{-1}{V}_l{\text{A}}_{l-1}^{-1}\right)$

Eigenvalue correction

We made an approximation earlier when we went from $E[a_{l-1}{a}_{l-1}^T\otimes D{s}_{l}D{s}_l^T]$ to $E\left[a_{l-1}{a}_{l-1}\right]\otimes E\left[D{s}_{l}D{s}_l^T\right]$

Using the eigendecompositions of $A$ and $S$:

$\text{A}={\text{Q}}_A{\Lambda }_A{\text{Q}}_A^T$

$\text{S}={\text{Q}}_S{\Lambda }_S{\text{Q}}_S^T$

we can write a more accurate expression for $\text{G}$:

$\text{G}\approx ({\text{Q}}_A\otimes {\text{Q}}_S)\Lambda ({\text{Q}}_A\otimes {\text{Q}}_S{)}^T$

where the diagonal matrix $\Lambda$ is defined as:

${\Lambda }_{ii}=E\left[\right(({\text{Q}}_A\otimes {\text{Q}}_S)D\theta {)}_i^2]=E[\left({\text{Q}}_S\text{vec}\right(D\theta \left){\text{Q}}_A^T{)}_i^2\right]$  

which "captures the variances of the pseudo-gradient projected onto each eigenvector of the K-FAC approximation".

We can get the damped inverse Gauss-Newton Hessian vector product approximation by adding $\lambda$ to the eigenvalues, obtaining:

$(G+\lambda I{)}^{-1}v\approx ({\text{Q}}_A\otimes {\text{Q}}_S\left)\right(\Lambda +\lambda I{)}^{-1}({\text{Q}}_A\otimes {\text{Q}}_S{)}^{T}v$

$=\text{vec}\left({\text{Q}}_S^T\right[\left({\text{Q}}_{S}V{\text{Q}}_A^T\right)\oslash \text{unvec}\left({\text{diag}}^{-1}\right(\Lambda +\lambda I\left)\right)\left]{\text{Q}}_A\right)$

Influence functions for autoregressive models

A few details change when we want to calculate $I_{f_{{\theta }^s}}\left(z_m\right)=-{\nabla }_{\theta }f\left({\theta }^s{)}^T\right(G+\lambda I{)}^{-1}{\nabla }_{\theta }L(z_m,{\theta }^s)$ for a Transformer language model trained with an autoregressive loss function. 

In this case, the property of interest $f$ considered (the thing we are calculating the influence on) is the log-likelihood of a particular token string completion $z_c$, given a token string prompt $z_p$^[4]:

$\text{log}p(z_c|z_p;\theta )$

The paper only considers measuring the influence on a subset of the Transformer's weights - only the MLP layers - so the MLP $G$ approximation derived above applies almost exactly. 

However, the parameter gradients are now summed over token indices:

$D{\theta }_l={\sum }_{t=1}^{T}D{\theta }_{l,t}={\sum }_{t=1}^T{a}_{l-1,t}\otimes D{s}_{l,t}$

Each diagonal block of $G$ is given by $E\left[D{\theta }_{l}D{\theta }_l^T\right]$, however we want to take into account how this second moment is affected by the inter-token correlations and so cannot as accurately directly approximate with  $E\left[a_{l-1}{a}_{l-1}\right]\otimes E\left[D{s}_{l}D{s}_l^T\right]$ as before. 

The paper presents the following middle-ground between efficiency and accuracy:

We first fit the covariance factors $\text{A}$ and $\text{S}$ as if the tokens were fully independent, and compute their respective eigendecompositions. Then, when fitting the diagonal matrix $\Lambda$, we use the exact pseudo-gradients $D{\theta }_l$ which are summed over tokens. This way, at least the estimated diagonal entries of the moments in the Kronecker eigenbasis are unbiased.

Implementing in PyTorch

As described above, the key ingredients for $I_f$ are:

1. The gradient of the property $f$ of interest with respect to the parameters $\theta$, evaluated at ${\theta }^s$
2. The inverse damped Gauss-Newton Hessian, which we can calculate from the expectations of the following quantities:
   1. $a_{l-1}$ - the MLP layer inputs
   2. $D{s}_l$ - the MLP pre-nonlinearity gradients (gradients of loss wrt output of linear transformation $W_l{a}_{l-1}$)
3. The gradient of the loss $L$ on the training data points (which we want to calculate the influence of) with respect to the parameters $\theta$, evaluated at ${\theta }^s$

We can get 1) and 3) by simply fetching parameter.grad^[5] after performing a backward pass of the loss on some input, target pair.

We can get 2a) using a forward hook that saves the input to a layer during the forward pass. We can get 2b) using a backward hook on the linear layer that saves the gradient wrt the linear layer's output. 

You can find my implementation attempt on GitHub here ^[6]- includes code applying influence functions analysis to a vanilla MLP trained on MNIST and a 2-layer transformer trained on a basic next character prediction task. 

Results of small experiment on MNIST

I trained an MLP on MNIST (with flattened images) and then used the influence function approximation code to extract influential training examples for particular predicted test set labels.

I found that influential training digits were usually more sloppy / unclear compared to the average MNIST digit, and shared some resemblance with the query image. Not all top influential images shared the same label as the query. I only searched a subset of the training corpus, for efficiency. 

Here are some examples, filtered by cases where the influence was non-negligible (some queries returned ~0 for all sampled training datapoints) (first image on left is query, followed by top most influential training images given by the approximation):

1. ^{^}

   The warm-start problem referenced by Bae et al. refers to the fact that for a not strictly convex objective, the influence of a training example in the neighborhood of a minimum ${\theta }^{\ast }$ may be different from the influence at a different initialization point. 

2. ^{^}

   The paper uses homogeneous vector notation to account for biases / affine transformations - you can assume there is a 1 appended to the activations $a$ and a bias vector appended to $W$ to cover this case.

3. ^{^}

   The paper refers to these as "pseudo-gradients" since they are sampled from the final output distribution and are distinct from gradients during training.

4. ^{^}

   The $z_p$, $z_c$ pair is referred to as the "query" in the paper, as we are "querying" which training examples were most influential for the model producing $z_c$ given $z_p$.  

5. ^{^}

   Specifically, concatenate a linear layer's .weight and .bias grads

6. ^{^}

   If you look through the code and find any bugs (quite possible) or performance improvements (definitely findable; e.g. more batching + splitting of GPU ops - WIP) I'd be super happy to merge PRs and/or hear from you! I hope to gradually improve this codebase and run larger experiments.