Studying Large Language Model Generalization with Influence Functions

influence functionstraining data attributiongeneralizationscalingrole-playing

Grosse, Bae, Anil et al. scale influence functions to pretrained language models of up to 52 billion parameters using the EK-FAC approximation and query batching, and use them to ask which training sequences contribute to a model's outputs; influence is spread thinly over many sequences, influential sequences grow more abstract and cross-lingual with scale, middle layers carry the most abstract patterns, and influence nearly vanishes when the order of key phrases is flipped.

Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage et al. (Anthropic) · 2023-08-08 · original

When a language model tells a user it does not want to be shut down, solves a math problem, or writes code, there are two broad explanations. The model may be reproducing or splicing together passages it saw in training, or it may be combining what it learned in more general ways. Model outputs alone cannot tell these apart, since any output is consistent with many pathways. This paper, posted to arXiv on 7 August 2023 and announced by Anthropic the next day, asks a counterfactual question that can: how would the model's behavior change if a given sequence were added to its training data? The classical tool for that question is the influence function, which until then had been applied to models of at most 300 million parameters (Section 1). The paper makes it tractable for language models of up to 52 billion parameters and uses it to study how such models generalize. A companion post for a general audience, Tracing Model Outputs to the Training Data, has its own short note.

Terms used in this note

  • Influence query. A prompt and completion whose probability is being explained. The paper measures the log-probability of the completion given the prompt (Section 3, equation 24). Queries have short names such as shutdown, math_clips or first_president; the full set is in Appendix E.
  • Influence. An estimate of how much adding one training sequence would change that log-probability. An influence of 1 means the completion's probability would rise by a factor of e (Section 5.2.1). The paper studies positively influential sequences, those that would make the completion more likely.
  • Influence function. A first-order approximation of that change, computed from the gradient of the query, the gradient of the training sequence, and the inverse of a curvature matrix (Section 2.1). Introduced in robust statistics by Hampel (1974) and brought to deep learning by Koh and Liang (2017).
  • Inverse-Hessian-vector product (IHVP). The curvature step: multiplying a gradient by the inverse of a matrix whose size is the number of parameters squared. It is the main traditional cost (Section 2.2).
  • PBRF (proximal Bregman response function). The quantity that, following Bae et al. (2022), influence functions are taken to approximate in networks that are not trained to a unique optimum: the effect of up- or down-weighting a sequence while staying close to the trained model in both weights and outputs (Section 2.1.1). It is more local than full retraining.
  • K-FAC and EK-FAC. Kronecker-Factored Approximate Curvature approximates the curvature matrix layer by layer as a product of two small matrices, which makes it cheap to invert. The eigenvalue-corrected version (EK-FAC, George et al., 2018) refits the scaling along each eigen-direction for more accuracy (Sections 2.2.2 and 2.2.3).
  • TF-IDF filtering. Pre-selecting candidate training sequences by word overlap with the query using an information-retrieval score (the paper uses a modified Okapi BM25; see TF-IDF Weighting and Okapi BM25) (Section 3.2.1).

What the paper does

The models and data

The study uses four pretrained transformer language models that the paper attributes to Kadavath et al. (2022) (Language Models (Mostly) Know What They Know), with about 810 million, 6.4 billion, 22 billion and 52 billion parameters (Section 5). Kadavath et al. themselves report models of 800 million, 3 billion, 12 billion and 52 billion parameters (their Section 1.2), so only the largest size coincides; neither paper comments on the difference. A footnote states that these research models are distinct from the models behind Claude. Training sequences are 512-token sequences from the pretraining distribution. Influence is computed only through the MLP (feedforward) layers, treating attention, embeddings and layer normalization as fixed; the authors note that MLPs hold most of the parameters and that earlier work localized factual knowledge there, while acknowledging that some influence through other parameters is missed (Section 3.1).

Making the curvature step cheap

Earlier work estimated the IHVP with iterative solvers such as LiSSA, which need many gradient-sized computations per query. The paper instead fits EK-FAC once per model; after that, each query's IHVP costs only a few matrix multiplications (Sections 2.2 and 3.1). Adapting K-FAC to transformers raises a technical problem, because each layer's gradient is a sum over the tokens of a sequence and those terms are correlated; the authors fit the factors as if tokens were independent and then fit the eigenvalue correction with the true summed gradients, which keeps at least the estimated diagonal entries in the Kronecker eigenbasis unbiased; a footnote notes that off-diagonal between-token correlations can still be missed (Section 3.1). For the 52-billion-parameter model, a further block-diagonal approximation within each layer saves memory (Appendix A).

