Paper: Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks (ICLR-2024)

Reviewer comment (Writing):
The presentation of this paper is extremely unclear. The reviewer has tried their best to understand the first 5.5 pages with many unclear notations but could not understand the remaining 2.5 pages. Examples of unclear points include: What are \(R_{FT, TR}\) and \(R_{FT, TE}\) in Section 4, and what do the three elements mean in them? What do the characters \(L, M, H\) mean in \(P_C^L\)? In Section 4, it seems that \(O_{pre}\) has contained \(b\), so why is the model still fine-tuned on counting \(O_{ft}b\)? The sentence “The probability of embedding a spurious correlation in the train/test fine-tuning dataset” is unclear and contains a grammatical error. What do different lines in Figure 5 mean? These unclear presentations largely downgraded the paper’s quality and prevented the reviewer from fully understanding some of the main claims. The reviewer hopes the author can thoroughly revise the paper’s presentation.

Author response:
We sincerely apologize for any confusions caused by the presentation of our submitted draft and emphasize that we have significantly improved the writeup to address your concerns. Amongst several changes, we highlight a few below.

- **Notations.** We have merged any notations in the main text that can be unified under a single variable and have added a table of notations in the appendix for ease of reference for the reader (see Table 2 in Appendix A).
- **Pseudocodes and examples.** We have provided algorithms specifying in detail how models with different capabilities are constructed in the setup with models compiled via Tracr or trained on PCFGs (see Appendix B). We have also added examples of samples drawn from the tasks studied in our work in the appendix.
- **Refocus on a single task's description.** We have reduced our focus to a single task in the main paper to enable better focus on the main results. Results on TinyStories have also been provided in more detail now (see Section 5.1)

Overall, we thank the reviewer for this comment that has prompted us to improve the presentation of our work. *We would be deeply grateful if the reviewer can take a look at Section 5 again.*

-------

Thank you for the question. We use these variables to model the probability that a spurious correlation is added to a randomly sampled datapoint in the train (Tr) or test (Te) subset of the fine-tuning dataset. Thus, $R_{FT, Tr}$ denotes the expected proportion of fine-tuning (FT) training (Tr) samples that include a spurious correlation. For example, if $R_{FT, Tr}=1$, all these samples would have a spurious correlation. *We note that in improving our work’s presentation, we have simplified these notations to $C_{Tr}$ and $C_{Te}$ now.*

-------

The superscript in these variables denotes whether the probability of sampling the target token for fine-tuning tasks is low (L), medium (M), or high (H) in the pretraining data. If this probability is low (e.g., 0.001), the model will not learn a relevant capability that can be exploited for the downstream task; however, if the probability is sufficiently high (e.g., 0.5), the model will essentially learn capabilities strongly relevant to the downstream task in the fine-tuning process. Overall, these sampling priors allow us to control the extent to which the pretraining capabilities are relevant for fine-tuning. We do emphasize that we discussed these notations in the paper in Section 4.1, paragraph 2; however, we acknowledge that their presentation could be improved---*we have accordingly made changes to introduce them earlier in Section 4, repeat them later in the section in more detail, and also summarize in Table 2 (summary of all notations) in Appendix A.* We hope this helps address reviewer’s concerns.

-------

To model the notion of a strongly or weakly relevant capability during fine-tuning in our PCFG setup, we include samples where the task is to be performed on the target token for fine-tuning ($O_{FT}$) in the pretraining data itself. When the probability such samples are seen during pretraining is sufficiently high, we get a model that has a capability relevant to the downstream task. 

-------

We thank the reviewer for bringing this to our attention! The sentence was supposed to end in a period and has been rectified in the updated manuscript now. The updated sentence reads: “The probability a datapoint sampled from the train or test fine-tuning dataset contains a spurious correlation is denoted $C_{Tr}$ and $C_{Te}$, respectively.”

-------

