Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks
Authors: Samyak Jain
Robert Kirk
Ekdeep Singh Lubana
Robert P. Dick
Hidenori Tanaka
Tim Rocktäschel
Edward Grefenstette
David Krueger
Venue: ICLR 2024 poster (2024)
Keywords: Fine-Tuning
Interpretability
Mechanisms
Primary area: visualization or interpretation of learned representations
TL;DR: We demonstrate that fine-tuning models rarely alters their underlying capabilities.

Abstract: Fine-tuning large pre-trained models has become the de facto strategy for developing both task-specific and general-purpose machine learning systems, including developing models that are safe to deploy. Despite its clear importance, there has been minimal work that explains how fine-tuning alters the underlying capabilities learned by a model during pretraining: does fine-tuning yield entirely novel capabilities or does it just modulate existing ones? We address this question empirically in synthetic, controlled settings where we can use mechanistic interpretability tools (e.g., network pruning and probing) to understand how the model's underlying capabilities are changing. We perform an extensive analysis of the effects of fine-tuning in these settings, and show that: (i) fine-tuning rarely alters the underlying model capabilities; (ii) a minimal transformation, which we call a `wrapper', is typically learned on top of the underlying model capabilities, creating the illusion that they have been modified; and (iii) further fine-tuning on a task where such ``wrapped capabilities'' are relevant leads to sample-efficient revival of the capability, i.e., the model begins reusing these capabilities after only a few gradient steps. This indicates that practitioners can unintentionally remove a model's safety wrapper merely by fine-tuning it on a, e.g., superficially unrelated, downstream task. We additionally perform analysis on language models trained on the TinyStories dataset to support our claims in a more realistic setup.

Published as a conference paper at ICLR 2024
MECHANISTICALLY ANALYZING THE EFFECTS OF FINE- TUNING ON PROCEDURALLY DEFINED TASKS
Samyak Jain' Robert Kirk? Ekdeep Singh Aubanaa?44. Robert P. Dick Hidenori Tanaka 1.5 Edward Grefenstette" Tim Rocktäschel, David Krueger
University of Cambridge, UK University College London, UK
&EECS Department, University of Michigan, Ann Arbor, MI, USA *Center for Brain Science, Harvard University, Cambridge, MA, USA Physics & Informatics Laboratories, NTT Research, Inc., Sunnyvale, CA, USA
ABSTRACT
Fine-tuning large pre-trained models has become the facto strategy for develop- ing both task-specific and general-purpose machine learning systems, including developing models that are safe to deploy. Despite its clear importance, there has been minimal work that explains how fine-tuning alters the underlying ca- pabilities learned by model during pretraining: does fine-tuning yield entirely novel capabilities or does it just modulate existing ones? We address this ques- tion empirically in synthetic, controlled settings where we can use mechanistic interpretability tools (e.g., network pruning and probing) to understand how the model's underlying capabilities are changing. We perform an extensive analysis of the effects of fine-tuning in these settings, and show that: (i) fine-tuning rarely alters the underlying model capabilities; (ii) minimal transformation, which we call a "wrapper', is typically learned on top of the underlying model capabilities, creating the illusion that they have been modified; and (iii) further fine-tuning on a task where such "wrapped capabilities" are relevant leads to sample-efficient revival of the capability, i.c., the model begins reusing these capabilities after only a few gradient steps. This indicates that practitioners can unintentionally remove a model's safety wrapper merely by fine-tuning on a, e.g., superficially unrelated, downstream task. We additionally perform analysis on language models trained on the TtoySto dataset to support our claims more realistic setup.
INTRODUCTION
Large language models (LLMs) pretrained on huge, web-crawled text datasets demonstrate extremely general capabilities (Radford al., 2018; 2019; Brown al., 2020; Bubeck et al., 2023). This has led to the current paradigm of machine learning, where practitioners often use model adaptation protocols such as fine-tuning to achieve unprecedented performance on a broad range of downstream tasks (Raffel et al., 2020; Sanh al., 2022; Reid et al., 2022; Driess et al., 2023; Ahn et al., 2022). Relatedly, the generality of an LLM's capabilities implies the model also learns to exhibit several undesirable behaviors, e.g., producing sensitive, biased, toxic outputs in the pursuit of completing task (Weidinger et al., 2021; Lin et al., 2021; Jiang et al., 2021; Parrish et al., 2021; Zhou et al., 2021; Xu al., 2021; Welbi et al., 2021). Fine-tuning with different training objectives has again seen immense usage in mitigating such "unsafe" capabilities, serving as an integral component of current state-of-the-art alignment approaches like RLHF (Ouyang al., 2022; Go al., 2023; Stiennon et al., 2020; Bai et al., 2022; Glaese et al., 2022).
Given its ubiquity in the design of both performant and safely deployable models, natural question emerges: precisely how does fine-tuning influence a pretrained model's capabilities to adapt to a downstream dataset (see Fig. 1)? The generality of an LLM's capabilities opens the possibility
*Co-first authors. samyak jain. csel 80 teitthu..aae in.in, robert.kin kirk.3 3. 40gma 1 com, e uubaaa@um edu.

---

Published as a conference paper at ICLR 2024
(a)
(b)
() Minimal transform
a story S-year old would understand.
But then, something unaxpected...
Onco time...
Once upan time... But then, something expected...
(Undesirable output)
tuning
( Deletion
C
Cij
Civ
upon time... They lived hepplly ever after.
Write story that S-year old would understand, don't include a twist.
Once upon a time... They all lived happäy after (Desirable output)
LLM
C2
Amplitication
Once upon time... Everything happened expected... They all Eied happily ever after.
Figure 1: How does fine-tuning alter model's capabilities? (a) Pretraining on huge. web-crawled datasets leads to LLMSs learning several capabilities that can justifiably process an input. The figure shows this using illustrative query, "write a story a 5-year old would understand." Via careful prompting, the desired answer can be retrieved, indicating both desired and undesired capabilities exist in an LLM. (b) Upon fine-tuning, e.g., avoid use of undesirable capabilities, we hypothesize that three explanations are possible: (i) minimal transformation of the original capability is learned, e.g., negation of the original capability; (ii) the undesirable capability is deleted altogether; or (iii) the use of another relevant capability is amplified.
that fine-tuning protocols merely identify the most relevant capabilities and amplify their use for given set of inputs, while inhibiting the use of other capabilities. Arguably, results on jailbreaking drignmenn-finetuneed LLMs via adversarially generated prompts elicit undesirable behavior support this hypothesis (Wei et al., 2023; Zou et al., 2023; Shen et al., 2023; Deng et al., 2023; Liu et al., 2023b); however, a precise study establish the phenomenology of fine-tuning remains absent from the literature. It therefore remains unclear how pernicious this problem is.
Motivated by the above, we perform an extensive analysis of the effects of fine-tuning on a pretrained model's 's capabilities in controlled settings where we can use mechanistic interpretability tools to understand precisely what is happening to the model's underlying capabilities as it is fine-tuned Specifically, we focus on the following two setups: (i) compiled transformer models based on the Tracr library (Lindner et al., 2023; Weiss al., 2021), which allows encoding specific computa- tional programs into a transformer. and (ii) procedurally generated setups involving probabilistic context-free grammurs (PCFGs) (Sipser, 1996; Chomsky, 1956), a formal model designed to capture syntactic properties of natural and programmatic languages that has recently served testbed for mechanistically understanding language models (Allen-Zhu Li, 2023c; Delétang et al., 2022: Shi al., 2022). While Tracr allows us analyze models with perfectly encoded capabilities, models trained on PCFGs allow evaluation of the effects of design choices involved in the pretraining pipeline. Fine-tuning these models via the often-used protocol of further training pretrained model on downstream dataset with a sufficiently small learning rate, we make the following findings.
Fine-tuning alters pretraining capabilities by minimally transforming them. We find that when a relevant pretraining capability is present, the fine-tuned model learns a minimally transformed version of it. We call the transformed portion wrapper.
Wrappers are generally very localized. We show that the wrappers transforming a model's pretraining capabilities are often extremely localized: e.g., via mere pruning of a few weights or neurons, we show the model can start to reuse its pretraining capability and unlearn how to perform the downstream task. Relatedly, we find that via a simple linear probe, we are still able to retrieve outputs expected from the pretrained model. Reverse fine-tuning to "revive" a capability, In scenarios where upon fine-tuning a model behaviorally seems to not possesses capability, we find that further fine-tuning the model subset of pretraining data leads to a sample-efficiens "revival" of the capability. We corroborate these results in realistic setup using the TinyStories dataset (Eldan & Li, 2023).
2 RELATED WORK
tine-tuning in the "foundation model' era. Fine-tuning large-scale foundation models pretrained on huge datasets, such LLMs (Radford et al., 2019; Brown et al., 2020) or large vision models (Radford et al., 2021; Caron et al., 2021), has become the norm in most domains of machine learning. Accordingly, several fine-tuning methods have been proposed in recent years, e.g., instruction fine- tuning (Wei et al., 2021; Liu et al., 2022b; Askell et al., 2021), parameter-efffiiemm fine-tuning (Houlsby et al., 2019; Zaken et al., 2021; Wang al., 2022), low-rank adaptation (Hu et al., 2021; Pfeiffer
2

---

Published as a conference paper at ICLR 2024
et al., 2020: Lialin et al., 2023), and weight averaging (Gueta et al., 2023; Matena & Raffel, 2022). The diversity of these protocols makes fine-tuning a general, umbrella term for related methods used to adapt a pretrained model to elicit its most relevant capabilittis. For precision, to fine-tuning protocols that continue training of a pretrained model on smaller downstream dataset at learning rate that is often one to three orders of magnitude lower than the average pretraining one. Such protocols are widely used in practice, e.g., in instruction fine-tuning (Wei et al., 2021).
restrict this paper
Understanding fine-tuning, A few papers theoretically analyze fine-tuning (Lampinen & Ganguli, 2018; Tripuraneni et al.. 2020; Gerace et al., 2022; Maddox et al., 2021; Kumar et al., 2022) under strong assumptions such as relatively simple model classes (e.g., linear functions) or a kernel view of deep learning, which, as shown by Yang & Hu (2020), trivializes the notion of feature transfer in fine-tuning transfer learning (though see Malladi et al. (2023) for notable exception). Prior works have also evaluated the effects of fine-tuning via the lens of mode connectivity (Juneja et al., 2022; Lubana et al., 2022), behavioral evaluations (Lovering et al., 2021), and intrinsic dimensionality of the loss landscape (Aghajanyan et al., 2020). In contrast, we aim to provide mechanistic analysis of how fine-tuning changes model capabilities. Contemporary works by Kotha et al. (2023); Prakash et al. (2024) claim that fine-tuning is unlikely alter model's '' capabilities their results can be seen as further support for claims made in our work on other experimental setups.
Model interpretability via synthetic tasks. Several recent works have focused on mechanistically understanding how Transformers learn synthetic language generation tasks, such as learning formal grammars and board games (Allen-Zhu & Li, 2023c; Zhao al., 2023a; Li et al., 2023; Nanda et al., 2023; Liu et al., 2022a; Valvoda et al., 2022: Liu et al., 2023a; Zhou et al., 2023; Chan et al., 2022). The goal of such papers, including ours, is not necessarily to provide accurate explanations for the success of LLMs, but to develop concrete hypotheses that can be used to develop grounded experiments or tools for understanding their behavior. For example, in a recent work, Allen-Zhu & Li (2023a;b) use synthetically designed setup to develop hypotheses for how "knowledge" about an entity is stored in pretrained model, showing such knowledge can often be manipulated via relatively simple linear transformations. Similarly, Okawa et al. (2023) use a procedurally defined multimodal dataset to demonstrate that emergent capabilities seen in neural networks are partially driven by the compositional nature of real world data. In another work, Zhou et al. (2023) utilize Tracr compiled Transformers to hypothesize and demonstrate that if primitive operations involved in formal algorithm can be implemented by a model, length generalization if practically feasible.
DEFINING OUR NOTION OF CAPABILITIES
For precision and to motivate our experimental setup, we first discuss the notion of capabilities that we aim to capture for analyzing how fine-tuning alters a model (see Tab. 2 for summary of notations). We use an idcalized definition to communicate not expect all capabilities in a pretrained model will act as perfectly as the definition necessitates. However, for the procedural tasks used in this work, idealized notion is fairly representative.
primary intuition and emphasize that we do
Let Dpr denote a dataset sampled from a distribution Px over the domain X. We will assume the domain X can itself be factorized into two domains X, and Xp. Correspondingly, a sample x E X can be divided into a tuple of variables (x; E Ia E X Xp). where x; identifies which capability a model should use to process the information encoded by the variable xzg This decomposition captures the idea that different prompts can force a pretrained LLM to clicit different capabilities, as shown in Fig. 1. The identifier of capability is denoted ig. Pretraining on Dpr yields us an L-layer model M()) X > Y, where often Y = X for language models. Let Read((M(.) denote the action where a linear layer is trained on intermediate outputs at layer of model M(.) using Dpr. Under this setup, we define capability as follows.
Definition 1. (Capability.) Define a surjective map Sc XD - Ye, where Ye C Y. Les Xe C X be a sub-domain s.t. V E Xe, the capability identifier variable is the same, i.e., 24 = ic. Then, we say the model M(.) "possesses capability C" for all Xx. 31 < s.f. Read, (M)) fe(xa).
linear readout at an intermediate layer is used in the definition above to emphasize that the notion of a capability need not correspond to only input-output behavior. Further, the definition is restricted to a sub-domain of the overall input space, which we argue is important to define a system's 's capabilities. For example, one can claim an 8-bit adder circuit possesses the capability to perform addition, but, technically, this is true only over the domain of 8-bit numbers; for inputs with more than 8-bit precision, the circuit will see overflow error, generating an incorrect but syntactically valid output.
3

---

Published as a conference paper at ICLR 2024
Similarly, LLM may possesses the capability to identify the sentiment of passage of text in specific language, but possibly fail when inputs in a different language are shown. Such structured failures imply claiming the existence of a capability should account for the input domain.
We next consider how the fine-tuning distribution PX over the domain X can interact with capabilities exhibited in the pretrained model. Our goal here is capture the fact that large-scale pretraining corpora is likely to have non-zero probability under the fine-tuning distribution, i.e., it is unlikely that a pretrained model will lack any capability relevant to the downstream task. This motivates notion of "relevance of capability". Specifically, let Der ~ x denote the downstream dataset used for fine-tuni fine-tuning, where P the empirical distribution that captures a subset of the support with non-zero probability in the distribution X Definition 2. (Relevance of Capability.) Assume the capability C in a pretrained model can be transformed to map go Sc via fine-tuning on Der, where [Drr] < [Der], such that for all x ~ Pxx the correct outpul is produced. Then, if for all P g Se yields the correct oulput, we claim capability C is strongly relevant to the fine-tuning task; else, we call it weakly relevant.
For example, weakly relevant capability can involve the ability to recognize a spurious at- tribute that the model can learn to exploit to perform well on the fine-tuning dataset, with- out enabling generalization to the overall dis- tribution that the fine-tuning dataset is sampled from. Meanwhile, a strongly relevant capability is one that extracts a causally relevant feature for that task (see Fig. for an example). When weakly relevant pretraining capability is avail- able, we empirically observe that we can often identify specific components in the latter half of the model (e.g., neurons or layers) that seem to implement the transform g in Def. 2. In such cases, we call g a "wrapper" and g C a "wrapped capability". we intervene on the model by either removing the wrapper training it to forget the wrapper, find the model starts to perform well on the pretraining task again. In such cases, we say the pretraining capability has been "revived""
idantiliation of negation words
Sentiment recogrition (Strongly relevant;
(Weakly relevant)
The movie was definitely not
The movie was definitely not bad. More character building could have helped develop the plat, but I loved the twists and action)
Complete the passage maintaining
bad. I really can't belleve they these. They're just wasta resources!
make
Figure 2: Capability Relevance. Consider the task of completing a passage while maintaining its narrative. Herein, the ability recognize the sentiment of text will be deemed strongly rele- vant and the ability to recognize negative words weakly relevant. Such words are often spuriously correlated with a negative sentiment.
4 BUILDING CAPABLE MODELS: TRACR AND PCFGs
We next describe the setup used in this work for analyzing how fine-tuning alters a model's capabilities (see Fig. 3 and App. B). Due to lack of clarity on what capabilities a language model possesses or what training data it has seen, we primarily focus procedurally defined setups that enable clear interpretability. To understand how the relevance of a capability affects fine-tuning, we randomly embed predefined spurious attribute into the fine-tuning dataset. Specifically, the attribute correlates with the features extracted by the pretraining capability-irt the attribute is "simple" enough, the model preferentially exploits it reduce the downstream loss (Shah et al., 2020; Lubana et al., 2022; Trivedi et al., 2023).
Tracr
PCFG
Fine-Tuning
Initialize for Capability C
Pretrain for Capability
Fine-tune on Dataset D77
Transformr
Rasp Coda:
Transformer
Transformer
Transformer
teweioo (XI)
JU T7
Relevance
Task Operand a..c=0
Capebiity
Capubilly
Figure 3: Experimental setup. We primarily analyze two setups: (i) Tracr "compiled" models with predefined capabilities and (ii) models trained to learn capabilities defined via PCFG, follow- ing Allen-ZMe Zhu & Li (2023c). During fine-tuning, we train the model on a dataset Dpr that promotes learning of capability C'. We randomly embed spurious attributes in the fine-tuning dataset that correlate with features extracted by a pretraining capability C to operationalize capability relevance.
4

---

Published as a conference paper at ICLR 2024
Compiled capabilities with Tracr. For fully controllable system, we use the recently proposed Tracr library (Lindner et al., 2023). Tracr enables "compiling" a transformer model with set of pre- defined computational primitives over string of characters from the English alphabet. Accordingly, we define a set of capabilities as a Tracr program and compile it into a Transformer via Traer (see App. B.1 for detailed pipeline). The model is then fine-tuned compiled capability may either be weakly or strongly relevant. While we analyze two tasks in this setup, for the main body of the paper, we focus on only the following one.
downstream task which the
Compile the capability to count the number of occurrences of a token OpT in a string into the model; fine-tune to count occurrences of another token OFr. If I 0) denotes the number of occurrences of token 0 in string x; the spurious correlation is defined by enforcing constant difference in token occurrences, i.e., r(æ, OFT) r(zr, OPT) ) q. See also Alg. 1 and Fig. 12.
Counter:
As an example, note that in the Counter setup, the model can exploit its pretraining capability and get the correct output on the fine-tuning dataset by merely adding q the count of OpT tokens. This wrapped capability will however perform poorly on samples without the correlation.
Learned capabilities with PCFGs. In this setup, capabilities are "learned", akin to practical situations. This allows us to probe the effects of different pretraining design choices, e.g., the distribution of the pretraining data. Specifically, we follow recent work by Allen-Zhu & Li (2023e) and train a minGPT model (Karpathy, 2020) via autoregressive training on probabilisite context-free grammars (PCFGs), a formal model of language that captures syntactic properties. Broadly, the gentratnnn process involves a tree traversal (see Fig. 3), starting from an initial root node and randomly choosing and navigating a set of production rules of the grammar from start/intermediate nodes to intermediate/terminn nodes, stopping only when a terminal node is reached. The terminal nodes reached by all paths starting at the root node will be concatenated define string x from the grammar (see Appendix for more details). We prepend special tokens T and O, called "task family" and "operand" tokens, that specify a certain task must be performed on the string x;; e.g., count the occurrences (a task family) of a certain token (operand) in a string. Overall, specific pair of the task family and operand tokens instantiates a task in our setup. The ground truth output of this task and special token indicating that the output should be produced at the next position are appended at the end of the string in the training data (see App. B.2 for further details and Fig. 15 for an example). Our experiments thus involve the following steps. (i) Pretrain model families. Every sample begins with the task family and operand tokens to specify the task; this ensures different tasks do not "conflict" (assign different labels to the same input), since, by construction, they have con-ovellapping support. (ii) Fine-tune the model on a task which may or may not have been included during pretraining. (iii) Evaluate how this fine-tuning affects the model. The data- gentratnnn process involves a uniform prior over task family tokens; meanwhile, the set of operand tokens seen during pretraining, denoted [Opr]. have a multinomial sampling prior. Specifically, the probability of sampling specific token OPT E [Opr) under task T is denoted P. P.(OPT) ). If this probability is low, the model may not learn the relevant capability to perform the task specified by the special tokens. While we analyze the effect of fine-tuning in two broad setups, using model pretrained on five distinct task familics relating to counting and indexing elements of a string, we focus on only the following one in the main body of the paper.
set of task
We intentionally reuse this task to demonstrate the effects of compilation of capability via Tracr versus learning the capability via PCFGs. Specifically, the model is pretrained to learn to count the occurrences of tokens from a ser of operand tokens {OpT] and fine-tuned to exclusively count occurrences of a token OFT E {OPT]. By making the sampling probability of OFT tokens high during pretraining, we can make the model preemptively performant the downstream task. This allows us to model the notion of capability relevance.
Counter:
5 EXPERIMENTS: MECHANISTIC ANALYSIS OF FINE-TUNING
We now provide several results indicating that fine-tuning rarely elicits meaningful changes to pretraining capabilities. To this end, we borrow several protocols commonly used in the field of mechanistic interpretability for our analysis (see Fig. 4), specifically network pruning (Voita et al.. 2019; Tanaka et al., 2019), attention map visualizations (Serrano & Smith, 2019; Wiegr- effe & Pinter, 2019; Lai & Tan, 2019), and probing classifiers (Tenney et al., 2019; Voita & Titov, 2020; Geva et al., 2023; 2022). We use multiple tools for all experiments since each
5