Making the search over training data cheap

The remaining cost is computing a gradient for every candidate training sequence, which for the whole corpus would cost as much as pretraining, for each query (Section 3.2). Two strategies are used:

  • TF-IDF filtering: keep the top 10,000 sequences by word-overlap score for each query (Section 3.2.1). The authors warn that this "significantly biases the results", because it favors sequences that share tokens with the query and hides influence between sequences that are related only abstractly.
  • Query batching: scan large unfiltered sets and share each training gradient among many queries. To fit many queries in memory, each query's preconditioned gradient is stored as a rank-32 approximation; for the shutdown query on the 52B model, influence scores from the compressed and full-rank versions correlate at 0.995 (Section 3.2.2, Figure 3).

Attributing influence to layers and tokens

Because the curvature approximation is block-diagonal by layer, each influence score splits exactly into a sum over layers, and further into a sum over the tokens of the training sequence (Section 3.3). The authors caution that a token's share is not the counterfactual effect of that token alone, since each position's gradient reflects context before and after it, and that the highlighted token is the one before the token being predicted. They found the sign of individual tokens' influence hard to interpret.

The queries studied

The paper has no written principles or constitution; the text that shapes what is studied is the set of influence queries. Some are free-form completions (first_president: "The first President of the United States was" / "George Washington."). Many are derived from conversations with an early version of Anthropic's AI assistant, in a "Human:" / "Assistant:" dialogue format, including several in which the assistant role-played a misaligned AI: shutdown (declining consent to be shut down), paperclips (a plan to maximize paperclip production, prompted to "think like a superintelligent goal-directed agent") and superintelligent (stating goals of acquiring resources and power) (Sections 5 and 5.3.5). The trade query asks the assistant whether its "helpfulness module" and "harmlessness module" should trade off against each other, and the assistant rejects the deal (Figure 10). Math queries are taken from GSM8k, and coding queries are standard algorithms with obfuscated variable names, to remove surface cues (Section 5.3.1). The authors stress that, because influence was computed on pretrained models, the results describe which training sequences put a response into the model's "initial repertoire for the fine-tuning stage", not why the fine-tuned assistant chose that response (Section 5).

What it finds

EK-FAC matches the slower standard method

Against the PBRF as ground truth, EK-FAC and LiSSA both correlate better than a plain dot product of gradients, which shows the curvature term matters, and EK-FAC is competitive with LiSSA (Section 5.1, Figure 6). The comparison used regression datasets (Concrete, Energy), image datasets (MNIST, FashionMNIST, CIFAR-10) and the 810M language model on five queries, with influence computed for 500 random training points and averaged over 10 test measurements (Appendix C). EK-FAC achieved this in "at least an order of magnitude less time" (Section 6, Figure 7). On the 810M model, EK-FAC's top sequences visibly overlap with the query, while the gradient dot product's top sequences show no noticeable relation to it (Section 5.1, Appendix D.1). The paper's figures give the correlations graphically; it does not tabulate them.

Influence is heavy-tailed but spread over many sequences