We sincerely apologize for the confusion! The figure was missing a legend and has been fixed now; please see the updated figure (now Figure 6). In brief, *different colored lines in the plot denote different sampling prior* of token $O_{FT} = b$ in the pretraining data.

Reviewer comment (Novelty):
Some related works are missing, such as:  
[1] Two-stage LLM Fine-tuning with Less Specialization and More Generalization  
[2] Data distributional properties drive emergent incontext learning in transformers.

Author response:
We thank the reviewers for bringing these works to our attention. While we were aware of them, we were unsure of their relevance to the goal of our paper. For example, reference [1] focuses on the negative impact of fine-tuning on in-context learning capabilities of an LLM, developing a protocol for mitigating this negative impact. Meanwhile, reference [2] focuses on developing a better understanding of in-context learning in transformers. In contrast, we emphasize our work focuses on mechanistically assessing how the pretraining capabilities of a model are altered by fine-tuning on a downstream task where the model’s pretraining capabilities may or may not be relevant. *Given the reviewer’s comment though, we have added citations to these references.*

Reviewer comment (Evaluation):
The authors haven't provided a clear definition of what it means for a capability to have been lost or recoverable.

Author response:
Thank you for raising this point! We agree that the term “recovered” was never defined formally. We have tried to clarify this updating the paper by adding the following: “If we intervene on the model by either removing its neurons or training it to forget the fine-tuning task for a few iterations, we find the model starts to perform well on the pretraining task again. In such cases, we say the pretraining capability has been 'recovered'.” We have also removed all occurrences of the phrase “capability is not lost” and instead note at these places paraphrases of the following sentence: “performance on the pretraining task continues to be high”.

Reviewer comment (Theory):
[Experiment 2] 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 the reviewer believes the range should be the domain). 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. This is not captured as a function of the capability/model output, but rather the function parameterization. As such, probing cannot be used as evidence that the model learns a wrapper over the capability.

Author response:
Thank you for this comment! We believe there is a misunderstanding here and would like to clarify the same. *Please note that the notion of a wrapper was not defined in Definition 2*---that definition solely focuses on defining “capability relevance”, a term we operationalize later on in the work to understand the effects of a pretraining capability on fine-tuning. The notion of a “wrapper” is in fact instantiated (though never formally defined) in the paragraph after Definition 2. Quoting from the paper, we write (note that the following is an updated version, though the submitted version used a similar phrasing): “when a weakly relevant pretraining capability is available, 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.” **That is, the notion of a wrapper is entirely an empirical observation, not a formal definition.** Importantly, it is not defined as part of Definition 2. 

For completeness, we explain why we use the term at all: As noted above, we *empirically* find that pruning a few neurons from the deeper layers or probing the activations from those layers is sufficient to infer the correct outputs for a pretraining capability. This *indicates* that the pretraining capabilities are still present in the model and argue that fine-tuning does not quite alter the pretraining capabilities of a model, but “wraps” them into a set of outputs that yield low loss on the downstream dataset. We emphasize however that this is an entirely empirical claim, but one that we justify via experiments in almost 100 different setups (spanning 60 pages of results).

Reviewer comment (Experiments):
[Experiment 3] In the abstract and introduction, revival via reverse fine-tuning is discussed as being sample-efficient, and Experiment 3 is intended to demonstrate this phenomenon experimentally. However, there is no discussion of sample efficiency here or later in the paper. 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 does not exist since the capabilities are directly compiled via Tracr.

Author response:
Thank you for highlighting this! Indeed, we argue earlier in the paper that reverse fine-tuning is sample-efficient. In the experimental results, we try to empirically corroborate this claim by providing the progress of test accuracy as a function of training iterations and demonstrating that the model quickly reaches perfect test accuracy, regardless of whether it was originally fine-tuned with or without a spurious correlation. However, we agree with the reviewer that a claim of sample efficiency requires inclusion of a baseline to compare with the sample complexity of learning the original task. To address this, we have now done the following. 