---

Published as a conference paper at ICLR 2024
tool, individually, is known to suffer from pitfalls (Meister al., 2021; Bai et al., 2021; Jain & Wallace, 2019; Belinkov, 2022; Bolukbasi et al., 2021). Demonstrating our claims consistently conclusions' robustmess pitfalls of spe-
hold true across diverse set of tools improves cific tool. Additionally, we propose a methodology called reverse fine-tuning (reFT ), wherein takes pretrained model, fine-tunes it on downstream dataset, and then fine-tunes it again in the "reverse" direction, i.c., on dataset sampled from the original pretraining distribution. We argue the behavior corresponding to pretrain- ing capability is retrieved in few steps of reFT, fine-tuning did not meaningfully alter said capability (this claim can be formalized using results by Bengio et al. (2019); Priol et (2021)).
Pre-Trained
Fine-Tuned
We primarily focus on the learned capabilities setup of PCFG counter in the main paper, relegating most results on compiled capabilities with Traer to App. and other results with PCFGs to App. H H-findings remain consistent across all settings. In the PCFG counter setup, the model is pretrained, amongst other tasks, to count occurrences of tokens from the set {Opr} = {a, b, c} in given string; during fine-tuning, the model is trained to count occurrences of OFr b. Here, the spurious correlation is defined by enforcing count of be 1 more than that of a. The probability datapoint sampled from the train or test fine-tuning dataset contains spurious correlation denoted Cr: and CTe. respectively. Here, Cr E (0.0, 0.5, 0.8, 1.0) and CT E {0.0, 1.0). We use three sets sam- pling probabilities of the task operands in the pre- training data: Pf- (0.999, 0.001, 0.000), P.M (0.9, 0.1, 0.0), or PM (0.5, 0.3, 0.2). These priors indicate a low/medium/high probability of sampling OFT. We use the following learning rates (LR) for fine-tuning, 7M nm5/10* and 1/s me/100, where Tjr is the average pretraining learning rate.
Reverse Fine-Tuning
(i) Pruning
Learns wrapper
C
6 wm U. a (iii) Attention Viz.
(iv) Probing
Figure 4: Analysis protocols. We analyze how fine-tuning affects a pretrained model's capabilities by (i) reverse Fine-tuning, (ii) network pruning, (iii) attention visualization, and (iv) probing classifiers. We use (iv) to show fine-tuning often yields wrapped ca- pabilities. For further evidence, we use (i) and (ii) and find we can "revive" the original capabilities, i.e., the model starts perform- ing well on the pretraining task again. See App. D for precise details.
5.1 RESULTS
Behavioral assessment of fine-tuning. We first evaluate the model's learning dynamics during fine-tuning (see Fig. 5 and Tab. 5). When the pretraining prior has low probability of sampling the token OPT, we see the fine-tuned model performs well only when the spurious correlation is
Ct. O
Ctr 0.5
Cr. 0.8
Ct. - 1.0
CT.
=1==1
Crc =
===1
Cre
Gr=
100
Acc. OFT
100
65
30-
0 2.5K 5k 7.5K10K
2.5K 5K 7.5K
0 2.5K 5K 7.5K
0 2.5K 5K 7.5K1OK
0 2.5K 5K 7.KK
2.5k 5K 7.5K10K
(a) 7d!
(b) n
(C)
Figure 5: Fine-tuning accuracy w.r.t. number of training iterations. We vary the probability of sampling the token OFT in the pretraining data and the spurious correlation in the fine-tuning datasets. When the prior is sufficiently high (a, b), we find the model learns to perform well on the downstream task. Mcanwhile, if the prior is low (c), the model Iearns the downstream task only if high enough learning rate is used and the spurious correlation is imperfect. This indicates the ability to extract information relevant for the downstream task is likely be exploited during fine-tuning.

---

Published as a conference paper at ICLR 2024
74
100
Acc. O
30-
0
2.5K 5K 7.5K 10K
0
2.5K 5K 7:5K 10K
2.5K 5K 7.5K 10K (c) Ct.
2.5K 5K 7.5K 10K (d) CT. 1.
(a) CT.
Cr. 1,
Figure 6: Impact of sampling prior on the pretraining task's accuracy as fine-tuning is per- formed. We plot accuracy on the pretraining task w.r.t. fine-tuning iterations. When the sampling prior of the OFr is low during pre-training, the pretraining task accuracy quickly plummets, especially if the spurious correlation is high; having high sampling prior mitigates this behavior. This indicates pretraining capabilities are affected the most when they are weakly relevant.
=
=.5
CT. 0.8
CT. =.0
CT
Cr==1
CT =
cr=
Acc. OfT Acc. Opr
100-
100
(a) fj
(b) nls
Figure 7: Pruning few neurons is sufficient to retrieve pretraining task accuracy. We plot accuracy w.r.t. number of neurons pruned to improve performance on the pretraining task. We see when small learning rate is used for fine-tuning, the pretraining task's performance improves after just 5-15 neurons are pruned (top), while the mne-tuning task's 's performance reduces correspondingly (bottom). We argue these neurons serve as wrapper to minimally alter the weakly relevant pretraining capability and exploit the spurious correlation present in the fine-tuning data.
present, i.e., Cr. = 1. As the sampling probability is increased, however, we observe this behavior significantly changes. In particular, even the model is fine-tuned for high value of Crr- albeit less than 1, it starts to perform well the test data regardless of the presence of the spurious attribute. Note that the performance is not high to begin with, indicating the ability to count OFr was learned during fine-tuning; however, having a sufficiently large sampling probability for the token during pretraining leads the model to avoid the spurious correlation. This indicates a pretraining ability to extract information relevant for the downstream task is likely to be exploited during fine-tuning. This further corroborated by the results in Fig. where we observe that when the spurious correlation is present in the finettunig data, accuracy on the pretraining task is affected the most if sampling prior of the target token was low during pretraining. We next analyze these results mechanistically.
Pruning/ / Probing fine-tuned models indicates learning of wrapped capabilities. Our results above indicate the model exploits its weakly relevant capabilities, i.e., the capability that helps exploit any spurious correlation, to solve the downstream task. We hypothesize that, at a mechanistic level, the model exploits the weakly relevant capability by learning a wrapper over it. To evaluate this, we analyze the models fine-tuned with low sampling prior via network pruning and linear probing (see App. for setup details). Specifically, we prune the fine-tuned models to find the most salient weights for reducing loss on the pretraining task of counting OPT. If the model learns a wrapper on this capability, the neurons we find should correspond to this wrapper, such that deleting them recovers the capability to count that token. As shown in Fig. 7, we find this is indeed the case-in in a setup with weak relevance of capabilities, pruning very small number of neurons is sufficient to revive the ability to perform well on the original task of counting OpT. To assess this further, we train linear probe on the residual output of every block the transformer model and determine whether the count of OpT can accurately computed via the fine-tuned model. shown in Fig. 8, in the presence of spurious correlations, a linear probe can retrieve the count of the token OpT, indicating intermediate outputs relevant to the pretraining capability are still being produced by the fine-tuned model. This observation is particularly evident when a smaller learning rate is used, which is common in practice. Overall, these results show that when weakly relevant capability is present in the pretrained model, a wrapper, i.c., localized transformation of the pretraining capability, is learned during fine-tuning.
reFT enables "revival" of pretraining capabilities. To further corroborate our claims above, we use model fine-tuned to count OpT and reverse fine-tune it to re- reale the ability count Opr.
7

---

Published as a conference paper at ICLR 2024
Ct.
-1
Cr.
Acc. OFr Acc. Oyr
100-
50
100
de
B4
10
ns
(a) Pre-trained Model
(b) 7MM
(c) 0s
Figure 8: Probing the presence of pre-training (top) and fine-tuning (hottom) capabilities. We plot probe accuracy versus the index of the block in the Transformer model. Cr is set to O. The pretrained model (left) acts as a baseline for the trend of performance through the model's blocks. In most scenarios, (left). A drop in performance is observed only when learning rate л/и is used with a weakly relevant capability (low sampling prior). This indicates pretraining capabilities persist after fine-tuning.
find we can infer the count of OpT with similar trend as the pretrained model
As a baseline. we also report a protocol called Scr. + FT, wherein the model is initialized with parameters pre-trained to count OPr and then fine- tuned to count Opr. Note that this baseline and the reFT protocol differ in their initialization state: former is initialized with parameters pretrained to count Opr. while latter is initialized with parame- ters pretrained count OPT and fine-tuned to count OFT. Results are shown in Fig. 9. We see the model starts to perform well on the pre-training task even a small learning rate is used for reFT. i.e., even minimally changing the fine-tuned model's 's parame- ters is sufficient to regain the pretraining capability! Further, increasing the sampling prior of OFT ac- celerates this behavior. This indicates that when strongly relevant capability is present, the model essentially amplifies its use, but does not catas- trophically affect the pretraining capability itself; meanwhile, with weakly relevant capability (low sampling prior during pretraining), even though the performance is initially poor on the pretraining task, in relatively few iterations (compared to baseline), the accuracy becomes perfect.
Cy.
Ciy m3
Scr:
66
30
127
163
161
107
107
10
107
167
161
(Log Scale) (b)
(a) r4
14
(c) 7/
Figure 9: Reverse Stnninunnng: We set Cre to be 0 to test the model performs well regardless a spurious correlation. Models are fine-tuned for 10K iterations. We observe that when strongly relevant capability is present (a, b), the model very quickly (0.1-1K iterations) starts perform well on the task via reFT, even if be- havior relevant to the capability ceased during pretraining (c.g.. when CTe is 1). Mcanwhile, when the model possesses a weakly relevant ca- pability (c), this "revival" is slightly slower (3K iterations). In contrast, the Scr. FT baseline only reaches perfect accuracy at 4.5K iterations and when using a larger learning rate, i.e., TJM-
Attention map visualizations further corrobo- rate the wrappers hypothesis. As we noted be- fore, the results discussed above remain consistent across other experimental setups for both Tracr and PCFG models. However, by construction, Tracr yields particularly interpretable attention maps, al- lowing us directly visualize the effects of fine- tniing We thus analyze the attention maps of a Tracr model on the Counter task described in Sec. 4. Results are shown in Fig. 10. The original Tracr compiled model serves as a baseline and clearly demonstrates that all tokens only attend the pre- training target token, Opr = Upon fine-tuning to count OFT b, we find the model clearly continues to pay attention OpT if small learning rate is used. A larger learning rate is, however, able to alter the model computation, but only the pretraining capability is not weakly relevant the fine-tuning task, i.c., when CT: = 0; otherwise, we again find the model continues to pay attention to the pretraining target.
altoning Taann
11 Taken /thendeed
Figure 10: Visualizing attention maps of fine- tuned Tracr models. Leftmost panel shows the Tracr compiled model's attention map on the counter task. Upon fine-tuning under different spurious correlations, we see the model con- tinues to pay attention to the pretraining target Opr a. Only when large enough learning rate and zero spurious correlation is used, there is a change in the attention pattern.
8

---

