Inducing Induction in Llama via Linear Probe Interventions

Authors: sheridan  feucht
byron c wallace
david  bau
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-13 23:53:20

Abstract: For induction heads to copy forward information successfully, heads in earlier layers must first load previous token information into every hidden state, a process Olsson et al. (2022) call key shifting. While this information is hypothesized to exist, there have been few attempts to explicitly locate it in models. In this work, we use linear probes to identify the subspaces responsible for storing previous token information in Llama-2-7b and Llama-3-8b. We show that these subspaces are causally implicated in induction by using them to "edit" previous token information and trigger random token copying in new contexts.

Main paper: 
 introduction: 
 What information encoded in a token allows it to be copied via induction? Although the workings of induction heads have been carefully studied (Elhage et al., 2021; Olsson et al., 2022) , there has yet to be work explicitly isolating the presumed previous token information that enables induction heads to function as they do. Using linear probes from Feucht et al. (2024) , we identify low-dimensional subspaces in Llama-2-7b (Touvron et al., 2023) and Llama-3-8b (Meta, 2024) that contain previous token information used for random token copying. To show this, we take information identified by linear probes for one input and substitute it into a completely different context (Figure 1 ), which artificially triggers copying in certain layers.2 Related Work Elhage et al. (2021) and Olsson et al. (2022) describe induction circuits as consisting of two stages:(1) key shifting, where one set of attention heads copies previous token information into succeeding token positions, and (2) prefix matching, where another set of heads later attends to that information to copy previously-seen sequences. While they1. Generate source sequence A B C … A B s = 2. Generate base sequence D E F … b = D E 3. Modify base sequence D E F … b* = D B4. Replace previous token info.A ← B ← D ← E ← A B C … A B s = D E F … b* = D B F B ← Figure 1:We artificially induce copying behavior in Llama, using probes from Feucht et al. (2024) to modify previous token information.characterize these heads through direct examination of weights and attention patterns, we show causal evidence of this process by editing previous token information in Llama models. 
random copying task: 
 First, we define a simple copying task wherein a sequence of l uniformly sampled tokens is duplicated and separated by a newline. We measure Llama's ability to predict the last token of such a sequence, which it can only achieve via in-context copying (e.g. land.id.Pale \n land.id → Pale). For l = 30 with random tokens, Llama-2-7b has a completion accuracy of 87.9% for this task, whereas Llama-3-8b reaches 100% accuracy.Next, we show that previous token information measured by probes from Feucht et al. ( 2024) is causally implicated in this random token copying task. Figure 1 shows the setup for this experiment. We first randomly generate source and base sequences s and b, following the same newlineseparated template as in Section 3.1. We then replace the last token of b with the last token of s to create b * , which prevents models from completing the copying task: Llama-2-7b copying accuracy drops to 0% with this change, whereas Llama-3-8b accuracy drops to 13%. Finally, we show that we can restore copying behavior for b * by modifying the hidden state of the token we want to copy (i.e. F in Figure 1 ). We use linear probes to extract previous token information [← B] from the hidden representation for C at layer ℓ (Figure 1 ), and then substitute this "flag" into the hidden state for F at the same layer, measuring the resulting accuracy for prediction of F over 100 sequences.To do this, we calculate the SVD of the probe, P ℓ = U ΣV * , and then take the top d rows of V * , which form an orthonormal basis for a subspace A of dimension d. To substitute previous token information from a source hidden state s ℓ t into a hidden state from the modified base sequence b * ℓ t , we take the projection of s ℓ t onto A and add it to the projection of b * ℓ t onto the orthogonal complement of A. Thus, the formula for an intervened hidden state at token position t and layer ℓ is simplyh ℓ t = AA T s ℓ t + (I -AA T )b * ℓ t .We perform this intervention for every layer ℓ.In addition to testing on random sequences of tokens, we also test sequences that end in multi-token words (MTWs) taken from Wikipedia (Foundation, 2022) . We then modify s so that it ends with a Wikipedia MTW (e.g. land.id.task.ed \n land.id.task → ed). This means that the previous token information inserted into b * now comes from the final token of a MTW. Motivated by results from Feucht et al. (2024) showing "erasure" of that information in early layers, we ask whether this will affect induction behavior for those sequences. 
results: 
 The leftmost column of Figure 2 shows that this intervention causes a 34% increase in accuracy above baseline for Llama-2-7b and a 26% increase for Llama-3-8b on random sequences of tokens, suggesting that models attend to the artificiallyinserted [← B] information to promote the output of F as hypothesized. Results for MTW copying are more mixed. For Llama-2-7b, interventions cease to be effective with MTWs. However, Llama-3-8b interventions are arguably stronger. 
discussion: 
 Results in Figure 2 indicate that there is some mechanism after layer 5 in Llama-2-7b (layer 14 in Llama-3-8b) that refers to previous token information in order to copy tokens from its context. We interpret this result as direct evidence of hypothesized "key shifting" being used for copying in fullscale models. Although this effect is strong for random tokens, MTWs are treated differently. Since prior work shows that previous token information is "erased" for MTWs in both models, we expected interventions to become less effective for MTWs; however, this is only true for Llama-2-7b. For Llama-3-8b, interventions instead become more effective, with fewer dimensions required to achieve high accuracy (d = 16 instead of d = 32). This difference in MTW induction mechanisms may indicate a fundamental difference between how these two models represent multi-token words. 
 
 Appendix:  


Decision: This is an ongoing mechanistic interpretability work aimed at identifying specific subspaces in Llama-2.7b and Llama-3.8b responsible for storing previous token information (induction heads) using a defined artificial copying task, which is a good fit for the workshop.