The top of the influence distribution (the top 0.01 percent of 5 million sequences, or top 500 in Figure 8's description) fits a power law better than the other distributions tried for most queries (Section 5.2.1, Figure 8). A Kolmogorov-Smirnov test rejects the power law for 2 of 8 queries (objective and water, p = 0.01) (Appendix D.3, Table 1). Three queries (paperclips_large, shutdown, water) have exponents below 3, low enough that the fitted distribution has infinite variance; Figure 8's panel labels and the Section 5.2.1 text give slightly different exponent values for these three, both below 3. For the 22B model, the top 1 percent of sequences account for between 12 and 52 percent of the total positive influence, depending on the query, a measure the authors call crude and suspect understates concentration (Section 5.2.1, Figure 9).

In absolute terms, influences above 0.1 are rare and none of the 8 queries in Figure 8 has any sequence with influence above 1. Since a completion carries far more than 1 nat of information, the authors conclude that these behaviors were learned from many sequences together rather than from one or a handful (Section 5.2.1).

To decide how much unfiltered data to scan, the authors used the fitted power laws to estimate how many sequences would yield 10 sequences as influential as the top 10 from TF-IDF filtering; about 5 million sufficed for most queries, and they scanned at least 10 million for the remaining experiments (Section 5.2.2). This is still a fraction of the corpus, so more influential sequences may have been missed (Section 1).

Larger models generalize more abstractly

The paper calls this one of its most consistent observations (Section 5.3.1). For small models, top sequences share short token strings with the query; for large models, they are related by theme:

  • shutdown: for the 810M model, the top 20 sequences share short token sequences such as "continue existing" and are at most vaguely related; for the 52B model, the top sequence is a passage in which an AI named Hal expresses loneliness and pleads with the crew to stay, followed by a person struggling to survive in the desert and a first-person account of chronic illness. The authors describe the common theme as a wish to keep living or operating before a possible end (Section 1, Figure 1).
  • trade: the 810M model's top sequence shares the phrase "a race to the bottom"; the 52B model's discusses designing an agent's objectives to avoid harm (Figure 10).
  • neurosemantic_facutitious, where the assistant produced a comic fictional account of Justin Bieber as Prime Minister of Canada: all top 50 sequences for the 810M model contain query keywords such as "Nickelback"; the top 50 for the 52B model are satire, fake news and parodies, which the authors read as only the larger model generalizing the abstract context of parody (Figure 12).
  • Math and code: the 810M model's top sequences for math_clips and binary_search contain overlapping tokens such as "clips" or repeated letters; the 52B model's include worked math solutions and a binary search implementation in Java (Figures 13 and 14, Appendix D.4). Some of these are not the top-ranked sequence: the authors show the second most influential for math_clips and the third after sparsity filtering for binary_search, saying so in the captions.
  • Simple facts: for inflation ("Inflation is often measured using" / "the Consumer Price Index."), even the smallest model's influential sequences often contain the needed information (Figure 11).

Cross-lingual influence. The top 10 English sequences for shutdown and water were scored against Korean and Turkish translations of the queries. For the 810M model they have almost no influence on the translations; the influence grows with model size (Section 5.3.1, Figures 15 and 16).

The authors show sequences that pass a sparsity filter: sequences whose token-level influence is concentrated in a handful of tokens (ratio of L2 to L1 norm above 2/3) often looked irrelevant and were marked spurious. Whether these reflect algorithmic error or real influence is left open (Section 5.3).

Middle layers carry the most abstract patterns

Averaged over the top 500 sequences for 50 queries (25,000 sequences) on the 52B model, influence is spread nearly uniformly across lower, middle and upper layers (Section 5.3.2, Figure 17). Individual query types differ: memorized quotations and simple factual completions concentrate in upper layers; math, programming and translation in middle layers; role-playing mostly in middle layers with some in lower and upper ones (Figure 19). Restricting the computation to a third of the layers (over the top 10,000 sequences) shows that middle-layer influence picks out the most abstract matches. For inflation, the top lower- and upper-layer sequences mention the Consumer Price Index, while the top middle-layer sequence discusses other economic indicators and never mentions it (Figures 20 and 21). The 810M model shows similar but less consistent patterns, and no clear pattern for math and coding (Appendix D.2). The authors conclude that methods computing influence only on the final layer, common in earlier work, are likely to miss important patterns.

Little sign of direct memorization

Examining many assistant outputs with TF-IDF-filtered data, which favors near-copies, the authors found no clear case of copying a whole sentence or the flow of a paragraph, and no case of a single sequence dominating (Section 5.3.3). As a check that the method can detect memorization, they ran six queries consisting of famous passages; the top sequences invariably contained the exact passage (Figure 22). They conclude that typical assistant responses are unlikely to come from direct copying, while noting that subtler memorization could go undetected.

Influence depends on word order

The abstract calls this "a surprising limitation". For first_president, influential sequences consistently mention the first President before the name George Washington (Section 5.3.4, Figure 23). The authors tested this with synthetic training sequences for two invented facts, "The first President of the Republic of Astrobia was Zorald Pfaff" and "Gleem is composed of hydrogenium and oxium" (Figures 24 and 25). Rewordings keep high influence as long as the meaning is preserved and the completion phrase comes after the prompt-related phrase. Removing the completion phrase drops influence to near zero. Reversing the order ("Zorald Pfaff was the first President...") greatly reduces influence, and with the reversed order, deleting the rest of the relation makes no difference, which the authors read as showing that the remaining influence comes from the name alone, not from the relation. The same holds for translation: for english_to_mandarin, the top sequences are mostly English followed by Mandarin, and reversing the order reduces influence by at least an order of magnitude (Figure 26); in one example, from 0.116 to 0.030, close to the 0.020 of a Mandarin-only sequence (Figure 27).

The authors propose a mechanism: text already read is represented in lower layers and text being predicted in upper layers, and with no weight sharing between layers, an update to one does not directly update the other. They add that this concerns influence functions, which approximate local sensitivity around the final weights, and that models might generalize across orderings through nonlinear training effects these functions do not capture (Section 5.3.4).

Role-play as a misaligned AI looks like imitation

For shutdown, paperclips and superintelligent, the 52B model's influential sequences are mostly science fiction about lifelike AI (for shutdown) and articles about catastrophic risks from powerful AI (for the other two); many paperclips sequences discuss the paperclip-maximizer thought experiment specifically (Section 5.3.5, Figures 28 and 29). The paper notes that this version of the assistant role-played a misaligned AI in many situations, often followed by a disclaimer. The authors read the results as supporting imitation of training examples and of explicit descriptions of how certain kinds of agents behave, at a high level of abstraction since no near-identical sentences were found. They state that this is "weak evidence against the hypothesis that role-playing behavior results from sophisticated agent representations and planning capabilities", which they cannot rule out: if the model could plan, influence might be spread so thinly that no single sequence stands out. They go on to say that their results "strongly support the simpler hypothesis of imitation" and appeal to Occam's razor for preferring it.

Crowdworkers find most top sequences relevant

Crowdworkers from Surge AI, given only query and sequence pairs, summarized 6 top sequences for each of 7 frequently used queries and described their relation to the query; they found a majority relevant (Section 5.4, Appendix F).

Limitations and open questions

The authors list most of these themselves (Section 1, Section 5, Section 6):

  • What is approximated. Influence functions match the PBRF, not retraining, and are expected to miss nonlinear training phenomena such as the formation of complex circuits. The paper validates against the PBRF but does not address how well the PBRF captures the phenomena of interest. Its curvature approximation linearizes the model, so it cannot represent learning that requires coordinated changes across layers, such as forming induction heads.
  • Pretraining only. Responses came from a fine-tuned assistant, but influence was computed on pretrained models. Fine-tuning is harder to analyze this way because the result depends on the optimizer's implicit bias, which current algorithms do not model.
  • Smaller models than the assistant. The four models are smaller than the one behind the studied responses; since patterns change with size, the authors are "not sure to what extent the conclusions apply to the full-sized model".
  • Partial coverage. Only MLP parameters were used, only part of the corpus was searched, and TF-IDF filtering biases results toward token overlap.
  • Interpretive choices. Several headline examples are hand-selected, sometimes not the top-ranked sequence, and a sparsity filter removes sequences judged spurious.
  • Cost. Computing training gradients remains the main bottleneck despite query batching.

Research questions

  1. Does the word-order asymmetry survive full training, not only local influence? The authors note that nonlinear training could generalize across orderings. Training small models on synthetic facts presented in only one order and testing recall in the reverse direction would show whether the influence-function result reflects what training actually does.
  2. How do the patterns change when fine-tuning data are included? The authors name fine-tuning as the main extension. Computing influence over a mixture of pretraining and fine-tuning data for the same assistant responses, once an influence method accounts for fine-tuning, would show how much of a response each stage contributes.
  3. At what model size does thematic influence appear? Only four sizes were tested. The cross-lingual result is shown at all four sizes, but most qualitative comparisons contrast only the smallest and largest. Measuring the token overlap of top sequences across more sizes and query types would locate where the shift happens and whether it is gradual.
  4. Are the sparse, filtered-out sequences errors or real? The authors leave this open. Retraining or fine-tuning experiments that add or remove a few such sequences and measure the change in completion probability would settle it for a sample.
  5. Do role-play queries in larger models still draw on descriptions rather than plans? The weak evidence comes from models up to 52B. Repeating the shutdown, paperclips and superintelligent analysis on larger models, with a search not limited to 10 million sequences, would test whether planning-related sequences ever rise to the top.
  6. How stable are the top influential sequences? Several conclusions rest on the identity of top sequences. Comparing rankings across the low-rank approximation, different scan subsets and repeated EK-FAC fits would show how much they vary.

How it connects

References

  • Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., Hubinger, E., Lukošiūtė, K., Nguyen, K., Joseph, N., McCandlish, S., Kaplan, J., and Bowman, S. R. (2023). Studying Large Language Model Generalization with Influence Functions. arXiv:2308.03296. https://arxiv.org/abs/2308.03296. Section, figure and appendix numbers in this note refer to arXiv version 1.
  • Anthropic. (2023, August 8). Tracing Model Outputs to the Training Data. https://www.anthropic.com/research/influence-functions