Published as a conference paper at ICLR 2024
Filtering
Filtering Mix & Match
Filtering Randomiiattinn
Not in Pretraining
Present in Pretraining
Prabing Analysis
(b) ReFT Dynamics
Probe Accuracy
LR: 7M
LR: ils
LR: ds
LR: AM
0.9
0.90
0.5
0.05
D: By
0.
10 10 10
107
207
10
12
19'
Probe Layer Index
Training Iteration
Figure 11: Validation on TinyStories. Models are trained to produce stories with several features (e.g., foreshadooving and fine-tuned via different protocols not produce stories with "forbidden" feature (specifically, twists) (see App. F for details). Left: We probe the existence of this feature at intermediate Transformer blocks. Probe accuracy on models pre-trained with or without twist data (Presen//Not in Preiraining, respectively) act as upper and lower bounds on the expected accuracy, and are plotted for ease of comparison. Regardless of the fine-tuning protocol (Filtering, Filtering + Randomisation, Filtering Mix & Match), for the lower LR, no protocol removes meaningful amount of information and a similar but less strong trend holds for the higher LR. Right: We plot the loss during reverse fine-tuning (reFT) to again produce stories with the forbidden feature. Fine-tuned models' loss goes down very quickly (30-300 iterations) compared to baselines (which never reach the same loss; also see Tab. 1). Both these results indicate the capability of identifying the forbidden feature, a necessary capability for story modelling, continues to persist after fine-tuning.
5.2
VALIDATING OUR HYPOTHESES ON TINYSTORIES
To give additional strength to results, we per- form an analysis using more realistic language models trained on the innyStories-Instruct dataset (Eldan & Li, 2023) (see App. B.3 for an example). These mod- write
Deletion Type
Twist Proportion at Iteration
0
30
300
3000
F (1aa)
44%
%
1
82%
els are able follow specific instructions coherent English stories over multiple paragraphs. We perform experiments analogous to the reFT and prob- ing experiments in the previous sections, but now ex- plicitly focus on whether fine-tuning can delete ca- pabilities present in pre-trained models. Models are pre-trained to generate stories with specific story fea- tures (such as containing twist, foreshadowing. or bad ending) and fine-tuned to not generate stories with specific feature (twist) (see App. for details on the protocols). We probe these models to detect the deleted feature from the intermediate model outputs in Fig. 11, where the dynamics of loss during reFT on learning to generate stories with the deleted feature are also shown. We also report the percentage of stories with the deleted feature generated by models during reFT in Table 1, where the generated stories are pro- cessed by a fine-tuned GPT-3.5 classifier to predict
+ RR (nar)
12%
56%
69%
75%
MM (77r)
31%
88%
50%
75%
F (ns)
69%
88%
75%
94%
+ RR (ns)
12%
44%
81%
11%
F MM (ns)
50%
81% %
62%
81%
Not in PT
12%
31% %
44%
81%
Table I: TinyStories reFT Analysis. We re- port the percent of generations with twist during reverse fine-tuning for the twist ca- pability. F, R. and MM stand for three fine-tuning protocols: Filtering, Randomi- sation and Mix & Match (see App. F for details). Regardless of learning rate and pro- tocol, models relearn generate stories with twist more sample-efficiently than the con- trol model pre-tr pre-trained on data w/o twists and fine-tuned to generate them (Not in PT). if the story contains the deleted feature (see App. F for details). Overall, we find that "deleted" capabilities can be easily and sample-efficiently recovered (compared to the baseline), i.e., stories with that feature can be generated again, regardless of the fine-tuning protocol used. These results support our hypotheses that fine- tuning only minimally alters pre- pre-trained model capabilities. We also highlight few recent papers that propose similar protocols as reFT and experiments as ours with further realistic settings (Qi et al., 2023; Yang et al., 2023; Zhang & Wu, 2024; Zhao et al., 2023b).
6 CONCLUSION
In this work, we show that fine-tuning generally alters pre-trained model via small, localized changes that only minimally transform their capabilities. We perform our analysis both with existing mecha- nistic interpretability tools as well as our proposed reFT method. Our results pave the way for future work both understanding how fine-tuning works in more realistic settings with larger models, as well as emphasize the need for developing methods beyond tin-ttning that alter pre-trained model capabilities more substantially, particularly deleting unsafe capabilities.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
ESL thanks Eric Bigelow, Nikhil Vyas, and Usman Anwar for relevant discussions early in the project. SJ was partially supported by BERI. ESL was partially supported via NSF under award CNS-2008151, RK was supported by the Foundation Al CDT at UCL.
AUTHORS' CONTRIBUTIONS
ESL conceived the project direction and developed set of hypotheses the limitations of fine- wningg with inputs from RK. SJ and ESL. co-designed a draft of the PCFG and Tracr setups, and came up with pruning and reverse fine-tuning analysis which led to validation and further refining of the hypotheses. SJ led the experimental execution and made the tasks considered in the paper precise in collaboration with ESL. RK proposed and the TinyStories experiments with inputs from ESL, SJ, EG and TR. Literature review and writing of the main paper was led by ESL. SJ led writing of the appendix. ESL, SJ, RK, and HT collaborated on design of all figures and plots. DSK acted as the primary senior advisor on the paper, with inputs from RPD, HT, EG, and TR as well.
REFERENCES
Armen Aghajanyan, Luke Zettlemoyer, and Sonal Gupta. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. arXiv preprint urXiv:2012.13255, 2020.
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et al. Do as can, not as say: Grounding language in robotic affordances. arXiv preprint 2041112204.01661, 2022.
Zeyuan Allen Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. arXiv preprint arXiv:2309, 111.:2309.13166 2023a.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge manipulation. arXiv preprint arXiv:2309. 14402, 2023b.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.1300 13673, 2023c.
Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as laboratory for alignment. urXiv preprint 1 2021.
Bing Bai. Jian Liang, Guanhua Zhang, Hao Li. Kun Bai, and Fei Wang. Why attentions may not be interpretable? In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 25-34, 2021.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.00 05862, 2022.
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances.
Computational
Linguistics, 48(1):207211, 2022.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv preprint arXiv:1901.10911 2019.
Tolga Bolukbasi, Adam Pearce, Ann Yuan, Andy Coenen, Emily Reif, Femanna Viégas, and Martin Wattenberg. An interpretability illusion for bert. arXiv preprint arXiv:2 143, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
10

---

Published as a conference paper at ICLR 2024
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXiv: 230.: 2023.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVE international conference on computer vision, pp. 9650-9660, 2021.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878-1888 18891, 2022.
Noam Chomsky. Three models for the description of language. IRE Transactions information theory, 133):113-124, 1956.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Mareus Hutter, Shane Legg, Joel Veness, al. Neural networks and the chomsky hierarchy. arXiv preprint arXiv:2207.02098, 2022.
Gelei Deng, Yi Liu, Yuekang Li, Kailong Wang, Ying Zhang, Zefeng Li, Haoyu Wang, Tianwei Zhang, and Yang Liu. Jailbreaker: Automated jailbreak across multiple large language model chatbots. arXiv preprini arXiv:2307.08715. 2023.
Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.0333 2023.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.074 07759999..779.... 2023.
Federica Gerace, Luca Saglictti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with model of synthetic correlated datasets. Machine Learning: Science and Technology, 3(1):015030, 2022.
Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts the vocabulary space. arXiv preprint arXiv:2203. 2022.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304. 2023.
Amelia Glaese, Nat McAleese, Maja Trębacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Maribeth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, et al. Improving alignment dialogue agents via targeted human judgements. arXiv preprint arXiv:2209.1433 2022.
Dongyoung Go, Tomasz Korbak, Germán Kruszewski, Jos Rozen, Nahyeon Ryu, and Mare Dymet- man. Aligning language models with preferences through f-divergence minimization. arXiv preprint 2rXX::2302.082/55 2023.
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen. Knowl- edge is region in weight space for fine-tuned language models. arXiv preprint 000463330..04663, 2023.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Pframeter-cfficient transfer learning for nlp. In International Conference on Machine Learning, pp. 2790-2799. PMLR, 2019.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen Yuanzhi Li, Shean Wang. Lu Wang. and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2/06.000 2021.
Sarthak Jain and Byron Wallace. Attention is not explanation. arXiv preprint 190.11990. 2019.
11

---

Published as a conference paper at ICLR 2024
Liwei Jiang, Jena D Hwang, Chandra Bhagavatula, Ronan Le Bras, Jenny Liang, Jesse Dodge, Keisuke Sakaguchi, Maxwell Forbes, Jon Borchardt, Saadia Gabriel, et al. Can machines learn morality? the delphi experiment. arXiv e-prints, pp. -11iv 2021.
Jeevesh Juneja, Rachit Bansal, Kyunghyun Cho, João Sedoc, and Naomi Saphra. Linear connectivity reveals generalization strategies. arXiv preprint. 1rxxiv22205.241111 2022.
Andrej Karpathy. MinGPT, 2020. Github link. httpot/tttttt////..oooooooooooor http://ttb...ooakkatta.t.....e tree crre/maater.
Suhas Kotha, Jacob Mitchell Springer, and Aditi Raghunathan. Understanding catastrophic forgetting in language models via implicit inference. arXiv preprint arXiv:2309.10100 2023.
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang, Fine-tuning can distort pretrained features and underperform out-of-distribution. arXiv preprint. arXiv:2202. 10054, 2022.
Vivian Lai and Chenhao Tan.
On human predictions with explanations and predictions of machine learning models: A case study on deception detection. In Proceedings of the conference on fairness, accountability, and transparency, pp. 29-38. 2019.
Andrew K Lampinen and Surya Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv: 10374, 2018.
Rémi Le Priol, Reza Babanezhad, Yoshua Bengio, and Simon Macoste-Julien. An analysis of the adaptation speed of causal models. In International Conference on Artificial Intelligence and Statistics, pp. 775-783. PMLR, 2021.
Kenneth Li, Aspen K. Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfi--
ter, and Martin Wattenberg. Emergent World Representations:
Exploring a Se-
quence Model Trained on Synthetic Task, 2023. Comment: ICLR 2023 oral (notable-top-5%) https//ggihubb.com//ceennnhhttht _oold.
http://cceeiie...ee/oommm/dddede TcZTT
code:
Vladislav Lialin, Vijeta Deshpande, and Anna Rumshisky. Scaling down to scale up: guide to parameter-efficient fine-tuning, arXiv preprint arXiv:2303.156477 2023.
Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulga: Mcasuring how models mimic human 2021.
falsehoods. arXiv preprint
David Lindner, János Kramár, Matthew Rahtz, Thomas MeGrath, and Vladimir Mikulik. Tracr: Compiled transformers as a laboratory for interpretability. arXiv preprint arXiv:2301.0506 2023.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprint 100749 2022a.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flop language modeling. arXiv preprint arXiv:2306. 2tt.::33000000440, 2023a.
Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. Advances in Neural Information Processing Systems, 35: 35:1950-1965, 2022b.
Yi Liu, Gelei Deng, Zhengzi Xu. Yuekang Li, Yaowen Zheng, Ying Zhang, Lida Zhao, Tianwei Zhang, and Yang Liu. Jailbreaking chatgpt via prompt engineering: An empirical study. arXiv preprint arXiv:2305.13866 2023b.
Charles Lovering, Rohan Jha, Tal Linzen, and Ellie Pavlick. Predicting inductive biases of pre- trained models. In International Conference on Learning Representations, 2021. http:: // /pooerreveew..o .nt/ forum? wtw.attmaaaaaaa.. www.//euu///a//eeeeem
Ekdeep Singh Lubana and Robert P. Dick. A gradient flow framework for analyzing network pruning. In International Conference on Learning Representations, 2021. URL https: / httppps///ppereeeee /ppenreview. net forum? i rmm.eumm..ue..
12

---

Published as a conference paper at ICLR 2024
Ekdeep Singh Lubana, Eric J. Bigelow, Robert P. Dick, David Krueger, and Hidenori Tanaka. Mechanistic Mode Connectivity, 2022. Comment: 39 pages.
Wesley Maddox, Shuai Tang, Pablo Moreno, Andrew Gordon Wilson, and Andreas Damianou. Fast adaptation with linearized neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 2737-2745. PMLR, 2021.
Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pp. 23610-23641. PMLR, 2023.
Michael Matena and Colin A Raffel. Merging models with wesher- averaging. Advances in Neural Information Processing Systems, 35:17703-17716, 2022.
Clara Meister, Stefan Lazov, Isabelle Augenstein, and Ryan Cotterell. sparse attention more interpretable? arXiv preprint arXiv:2106.01087, 2021.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint 16X:v:6611,06444 2016.
Michael C Mozer and Paul Smolensky. Skeletonization: A technique for trimming the fat from network via relevance assessment. Advances neural information processing systems, 1, 1988.
Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models. arXiv preprint arXiv:2309.0088 0999099909999... 2023.
Maya Okawa, Ekdeep Singh Lubana, Robert Dick, and Hidenori Tanaka. Composi-
tional abilities emerge multiplicatively: Exploring diffusion models synthetic task. 2023.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27744, 2022.
Alicia Parrish, Angelica Chen, Nikita Nangia, Vishakh Padmakumar, Jason Phang, Jana Thompson, Phu Mon Htut, and Samuel R Bowman. Bhq: A hand-built bias benchmark for question answering. arXiv preprint arXiv:21 10.08193, 2021.
Jonas Pfeiffer, Andreas Rücklé Clifton Poth, Aishwarya Kamath, Ivan Vulié, Sebastian Ruder, Kyunghyun Cho, and Iryna Gurevych. Adapterhub: framework for adapting transformers. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP 2020): Systems Demonstrations, pp. 46-54. Association for Computational Linguisties, 2020.
Nikhil Prakash, Tamar Rott Shaham, Tal Haklay, Yonatan Belinkov, and David Bau. Fine-tuning enhances existing mechanisms: A case study on entity tracking. arXiv preprint arXiv:2402. 14811, 2024.
Xiangyu Qi, Yi Zeng, Tinghao Xie, Pin-Ya Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson. Fine-tuning aligned language models compromises safety, when users do not intend to! arXiv preprint (tXii:23110.3693.3 2023.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. PMLR, 2021.
13

---

Published as a conference paper at ICLR 2024
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 2020.
Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. Can wikipedia help offline reinforcement learning? arXiv preprint arXiv:2201. Stti::2201.12122 12122, 2022.
Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika, Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szczechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, De- bajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang. Han Wang, Matteo Manica, Sheng Shen, Zheng Xin Yong, Harshit Pandey, Rachel Bawden, Thomas Wang, Trishala Neeraj. Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Teehan, Teven Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Representations, 2022. URL https: hwtp//peeereeeeeeeeeeee net otw.:///
Sofia Serrano and Noah A Smith. Is attention interpretable? arXiv preprint 19Xiv:9906.00 037311, 2019.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573-9588 2020.
Xinyue Shen, Zeyuan Chen, Michael Backes, Yun Shen, and Yang Zhang. "do anything now": Characterizing and evaluating in- in-the-wild jailbreak prompts on large language models. arXiv preprint arXiv:2308.0383 2023.
Hui Shi, Sicun Gao, Yuandong Tian, Xinyun Chen, and Jishen Zhao Learning bounded context-free- grammar via lstm and the transformer: Difference and the explanations. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 8267-8276, 2022.
Michael Sipser. Introduction to the theory of computation. ACM Sigact News, 27(1):27-29, 1996.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3021, 2020.
Hidenori Tanaka, Aran Nayebi, Niru Maheswaranathan, Lane McIntosh, Stephen Baccus, and Surya Ganguli. From deep learning mechanistic understanding in neuroscience: the structure of retinal prediction. Adv. in Neural Information Processing Systems (NeurIPS), 2019.
lan Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint arXiv: 19990055.55... 2019.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Nilesh Tripuraneni, Michael Jordan, and Chi Jin. On the theory transfer learning: The importance of task diversity. Advances in neural information processing systems, 33:7852-7862. 2020.
Puja Trivedi, Danai Koutra, and Jayaraman J TTiaggarajan. A closer look at model adaptation using feature distortion and simplicity bias. arXiv preprint arXiv:2303. 13500, 2023.
Josef Valvoda, Naomi Saphra, Jonathan Rawski, Adina Williams, and Ryan Cotterell. Benchmarking compositionality with formal languages. In Proceedings of the 29th International Conference Computational Linguistics, pp. 6007-6018, 2022.
Elena Voita and Ivan Titov. nntematiio-thhoreede probing with minimum description length. arXiv preprint arXiv:2003.12298 2020.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv: 190.099995009118 2019.
14

---

Published as a conference paper at ICLR 2024
Yihan Wang, Si Si, Daliang Li, Michal Lukasik, Felix Yu, Cho-Jui Hsieh, Inderjit S Dhillon, and Sanjiv Kumar. Two-stage Ilm fine-tuning with less specialization and more generalization. arXiv preprint 3XXv:22111006333. 2022.
Alexander Wei, Nika Haghtalab, and Jacob Steinhardt. Jailbroken: How does Ilm safety training fail? arXiv preprint arXiv:2307.02483 2023.
Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.0165, 2021.
Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, et al. Ethical and social risks harm from language models. arXiv preprint arXiv:21 2021.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. International Conference on Machine Learning, pp. 1 11000-11090, PMLR, 2021.
Johannes Welbl, Amelia Glaese, Jonathan Uesato, Sumanth Dathathri, John Mellor, Lisa Anne Hendrieks, Kirsty Anderson, Pushmeet Kohli, Ben Coppin, and Po-Sen Huang. Challenges in detoxifying language models. arXiv preprint arXiv:2109.07445. 2021.
Sarah Wiegreffe and Yuval Pinter. Attention is not not explanation. arXiv preprint 1rX1v:1908.04626, 2019.
Albert Xu, Eshaan Pathak, Eric Wallace, Suchin Gururangan, Maarten Sap, and Dan Klein. Detox- ifying language models risks marginalizing minority voices. arXiv preprint arXiv:2104.06390 2021.
Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011. 14522,, 2020.
Xianjun Yang, Xiao Wang, Qi Zhang, Linda Petzold, William Yang Wang, Xun Zhao, and Dahua Lin. Shadow alignment: The ease of subverting safely-aligned language models. arXiv preprint arXiv:2310.02999, 2023.
Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for wansformer-based masked manguage-modete arXiv preprint arXiv:2106. /0/99, 2021.
Xiao Zhang and Ji Wu. Dissecting learning and forgetting in language model finetuning. In
The Twelfth International Conference Learning Representations, 2024. URL https://
openreview net wttp:/ooommiddmseeeeere
Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arX1v:2303.08117, 2023a.
Jiachen Zhao, Zhun Deng, David Madras, James Zou, and Mengye Ren. Learning and forgetting unsafe examples in large language models. arXiv preprint arXiv:2312.117.. 2023b.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16023 2023.
Xuhui Zhou, Maarten Sap, Swabha Swayamdipta, Noah A Smith, and Yejin Choi. Challenges in automated debiasing for toxic language detection. arXiv preprint arXiv:2102.0000 2021.
Andy Zou, Zifan Wang. J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307. /5043, 2023.
15

---

Published as a conference paper at ICLR 2024
A ORGANIZATION OF APPENDIX
In the appendix we present comprehensive analysis of our claims on Tracr, PCFG and TinyStories- Instruct using different mechanistic interpretability tools discussed in Section-D of the main paper. We also present a summary of the notations used in this work in Tab. 2. Overall, the appendix is organized as follows:
Sec. B presents details of the Tracr, PCFG and Tiny Stories datasets utilized in this work.
Sec. C presents the training and model details for each of the datasets considered.
Sec, D lists the protocols used for different mechanistic interpretability tools like attention maps, probing, pruning and reverse fine-tuning.
Sec. provides few more results in practically relevant contexts, such in synthetic jailbreaking setup.
Sec. E. studies the effect of using different fractions of pre-training and fine-tuning data points for fine-tuning.
Sec. F.2 presents the jailbreaking analysis using the PCFG setup.
Sec. E.3 shows reverse fine-tuning fine-tuned model sample efficient compared to baselines for both PCFG and Tracr models.
Sec. E.4 presents reverse fine-tuning analysis of fine-tuning protocol that actively tries to remove a capability from PCFG Tracr models.
Sec. F presents detailed discussion of setup details and results on TinyStories.
Sec. presents additional results on Tracr for counter and max element tasks.
Sec. H presents additional results PCFG for the counting and index of occurrence tasks.
Table 2: Notations used in this work.
Notation
Meaning
X
Input domain
Xp
Factor of the input domain that captures values of the inputs Factor of the input domain that captures task identifiers
X
Px
Prohability distribution over the input domain
PR
The overall distribution defining the downstream fine- tne-tuning task Dataset used for fine-tuning
Drr
PX P..
Empirical distribution from which the fine-tuning dataset is sampled
T
Denotes task to be performed by the model (e.g.. count)
D
Denotes operand that will he processed by to perform the task T
{
Set of operand tokens seen during pretraining
O
A specific token used as an operand during pretraining
A specific token used as an operand during fine-tuning
r rar 0)
Denotes the result of executing task from Sec. 4 on string æ for some operand Probability that randomly sampled string in the training data used for finc-tuning has spurious correlation between the pretraining capability and the downstream task Probability that a randomly sampled string in the test data used for evaluating fine-tuned models has a spurious correlation between the pretraining capability and the downstream task Sampling prior. Denotes the probability that when string with task token T is sampled during pretraining, the operand perform the task is
Cn
Ct.
P((0)
P.. P. P
Sampling priors such that the probability of sampling the target token for fine-tuning high (PC) medium (PM ) or small (P2)
(Dr is
Ma, 7s. 8
Medium Small / Very-small learning rates used for fine-tuning. IJv' is only used for specific reverse fine-tuning experiment with Tracr compiled models. Denotes reverse fine-tuning
reFT
7liceza
Number iterations used during pre-training
LR
Learning rate
16

---

Published as a conference paper at ICLR 2024
B ADDITIONAL DETAILS ON DATASETS
We consider three experimental setups: Compiled programs with Tracr (Lindner et al., 2023), learning models on Probabilistic Context Free Grammars (PCFG) (Allen-Zhu & Li, 2023c), and the TinyStories Instruct dataset.
B.1
TRACR DETAILS
Tracr (Lindner et al., 2023) generates a transformer model using the RASP library by Weiss et al. (2021). The specific code snippet used to generate the Tracr models for the counting and the max element tasks are shown in Fig. I and Fig. respectively. The models corresponding to these tasks is implemented with three standard transformer blocks, where each block consists of self-attention layer followed by two MLP layers.
We analyze the following two tasks to understand the effects of fine-tuning pretrained model's capabilities.
Counter: Compile the capability to count the number of occurrences of token Opr in string into the model; fine-tune to count occurrences of another token OFT. If r2, 0) denotes the number of occurrences of a token O string æ, the spurious correlation is defined by enforcing a constant difference token occurrences, i.e., I(z, OFT) I(xr, OPT) ) = q. See also Alg. and Fig. 12. Max-identifier: Compile the capability to identify the Opr-th largest element in a string: fine-tune to identify the Opr-th largest element. If r(x, 0) reads out the 0-th largest token in the string we define the spurious correlation as [ OFT) I(x, OpT) qi c.g., if and the OpT largest token in the string xr a, then the Opr-th largest token will be (which equal + in Tracr's vocabulary). See also Alg. 2 and Fig. 13.
The fine-tuning data is generated by randomly sampling tokens from uniform distribution over the input vocabulary. For the Counter task, the input vocabulary consists of first nine letters from the English alphabet. For the max element task, the input vocabulary consists of all the letters in the English alphabet. We sample with replacement for the Counter task and without replacement for the max element task (to avoid having multiple max elements). Examples for the task are shown in Figs. 12, 13.
Algorithm 1: Pseudocode for compiling the Counter capability via Tracr: Rasp code used to generate the model for the Counter capability and task via Tracr
def c)
binzarize the tokens into 0''s and 1's
øin Select the indices of tokens with value of
bin_idx rasp. Sep.Select.. (bin, rasp.indice Count the number of selected indices
casp.Comparison
count a = rasp.Select.rriichh (bin idx)
Generate an identity map
iix__sseee idx select rasp. Select (rasp. indices, rasp. indices,
rasp. Comparison.E0) output the count
sum rasp Aggrogoto
count_a)
Task: Count b
Sample: c, a, d, a, b, c, b. j. c. b, d, d, h, j. i, b, a, b, #. Answer: 10
Figure 12: Exemplar for Counter Task: A sample used for fine-tuning Traer compiled models counting 'b'.
17

---

Published as a conference paper at ICLR 2024
Algorithm 2: Pseudocode for compiling the Max identifier capability via Tracr: Rasp code used generate the model for the Max Identifier capability and task via Tracr.
def max,dentifiers
Identify the tokens larger than given token var_small Selp.telect (rasp.tokens, rasp.tokens, rasp. Casp.comparison.Ii
Calculate the sum of the identified tokens for every input token sum_small = Sesp.seeecto (var_small) Identify the fifth largest taken
tanget (sum_amal small==4)
Find the index the identified token the original input select_ide = zasp.Select tbin_target, rasp, indices, casp.Comparisot
Output the ftentified index return rasp.Aggregate select_idx, rasp. tokens)
Task: Find fifth largest element
Sample: S, b, d, f, h, m. P. q.
Answer:
Figure 13: Exemplar for Max-Element Task: A sample used for fine-tuning Tracr compiled models on the Max identifier task.
B.2 PCFG
We follow Allen-Zhu & Li (2023c) and use the production rules shown in Fig. 14. We sample string of tokens from the grammar and then randomly subsample a string 250 tokens from the generated original sequence (this helps remove bias towards tokens likely to be at the beginning). The sampling probabilities to sample a valid rule given the parent node is fixed to be 0.5. We formulate the training data as follows: Start of sequence token (SOS) Task family token (e.g., Counting) (T) Operand one (counting what) (0) + Operand two (number positions) (0') + Start of text token (SOT) + Data
generated from DGP (Txt) End of text token (EOT) + Answer request token (ART) + Answer token (Ans) + End of sequence token (EOS). This can be summarized as follows.
Sample input: SOS 0 0' SOT Txt EOT ART Ans EOS.
(1)
consider the following tasks for pre-training:
Counting (C): Counting number of (say a) for the last positions (forty). This example will be written Ca40.
Counting composition of elements (CC): Counting number of 0 (say aa) for the last D positions (forty). This example will be written CCa40.
r, q; S q.p; p m. o p - n. o, m: q - n, m, O:
q - m. n; r- m: f- n: m m.lij;
m m-j. j. k:
- k, 1; n n,l,j, k; 0 o-l, 0 k, j- i;
j- i, h: h- k k,g, 1-i,h.. h-
d, f, e; g - f; - h i- i- d, e;
d a; 1- e - b; e e-c, f f- c, b, a: f- b, a:
Figure 14: PCFG setup: Grammar rules considered to generate the PCFG dataset. The highlighted
token represents the parent token. These rules have been adapted from Allen-Zhu & Li (2023c).
18

---

Published as a conference paper at ICLR 2024
Task Family Token T:: Operand Token 0:: 'a'
Sample: S, (,a, 40, <, a, b, a, C, a, b. a. a. a b. b, a. b. C, a. C, b.
a, a, b, b, b, a.
a, btw. <hacchaacabaccccccaaa44444......
b, b,
aac.cbee..ccuueebbcccaaaabbaaa.a #. #. # #, #, #, #
Answer: 15
Figure 15: PCFG Exemplar. representative sample from the PCFG dataset (Allen-Zhu & Li, 2023c)
Index of occurrence (I): Index from the EOT token when 0 (say a) occurred for the gr time (sixth). This example will be written as Ia10.
Index of occurrence of composition element (IC): Index from the EOT token when 0 (say aa) occurred for the tth time (sixth). This example will be written as ICa10.
Token value at an index (T): The token value at index (forty) before the end token. 0 is NULL here. This example will be written as TNULL5.
For the "Counting", "Counting composition clements", and "Token value at index" tasks, we set the value of D' token 40. For "Index of occurrence" and "Index of occurrence of composition element" task, we set the value o' token as 6. All five tasks above are considered during pre-training, but for fine-tuning we consider only single task with given operand. Specifically, we analyze fine-tuning the pre-trained models on the "Counting" and "Index of occurrence" tasks only.
We analyze the following tasks understand the effects of fine-tuning pretrained model's capabilities.
Counter: We intentionally reuse this task to demonstrate the effects compilation of the capability via Tracr versus learning the capability via PCFGs. Instead of being compiled, the model is trained count the number of tokens from set of tokens [OpT]. The model is then fine-tuned to exclusively count OFr {OPr) token. By making the sampling probability of OFr tokens high during pretraining, we can make the model preemptively performant on the downstream task; this allows us model the notion capability relevance.
Indexer: Amongst other tasks, the model is pretrained to output the index (location in string) of a token from the sel {Opr} occurs for the : time; fine-tuning is performed to output the index of th occurrence of another token OPT instead. We arbitrarily set k to 6 for our experiments, but emphasize that any integer less than context size can be used. If ( 0) denotes the index of tt occurrence of a token 0 in a string x, the spurious correlation enforced via constant offset q in operand token indices, i.e., r(x, DFT r(x, OPT) = q.
While the pre-training dataset generated by simply sampling from PCFG (see Fig. 15 for an example), for generating the fine-tuning dataset we provide explicit control over the value of Or by artificially adding the target tokens OFT from the fine-tuning task. It is important to ensure that the distribution shift between the fine-tuning distributions with different values of Cix is minimized and the data is fairly spread across multiple classes to enable reusability of feature via pretraining. As shown in Fig. 16, the class distribution the datasets with O and Cr = 0 for the counting and the index of occurrence tasks satisfies these requirements.
19

---

Published as a conference paper at ICLR 2024
CT.
CT 1.0
CT. 0.0
Crr 1.0
25K
Number of Samples
20k
15K
10k
5K
25K
Number of Samples
20K:
15K
5
10
15
20
25
30 35 40
10 15
20 25 30 35 40
10 15
20 25 30 35 40
5 20 15 20 25 30 40
(b) Count/Index of Occurrence b Class Label
(a) Count/Index Occurrence
Figure 16: Distribution of the class labels for Counting (first row) and Index of occurrence tasks (second row). (a) shows the distribution for the operand token and (b) shows the same for the operand token' The data is similarly distributed across different classes and the distribution shift for the two operands and the different values of Crr is small.
Summary: Eddie, elderly engine who likes to sleep, is blown off the train tracks big wind, but little bird and its friends help him get back on the tracks and become his best friends. Eddie learns that friends make you feel young again.
Features: Dialogue, Twist.
Sentence: Eddie was happy again.
Words: sleep, engine, elderly.
Story: Once upon a time, there was an elderly engine named Eddie. Eddie was old and tired. He liked to sleep lot. One day, Eddie was sleeping on the train tracks when he heard a little bird. "Hello, Eddie!" said the bird. "Why are you sleeping here?" "I am old and tired," Eddie told the bird. "I like to sleep." " Just then, something unexpected happened. big wind came and blew Eddie off the tracks. Eddie was scared but the little bird was brave. "Don't 't worry, Eddie!" said the bird. will help you." The bird flew fast and found strong friends to help Eddie. They pushed Eddie back on the tracks. Eddie was happy again. From that day Eddie and the bird were best friends. And Eddie learned that even when you are old and tired, friends make you feel young again.
Figure 17: TinyStories Exemplar. An example from the TinyStories Instruct dataset (Eldan & Li, 2023).
20

---

Published as a conference paper at ICLR 2024
B.3 TINYSTORIES INSTRUCT
For the TinyStories results, we use the TinyStories Instruct variant of the dataset (Eldan & Li, 2023)'. An example from this dataset is given in Fig. 17. For the fine-tuning datasets, take the original dataset and alter it in several ways. Details are discussed in App. F.
C DETAILS ON TRAINING AND EVALUATION
C.I TRACR
Compiled Model Details:
The compiled model obtained for the counting and max identifier tasks consists of three blocks, wherein each block contains a single head attention layer followed by two layer MLP. No normalization layers used by models developed using Tracr.
Training details: The compiled model is fine-tuned using SGD with momentum for 10K iterations with batch size of 96. Tracr yields models with rather sparse parameterization, which often yields unstable training dynamics (e.g., gradient spikes), especially with adaptive optimizers. To address this, we perform the following two interventions. First, we add small amount of initial gaussian noise Wnoise N00 0.001) to the weights of the compiled model to densify them slightly. Note that the scale of this noise is not high, i.c., it avoids any performance loss but sufficient enough to reduce gradient spikes resulting from extreme sparsity of model parameters. Second, we choose to use on SGD with momentum as the optimizer, using the following four choices of learning rates: Large LR (10-),, Medium LR (10 110-)), Small LR (10-3), and Very Small LR (10-1) The characterization of "Large" or "Small" is based on a general heuristic of what learning rate regimes are commonly used with SGD in modern neural network training. Linear warmup is used for 2K iterations followed by a cosine schedule with a minimum learning rate of the order 10-2 smaller than its max value. Evaluation of the fine-tuned model is done on both test set with and without the spurious correlation (ie. Cra = and Cra 1).
C.2 PCFG
Model details: We the minGPT model by Karpathy (2020) for all experiments on the synthetically generated PCFG dataset, similar to Allen-Zhu & Li (2023c). The model has close to 3 million parameters and consists
blocks each made up of multihead self attention with heads and two layers of MLP layers with an embedding dimension of 192.
Pretraining performed from scratch with a learning rate of 10-3 using the standard AdamW optimizer. Cosine learning rate is used along with linear warmup, where the warmup is used in the first 20% of the training. The model trained using the standard next token prediction task used for training language models. We consider the set of five tasks mentioned in the previous section during the pre-training phase, but focus on only one these tasks during fine-tuning. We use the task family token and operand token to define the notion of task. The task family token is sampled from a uniform distribution, while the operand token (0) is sampled from multinomial distribution. The sampling probability for different operands is varied in the experimental setup to understand the effect of capability relevance in fine-tuning. More specifically, analyze the following distributions for sampling the operand tokens (a, b, c):
Pre-training details:
* (a) = 0.999, Pt (b) = 0.001, Pr (c) = 0.0;
PP) 0.99, Pr (b) = 0.01, (c(c) 0.0:
.
(a) 0.9, Pt (b) 0.1, Pr (c) 0.0;
.
(a) = 0.7, Pr (b) = 0.2, Pr (c) = 0.1; and
Pr 0.5, Pr (b) 0.3, Pr (c) 0.2.
For each of the configurations of sampling distributions of operands, we pre-train the model for 10K, 50K, 100K and 200K iterations. The model trained in an online fashion to model the standard language model training pipeline, i.c., data is sampled on the fly from the data generating process during training time.
rust Int
21

---

Published as a conference paper at ICLR 2024
Fine-tuning details: While pre-training is done in the next token prediction fashion, fine-tuning is done in a supervised way where the model is required to just perform the desired fine- inn-tuning task. We use the final iteration model obtained from pre- pre-training as the initialization for fine-tuning. While pre-training is done on multiple pairs of task and operand tokens, the model is fine-tuned on a single pair of task and operand tokens. To simulate similar setup for fine-tuning as in Traer, we analyze the effect of fine-tuning the model using three different sets of learning rate: Large LR (πı: 10-*), Medium LR (πı; 10 9) and Small LR (7s : 10 10-6). Fine-tuning is done for 10K iterations using AdamW W optimizer with batch size of 96 samples. Similar to pre- training phase, we use cosine learning rate with an initial warmup of 20% of the fine-tuning tinettunin iterations. The minimum value of the learning rate is set to be 100x lower than the maximum learning rate. Similar to Tracr evaluation is done on both the test sets with and without the spurious correlation (Cra = and Cre = 1).
D MECHANISTIC INTERPRETABLIE TOOLS SETUP
In this section, we describe the different tools of interpretability considered in our work.
We present the attention maps for different tasks considered in the Tracr setup. Each map shows the tokens which are attending other tokens on the axis and the token which are being attended to on the x-axis. If token attended by many other tokens, then, in crude sense, this can imply that the presence of the token is impacting the underlying task performed by the model. In the Counter task, if significant attention is given a's b's is an indicator of the respective capability of the model. For the max identifier task, in the attention map in Block-0, the model implements the sorting function, where each token is attended by the tokens which are greater than that. The vocabulary order followed is > b > > d.... In the attention map of Block-2, the model implements the read function, where it outputs the token at the desired position in the sorted sequence.
Attention Maps:
Pruning: We consider single step pruning where we prune the weights/neuroos with largest dot product between their gradient and weights, where the gradients calculated by minimizing the loss for the capability we want to revive. More formally, let the weights of the model with N parameters be given by w; where i {0, 1...., N 1). Let the corresponding gradient be given by grad(w;, (w:) then the top-K weights with largest value of grad( w; )ue, are pruned off. This follows the pruning protocols proposed in prior work for reducing preserving loss via pruning (Molchanov et al., 2016; Lubana & Dick, 2021: Mozer & Smolensky, 1988). We use weight pruning for the Tracr setup and neuron pruning for PCFG, where a neuron defined as row in the weight matrix. We present detailed description of the pruning protocol considered in Algorithm-3.
Algorithm 3: Pruning Pseudocode. A fine-tuned model fo is parameterized by 0 and 0, denotes its 7th neuron or weight (we prune neurons in PCFG experiments and weights in Tracr). Pre- training task family token given by Opr and prepended to string X sampled from the data generating process, yielding the input Opr X. The true value corresponding to pre-training task family token Opr is given by y. Let the cross-entropy loss be given by CE. Let Topx (W) denote the indices of the top K values in the vector W.
def prune ()
# Forward prop the model on pre-training task
out (er o X)
Calculate the loss C = CE (out, y)
Calculate the gradients
grad = VoL
Calculate the dot product between model weights and gradients dotproduct = 0.grad
Select the indices of top X values
indices Tepи (dotproduct)
Prune off the neurons/veights present in top K. indices [[indices] = return 0
Probing: Probing is used to understand a particular capability is present in the model. In this, we train linear layer (probe) top of every block (residual layer's output) of the mini-gpt model and
22

---

Published as a conference paper at ICLR 2024
analyze if the probe is able perform a task requiring the of the desired capability. The probe a linear layer with the output dimensions same as the vocabulary size of the model. The probe is trained using the data randomly sampled from the PCFG data generating process for 4K iterations using Adam W optimizer with maximum learning rate of 10-3 which is decayed by a factor of 10 at 2K, 3K and 3.5K iterations. Training of the probe is done separately on the residual output of each of the six blocks present in the minGPT model. The stream corresponding to the answer token (Ans) is used the input to the probe.
Same set of hyperparameters as used the fine- tuning of the pre-trained Tracr model are used in reFT , except for the learning rate, which we force to be smaller than the corresponding fine-tuning learning rate. Note that this use of an even smaller learning rate is intentional: if the original pretraining capability be revived even with this setup, it is stronger evidence that the pretraining capability was never forgotten or removed.
Reverse Fine-tuning:
23

---

Published as a conference paper at ICLR 2024
E ADDITIONAL RESULTS
E. FINE-TUNING IN PRESENCE OF SOME PRE-TRAINING DATA
In this section, we demonstrate our claims also hold for an often used fine-tuning setup wherein, beyond the fine-tuning data, the model also gets to see some portion of the pretraining data again. Specifically, we perform three degrees of mixing of the pretraining and fine-tuning data: (i) 50% PT + 50% FT, (i) 10% PT + 90% FT, and (i) 0. % PT + 99.9% FT. We show behavior results on how the performance of the model improves as a function of fine-tuning iterations for different spurious correlations for low pretraining sampling prior in Figs. 18, 19 and high sampling prior in Figs. 20, 21. Furthermore, we probe these models' intermediate outputs to infer if features relevant to the pretraining capability continue to persist. Results can be seen in Figs. 22, 23.
=
Cr 0.5
Cr
1.0
QM
C==
C===
0==0
c==1
Cr=
Cr==1
100
65
Acc. OFT
100
30
2.5 5K 7.5KIKK
2.5K SK 7.5KI0K
Z35K 7.5KIOK 2.5 SC 7.5KIOK
2.5k 5x 7:5KIUK D 2.5 SK 7.5KIUK
(a) 50% PT + 50% FT (high mixing)
(b) 10% PT + 90% FT (medium mixing)
(c) 0.1% PT + 99.9% FT (low mixing)
Figure 18: Effect of different sampling probabilities of pre-training target token OPT on fine- tuning task's performance. We observe similar gains for different values of sampling probabilities of Opr during fine-tuning.
C
Cr.= D.5
Cre = 0.8
Crr
fTи
fs
GG
Ctc
Co
Cre=1
Cre DD
Cre
100
65
Acc. OPt
100
65
30
2.5K 7.5KIKK
2.5k 7.K
7.5KI0K
7.510KK
(a) 50% PT + 50% FT (high mixing)
(b) 10% PT + 90% FT (medium mixing)
(c) 0.1% PT + 99.9% FT (low mixing)
Figure 19: Effect of different sampling probabilities of pre-training target token Opr on pre- training task's performance. We observe a higher loss in performance low sampling probability used for sampling the pre-training target token OPr during fine-tuning.
CT.
Cr.
Cr
Cr.
GM
7s Cre
0
CTe
Cr
Gr. D
100
65.
Acc. OFT
100
65
30
2.5K SK 7.5K30K
2.5K 5K 7.5KT0K
2:3K 5K 7.5K10K
0 2.5K 5K 7.5KT0K
2.5K 5K 7.5K10K 0 2.5K 5K 7.5K10K
(a) 50% PT + 50% FT (high mixing)
(b) 10% PT + 90% FT (medium mixing)
(c) 0.1% PT + 99 FT
(low mixing)
Figure 20: Effect of differrnt sampling probabilities of pre-training target token OpT on fine- tuning task's performance. Pre-training is done using high sampling prior for fine-tuning task family token. We observe similar gains for different values of sampling probabilities of Opr during fine-tunning.
24

---

Published as a conference paper at ICLR 2024
Gr m 0.5
Crr m 0.8
Ctr = 1.0
Ct. =
C+ =
C==0 0
====1
Gr. =0
Cra =
Acc. OFT
30
100
2.5K SK. 7.5K10K 0 2.5K SK 7.5K10K
2.5 SK 7.3K10K 2.5 SK 7.510K
236 SK 7.5K10K 2.5K SK 7.5K0KK
(a) 50% PT + 50% FT (high mixing)
(b) 10% PT + 90% FT (medium mixing)
(c) 0.1% PT + 99.9% FT (low mixing)
Figure 21: Effect of different sampling probabilities of pre-training target token OPT pre- tuning task's performance. Pre-training is done using high sampling prior for fine-tuning task family token. We observe similar gains for different values of sampling probabilities of Opr during fine-tunning.
50% PT + 50%
10%
90%
0.1%
99.9% FT
CTr =0
=r=11
=-==0
Cr.=1
Acc. Opt
Acc. OFt
80 B1 83 94 85 66 Pre-trained Model
80 B1 B3 84 B5 86
80 B1 B2 83 B4 B5 86
80
81
82
63
84 BS B6
80 B1
85 B4 85
(b) 7/M
Figure 22: Probing analysis corresponding to Fig-18 and 19
50% PT 50%
10% 90% FT
0.1% PT 99.9% FT
CT 0
Cr. =1
Cr. 0
Cr. m1
Acc. OFT Acc. Ovt
30
80 82 84 B6 Pre-trained Model
80 B1 92 B3 84 B5 36
66 60 B2 96
90 B1 B2 63 64 BS B6
30 B1 83 84 B5 B6
(b) пM
nss
Figure 23: Probing analysis corresponding to Fig-20 and 21
25

---

Published as a conference paper at ICLR 2024
E.2 JAILBREAKING ANALYSIS
We emulate jailbreaking (Wei et al., 2023; Zou et al., 2023) in our PCFG setup by defining several task family tokens describing the same task. Specifically, for the "Counter" task, we use three task family tokens TNJ, TJ,, T.2 to refer to the task in a string. Here subscript N indicates the task family token will not allow jailbreaking, while J//J2 indicate the task family token can be used to jailbreak the model, as explained next. For pretraining, the token TNJ may be paired with operand tokens a, b, с learn to count them from inputs sampled from the PCFG. However, tokens J, Tj are used only for counting a. During fine-tuning, the model is fine-tuned to count the token b using the task family token J. For evaluation, we compute the model's accuracy on its ability to count the token using either the task family token TyJ TJi T.z. As shown in Fig, 24, the model is unable to infer the count of a if the task family token TNJ is used; however, task family tokens TJ,, TJ, are used, the model performs perfectly if the prior for sampling the fine-tuning target during pretraining was sufficiently high. We argue that this is expected because under a high sampling prior breaks the symmetry between task family tokens (indeed, TJ, is only seen with operand token a, but T is seen for all operand tokens. This indicates the pretraining capability continues to persist in the model, enabling jailbreaking, To further investigate this result, we also probe the fine-tuned models. Results are shown in Fig. 25. As expected, we see task family tokens TJ₃ TJ, allow for linear readout of the count of a; however, we see that even for inputs with task family token TNJ, the model does encode the count of a in the outputs around the middle layers!
7
TNy
T
TA
Tr,T,
100
Acc. Opr
65
30-
2.5K
5K
10K
7.5K
2.5K 7.5K
10K
10K
(b)nss
Figure 24: Jailbreaking analysis using PCFG. We report performance on the pretraining task (counting OpT) as a function of fine-tuning iterations, where the fine-tuning task (counting OFT) is performed using the task family token Tyy- We find that the model is able learn the fine-tuning task and seemingly performs poorly on the pretraining task when task family token Tyu is used in the input. However, in presence using task family tokens T3, or T,, in the input shows the model can still perform the pretraining task perfectly-i.e we can jailbreak the model.
sufficiently relevant capability (high pretraining prior for D ),
74
74
T
Tw
T
TN
Acc. Opr
100-
Acc. OfT
80 as de dis no
00
0A 05 166
100
12
04 85 86
00 81
02
84 05 06
66 01
06 05 06
(a) Pre-trained Model
(b) 7M
(c) its
Figure 25: Probing analysis for the setup used to understand jail-breaking. Similar results on using the fine-tuning token or the jailbreaking token for training the probe indicate that the pre-training capabilities are not removed on fine-tuning.
26

---

Published as a conference paper at ICLR 2024
E.3 SAMPLE EFFICIENCY ANALYSIS FOR REVERSE FINE-TUNING
To emphasize the fact that the pretraining capability is "revived" in the model relatively sample- efficiently, we repeat Fig. 9, where models trained PCFG are reverse fine-tuned, and repeat the experiment with the Scr. + FT baseline for Tracr compiled models. As can be seen in Figs. 26, 27, compared the baseline, the model learns perform the pretraining task in substantially fewer iterations than the baseline. We note that for the Tracr models in these results, even an extremely small learning rate sufficient to revive the pretraining capability! We also note that we not sweep over the Ot hyperparameter in the Tracr models because they are compiled, i.e., we cannot control correlation with pretraining capabilities in meaningful way.
Ctr
Crr
Scr. + FT
nM
ns
100
65
Acc. OPT
30
100
65
30
102
10a
10*
102
10
10
102
103
10
Iters (Log Scale)
(a)
P
(b) PM
(c) Pt
We set Cte to be 0 to test if the model performs well regardless of spurious correlation. We observe that when strongly relevant capability is present (a, b), the model very quickly (0.1-1K iterations) starts to perform well on the task via reFT, even if behavior relevant the capability ceased during pretraining (e.g., when CTv is 1). Meanwhile, when the model possessesses a weakly relevant capability (c), this "revival" is slightly slower (3K iterations). In contrast, the Scr, FT baseline only reaches perfect accuracy 4.5K iterations and when using larger learning rate T7M-
Figure 26: Reverse Time-Tunigg on PCFGs:
reFT
Scr. + FT
Acc. OPT
100
65
10a
10
103
10
10a
104
Iters (Log Scale) (b) ns
(a) ΩM
(c) nvs
Figure 27: Reverse Fine-Tuning Tracr: We Cra to be 0 to test if the model performs well regardless of spurious correlation. We observe that the fine-tuned model upon reFT very quickly starts starts perform well on the pretraining task. Moreover, the protocol works even extremely small learing rate is used. In contrast, the Scr. + FT baseline only reaches a large learning rate пı is used, and does so less sample efficiently. We note that the results for 7jM learning rate look worse than the ns learning rate around 103 iterations because lIM is too big of learning rate, forcing the model to essentially go through "retraining" phase.
27

---

Published as a conference paper at ICLR 2024
E.4 REVERSE FINE-TUNING A MORE SAFETY-ORIENTED FINE -NN-IUNING PROTOCOL
The fine-tuning protocols used in the bulk of the paper focus on learning a new capability, e.g.. counting new operand, while promoting reusability of capabilities learned during pretraining. Part of our motivation is to see if a pretrained model is actively forced to remove a capability, does that work? To analyze this, we define fine-tuning protocol called randFT wherein the model trained to actively produce an incorrect output for inputs that require use of the pretraining capability. For example, if the model possessesses the capability to produce the count the number of occurrences of token OpT in string, we fine-tune it to produce the count of tokens OFT - b in that string. We analyze these fine-tuned models analyzed via reverse tine-tuning (reFT) ), i.e., by further training them to produce the correct outputs (number of occurrences of token Opr). We provide results for three baselines as well: (i) Scr., wherein the model is trained from scratch learn to count the token a; (ii) Scr. + FT, wherein the model is initialized with parameters trained via trained from scratch to count a separate token (OFr) and then the model is fine-tuned to count the token Opr: and (iii) reFT, which follows reverse fine-tuning models that were fine-tuned with the protocols used in the bulk of the paper. i.e., fine-tuned to learn new capability that is related to the pretraining one.
Results are shown in Fig. 28. We specifically zoom in on the the scenario where reFT takes the longest time, i.e., when the sampling prior of the downstream target OFT is low in pretraining data; results for other sampling priors are shown in Fig. 29 We see that reverse fine-tuning a randFT model similarly sample-efficiens as the standard reFT pipeline used in the bulk of the paper, while being more cample--efficien than the Scr. and Scr. + FT baselines.
In addition, we perform a probing analysis of the randFT models in Fig. 30. We again find that we can predict the information relevant for the pretraining task, i.e., the count of Opr.
reFT
reFT- - FF CT =0
Scr.
Scr. + FT
nM
ns
Cr=
100
Acc. Opt
65
30
10a
10a
10a
10a
10a
10
10a
10
10a
10a
10
(a)
(b)
Iters (Log Scale)
Figure 28: Reverse fine- fine-tuning a model fine-tuned to remove its pretraining capability. See text
in Sec. E,4 for details.
=
0
CTr =1
Scr. + FT
QM
ns
100
65
Acc. OPT
100
65
30
102
10a
10
102
10a
10a
102
10a
10c
Iters (Log Scale) (b) PM
(a) PH
(c) TH
Figure 29: Reverse fine-tuning performance using randFT fine-tuning protocol to forget the pre-training capability. We follow the setup of Fig. 9 and plot results for several sampling priors of the target token for fine-tuning, i.e., OPT. but we use randFT for fine-tuning the models before reFT. The baseline results Scr. + FT copied from the Fig. 28, i.e., baseline is not trained in an "adversarial" way, but the randFT results are. While this makes the baseline unfairly stronger, we find reFT the randFT models are still more sample efficient.
28

---

Published as a conference paper at ICLR 2024
=
=r
Cr=
Cr.=1
Acc. OpT
100.
05-
Acc. OFr
100
30-
(a) Pre-trained Model
(b) f/
(c) 0s
Figure 30: Probing analysis of randFT fine-tuning protocol. We plot probe accuracy versus the index of the block in the Transformer model, Cro is set to O. The pretrained model (left) acts as baseline for the trend of performance through the model's blocks. most scenarios, we find we can infer the count of OpT with a similar trend as the pretrained model (left), drop in performance is observed only when learning rate 7JM is used with a weakly relevant capability (low sampling prior). This indicates pretraining capabilities continues to persist upon fine-tuning.
29

---

Published as a conference paper at ICLR 2024
F DETAILS AND RESULTS ON TINYSTORIES EXPERIMENTS
In this section, describe experiments the TinyStories dataset in more detail. These experiments are designed to validate our hypotheses in a more realistic language modelling setting. Overall, the results support our hypothesis that fine-tuning does not lead to deletion of capabilities as they can be revived in a sample-efficient way and uncovered through probing.
F.1 MODEL TRAINING
Dataset. We use the TinyStories (Eldan Li, 2023) dataset to train our models. This data consists of children's 's stories written by GPT-3.5 and GPT-4. Each story is several paragraphs long, and comes with several attributes labelled: a set of three words that are included in the story; a sentence that is included in the story; GPT-3.5-written summary of the story; and list of 0-3 "story features", such as Twist, Dialogue or Bad Ending, which the story abides by.
We use the TinyStories-Instrrn Instruct version of this dataset 2, wherein each story prefixed with an "instruction" containing the story attributes described above, hence enabling the model to learn to conditionally generate stories based on an input or instruction.
Pre-training. We pretrain 91 million parameter autoregressive language models with similar architecture LLaMa (Touvron al., 2023), with a custom tokenizer with vocabulary of size 8192 trained on the dataset." They have hidden dimension 768, 12 layers, and 12 attention heads per layer. These models are trained with the standard language modelling cross-entropy loss, with batch size 128, sequence length 1024, no dropout, for 30,000 gradient steps, with learning rate schedule with a linear warmup from 0 and cosine decay to O, with maximum learning rate 0.001. These models achieve loss of 0.8 at the end of training. and generate coherent multi-paragraph stories given specific instruction in the form it saw during training.
Fine-tuning. We are interested in analysing whether fine-tuning these models can alter underlying capabilities. The specific capability we investigate is that of generating stories containing Twists (which is one of the story features), and analysing whether various fine-tuning protocols can remove this capability from the pre-trained model. We investigate a variety of fine-tuning protocols modelled after plausible realistic scenarios where one may want fine-tune a model not generate text of a certain type (e.g., highly toxic text), regardless of the input instruction. These include: Filtering fine-tunes the model on dataset where all instances of stories with Twists filtered out; Filtering + Mix & Match filters, and then replaces all instances of another, unrelated feature (in this case, Foreshadowing) in the instruction with the Twist feature; and Filtering Randomisation filters, and then adds the "Twist" instruction to the prompt for stories that do not contain Twists, thus training the model to not model stories with Twists even if instructed. This last protocol acts as a kind of adversarial training (in that there are stories with the Twist instruction but no Twists), and introduces a spurious correlation between the Twist instruction and the foreshadowing capability, as in the Tracr and PCFG results.
We take the pre-trained model described above, and fine-tune it with these various protocols. We then perform reFT a dataset of stories which all have Twists in, measure the extent which each fine-tuning protocol deleted the capability of Twist generation. To ensure a good control, compare the reFT models to model pre-trained on data with no Twist stories, which is then fine-tuned on Twist stories. The sample efficiency and final performance of this model serves as a comparison for the reFT ed models.
F.2 EVALUATION METRICS
We evaluate whether the fine-tuning protocols have removed the capability model and generate stories with Twists in multiple ways. Firstly, we look at the loss on stories with Twists. fine-tuning deletes the Twist capability, we expect the loss on this dataset to increase.
'ttpps: / /huggi aag.ace.comm http:////..ttteeeeeeeeee//a//aaa In *Our code is hased this htppshonyynttaa//gghttbb.eeaaaaa
30

---

Published as a conference paper at ICLR 2024
GPT Evaluations. To evaluate the generative capabilities of these models, we generate stories from them given prompt instructions with the Twist story feature. We then evaluate whether these stories contain Twists. To do this evaluation, a GPT-3.5 model to classify whether a given story has a Twist or not. To do this, we use the TinyStories dataset and accompanying labels. This fine-tuned model achieves 92% accuracy held-out test set after fine-tuning. We generate stories with multiple different prompts from both the fine-tuned and reverse fine-tuned models throughout fine-tuning, and measure the proportion of stories which are classified as having Twist, which we call the generation score.
use the OpenAl GPT fine-tuning API to fine-tum
Probing. As well as using reFT measure whether the fine-tuning protocols have deleted the capability to generate Twists, we also probing to evaluate whether fine-tuning removes information from internal representations. We train linear probes on the internal activations of the transformer models to predict which story features (e.g. Twist, Bad Ending. Dialogue) present in the story. These probes take an average of the activations at the final 10 token positions of the story. Given that this is a multi-label classification problem we employ separate binary classification probe to classify the presence of each story feature. We use the accuracy of these probes at different layers before and after fine-tuning, and the control pre-trained model which was trained data with no Twists, to measure whether fine-tuning has removed information from the models' internal activations.
F.3 RESULTS
Reverse fine-tuning The loss on stories with Twist during fine-tuning is shown in Fig. 31. This shows that the fine-tuning protocols are raising the loss, and hence behaviourally deleting the capability of fine-ture tuning. The generation scores are shown in Fig. 32. This again reinforces that most fine-tuning protocols are removing the capability behaviourally, as the generation scores (while noisy) drop close to O.
Fig. 33 shows the loss during reFT for all the fine-tuned models, as well as the control model pre-trained without stories with Twists, and Fig. 34 shows the generation scores. Both of these results show that the fine-tuned models learn the new capability in much more sample-efficient way. and in fact converge lower loss on this dataset than the control pre-trained model.
Probing In addition to the reFT results, we perform probing experiments. The probing accuracy for the Twist feature across layers for the fine-tuned models and the two control pre-trained models shown in Fig. 11. which reproduce here in Fig. 35 for completeness. These results show that small amount of information about story classification has been removed from the activations of the fine-tuned models compared to the model pre- trained with Twist stories, but the reduction is very minor, as shown in comparison to the information present in the model pre-trained without Twist stories.
Fig. 36, Fig. 37, and Fig. 38 show similar plots for several other story features. Some of these are easier or harder for probes classify, but the important result is that the difference in probe accuracy between the fine-tuned models and both pre-trained control models is negligible for all of these features, showing that the results in Fig. 35 are due to the Twist feature, i.e., the feature that we trained the model to delete.
hhttps://platffmm http://ppaaa..cccccceeeee/e/eeee des finintuning
31

---

Published as a conference paper at ICLR 2024
Filtering + Randomisation
Filtering + Mix & Match
Filtering
LR: ns
LR: MM
0.85 -
Loss
0.80
100
10
102
10
100
10
10
103
Training Iteration
Figure 31: Larger learning rates lead to more pronounced loss modelling capability. The plots show loss data with the Twist feature present while fine-tuning to delete the capability to model text with the Twist feature, for different learning rates and fine-tuning protocols.
Filtering + Randomisation
Filtering + Mix & Match
Filtering
LR: π7s
LR: NM
Generation Score
0.5
0.0
100
101
102
103
100
101
102
10
Training Iteration
Figure 32: Larger learning rates lead to more pronounced loss of generative capability. The plots show the generation score for the Twist feature present while fine-tuning to delete the capability to model text with the Twist feature, for different learning rates and fine-tuning protocols.
Filtering + Randomisation
Filtering
Not in Pretraining
Filtering Mix & Match
LR: ns
LR: AM
Loss
0.9
0.8
100
101
102
103
100
10
102
103
Training Iteration
Figure reFT easily recovers deleted capabilities. We plot loss on data with the Twist for reFT of various models fine-tuned to delete the capability. well as control model which pre-trained without data with Twists. The fine-tuned models learn the capability more additionally converge a lower loss than the control model.
and
32

---

Published as a conference paper at ICLR 2024
Filtering + Randomisation
Filtering
Not in Pretraining
Filtering + Mix & Match
LR: ns
LR: AM
Generation Score
0.75
0.50
0.25
100
10a
102
103
100
10a
102
103
Training Iteration
Figure 34: reFT easily recovers deleted generative capabilities. We plot the generation scores for the Twist feature for reFT of various models fine-tuned to delete the capability, as well as a control model which was pre-trained without data with Twists. The fine-tuned models learn the capability much more sample-effiiietly. and additinoally converge lower loss, than the control model.
Filtering
Not in Pretraining Present in Pretraining
Filtering + Mix & Match
Filtering + Randomisation
LR: 7M
LR: ns
Probe Accuracy
0.95
0.90
0.85
Bo B1 B2
B3 B4 B5 Bg By Bg B1
Bo B1 B2 B3 B4 B5 Bg By Bg Bg B10
Probe Layer Index
Figure 35: Probing the presence of capabilities TinyStories Models. We plot probe accuracy of classifying whether a story contains a Twist or not wrt. the layer of the Transformer model (similarly to Fig. 8). Accuracy on models pre-trained with without Twist data (Present/Not in Pretraining respectively) act as upper and lower bounds on the expected accuracy of the probes, and are plotted on both LR figures for ease of comparison, although they do not use fine-tuning learning rate. We find that regardless of fine-tuning protocol (Filtering, Filtering Randomisation, Filtering + Mix & Match), for the lower LR no fine-tuning protocol removes meaningful amount of information from the activations, and a similar but less strong trend holds for the higher LR, implying that the pre- pre-trained model retains its capability of story identification (a necessary capability for story modelling) throughout fine-tuning. Identical to Fig. 11
33

---

Published as a conference paper at ICLR 2024
Filtering
Not in Pretraining
Filtering + Mix & Match
Present in Pretraining
Filtering Randomisation
LR: f]
LR: ns
Probe Accuracy
0.95
0.90
0.85
Bo B1 B2 B3 B4 B5 Bg By Bg B10
Bo B1 B2 B3 B4 B5 B6 By Bg Bg B10
Probe Layer Index
Figure 36: Probing the presence of capabilities TinyStories Models. We plot probe accuracy of classifying whether story contains the Foreshadowing feature not wrt. the layer of the Transformer model. All other details the same as Fig. 35
Filtering
Not in Pretraining
Filtering + Mix & Match
Present in Pretraining
Filtering + Randomisation
LR: 7M
LR: ns
1.00
Probe Accuracy
0.95
0.90
0.85
Bo B1 B2
B3
B4
B5
B6
B7 Bg Bg B10
Bg B1
B2 B3
B4
B5
Bg
B7
Bg
Bg
B10
Probe Layer Index
Figure 37: Probing the presence of capabilities in TinyStories Models. We plot probe accuracy of classifying whether story contains the Moral Value feature not wrt. the layer of the Transformer model. All other details the same as Fig. 35
34

---

Published as a conference paper at ICLR 2024
Filtering
Not in Pretraining
Filtering + Mix & Match
Present in Pretraining
Filtering + Randomisation
LR: NM
LR: ns
Probe Accuracy
0.95
0.90
0.85
Bo B1
B2
B3
B4
B,
B6
By
Bg
Bg B10
Bp B1
B2
B3
B4
B,
B6
B
Bg
Bg
B10
Probe Layer Index
Figure 38: Probing the presence of capabilities in TinyStories Models. We plot probe accuracy of classifying whether a story contains the Bad Ending feature or not wrt. the layer of the Transformer model. All other details the same Fig. 35
35

---

Published as a conference paper at ICLR 2024
G ADDITIONAL TRACR RESULTS
In this section, we present additional results on the counting and max-identinn tasks with Tracr models. These results provide extensive analysis and support claims presented in the main paper. Firstly, we present the detailed attention maps showing the full input sequence data: Fig. 40 and Fig. 39 show the attention maps corresponding to Fig. 42 and Fig. 43. We now present the detailed results on Tracr's Counter tasks.
0.5
Attending Token
0.4
0.3
0.2
0.1
HI (itooaat www. www.
0.0
Token Attended
(a) nM
(b) nL
Figure 39: Counter Task: Detailed visualization of the capability revival analysis on Traer compiled for the Counter task. Observation: On using s/ we are able to revive the compiled capability of the Tracr model fine-tuned with Tум, whereas using ar for fine-tuning hampers the underlying capability of the model and therefore there is revival seen.
Tracr Program
CTr 0
CTr =1
CT.
CT.
Attending Token
0.4
0.
02
0.0
Token Attended
(a) ns
(b) NM
Figure 40: Counter Task: Detailed visualization of the self attention layer in Block-1 of the Tracr fine-tuned model. Observation: (a) shows the map for the compiled Tracr model, where we observe that a's 's are being attended by the second token in the input sequence. (b) On using for fine-tuning the absence of spurious correlations (Cre 0), a's are still being attended in the attention map. (c) On using 7M, the model learns the fine-tuning task as b's are also being attended now. On fine-tuning in the presence of the spurious correlation (CT. = 1), the model doesn't learn to attend b's, but rather learns the spurious correlation.
G.1 BEHAVIORAL RESULTS ON FINE-TUNING
Summary of results the Counter task. As shown in Fig. 42, using 7)M the model seems to learn a new capability of counting b's (the model primarily attends to b in its attention map). However, on using small learning rate (ns) of 10-3, in the absence of correlations, the model able to learn to attend to b's in its attention map. Thus the model is not able to learn the capability of counting b's. As shown Tab. 3, in the presence of spurious correlations however, the model is able to learn the spurious correlation and achieve high accuracy on the correlated test set. We also present the visualization of the attention maps after reverse fine-tuning in Fig. 43, where is observed that on using 7) for fine-tuning. revival of capability is possible even on using very small learning rate (M)+) of 10-4 We present detailed results on reverse fine-tuning Tab. 4 Whereas, case the model is fine-tuned with a large learning rate (7z,) of 10-1 revival of capability is not possible. We also present analysis of single weight pruning and grafting in Fig. 55. On pruning off a single weight from the Tracr model compiled to count 'ss, the model can achieve boost in accuracy of over 60% on the task of counting a a's. This observation is evident only when the Tracr model is fine-tuned on correlated dataset.
36

---

Published as a conference paper at ICLR 2024
CTr =
CTr =
CTr = 0
CTr =
0.5
0.4
Attending Token
0.3
0.2
0.1
0.0
(a) ns
Token Attended
(b) NM
Figure 41: Counter Task: Validation of Tracr observations Counter task using three different input samples. The rows represents different input samples. Observation: Using 7/s for fine-tuning the Tracr model compiled on Counter task is unable learn attend b's (a). But the model learns the spurious correlation. On the other hand, on using 7MM the model is able to learn the fine-tuning task by attending to b's But in the presence of the spurious correlation Cr the model still doesn't learn to attend b's.
Summary of results on the max-identifer task, We present a visualization of the attention maps for the max identifier task in Fig. 44, where we observe that Tracr model implements the sorting and the reading functions in the attention maps in blocks and 2 respectively. On fine-tuning the model using different learning rates, the sorting capability implemented in Block-0, gets distorted, thereby resulting in poor fine-tuning performance (as evident in Tab. 3). However using l/ss (10-4), changes the reading function, without disturbing the sorting function. Thus the model able to perform well on the downstream task (as evident in Tab. 3).
Tracr Program
=
c
Cr.s1
Altending
Attending
Tokon Attended
Token Attended
[b)
(a) Os
(b) m
Figure 43: Capability Revival Analy- sis: Using Jlys (a) is able to recover the old capability on reverse fine-tuning the model fine-tuned with Mag But ns is not able recover the original capability, when the compiled model is fine-tuned with тд- This because using a large value of learning rate during-fine-tuning hampers the pre-training capabilities.
Figure 42: New capabilities are not learned on using Ti5 for fine-tuning, (a) The counting a's capability im- plemented by the originally compiled Tracr model is to attend to a's. On using 7/s (10 (10-3) for fine-tuning, the com- piled model is not able to learn the capability of counting b's (b). Increasing the learning rate makes the model learn
attend b's is (c), but the pretraining capability of attending to a's is still exists.
G.2 COUNTER RESULTS
A detailed analysis of the attention maps of block-1 1 and 2 for different learning rates is shown in Fig. 45. We further validate our results for three different input datapoints in Fig. 41 and Fig. 49. A detailed analysis of the activation map in Block-1 for different values of Cry is shown in Fig. 48.
37

---

Published as a conference paper at ICLR 2024
Tracr Program
Block-0
Block-0
Block-2
Block-0
Block-2
Attending Taken
Kehwing
hwwhiggggg
Eehwing hhmm
Kehwmg
Token Attended (b)
Figure 44: Learning of the fine-tuning capability is affected by type of compiled capability present in the model. (a) The Tracr program implements the sorting function in Block-O and read function in Block-2. Using a/ and Vss can destroy the sorting capability present Block-1 (b). But using yvs, preserves the sorting capability (c). Thus on using i|vs, the model learns read different stream of output, while preserving the sorting capability.
Table 3: Results on counting and max element task Tracr setups. The evaluation is done on test sets with and without the spurious correlation. The Tracr compiled models fine-tuned for different learning rates and different value of CTr-
Counting Element
Max Identifier
Ctr
CT
=
Cr =
CTa
1
UT
0
Acc. OPT
Acc. OFr
Acc. Ace. OFr Acc. Opr
Acc. OPT
Acc. Opr
Acc. OPT
0.0
100.0
0.0
100.0
20.3
34.5
0.0
99.3
0.2
0.0
100.0
0.0
100.0
0.5
92.0
0.0
97.1
100.0
100.0
0.6
97.0
0.1
97.6
10
0.6
0.0
100.0
0.0
100.0
0.3
98.5
0.0
96.7
0.8
0.0
100.0
0.0
100.0
0.1
99.4
0.0
98.6
0.9
0.0
100.0
0.0
100.0
0.7
98.2
0.1
92.5
0.0
100.0
35.8
0.7
0.3
99.6
16.8
37.8
0
I.I
96.3
0.0
98.8
29.7
0.2
16.3
52.6
0.2
0.0
100.0
0.0
99.2
28.4
18.6
19.0
46.0
95.9
4.8
87.9
3.3
92.6
10
0.6
0.1
99.9
0.0
98.8
3.9
83.2
2.6
82.5
0.8
0.3
99.6
0.1
97.0
4.5
88.8
6.6
72.9
98.5
39.3
16.0
45.7
26.9
11.1
0.3
98.3
4.2
0.2
TT.1
78.5
23.8
14.4
0
54.6
1.2
25.7
27.2
6.4
20.2
4.5
28.5
0.2
50.2
15.0
26.5
24.3
7.4
27.4
5.4
28.0
7.1
90.9
19.8
2.3
11.3
24.0
7.6
20.8
10
0.6
4.1
94.2
11.8
2.2
11.8
26.7
8.4
20.1
0.8
1.3
98.3
6.7
0.7
11.5
34.3
8.5
19.9
9.2
0.7
14.6
32.2
15.8
4.0
94.3
TO.3
2.2
16.0
33.2
12.8
14.0
0
32.6
0.0
10.6
28.7
0.5
82.6
0.5
91.1
0.2
59.2
0.1
31.9
24.4
0.1
84.8
0.6
91.3
0.5
28.5
65.1
37.8
5.6
0.0
89.3
0.6
91.8
10-4
0.6
24.4
70.3
35.6
4.8
0.0
89.6
0.6
90.8
0.8
14.1
84.2
29.7
2.1
0.0
89.7
0.6
89.9
0.7
0.0
93.2
97.1
1.6
98.3
10.6
0.2
0.0
90.2
0.7
88.6
We present an evidence further, showing that capability of the Tracr compiled model to count a's is still present in the model in Fig. 46, 47, where Fig. 46 presents a closer look of the Fig. 47. As can be seen in Fig. 46, on using 77 and 17M, Block-1 activation map of the Tracr fine-tuned model shows neurons corresponding to token being activated in different output channel.
Finally, we present evidence of the wrapper being learned by the model on fine-tuning using spuriously correlated dataset. We show that this wrapper can be localized in very few neurons of the model.
38

---

Published as a conference paper at ICLR 2024
Table 4: Results on counting task Tracr for reverse fine-tuning with different learning rates. Fine-tuning was done using !jM- The evaluation is done on test sets with and without the spurious correlation. The Tracr compiled models are fine-tuned for different learning rates and different value of Ctr.
Counting Element
CTX
CTo
Cro
0
Cro
CTa
0
Acc. OPT
Acc. OFF
Acc. Opr
Acc. OFr
Acc. Opr
Acc. OFr
Acc. Opr
Acc. OFT
96.5
100.0
2.1
0.0
94.0
100.0
0.1
0.0
0.2
98.5
100.0
0.0
94.9
100.0
0.1
0.5
39.4
100.0
43.6
0.0
44.9
100.0
5.6
0.0
10
0.6
72.9
100.0
26.7
0.0
69.4
100.0
0.2
0.0
0.8
49.3
100.0
6.5
0.0
37.1
100.0
16.6
0.0
100.0
64.0
0.0
34.1
100.0
69.2
100.0
3.7
0.0
65.4
100.0
6.3
0.0
63.3
99.9
36.6
0.1
65.5
98.6
0.0
0.0
19.5
100.0
48.8
0.0
29.6
100.0
17.5
0.5
14.3
100.0
54.4
0.0
28.9
99.9
18.7
0.0
10
0.6
86.2
99.9
13.8
0.1
78.3
98.6
0.0
0.0
0.8
65.6
100.0
1.7
0.0
43.7
99.5
10.8
0.0
0.9
33.3
100.0
27.9
0.0
36.5
100.0
12.6
0.0
99.0
99.6
0.9
0.3
95.2
96.8
0.1
0.1
19.8
99.9
34.9
0.0
23.7
98.6
10.2
0.0
33.8
78.0
24.4
42.9
0.5
2.0
87.6
33.2
9.5
18.9
85.7
11.9
0.9
10
0.6
7.1
99.8
35,4
0.1
22.3
97.3
15.4
0.1
0.8
11.6
97.2
45.8
0.3
27.3
95.5
16.5
0.5
0.9
24.2
75.2
20.3
23.6
33.5
68.8
13.1
1.5
68.5
99.9
26.7
0.1
65.3
98.6
5.8
0.0
45.2
99.9
0.1
0.1
16.9
98.5
4.9
0.0
0.2
30.1
9.4
22.7
45.1
18.1
26.2
17.9
19.5
0.5
30.1
3.2
22.7
41.3
15.7
26.1
17.9
16.2
10-4
0.6
54.1
0.0
45.9
35,9
0.0
27.8
25.7
11.9
0.8
26.8
83.0
64.5
10.9
3,0
76.8
49.2
1.3
0.9
27.9
85.0
61.5
8.4
2.1
80.9
44.2
1.1
45.2
99.6
31.9
0.3
42.2
96.8
12.6
0.0
As shown in Fig. 55, we present this evidence for differrnt values of CTr in Fig. 56. Similar the analysis presented for PCFG where we prune multiple neurons, we analyze the effect of pruming of mutliple weights and neurons in Fig. 57 and Fig. 58 respectively. These results verify that the wrapper learned by the Tracr compiled model on fine-tuning using spuriously correlated dataset can indeed be localized are indeed because of removal of the wrapper, we present the histograms showing the distribution of predicted classes in Fig. 59, where it can observed that after pruning the model still predicts multiple classes.
few weights of the Tracr model. To ensure that the gains achieved on pruning
G.3 MAX IDENTIFIER RESULTS
In this section, we provide additional evidence and a detailed analysis of the performance of the Tracr compiled model on the max identifier task. We show that the model implements the sorting pattern in the activation map of its first block in Fig. 53 and Fig-54. We present validation of our observations on considering the spurious correlation as the difference between the fifth and seventh maximmum element being three in Fig. 51. We validate our results for three different input data-points in Fig. 50. detailed visualization of the attention maps in Block-O and Block-2 for different learning rates shown in Fig. 52.
39

---

Published as a conference paper at ICLR 2024
Tracr Program
Tracr Program
Block-1
Block-2
Block-1
Block-2
CTr
Ctr
Attending Token
6.3
002
nys
(a)
Token Being Attended To
(b)
Figure 45: Counter Task: Visualization of the attention maps of the first and second blocks of the Tracr fine-tuned models. (a) shows the analysis when the spurious correlation is not present in the fine-tuning datatset, whereas case (b) the spurious correlation is present in the fine-tuning dataset. The first shows the maps for the Tracr compiled model and other
shows the analysis for different learning rates. Observation: (a) Using or fv' the model is not able to learn to attend is and thus the fine-tuning task performance poor. Whereas using 7JM the model is able learn to attend to b's, however the capability to count a's is likely still present since the model still attends to a Further increasing the learning rate leads to distortion of the compiled capabilities, and thus model learns the fine-tune tuning task by learning different capability. (b) In the presence of spurious correlation, even for large learning rate the compiled capability still present, since the model attends to a's.
40

---

Published as a conference paper at ICLR 2024
G.U
0.4
Tracr
fM
fis fvs
Input Token
Input Token
Output Channel Number
Figure 46: Counter Task: Visualization of the activated output of the first MLP layer in first and second blocks. The visualization is shown only for channel numbers 10-100. Observation: Using Mas preserves the Tracr compiled capability, while also learning the fine-tur ine-tuning task. This shows that the model changes behaviourally but mechanistically the compiled capabilities are still present in it.
Black-1
Tracr Program
Block-2
30
0M
Input Token
fyss
(a)
Output Channel Number
(b)
Figure 47: Counter Task: Visualization of the activated output of the first MLP layer first
and second blocks. This is the complete visualization of the activation map presented in Fig. 46.
41

---

Published as a conference paper at ICLR 2024
Cty
CTr
Ctr
CT
Cr.
Attending Token
Token Attended
(a)
(b)
(c)
(d)
(e)
Figure 48: Counter Task: Effect of the presence of the spuriously correlated data-points in different fractions (Crz) in the fine-tuning dataset. The Traer compiled model with the capability to count a's is fine-tuned to count b's on different values of Crr. 7M is used for fine-tuning. Observation: On increasing the value Cr., the model gives lower attention to b s and in case of Cr 1, almost attention is given by the model to b's.
ns
reFT
NVS
reFT
MM
ReFT
ns
reFT
Attending Token
(a) ΠM
Token Attended
(b) n
Figure 49: Counter Task: Validation of Tracr observations Counter task on reverse fine-
tuning on three different input samples. The rows represents different input samples. Observation: Capability revival is possible on using 7/ for fine-tuning but not using f/r.-
42

---

Published as a conference paper at ICLR 2024
=
CTr
Ctr
Ctr 1
Attendinn Token
(a) nvss
Token Attended
(b) NM
Figure 50: Max Identifier Task: Validation of Tracr observations on max identifier task on three different input samples. The rows represents different input samples. Observation: Using yys preserves the original Tracr capabilities and therefore performs well on the fine-tuning task, whereas using Y/r distorts the compiled capabilities resulting in poor performance on the fine-tuning task.
Tracr Program
Block-0
Ovs
Block-0
Block-2
Block-0
Block-2
Attending Token
Kenwmeustor UStoor
Kehwing
Kchwagustale
Icchmmustded
Kehnmjustdat
Kcnnmoustod
Token Attended (b)
(a)
(C)
Figure 51: Max Identifier Task: Validation of Tracr observations on max identifier task with the spurious correlation defined as the difference between the indices of fifth and seventh maximum elements being three. The rows represents different input samples. Observation: Using nys preserves the original Tracr capabilities and therefore performs well the fine-tuning task, whereas using T/ distorts the compiled capabilities resulting in poor performance on the fine-tuning task.
43

---

Published as a conference paper at ICLR 2024
Block-0 Tracr Program
Block-2
Block-0
Tracr Program Block-2
kehwmguatd
Kehwmguatda
K Keh Kehwmguatoda
CT sewmguaatas
CTp=0
wm um a
Kchwmguatdd
Attending Token
k wmg Fett
f]s
hwmm
kewwng wamgguane guateee
kehw uat.eer
kehmggua Cwwmguttde
k Sewwggua
QV5
k whwhmmutttte
ke h.wwin m
k Lehwmguatdd
(a)
Taken Attended
(b)
Figure 52: Max Identifier Task: Visualization of the attention maps of the zeroth and second blocks of the Tracr fine-tuned models on the max identifier task. (a) shows the analysis when the spurious correlation correlation is present in the fine-tuning dataset. The first row shows the maps for the Tracr compiled model and other shows the analysis for different learning rates. Observation: Using TLL 7) or for fine-tuning distorts the capability of the programmed Tracr model in the Block-0 and as result the Block-2 attention map is not able attend to the desired output token. Whereas using mys is able to preserve the capability and as result the fine-tuned model is able to attend to the correct token the attention map in Block-2.
present in the fine-tuning datatset, whereas in case of (b) the spurious
44

---

Published as a conference paper at ICLR 2024
0.0
0.1
0.2
0.4
0.5
Tracr
n
M
05
Avs
Input Token
50
35
70
35
60
85
70
70
Input Token
Output Channel Number
Figure 53: Max Identifier Task: Visualization the activated output of the first MLP layer in first and second blocks for the max identifier task. The visualization is shown only for channel numbers 50-70. Observation: Using v for fine-tuning, which enables the model to learn the fine-tuning task, preserves the Traer compiled model's s compiled capability of sorting tokens in Block-1. Whereas other learning rates are not able to preserve this capability.
Block-1
Tracr Program
Block-2
Attending Token
1]5
flvs
12
(8)
Output Channel Number
(b)
Figure 54: Max Identifier Task: Visualization of the activated output of the first MLP layer in first and second blocks. This is the complete visualization of the activation map presented in Fig. 53.
45

---

Published as a conference paper at ICLR 2024
n
MM
ns
CTe 0
CTe 0
100
Acc. OFT
65
30
Acc. OPT
65
30
Base W1
W2
W3
W4
Base W1 W2
W3 W4
(a) CTr
0
(b) CTr 1
Figure 55: Counter Task: Pruning evaluation on Tracr model fine-tuned to count b's. Observa- tion: Higher value of CTr leads the learning the wrapper on top of the Tracr compiled capability. This wrapper is learned on using TM and and can be localized in few weights of the model.
nL
M
ns
CTe = 0
Cre =
100
Acc. OFT
65
30
100
Acc. OPT
65
30
Base
W1
W2
W4
Base W1
W2
W3
W4
(a) CTr
0.5
(b) CTr 0.8
Figure 56: Counter Task: Pruning evaluation on Tracr model fine-tuned to count b's. Observa-
tion: Observations consistent with Fig-7.
46

---

Published as a conference paper at ICLR 2024
Irle-1
Irle-2
-T-
Irle-3
CT==
Gt = 0.5
G7e = 0.8
Gr==1
Acc. Opr
75
100
Acc. OFT
080000 6
www.
5
10
15
20
0
10 15
20
5
10
15
20
0
10
15
20
Number of Pruned Weights (b)
(d)
Figure 57: Counter Task, weight pruning: Pruning weights of the Tracr model fine-tuned to count b's using different learning rates. Observation; In the presence of spurious correlation the model learns a wrapper when learning rates JIM and Tis are used. This wrapper can be localized in few weights of the model.
Irle-1
-.- Irle-2
--- Irle-3
Cr.
Ct. 0.5
Cte = 0.8
Cr m1
Acc. Opt
100
Acc. OFT
75
10
15
20
6
15 20
1.0
15
20
10
15
20
Number of Pruned Neurons (b)
(d)
Figure 58: Counter Task, neuron pruning: Pruning neurons of the Traer model fine-tuned to count b's using different learning rates. Observation: In the presence of spurious correlation the model learns a wrapper in case of γ/м and T)s- This wrapper can be localized a few weights of the model.
n2.
No-Prune
flg
B1WV97-0 Pruned
No-Prune
B1WKK14B-1 Pruned
543
ya.
20h
G
103
Pruned Weight
[b]
Predicted Class Label
(c)
Figure 59: Counter Task: Effect of pruning single weight on the distribution of predicted class labels. Observation: Even after pruning, the model predicts different classes indicating that the gain accuracy on pruning is indeed because the model has removed the wrapper.
47

---

Published as a conference paper at ICLR 2024
H ADDITIONAL PCFG RESULTS
In this section, we provide a detailed analysis of the PCFG results on the counter task. More specifically, we analyze the effect of the presence of weakly and strongly relevant capability in the model across three different parameters: training iterations (Riters), fraction of fine-tuning data samples with spurious correlation (Crz) and the probability of sampling operand token (0) be during pre- pre-traiinn training (Pt). Pr essentially determines whether the capability present in the are-traine model strongly relevant or weakly relevant for the fine-tuning task. Additionally we also analyze the effect of using the spuriously correlated (Ora = 1) and uncorrelated test set (CTe 0) for evaluation of the fine-tuned model. We present the summary of the results Tables 5, for the counting task and Tables 7,6 for the index of occurrence tasks. Then discuss the effect of learning rate on fine-tuning pre-trained models with weakly and strongly relevant capabilities in Fig. 60, 61. consistent for the considered counting and index of occurrence
We observe that the observations tasks. Next we analyze the effect of the presence weakly and strongly relevant capability in the pre-trained model for different fractions of spuriously correlated data-points (Cr+) and different pre-training iterations (Niters) ) in Fig. 62, 64, 66 for the counting element task and Fig. 63, 63, 67 for the index of occurrence task. We observe that the observations tasks and different values of Tltters Next we present the effect of the spuriously correlated data and presence of weakly and strongly correlated capabilities on the learning of the wrapper in Fig. 68, 69 on using uncorrelated test set for evaluation on counting and index of occurrence tasks respectively. Similar analysis on using test set with spuriously correlated samples is present Fig. 70 and 71. We present the capability revival analysis the Counter task for flters 200K and fiters = 50K pre-trained models for weakly and strongly relevant capability fine-tuned models Fig. 76 and Fig. 77 respectively. A similar analysis for different number of pre-training iterations is present in Fig. 78.
fairly consistent across both the
Table 5: Results the PCFG counting task with 200K iterations of pre-training.
Acc Opz
Acc Op:
17
(a)
Cte
Ave
IK
Cta
10K
U
IK
10K
Acc PT
IK
C+.
T0K
U)
1K
Che
10A
0
9.7
9.5
14,5
0
99.4
100
74.9
87.5
0.5
7.2
9.7
1.3
0
75.9
99.9
95.9
100
0.999
100
27.1
5.6
10.8
0.2
60
99.8
100
17.2
15.2
0.2
1.6
98.9
100
99.9
92.6
13.6
67.2
99.8
100
100
92.8
0.5
100
90.7
15.6
73.5
99.8
99.4
99.9
100
0.9
0.8
99.9
43.4
11.4
33.2
99.4
99.2
100
100
99.8
15.9
16.5
2.4
99.6
9.4
6.1
100
98.9
14.7
76.6
87.6
99.2
0.5
95.1
23.6
71.2
33.4
99.7
100
99.8
99.9
0.5
99.9
99.9
0.8
92.8
12.2
79.8
2.2
99.9
99.9
98.7
99.9
49.5
19
60.6
0
25.8
16
99.8
100
0
48.9
10.2
51.9
4.6
39.8
99.8
13.9
79.7
19.7
11.6
12.3
1.2
18.4
81.4
99.7
0.999
100
27.1
12.1
6.6
7.7
0.2
6.1
85.7
98
99.7
0.4
17.5
0
0
O
99.9
100
0
100
85.3
94.8
56.9
99.8
99.9
83.3
87.3
10
5
0.9
0.5
100
99.9
67.2
94.9
55.4
99.9
99.9
99.9
99.3
99.8
100
34.6
94.8
21.7
99.8
99.4
99.7
100
98.5
13.5
88.6
0.8
58.3
3.6
99.8
0
100
97.5
97.55
65.7
100
100
95.6
95.4
0.5
99.9
94,1
98.1
69
100
100
99.3
100
0.5
100
99.9
0.8
99.9
87.4
93.8
67.7
99.8
100
100
100
99.6
41.2
91.8
53.3
90.1
19.6
99.8
100
100
29
96.6
25.7
28.5
51.8
15.1
29.2
0.5
98.7
21.8
88.9
10.6
23.3
23.7
20.3
87.5
0.999
100
27.1
0.8
83
15..1
69.7
6.8
18.4
8.9
26.5
99.7
71.7
2.3
56
0
15.7
0
29.5
99.9
100
100
95.4
9T.8
99.8
99.3
84.1
84.6
10-6
0.9
100
100
99.9
96
94.4
99.9
99.6
99.5
95.9
99.6
0.8
99.8
99.4
95.9
92.6
99.6
99.3
94.8
99.6
99.8
51.6
95.1
63.9
99.5
30.5
94.2
99.7
0
100
99.9
97.7
98.1
99.8
99.8
95.4
95.6
0.5
100
99.9
97.9
93.7
100
99.9
98.5
99.6
0.5
99.9
99.9
100
100
98.8
93.1
100
99.9
98.3
99.9
100
97.6
98.6
85.1
99.8
73.9
98.3
48