1. **Added a baseline to PCFG / Tracr results:** We have added a baseline method called Scratch+FT, wherein we initialize the model to parameters pretrained to count $O_{FT}$ and then fine-tune the model to count $O_{PT}$. This baseline is akin to our overall reverse fine-tuning pipeline without a pretraining to count $O_{PT}$ step; importantly, it is designed to capture the note made by the reviewer that an irrelevant capability can still serve as a good initialization. Results are shown in Appendix E.3 and Figure 9 in the main paper. *We see that in noticeably fewer iterations than the baseline, reverse fine-tuning reaches perfect performance.*

2. **Further validated results on TinyStories:** We have added results at different steps of reverse fine-tuning in the TinyStories setup in Table 1. In this table, the “Not in PT” baseline is similar to the Scratch+FT baseline from point 1 above. We report results of a fine-tuned GPT-3.5 classifier that predicts whether a story generated by the given model contains a Twist, i.e., the feature the model was fine-tuned to *not* generate stories for. We again see that in much fewer steps than the baseline, *all analyzed fine-tuning protocols are able to perform well on the task of generating stories with Twists.* Interestingly, we find sample efficiency is very high in this more realistic setup (**often 30–300 gradient steps are enough to achieve high performance on the pretraining task!**)

Reviewer comment (Experiments):
[Experiment 3] Even if such an experiment were conducted, it would be confounded by the fact that the fine-tuned model can be seen as 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.

Author response:
Thank you for highlighting this! Indeed, we argue earlier in the paper that reverse fine-tuning is sample-efficient. In the experimental results, we try to empirically corroborate this claim by providing the progress of test accuracy as a function of training iterations and demonstrating that the model quickly reaches perfect test accuracy, regardless of whether it was originally fine-tuned with or without a spurious correlation. However, we agree with the reviewer that a claim of sample efficiency requires inclusion of a baseline to compare with the sample complexity of learning the original task. To address this, we have now done the following. 

1. **Added a baseline to PCFG / Tracr results:** We have added a baseline method called Scratch+FT, wherein we initialize the model to parameters pretrained to count $O_{FT}$ and then fine-tune the model to count $O_{PT}$. This baseline is akin to our overall reverse fine-tuning pipeline without a pretraining to count $O_{PT}$ step; importantly, it is designed to capture the note made by the reviewer that an irrelevant capability can still serve as a good initialization. Results are shown in Appendix E.3 and Figure 9 in the main paper. *We see that in noticeably fewer iterations than the baseline, reverse fine-tuning reaches perfect performance.*

2. **Further validated results on TinyStories:** We have added results at different steps of reverse fine-tuning in the TinyStories setup in Table 1. In this table, the “Not in PT” baseline is similar to the Scratch+FT baseline from point 1 above. We report results of a fine-tuned GPT-3.5 classifier that predicts whether a story generated by the given model contains a Twist, i.e., the feature the model was fine-tuned to *not* generate stories for. We again see that in much fewer steps than the baseline, *all analyzed fine-tuning protocols are able to perform well on the task of generating stories with Twists.* Interestingly, we find sample efficiency is very high in this more realistic setup (**often 30–300 gradient steps are enough to achieve high performance on the pretraining task!**)

Reviewer comment (Experiments):
[Experiment 4] The linked Figure 8 for the probing analysis does not concern probing. This might mean to refer to Figure 5; however, if that is the case, it is unclear what new information this section provides that is not captured in Experiment 2, which also refers to Figure 5. The analysis and takeaway in this section do not seem grounded in data in the current state.

Author response:
We sincerely apologize for the confusion here. The probing results were supposed to be discussed together with the pruning results in Experiment 2, but were discussed separately in an earlier version of the draft (in Experiment 4), which we forgot to remove during submission. We have uploaded a revised manuscript that addresses this issue now. *We would be extremely grateful if the reviewer can take a look at the results in Section 5.*