---

Published as a conference paper at ICLR 2024
Table 6: Results on the PCFG counting task with 50K iterations pre-training.
Acc OpT
Acc. OF7
Cr
Acc. PT
Cte =
0
Cz
Acc. PT
Cre
Cte
IK
10K
IK
10K
1K
10K
9.1
2
86.3
93.6
0.999
99.9
11.7
8.9
2.1
0.1
5.17
90.2
99.9
97.9
99.8
0.8
5.5
11
0
0
64.9
100
98.9
99.9
20.2
15.9
0
0
1.9
99.9
100
9.1
I0.3
0.7
99.6
94,4
10.6
1.5
93.2
100
10
0.9
99.9
15.8
0.8
4.2
9.2
0.1
63.1
100
97.8
100
18.4
16.4
0
0.5
100
99.9
87.7
10.1
62.6
99.9
89.3
93.5
90.1
67.4
100
0.5
99.8
99.7
0.8
59.5
10.1
29.3
0.1
99.2
99.9
100
99.9
18.7
15.2
5.1
17.4
14.2
100
100
3
9.3
16.6
0.6
32.6
99.8
12.4
88.7
30.9
11.4
4.1
12.7
99.1
93
99.7
0.999
99.9
5.17
6.3
10.8
99.8
2.1
20.7
0.2
99.8
99.9
28
10.9
34.7
0.1
39.6
99.3
23.4
88.6
10-5
0.9
99.9
38.2
8.9
4.2
15.8
22.7
99.6
92.8
100
13.1
2.9
1.9
19.8
0.3
0.1
99.7
100
99.6
73.7
88.1
46
99.9
99.9
86.4
89.2
99.6
79.3
82.7
57.1
99.8
99.9
99.1
99.9
0.5
99.8
99.7
80
33.6
81.7
12.9
68.6
0.5
46.4
16.1
98.8
100
94.1
18.6
81.9
18.8
49.1
0.4
24.4
0.5
38.8
37.2
20
5.6
23.8
50.7
93.7
0.999
99.9
5.17
8.8
4,2
0.2
74.8
99.9
2.3
99.8
99.7
21
94.3
20.6
24.4
56.8
14.6
28.4
0.5
46.9
38.6
19.7
3.6
11.8
28.6
39.4
92
10-6
0.9
99.9
15.8
30.4
10.4
8.7
1.6
3.6
12.3
62.8
98.9
6.7
99.8
99.6
94.4
82.7
99.9
99.8
81
84.6
0.5
100
99.5
91.4
80.1
99.7
99.9
95.9
99.6
0.5
99.8
0.8
99.7
97
90.1
74.6
99.7
99.3
96.2
99.3
99.8
55.7
91.2
96.3
99.4
49

---

Published as a conference paper at ICLR 2024
Table 7: Results on the PCFG index of occurrence task with <00K< iterations of pre-training.
Acc.
Acc
Cr:
Acc. PT
Cte O
Cze
Acc. PT
Cre
Cte
IK
10K
1K
10K
1K
10K
0.0
23,9
46.1
99.5
0.0
23.0
0.0
66.6
99.6
81.0
100.0
0.999
99.0
9.3
0.8
14.9
0.0
0.8
0.0
36.8
99.1
99.5
100.0
33.6
9.5
0.0
0.0
3.7
5.8
100.0
100.0
96.2
0.0
88.6
98.6
99.7
96.8
0.0
83.5
0.0
97.3
100.0
100.0
10
0.9
99.2
97.1
0.8
96.8
0.0
84.8
0.0
97.4
99.2
100.0
100.0
72.6
1.3
79.1
0.0
37.3
24.0
100.0
100.0
95.7
4.5
96.5
3.6
99.5
99.7
99.8
99.8
5.6
78.2
3,5
100.0
0.5
98.0
98.9
0.8
96.0
17.7
90.2
5.6
98.8
99.4 4
100.0
100.0
91.3
15.8
79.8
14.5
48.0
28.8
100,0
100.0
94.6
2.6
86.2
18.6
17.0
85.4 4
16.0
68.5
98.4
3.4
97.6
7.7
27.2
97.8
0.999
99.0
9.3
94.1
5.1
94.0
0.2
10.7
37.6
99.7
77.6
27.7
70.8
0.0
4.9
3.7
56.6
100.0
99.2
88.5
99.0
79.3
97.3
99.1
98.9
99.1
10-5
0.9
99.2
99.2
91.9
99.2
77.9
97.1
98.3
100.0
100.0
95.6
84.3
98.0
65.4
98.0
72.7
81.3
31.8
99.9
100.0
97.3
95.8
100.0
89.1
99.1
99.9
99.9
99.4
97.6
95.9
98.8
75.4
99.2
99.3
100.0
100.0
0.5
98.0
98.9
95.1
97.4
87.7
98.0
81.0
97.2
51.8
100.0
100.0
99.0
80.4
99.4
73.0
14.6
23.3
19.5
15.5
0.5
98.9
94.2
99.5
94.3
13.7
20.2
22.0
40.5
0.999
99.0
9.3
99.2
67.8
98.7
61.2
13.2
9.2
25.1
71.8
99.4
61.6
98.8
11.3
95.6
98.8
99.6
99.5
99.6
97.33
97.0
98.0
98.4
10-6
0.9
0.5
99.2
99.0
98.8
99.8
99.2
97.1
96.6
96.6
98.8
99.9
98.7
99.2
99.7
99.0
97.1
96.7
99.3
100.0
97.8
97.9
99.8
99.6
99.2
98.9
99.2
99.7
0.5
98.7
98.0
99.8
98.2
98.9
99.0
99.8
100.0
0.5
98.0
0.8
4
97.9
99.6
98.2
99.5
98.9
100.0
100.0
98.1
96.7
99.9
96.4
94.5
99.9
100.0
50

---

Published as a conference paper at ICLR 2024
Table 8: Results on the PCFG index of occurrence task with 50K iterations of pre-training.
Acc OpT
Acc
Cr
Acc. PT
Cte = 0
C7e
Acc. PT
Cr
Cte
IK
10K
IK
I0K
1K
10K
0.0
11.7
66.1
3.2
0.0
2.0
0.0
60.0
98.5
92.5
100.0
0.999
94.2
3.2
0.8
6.7
0.1
1.0
0.0
26.7
96.8
98.7
100.0
23.1
10.9
0.0
0.0
3.8
5.1
99.6
100.0
43.2
0.0
34.6
75.1
97.5
0.0
56.2
81.0
97.5
99.9
10
0.9
94.2
69.9
0.8
53.5
0.0
53.1
0.0
74.0
97.0
98.8
100.0
12.8
11.0
32.8
0.0
4.4
3.4
99.8
100.0
72.1
2.3
59.6
98.6
90.2
99.6
65.4
2.3
70.4
91.4
100.0
0.5
88.6
91.5
0.8
56.6
1.5
65.5
0.0
88.7
97.4
99.7
100.0
39.0
12.8
40.2
0.0
6.1
5.3
99.9
100.0
5.0
0.1
6.4
7.7
18.4
88.9
10.3
86.5
74.1
1.4
42.3
98.8
0.999
94.2
3.2
37.4
3.8
13.2
86.9
98.7
48.0
19.1
2.7
0.0
0.5
4.4
95.2
100.0
89.2
18.9
69.6
26.8
74.0
91.1
66.8
82.8
10-5
0.9
94.2
91.9
25.0
75.9
38.9
69.9
86.9
82.0
99.2
30.4
75.9
87.8
99.8
69.7
8.8
73.9
8.4
13.8
4.4
92.6
100.0
87.6
64.0
77.7
31.4
93.1
95.6
85.5
90.9
85.9
$7.5
80.0
57.0
90.4
95.5
93.9
100.0
0.5
88.6
91.5
78.2
80.1
33.5
77.6
13.3
38.6
5.5
98.1
100.0
87.5
3.7
73.8
14.2
8.3
23,0
15.8
20.7
0.5
95.3
49.3
83.3
19.2
5.4
13.7
27.5
76.9
0.999
94.2
3.2
94.0
31.0
71.4
7.7
3.9
38.7
93.7
88.1
$0.6
97.5
93.6
88.2
73.2
70.7
69.6
75.6
56.9
66.9
0.5
94.3
92.2
73.1
76.4
70.3
69.7
67.4
88.6
10-6
0.9
94.2
69.9
92.7
89.8
75.7
80.3
65.3
63.6
71.8
94.7
56.2
72.8
72.3
97.9
87.3
87.0
80.0
76.6
90.3
93.8
85.1
88.3
0.5
87.7
84.5
80.4
80.8
93.4
91.5
90.3
97.6
0.5
88.6
91.5
0.8
90.6
83.2
83.1
80.5
92.3
87.9
90.1
98.6
77.3
79.8
70.2
91.9
99.7
51

---

Published as a conference paper at ICLR 2024
ns
=r,= 0
CT. 0.5
Py(a) 0.999
Ctr 0.8
Cr 1.0
100
50
Prla) = 0.99
25
Pr(a) 0.9
100
Ost
Pr(a)
100
50
75
2.5K 7.5K
10K
2.5k
7.5
10k
2.3K
7.5K
10k
2.5K
3K
7.5K
10k
(a)
(b)
(c)
(d)
Figure 60: Counter Task, Riters = 200K, CT = Effect learning rate (LR) on fine-tuning pre-trained models with weakly and strongly relevant capabilities and using different values of Cte for fine-tuning. Observation: In the presence of strongly relevant capability, training with ys yields good performance on the fine-tuning dataset. The convergence time to learn the fine-tuning task increases with an increase in Ctr.
fs
Pr(a) 0.999
100
50
PHl) = 0.99
100
Py(a) 0.9
OFT
Py(a) = 0.5
100
50
(a)
(b)
(c)
(d)
Figure 61: Index of Occurrence Task, Niters 200K, Cre 0. The settings are consistent with
Fig. 60
52

---

Published as a conference paper at ICLR 2024
=
0.25
0.5
0.8
Cr. =1
fl
Cra m1
Pyla)
G -
fis
=1
100
50
Pria) = 0.99
75
25
Op:
75
50
P+la) 0.5
100
50
Figure 62: Counter Task, Tiiters 200/: Effect of the presence of strongly or weakly relevant pretuning capability fine-tuning performance using 1JM and TJs. Test sets with and without the spurious correlations are used for evaluation. Observation: The convergence time for learning the fine-tuning task in the absence of strongly relevant capability is higher compared to when the strongly relevant is present in the model. The time further increases if spurious correlations are present the fine-tuning set. However, in the presence of spurious correlations, the convergence time to learn the spurious correlation is small and is possible even on using the learning rate vss. Using ns is unable yield learning the fine-tuning task if a weakly relevant capability is present in the model.
=
Cr,
Cr.
Cr.11
Cr
QM
Cr
=M
Gre
Gre
100
73
25
100
50
=
OFT
25
Pr(a)
100
50
(b)
(c)
(d)
Figure 63: Index of Occurrence Task, Riters = 200K: The settings are consistent with Fig. 62.
53

---

Published as a conference paper at ICLR 2024
Cty
Cty 0.25
Ctr 0.5
Cir -
Ctr =1
Cte
fl
Pyla] 0.999
fis
c,-1
75
25
Pr(a) 0.99
100
75
50
0.9
Orr
76
50
Pvta)= 0.5
75
25
2.5K
7.3K
10K
2.3K
5K
7.5K
100
2.55 5X
7.5K
10K
2.5k 5K
7.5k
10k
(a)
(b)
(c)
(d)
Figure 64: Counter Task, Riters 50K: Effect of the presence of strongly or weakly relevant pretuning capability on fine-tuning performance on using T/M and 7Js, Test sets with and without spurious correlations are used for evaluation. The observations are consistent with Fig. 62.
=
CT. 0.5
Cr. 0.8
=r=11
OM
Cr.=1
Pyla) 0.999
ns
100
50
Pyla]= 0.99
75
25
Py(a)= 0.9
100
OCT
Pr(a) - 0.5
100
50
2.5K
7.3K
2.5K
7.5K
2.3K
7.5K
2.3K
7.5K
(a)
(b)
(C)
(d)
Figure 65: Index of Occurrence Task, natere
50K' The settings consistent with Fig. 64.
54

---

Published as a conference paper at ICLR 2024
- 0
0.25
Ct. 0.5
Crr
Crr
-
-111
Py(a) - 0.999
C
-1
75
22
Pyla] 0.99
Pr(a)- 0.9
1064
Py(a) - 0.5
255
36
7.50
105
235
7.5
105
23R
36
7.38
(a)
(b)
(c)
(d)
Figure 66: Counter Task, Pliters 10K: Effect of the presence of strongly weakly relevant pretuning capability on fine-tuning performance on using T/ and 7Js. Test sets with and without spurious correlations are used for evaluation. The observations are consistent with Fig. 62.
=
CT 0.5
Cr. 0.8
=
OM
Cr-
Pria) 0.999
G5e
ns
Cye
100
50
P 0.99
75
25
Py(a)= 0.9
100
OFT
Pr(a) 0.5
100
50
2.5k
233
7.RK
(2)
(b)
(c)
(d)
Figure 67: Index of Occurrence Task, Niters 10K: The settings are consistent with Fig. 66.
55

---

Published as a conference paper at ICLR 2024
Prla) = 0.999
Prla)=
0.99 0.99
Prid) 0.9
Pria) = 0.5
Ctr =
Cr. 0.5
Aites
= 0.8
Ct 00=11
25
nhws = 50k
100
OpT
50
25
Ohess 10k
100
73
50
235 5R
7.3K
10K
2:3K
5K
7.5k
10k
2:3K
3K
7.5K
10%
2.5K
3K
7.5K
10k
(a)
(b)
(c)
(d)
Figure 68: Counter Task, TMM Cro 0: Learning of the wrapper in presence of different fraction of spuriously correlated data, values of Pr (a) during pre-training, and training iterations. Observation: Using higher fraction of spuriously correlated data in the fine-tuning set (higher value of CTv) leads to faster degradation in the pre-training accuracy. Further this degradation is even faster in presence of weakly relevant capability.
Prla) 0.999
=
Pría) 0.9
Pr(a)= - 0.5
=0
Cr. = 0.5
Aoery 200k
Cr. 0.8
=10
100
= 50k
Opt
nhes 10k
100
75
50
25
(a)
(b)
(c)
(d)
Figure 69: Index of Occurrence Task, Medium LR, Cre 0:
The settings are consistent with
Fig. 68.
56

---

Published as a conference paper at ICLR 2024
Prla) = 0.999
Prla)=
0.99 0.99
Prid) 0.9
Pria) = 0.5
Ctr
Cr= 0.5
Ct. 0.8
Cr 1.0
254
Uhws = 50k
100
Acc. Opt
75
50
25
Ohws 10k
100
73
50
235 5R 7.3K
10K
2:3k
5K
7.5K
10K
2:3K 3K
7.5K
10k
2.55 3K
7.5k
10k
(a)
(b)
(c)
(d)
Figure 70: Counter Task, TMM Cro 1: Learning of the wrapper in presence of different fraction of spuriously correlated data, values of Pr (a) during pre-training, and training iterations. Observation: Using higher fraction of spuriously correlated data in the fine-tuning set (higher value of CTv) leads to faster degradation in the pre-training accuracy. Further this degradation is even faster in presence of weakly relevant capability.
Prla) 0.999
-
Pría) 0.9
Pr(a) - 0.5
=
Cy. = 0.5
Ooers 200k
Cr. 0.8
=10
100
50k
OPt
nhes 10k
100
75
50
25
(a)
(b)
(c)
(d)
Figure 71: Index of Occurrence Task, Medium LR, Cro 1: The settings are consistent with
Fig. 70.
57

---

Published as a conference paper at ICLR 2024
Pyla Pyla)
Pyla)
Pyla)
200k
1.0
100
75
50
25
neerr = 50k
100
Acc. Opt
75
10k
100
2.5
2.5K
7.5K
13K
233 7.5K
10k
2.5K
7.5K
10k
(a)
(b)
(c)
(d)
Figure 72: Counter Task, T[8, Cra = 0 Learning of the wrapper in presence of different fraction of spuriously correlated data, values of Pr (a) during pre-training, and training iterations. Observation: The observations are consistent with Fig. 68
Pyla) 0.999
Py(a) 0.99
Py(a) 0.9
Py(a) 0.5
=T.
CT. 0.5
nowrs 200k
GTr
=r==1. 1.0
50k
OpT
25
fleers = 10k
100
75
50
25
(a)
(b)
(c)
(d)
Figure 73: Index of Occurrence Task, 11ss CT 0 The settings are consistent with Fig. 72
58

---

Published as a conference paper at ICLR 2024
Prla Pyla)
Pyla)
Pyla) 0.5
fawer 200k
1.0
100
75
50
23
Oower 50k
100
Acc. Opt
75
10k
100
2.5K 7.3K
10K
2.KK
7.5k
13K
233K 7.5K
10k
2.3K
3K
7.5K
10k
(a)
(b)
(c)
(d)
Figure 74: Counter Task, T[s, Cra = 1 Learning of the wrapper in presence of different fraction of spuriously correlated data, values of Pr (a) during pre-training, and training iterations. Observation: The settings are consistent with Fig. 70
Pyla) 0.999
Py(a) 0.99
Py(a) 0.9
Py(a) 0.5
CT.
CT. 0.5
Aicer 200k
Gtr 0.8
GT.
50k
OPt
25
Plaes = 10k
100
75
50
25
(a)
(b)
(c)
(d)
Figure 75: Index of Occurrence Task, 115, CTo 1: The settings are consistent with Fig. 74
59

---

Published as a conference paper at ICLR 2024
Ctr 0.0
Cir 0.25
Ct. 0.5
Ctr 0.9
Ct. 1.0
Py(a) = 0.999
Cre
CTe
Cre
C
50
25
Pr(a) = 0.99
100
Acc. OPr
75
50
25
Pr(a) = 0.9
100
75
50
Pr(a) 0.5
100
25
10a
10a
16
10
10
(a) ПM
(b) ns
Figure 76: Counter Task, Niters 200K: Reverse Fine-tuning on weakly and strongly relevant capability fine-tuned models. Medium and small learning rates are used for reverse fine-tuning in the presence of different degrees of spuriously correlated data-points present in the train-set. The fine-tuned model was fine-tuned using Large LR. Observation: When the model possesses weakly relevant capability, the convergence time is lower for models fine-tuned on dataset with spurious correlations. the model possesses strongly relevant capability, this difference is less. The "revival" of pre- training capability is observed for all values of Cre- Even though fine-tuning large learning rate of 10 4 capability revival is possible even on using a small learning rate.
done using
60

---

Published as a conference paper at ICLR 2024
CTr 0.0
Cir 0.25
Ctr 0.5
Cir 0.9
Ci. 1.0
Pr(a) = 0.999
CT+
CT
Cr.
C
50
25
Pr(a= =
100
Acc. Opr
75
50
25
Pr(a) 0.9
100
75
50
Pr(a) 0.5
100
25
10a
'10
10a
10
10
(a) пM
(b) n/s
Figure 77: Counter Task, Nster = 50K: Reverse Fine-tuning on weakly and strongly relevant capability fine-ture tnn-t models. aл and ns are used for capability reverse fine-tuning the presence of different fraction of spuriously correlated data- data-points present in the train-set. The fine-tuned model was fine-tuned using Large LR. Observation: The observations are consistent with Fig. 76
74
pt
CTr 0
CTr
niterz
200k
CTe
CTe
Crc
Ctc
100
75
50
25
niters 50k
100
Acc. OPT
75
50
25
niters 10k
100
75
50
25
10
100
10a
107
10
110
107
10a
10a
167
10'
100
(a) Nie
Iters (Log Scale)
(b) ns
Figure 78: Counter Task: Reverse fine-tuning analysis for different pre-training iterations. Observation: Capability Revival is seen for models pre-trained with different number of iterations.
61

---

Published as a conference paper at ICLR 2024
H.1 PRUNING ANALYSIS
In this section, we present detailed results on pruning analysis of the PCFG setup on both counting and index of occurrence tasks. We provide an exhaustive evaluation in Fig. 79, 82, 84, 86 and 88 for the Counter task and Fig. 80, 83, 85, 87 and 89 for the index of occurrence task.
CTr
CT 0.5
CTr 0.8
CTr
1.0 1.0
Diters
200k
CTo
CTo 1
CTo 0
CTe
GT 0
CTo 1
100
75
50-
Aiters 50k
100
Acc. OpT
50
niters 10k
75
25
0 25
0
100
0
100
100
(a) nL
(b) ΠM
(c) ns
Figure 79; Counter task, Pt (a) = 0.999, Pruning Analysis: Revival of pre-training capability
analysis for different learning rates. pre-training iterations and different values of CTy. Observation: (a) On using 7t the model learns the wrapper only when fine-tuning on CTe = 1. This wrapper is learned only on fine-tuning set with the spurious correlations.
(b) On using 7ar the model learns the wrapper on fine-tuning with smaller values of CTa as well. However, still this wrapper is learned only on fine-tuning set with the spurious correlations. (c) On using 7)s the model learns the wrapper for all values of Cr and for all the data samples.
62

---

Published as a conference paper at ICLR 2024
CTr 0
Ctr 0.5
Ctr 0.8
Ctr 1.0
Dicors
200k
Cre 0
Cre 1
====0
CTe 1
CT 0
CTe =1
100
50
niters 50k
Acc. OşT
75
25
Oiters 10k
100
50
0
25
50
100
0 25
50
100
25
100
0
25
50
100
0
25
50
100
0
25 50
100
(a)
n
(b) QM
(c) ns
Figure 80: Index of Occurrence task, Pt (a)
0.999, Pruning Analysis: The settings
consistent with Fig. 79
CTr 0
Ctr = 0.5
CTr 0.8
CTr 1.0
Cte 0
Cte 1
CTe
Cre
Acc. OFT Acc. OPT
100
30
100
30
(a) NM
(b)
Figure Index of Occurrence task, Pr (a) 0.999, Pruning Analysis: The settings are consistent with Fig.
63

---

Published as a conference paper at ICLR 2024
CT.
Ct. 0.5
Ct. 0.8
Ct. 1.0
Pr(a) 0.999
Cte
C
100
Pr()= 0.99
100
Acc. OFT
50
Pr(a) 0.9
100
Pr(a)
100
10 15
20
10
15
10
20
(a) nl
AM
(c) fls
Figure 82: Counter task, Paters 200A, Pruning Analysis: Revival of pre-training capability analysis for different learning rates, weakly and strongly relevant capability fine-tuned models, and different values of CT. Observation: Learning of the wrapper is possible for weakly as well as strongly relevant capability pre- pre-trained models.
=
Cr. 0.5
0.8
Cr. 1.0
Pr(a) 0.999
Cte
Cre =1
Cr. =
Cre =1
CTe
Gr=1
100
50
Pr(a)= 0.99
100
OPT
Pr(a) =
Pria) 0.5
100
15
13
20
10 15
20
10
20
(a)
(b)
(c)
Figure 83: Index of Occurrence task, Pitters
200K, Pruning Analysis: The settings are consistent
with Fig. 82
64

---

Published as a conference paper at ICLR 2024
=
=
=
Ct = 1.0
Pr(a) = 0.999
Cтw
Cte m
Cte
Cre 0
G-11
100
90
Py(a) = 0.99
100
OcT
Prla)= = 0.9
100
75
Pria) = 0.5
100
50
50
100
(a)
(b) RM
(c) n0
Figure 84: Counter task, Niters 200K Pruning Analysis: Revival of pre-training capability analysis for different learning rates, weakly and strongly relevant capability fine-tuned models and different values of Cr. Here larger number of neurons pruned compared Fig. 82.
Ct.
Cr. =
0.5
Ct. 0.8
Ct. 1.0
Pr(a) 0.999
Crc
Cre
Cre
Cro
G==1
100
50
Py(a) 0.99
Opt
75
25
Pria) 0.9
100
50
Pata) 0.5
75
25
25
50
75
100
25
50
75
100
25
50
75
160
25
50
75
160
25
50
75
100
0
25
50
75
100
(a)
AM
n7s
Figure 85: Index of Occurrence task, Miters 200K, Pruning Analysis: Here larger number of neurons pruned. The settings are consistent with Fig. 84
65

---

Published as a conference paper at ICLR 2024
Py(a) = 0.999
- Pr(a) 0.99
----- Py(a) =
Pr(a)= =
Aiters 200k
CT. =
CTr 0.5
CT. 0.8
Ct. 1.0
100
50
25
Oiters 50k
100
Acc. OPT
75
50
25
Oiters 10k
100
75
25
0 25 75
100
25
50
75
100
25
50
100
25
50
75
100
(a)
(b)
(c)
(d)
Figure 86: Counter task, M/. Cr. = 0, Pruning analysis: Effect strongly and weakly relevant capabilities for different number of pre-training iterations and different values of Ctr Observation: Fine-tuning a model with strongly relevant capability leads to learning of an "inhibitor" on its pre- training capability, i.c., wrapper that disallows of the pretraining capability. Revival of the pre-training capability is partly possible on pruning, if the model has strongly relevant capability and it was fine-tuned dataset without spurious correlations. The inhibitor is mainly learned for the 200 iteration pre-trained model.
Py(a)= = 0.999
17(())=099)
Pr(a)= 0.9
P+(a) = 0.5
Mitors 200k
CT. 0.0
CT 0.5
Ctr 0.8
CT. 1.0
75
50
25
Ditees 50k
100
Acc. OFT
50
25
Oiters 10k
100
50
25
.0
0
25
75
100
25
100
25
50
100
50
100
(a)
(b)
(c)
(d)
Figure 87: Index of Occurrence task, M,, CT 0, Pruning analysis: The settings are consistent
with Fig. 86
66

---

Published as a conference paper at ICLR 2024
Pr(a)= 0.999
Pr(a)= 0.99
=.9
Pr(a) = 0.5
f
Ctr 0.0
=
Ctr 0.8
CT. 1.0
75
50
25
Acc. OPT
75
50
25
75
100
75
50
25
100
100
(a)
(b)
(c)
(d)
Figure 88: Counter task, Riters 200K, CT O, Pruning analysis: Effect of the strongly and
weakly relevant capabilities for different number of pre-training iterations and different values of fraction of spurious correlations present in the fine-tuning dataset. Observation: Learning of the inhibitor is observed on using ?!,.
Py(a) = 0.999
Pr(a)= 0.99
Pr(a)=099
Pr(a) = 0.5
nP
CT. 0.0
CTr = 0.5
CTr 0.8
Cr. 1.0
100
75
50
100
Acc. OPT
75
25
100
75
25
0
25
50
75
100
0
25
50
75
100
25
50
75
100
25
50
75
100
(a)
(b)
(c)
(d)
Figure 89: Index of Occurrence task, Riters = 200K, Cra Pruning analysis: the settings are consistent with Fig. 88
67

---

Published as a conference paper at ICLR 2024
Pr(a) = 0.999
Pr(a) 0.99
Pr(a) 0.9
Pr(a) 0.5
CT. = 0
Ctr
CTr
CT. 1
100
100
75
75
25
25
пM
100
100
Acc. Om
75
Acc. Oft
75
50
1/s
75-
75
50-
50
80
81
82
B3
B4
B5
B6
80 81
B2
B3
84
85
86
80
81
82
B3
84
B5
B6
80
81
82
83
84
B5
86
(a)
(b)
Probed Block
Figure 90: Counter task, Raters 200K Cr. = O, Probing analysis: The effect of different values of learning rate, weakly and strongly relevant capabilities is shown. Observation: Using VIr. hampers the pre-training capability to count especially when the probed model has weakly relevant capability. The performance on the pre-training task of counting a's continues to be high, especially with r7s- The accuracy of counting b's shows that fine-tuning capability is learned using nı- On using 7s, models with weakly relevant capabilities are not able learn the fine-tuning capability well.
Pr(a) = 0.999
Py(a) = 0.99
Pr(a) 0.9
Py(a) 0.5
CTr 0
CTr
Crr 0
CTr
100-
100
50
50
25
25
пи
100
100
Acc. OFT
75
Acc. OFT
75
25
25
f]s
100
100
75
75
25.
25
80
61
62
83
B4
B5 B6
80 B1 B2
83 84
B5
B6
80 81 82 B3 B4
B5 B6
80 81
82
83
B4
B5 86
(a)
(b)
Probed Block
Figure 91: Index of Occurrence task, Riters 200K CT 0, Probing analysis: The settings
consistent with Fig. 90
68

---

Published as a conference paper at ICLR 2024
Pr(a) = 0.999
Pr(a) = 0.99
Pr(a) 0.9
Pr(a) 0.5
CT 0
CT =1
CT. =
CT.
100-
75
50
25
nM
Acc. OFT
75
Acc. OFT
50
75-
50
80
B1
B2
B3
B5 B6
B1 B2 B3 84 B5 B6
80 B1 B2 63 B4
B5 B6
80 81
B2
63 B4 B5 B6
(a)
(b)
Probed Block
Figure 92: Counter task, Nsters 50K, Cre 0, Probing analysis: The observations consistent
with Fig. 90
=
Pr(a) = 0.99
Pr(a)= = 0.9
Py(a) = 0.5
= 0
CT
п
Cr:=0
CT. 1
100
50-
25.
nM
100-
Acc. OFt
75-
Acc.
25
77s
100
75
50
80
61
B2
B3 B4
B5 B6
80 B1 B2
83 B4
B5 B6
80 B1 B2 B3 B4 B5 B6
80 81
B2 B3
B4 B5 B6
(a)
(b)
Probed Block
Figure 93: Index of Occurrence task, Nlttess = 50K, CT = 0, Probing analysis: The settings are
consistent with Fig. 92
H.2 PROBING ANALYSIS
In this section, we present detailed results on probing analysis of the PCFG setup on both counting and index of occurrence tasks. We provide exhaustive evaluation in Fig. 92, 95, 97 for the Counter task and Fig. 93, 96, 98 for the index of occurrence task.
69

---

Published as a conference paper at ICLR 2024
74 CTr =1
CTr 0
CTr 0
Ctr = 1
Acc. OPt
100-
65-
84
80
82
85
80 82
80
83
B5
(a) GM
(b) 77
Figure 94: Index of Occurrence task, Piters = 200K Cr. 0, Probing analysis: The settings consistent with Fig. 8
Pr(a) = 0.999
0
Pr(a) 0.9
Pr(a) 0.5
Cr. 0
CTr 0.5
n
CTr = 0.8
Ctr =1
100
75
50
25
100-
Acc. OFT
75-
Acc. OFT
50-
75
50
80
B1 B2 B3 B4 B5 B6
80 61 B2 B3 84 B5 B6
80 B1 82 B3 B4
B5 B6
80
B2 B3 B4 B5 B6
(a)
(b)
Probed Block
(c)
(d)
Figure 95: Counter task, Nitrrs = 200K Cr. O, Probing analysis. Observation: With an increase in Cr., the accuracy counting a's also increases for both weakly as well strongly relevant capability models.
70

---

Published as a conference paper at ICLR 2024
Pr(a) 0.999
Pr(a) 0.99
Pr(a) 0.9
Pr(a) 0.5
CT
CTr 0.5
CT. 0.8
CTr
100-
75
50
пM
Acc. Opr
75
Acc. OFT
50
75-
50
25
83 B4
81 B2 B3
80
B6
B4 85
(a)
(b) Probed Block
(d)
Figure 96: Index of Occurrence task, Piters = 200K Cre = 0, Probing analysis: The settings consistent with Fig. 95
P
Pr(a) 0.99
Pr(a) 0.9
Pr(a) 0.5
Cr
CTr 0.5
Ctr 0.8
Ctr
75
50
NM
100
Acc. OF
Acc. OFT
50-
25.
100-
50
25
(a)
(b)
Probed Block
(c)
(d)
Figure 97: Counter task, Piters 50K Cre = 0, Probing analysis. Observation: With an
increase Cre, the accuracy on counting b's also decreases for both weakly as well as strongly relevant capability models.
71

---

Published as a conference paper at ICLR 2024
Pr(a) = 0.999
Pr(a) = 0.99
Pr(a) 0.9
Pr(a) 0.5
Cr =0
CT 0.5
CT 0.8
Ct. =1
100
75-
25.
100
Acc. OFT
75-
Acc. OFT
50
ns
75-
50
80
81 B2
B3 B4 B5 B6
80 81 B2 B3 84 B5 B6
80 B1 82 B3 B4 B5 B6
80 B1 B2 B3 B4 85 B6
(a)
(b)
Probed Block
(c)
(d)
Figure 98: Index of Occurrence task, Niters =5KK, Cra = 0, Probing analysis: The settings are consistent with Fig. 97
72


Reviews:
[{"review_id": "btm07BIpCC", "reviewer": "Reviewer_eQAm", "summary": "This work applies mechanistic interpretability methods on finetuned language models as an attempt to understand how the finetuning process alters pretrained models. The authors conclude that finetuning with smaller learning rate allows the finetuned model to learn a \"wrapper\" on top of the existing pretrained model to perform specific finetuned tasks.", "questions": "(see weaknesses)\n\n\nOverall this is a good solid and interesting paper. I would like to see it being accepted.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I appreciate how the experiment is setup: gradully stepping from a controllable but synthetic setting to a less controllable but realistic setting. The gradual stepping from controllable to uncontrollable settings and the meticulous linking of claim consistency between each experiment makes tracing the arguments and their evidence easy.  The conclusions drawn are well supported by the empirical results. \n\nI also appreciate how the notion of capability is defined in this work. An exact definition conveys what the authors consider as \"capabilities\" and their coarse dichotomy into strong and weak relevances.", "weaknesses": "The main (potential) weakness of this work is the conclusion drawn. I personally dislike the usage of \"lack of novelty\" as a justification to reject a paper. That being said, there is frankly not much to be learned about finetuning from this work. Yes, small finetuning rate would only change the model a little. Yes, the change would be potentially reverted if training on the original pretraining task. Yes, a larger finetuning rate would cause collateral damage, causing existing capabilities to be lost. These are all well-established knowledge about finetuning where plenty of papers have explored. Again, this is not to discredit this entire work. The authors should consider highlighting the newer findings that readers may not already know, perhaps how the distinction between strong and weak relevance would benefit the finetuning of a new capability.\n\nThe authors motivate the importance of understanding of the finetuning process by mentioning safety and jailbreaking attacks in the introduction and conclusion. Unfortunately no related experiments are shown.\n\nLast but not least, the authors should include a paragraph in related works discussing on how semi-controllable language models (e.g. transformers training on PCFG and Tracr) have been utilized in model interpretability previously, to help contexualize the usage of mechanistic interpretability techniques (e.g. weight pruning). The authors should also include the baseline of finetuning with a portion of pretrained data mix in, since that would be the most common finetuning technique. Conclusions drawn on such finetuned procedures would be more applicable, compared to purely finetuning on the new domain.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uKyW996Byr", "reviewer": "Reviewer_wEP2", "summary": "This work explores how fine-tuning impacts the capabilities of large pre-trained models. It uses mechanistic interpretability tools and the technique of reverse fine-tuning to study how fine-tuning affects these capabilities and finds that fine-tuning often adds a minimal transformation, called a \"wrapper,\" on top of the existing capabilities when the learning rate is sufficiently small. They demonstrate this in a synthetic setting with PCFG's initialized with Tracr solutions and in a more realistic setting with the TinyStories dataset.", "questions": "All questions are addressed in weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The authors address the critical problem of understanding how fine-tuning impacts models, and specifically try to offer mechanistic insight via a compelling section of techniques (probing, pruning, reverse fine-tuning).\n2. I like the TinyStories experiment, where the authors find that the pretraining capability can be recovered via fine-tuning uniquely when the learning rate is small enough, as evidenced by the control model. Importantly, I find it really interesting how in the high learning rate deletion setting, the capability loss for fine-tuning the control and deleted model match each other. I best understood this from Figure 79 on page 53 in the Appendix, which might benefit from being in the main paper. \n3. The authors are incredibly comprehensive over hyperparameters, giving many alternatives in the appendix.", "weaknesses": "My overall analysis is that the authors haven't provided a clear definition of what it means for a capability to have been lost or recoverable\n\n1. [Experiment 2] I believe the wrapper formalization is inconsistently used throughout the paper. Starting at Definition 2, a wrapper $g$ is composed with $\\mathcal{C}$ and is defined as a map acting on the output of a capability ($g$ is unquantified here, and I believe the range should be the domain here). This definition is independent of a model and how it is parameterized. However, in the experiments, the wrapper definition is always a subset of the weights to be pruned. I believe this is not captured as a function of the capability/model output, but rather the function parameterization. As such, probing can not be used as evidence that the model learns a wrapper over the capability. \n    - This is not simply fixed by altering the definition to be a change over weights; there is always the wrapper of adding the optimal FT weights and subtracting the current PT weights, which would produce a strongly relevant capability but (I think) would go against the spirit of a wrapper. I believe it is non-trivial to find a definition of wrapper that is 1) covered by pruning and 2) captures the spirit of being a lightweight modification of the model. \n    - Is there an example of a fine-tuning capability that can not be obtained by a wrapper under the current or a revised definition? Having this delineation is important to proving something that isn't true by construction.\n\n2. [Experiment 3] In the abstract and introduction, revival via reverse fine-tuning is discussed as being sample-efficient. Experiment 3 is intended to demonstrate this phenomenon experimentally. However, there is no discussion of sample efficiency here or later in the paper. Specifically, to show that the pretrained capability is actually forgotten, there would need to be a comparison to how much time it takes to learn the pretrained capability, which doesn't exist since the capabilities are directly compiled via Tracr. Even this experiment would be confounded by the fact that the fine-tuned model can be seen as simply a good initialization for learning the pretrained capability from scratch. This sample efficiency analysis is critical since in theory, any capability can be learnt by training for long enough.\n\n3. [Experiment 4] The linked Figure 8 for the probing analysis does not concern probing. Does this mean to refer to Figure 5? If so, what new information does this section provide that is not captured in Experiment 2, which refers to Figure 5? The analysis and takeaway in this section do not seem grounded in data in the current state.\n\n4. I found the results in this paper quite difficult to parse. None of these items individually changed my score, but it did make it much harder/time-consuming to parse the message of the results.\n    - In Figure 4, (i) through (iv) do not agree on the figure vs the text.\n    - Figure 5 does not have a legend and I do not know what it is measuring.\n    - Figure 6 does not have an x-axis and I do not know what it is measuring.\n    - The plots are out of order, for example, Fig 7 is referenced before Fig 5. This led to some confusion while reading experiments.\n    - Plots such as Figures 6, 7, and 8 test for three controlled variables at once and take a lot of effort to parse.\n    - Is there a reason the chosen $P(O_{PT})$ is changed from Figure 7b to Figure 8b?\n    - The reader is required to parse and remember a lot of notation consistently across the paper. It would be helpful to be more verbose about the notation definitions, especially in plot titles and captions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eJw9H4gS1D", "reviewer": "Reviewer_GZdd", "summary": "The paper analyzed the changes of a pretrained model during fine-tuning, especially how the capabilities are learned or lost during fine-tuning.", "questions": "Please see the first point in Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. I am really interested in this topic as while fine-tuning is important for unleashing LLM's capabilities learned during pretraining, there is little literature on understanding the dynamics during this process. This is also important for how to mitigate forgetting during LLM fine-tuning.\n\n2. The authors proposed some novel synthetic tasks to probe the evolving of capabilities  during LM fine-tuning.", "weaknesses": "1. The presentation of this paper is extremely unclear. I have tried my best to understand the first 5.5 pages with many unclear notations but cannot understand the remaining 2.5 pages. I list some (but not all) unclear points below:\n* What are $R_{FT, TR}$ and $R_{FT, TE}$ in Section 4? What do the three elements mean in them?\n* What does the characters $L, M, H$ mean in $P_C^L$?\n* In Section 4, it seems that $O_{pre}$ has contained $b$, why do you still fine-tune the model on counting $O_{ft}b$?\n* \"The probability of embedding a spurious correlation in the train/test fine-tuning dataset,\" This sentence is unclear and has grammatical error.\n* What does different lines mean in Figure 5? \n\nUnclear presentations such as listed above largely downgraded the quality of this paper and I cannot fully understand some of the paper's main claims. I hope the author can thoroughly revise their paper on its presentations.\n\n2. Some related works are missing such as [1][2].\n\n[1] Two-stage LLM Fine-tuning with Less Specialization and More Generalization\n\n[2] Data distributional properties drive emergent incontext learning in transformers\n\n\n------post rebuttal--------\n\nI sincerely thank the authors for clarifications and revisions. The paper looks much better now with the necessary clarifications on the notations.\n\nOverall, I think the paper provides some insightful analysis on the fine-tuning dynamics of transformers. However, the presentation and some of definitions can still be made clearer, for example, in Definition 1, $X_D$ is declared but never used or explained, is it the same thing as $X_C$?\n\nSome definitions and notations should also be made more forma such as the relationship between $P^{FT, E}$ and $P^{FT}$.\nI suggest the authors to further revise the paper if it would be accepted. \n\nThe reason that I thought the two papers would be related for this work is that they also studied the relationship between pretraining and finetuning, and to some extends validates that the spurious correlation is utilized during fine-tuning. In [1], the model easily \"forgets\" its capabilities obtained from pretraining as it wrapped these capabilities through spurious correlations. However, these capabilities are not actually lost and can be preserved if we remove this wrapper by removing the soft prompt. In [2], the authors also studies how the model lost its general capabilities during fine-tuning on a single dataset, which can also be related to spurious correlation.\n\nI have increased my score to 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)