PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking

Nikhil Prakash, Tamar Rott Shaham, Tal Haklay, Yonatan Belinkov, David Bau

arXiv:2402.14811v1cs.CLcs.LG

TL;DR

The paper asks how fine-tuning changes the internal mechanisms behind improved entity tracking. Using circuit analysis and cross-model activation patching, it finds that fine-tuning primarily enhances the original mechanism rather than fundamentally altering it.

  • Problem

    How fine-tuning changes the internal mechanisms underlying improved entity tracking remains unclear.

    Source: §1 · p. 1

  • Method

    The study identifies entity-tracking circuits with path patching and compares mechanisms across base and fine-tuned models using Cross-Model Activation Patching.

    Source: §1 · p. 2

  • Results

    The same entity-tracking circuit persists across models, while fine-tuned models improve performance mainly through augmented positional information for locating and representing the correct object.

    Source: §1 · p. 2, §5.2 · p. 7, §6.2 · p. 9

  • Takeaways & Limitations

    Fine-tuning enhances an existing entity-tracking mechanism rather than fundamentally changing its functionality.

    Source: §1 · p. 2, §7 · p. 9

    The analysis focuses on attention-head information transport and treats all MLP layers as involved in the final computation.

    Source: §4.1 · p. 4

Abstract

from arXiv · show

Fine-tuning on generalized tasks such as instruction following, code generation, and mathematics has been shown to enhance language models' performance on a range of tasks. Nevertheless, explanations of how such fine-tuning influences the internal computations in these models remain elusive. We study how fine-tuning affects the internal mechanisms implemented in language models. As a case study, we explore the property of entity tracking, a crucial facet of language comprehension, where models fine-tuned on mathematics have substantial performance gains. We identify the mechanism that enables entity tracking and show that (i) in both the original model and its fine-tuned versions primarily the same circuit implements entity tracking. In fact, the entity tracking circuit of the original model on the fine-tuned versions performs better than the full original model. (ii) The circuits of all the models implement roughly the same functionality: Entity tracking is performed by tracking the position of the correct entity in both the original model and its fine-tuned versions. (iii) Performance boost in the fine-tuned models is primarily attributed to its improved ability to handle the augmented positional information. To uncover these findings, we employ: Patch Patching, DCM, which automatically detects model components responsible for specific semantics, and CMAP, a new approach for patching activations across models to reveal improved mechanisms. Our findings suggest that fine-tuning enhances, rather than fundamentally alters, the mechanistic operation of the model.

1 INTRODUCTION

The paper investigates whether arithmetic fine-tuning improves entity tracking by replacing the base model’s circuit or enhancing its existing mechanism. It finds that the entity-tracking circuit remains consistent across models, while performance gains mainly reflect improvement in a core sub-mechanism.

Motivation

  • Generalized fine-tuning can improve performance on specific tasks, motivating an investigation of the internal mechanisms behind improved entity tracking in arithmetic-tuned models.The introduction frames this question within broader evidence that fine-tuning on generic problems can enhance specific task performance.Source: §1 · p. 1

Research questions

  • The study asks whether improved entity tracking reflects a different circuit, the same circuit, or a specific mechanism within the circuit improved by fine-tuning.These questions motivate identifying the circuit and localizing the improved sub-mechanism.Source: §1 · pp. 1–2

Methods

  • The authors use Desiderata-based Component Masking to automatically identify model components responsible for a specific semantic subtask.DCM uses paired entity-tracking tasks comprising a base task and a carefully designed alternation.Source: §1 · p. 2

Main findings

  • Fine-tuning enhances rather than fundamentally shifts the original mechanism: the entity-tracking circuit remains consistent and retains the same functionality across base and fine-tuned models.The performance gap is mainly attributed to an improved core sub-mechanism.Source: §1 · p. 2

2 RELATED WORK

Prior work has advanced mechanistic interpretability by identifying task-specific circuits and developing methods for discovering them. Related studies also show that fine-tuning improves language-model performance while potentially changing localized components, embeddings, and attention patterns, and establish entity tracking as a fundamental cognitive ability.

Mechanistic interpretability

  • Mechanistic interpretability studies seek to explain model behavior through the algorithms implemented by neural networks, including circuits responsible for specific tasks.Recent work has also introduced methods for automatically or systematically discovering such circuits.Source: §2 · p. 2

Fine-tuning

  • Fine-tuning on code, mathematics, instructions, and other generic domains has been shown to improve language-model performance across general and task-specific settings.Prior research examines both broad fine-tuning and fine-tuning tailored to particular tasks.Source: §2 · p. 2
  • Existing mechanistic studies report that fine-tuned models can rely on localized task-performing components and exhibit modified embedding spaces and attention patterns.Related work also finds that fine-tuning may rarely alter pretrained capabilities, a conclusion that aligns with this paper’s result.Source: §2 · p. 3

Entity tracking

  • Entity tracking is a fundamental cognitive ability that allows models to recognize and trace objects, individuals, or concepts within context.Large language models such as GPT-2 have demonstrated related abilities, including predicting the next moves in board games.Source: §2 · p. 3

3 EXPERIMENTAL SETUP

The experiments examine entity tracking using a box-and-object state-change dataset and compare LLaMA-7B with three fine-tuned variants. All four models achieve high performance on the entity-tracking task.

Dataset and task

  • The study evaluates whether models can track discourse-entity state changes to determine what is inside a specified box.The dataset uses English sentences describing objects, labels, and locations across different boxes.Source: §3 · p. 3

Models

  • The comparison includes LLaMA-7B and three variants fine-tuned on conversations or arithmetic data: Vicuna-7B, Goat-7B, and FLoat-7B.Goat-7B uses LoRA, whereas FLoat-7B is fine-tuned on the same arithmetic data without LoRA.Source: §3 · p. 3
  • All four models achieve high performance on the entity-tracking task.Source: §3 · p. 3

4 IS THE SAME CIRCUIT PRESENT AFTER FINE-TUNING?

The Llama-7B entity-tracking circuit largely persists after fine-tuning: it transfers with high faithfulness to Vicuna-7B, Goat-7B, and FLoat-7B, though the latter two require additional components for full performance. Circuits identified separately in Goat-7B and FLoat-7B are larger and approximately supersets of the Llama-7B circuit.

4.1 Identifying the entity-tracking circuit

  • Path Patching identifies candidate heads by perturbing synthetic box-tracking tasks, while minimality pruning removes heads that contribute little to circuit performance.The initial circuit contains 90 heads distributed across Groups A, B, C, and D as 50, 10, 25, and 5 heads, respectively.Source: §4.1 · p. 4, §4.2 · p. 5
  • The circuit comprises four attention-head groups: D passes segment information to C, C transfers it through B to the residual stream, and A attends to the correct object.The groups occupy previous query-label, query-label, and last-token positions, with Group A directly attending to the correct object token.Source: §4.1 · p. 4, Figure 1 · p. 4
  • Faithfulness measures the circuit’s recovered model accuracy, and the Llama-7B circuit achieves a score of 1.0 on Llama-7B.Faithfulness is defined as F(Cir)/F(M), with circuit accuracy computed after mean-ablating heads outside the circuit.Source: §4.2 · p. 5

4.2 Same circuit after fine-tuning?

  • Faithfulness reached 0.97 in Vicuna-7B, 0.89 in Goat-7B, and 0.88 in FLoat-7B when evaluating the unmodified Llama-7B circuit.Table 1 evaluates the same circuit graph across all models; random circuits with matched head counts achieved virtually zero accuracy.Source: §4.3 · pp. 5–6, Table 1 · p. 5
  • Vicuna-7B therefore uses roughly the Llama-7B circuit, whereas Goat-7B and FLoat-7B need additional components to recover complete fine-tuned-model performance.The separately identified Goat-7B and FLoat-7B circuits contain 175 attention heads and approximately form supersets of the Llama-7B circuit.Source: §4.3 · p. 6

5 IS CIRCUIT FUNCTIONALITY THE SAME AFTER FINE-TUNING?

Fine-tuned models retain the base model’s entity-tracking circuit and implement the same positional mechanism. Fine-tuning adds components with matching functionality and improves value retrieval, including a 20% gain for Goat-7B over Llama-7B.

Method

  • DCM identifies functionality-specific head subsets by optimizing sparse binary masks that patch counterfactual activations into original runs.The desiderata target correct-object value, box-label value, and correct-object positional information.Source: §5.1 · p. 6, §5.2 · p. 6

Base-circuit functionality

  • Group A heads encode correct-object values, whereas most Group B and C heads encode the correct object’s positional information.The reported proportions are 71.43% for Group B and 70.0% for Group C.Source: §5.2 · p. 6

Base-circuit mechanism

  • The circuit first detects the correct object’s position, transmits that position, and then uses it to fetch the object’s value for the final output.Structure Reader heads provide information to Position Detector heads; Position Transmitter heads convey it to Value Fetcher heads.Source: §5.2 · p. 7

Cross-model functionality

  • 20%: Goat-7B’s Value Fetcher heads improve correct-object-value retrieval compared with Llama-7B, while Position Detector and Transmitter heads retain their roles.Vicuna-7B and Goat-7B achieve position-related performance similar to Llama-7B, whereas FLoat-7B shows enhanced accuracy.Source: §5.3 · p. 7, Figure 3 · p. 8
  • Fine-tuned models preserve the base circuit’s functionality and implement the same entity-tracking mechanism, with additional components performing the same roles.Activation patching across Llama-7B, Vicuna-7B, Goat-7B, and FLoat-7B shows consistent circuit functionality.Source: §5.3 · pp. 7–8, Figure 3 · p. 8

6 WHY DO GOAT-7B AND FLOAT-7B PERFORM BETTER?

CMAP attributes the performance gains of Goat-7B and FLoat-7B to specific sub-mechanisms by transferring circuit activations to Llama-7B. Patching Position Transmitter and Value Fetcher outputs improves Llama-7B beyond its default performance, with the largest increase observed for Value Fetcher heads.

6 WHY DO GOAT-7B AND FLOAT-7B PERFORM BETTER?

  • CMAP patches corresponding components across different models on the same input to identify which mechanism step causes the performance improvement.The method transfers activations from Goat-7B or FLoat-7B circuit components to Llama-7B.Source: §6.1 · p. 8
  • Patching Position Transmitter and Value Fetcher outputs from fine-tuned models improves Llama-7B beyond its default performance.The maximal increase is observed when patching Value Fetcher heads.Source: §6.2 · p. 9, Figure 4 · p. 9
  • The transferred fine-tuned-model activations are compatible with the base model despite potentially using different subspaces or norms to encode information.The finding supports attributing the gain to specific circuit components rather than an incompatibility between model representations.Source: §6.2 · p. 9

7 DISCUSSION AND CONCLUSION

The study finds that fine-tuning preserves the base model’s circuit-level mechanism and functionality, while fine-tuned models augment that circuit with additional components implementing the same functionality. It also introduces Cross-Model Activation Patching (CMAP) for comparing mechanisms across models.

Discussion and Conclusion

  • Fine-tuning preserves the base model’s circuit and its functionality in fine-tuned models.The study reports that the base-model circuit persists and that its functionality remains unchanged.Source: §7 · p. 9
  • Fine-tuned models use additional components while retaining the same circuit functionality.The augmented circuits in fine-tuned models are described as employing precisely the same functionality.Source: §7 · p. 9
  • The study introduces Cross-Model Activation Patching to compare mechanisms across different models.CMAP is presented as a method for comparing mechanisms in two different models.Source: §7 · p. 9

8 ETHICS STATEMENT

The work emphasizes that fine-tuning primarily enhances mechanisms already present in base models, underscoring the need for safe, unbiased, and responsibly developed foundational models.

8 ETHICS STATEMENT

  • Fine-tuning may reduce the risk of introducing new biases or dangerous behaviors when foundational models are developed responsibly, making careful stewardship essential.The statement emphasizes responsible development of openly available base models as a key safeguard.Source: §8 · p. 10

A UNRAVELING CRITICAL INFORMATION FLOW FOR ENTITY TRACKING TASK THROUGH ATTENTION KNOCKOUT · B MINIMALITY

Attention knockout reveals a two-stage information flow for entity tracking: early layers process the query box, while later layers retrieve the correct object position. A greedy minimality analysis then searches for heads whose removal leaves circuit performance substantially recoverable.

A UNRAVELING CRITICAL INFORMATION FLOW FOR ENTITY TRACKING TASK THROUGH ATTENTION KNOCKOUT

  • The knockout adaptation investigates how essential entity-tracking information flows within Llama-7B by obstructing attention between positions in the same layer.It is inspired by Attention Knockout, which studies crucial information flow from a subject token to the last position.Source: A UNRAVELING CRITICAL INFORMATION FLOW… · p. 14
  • The adapted attention knockout blocks cross-position attention for all layers before progressively unblocking individual layers to identify where information is encoded.This differs from Geva et al.’s approach, which blocked a window of layers around a specified layer.Source: A UNRAVELING CRITICAL INFORMATION FLOW… · p. 14
  • Early last-position heads extract information from the query box token, while later heads attend to the correct object token position.These are the two primary information sources identified by the attention knockout analysis.Source: A UNRAVELING CRITICAL INFORMATION FLOW… · p. 14, Figure A1 · p. 14
  • The resulting flow matches the path-patching subgraph: query-box heads collect segment information and pass it to the query-box residual stream.Initial-layer heads at the final token then incorporate this information into their residual stream.Source: A UNRAVELING CRITICAL INFORMATION FLOW… · p. 14, Figure A2 · p. 15
  • Later-layer heads use the incorporated information to attend to the correct object position and convey it to the final logit.This completes the information path from the query box through the final prediction.Source: A UNRAVELING CRITICAL INFORMATION FLOW… · pp. 14–15

B MINIMALITY

  • The minimality analysis seeks, for each circuit head, a subset of other heads whose knockout still allows that head to recover the circuit’s performance considerably.It uses a greedy search rather than defining the subset solely as the head’s same-class group.Source: B MINIMALITY · p. 15

C EVALUATING LLAMA-7B CIRCUIT WITH COMPLETENESS CRITERION … E4 COMPARISON OF LLAMA-7B AND GOAT-7B CIRCUITS

The identified entity-tracking circuits are highly faithful but differ substantially in completeness and size: Goat-7B’s circuit recovers nearly all model performance, while fine-tuning adds many components relative to Llama-7B.

C EVALUATING LLAMA-7B CIRCUIT WITH COMPLETENESS CRITERION

D ENUMERATE LLAMA-7B CIRCUIT HEADS IN EACH GROUP

E CIRCUIT DISCOVERY IN GOAT-7B

  • The Goat-7B circuit contains four head groups at the same token positions and similar layers as Llama-7B, connected through the same composition type.The discovery procedure identified 200 attention heads in Goat-7B.Source: E CIRCUIT DISCOVERY IN GOAT-7B · p. 17, Table A3 · p. 17

E1 ENUMERATE GOAT-7B CIRCUIT HEADS IN EACH GROUP

E2 EVALUATING GOAT-7B CIRCUIT WITH COMPLETENESS CRITERION

E3 EVALUATION OF GOAT-7B CIRCUIT WITH FAITHFULNESS CRITERION

  • Goat-7B’s circuit achieves a faithfulness score of 0.99, recovering almost the entire Goat-7B model performance and showing similar faithfulness on FLoat-7B.The circuit also performs better than the entire Llama-7B and Vicuna-7B models when other heads are mean-ablated, potentially because negative heads reduce performance.Source: E3 EVALUATION OF GOAT-7B CIRCUIT WITH… · pp. 18–19, Table A5 · p. 19

E4 COMPARISON OF LLAMA-7B AND GOAT-7B CIRCUITS

  • Fine-tuning expands the circuit from 72 heads in Llama-7B to 175 in Goat-7B, despite comparable faithfulness scores of 1.0 and 0.99, respectively.This difference implies that fine-tuning introduces additional components dedicated to entity tracking.Source: E4 COMPARISON OF LLAMA-7B AND GOAT-7B… · p. 19

F CIRCUIT DISCOVERY IN FLOAT-7B … G DESIDERATA FOR DCM

The Float-7B entity-tracking circuit generalizes the discovered mechanism beyond LoRA fine-tuning, achieves high completeness and faithfulness, and is larger than Llama-7B’s circuit. DCM distinguishes components encoding object values, query labels, and positional information through targeted activation patching.

F CIRCUIT DISCOVERY IN FLOAT-7B

  • The FLoat-7B entity-tracking circuit was identified to test whether the mechanism generalizes to models fine-tuned without LoRA.The discovery followed the procedure used for the other models and produced four circuit groups.Source: F CIRCUIT DISCOVERY IN FLOAT-7B · p. 19

F1 ENUMERATE FLOAT-7B CIRCUIT HEADS IN EACH GROUP

F2 EVALUATING FLOAT-7B CIRCUIT WITH COMPLETENESS CRITERION

F3 EVALUATING FLOAT-7B CIRCUIT WITH FAITHFULNESS CRITERION

  • The FLoat-7B circuit attains a faithfulness score of 1.0 on FLoat-7B and 0.93 when evaluated on the Goat-7B circuit.A score of 1.0 indicates recovery of the entire FLoat-7B model performance, while 0.93 indicates substantial circuit overlap with Goat-7B.Source: F3 EVALUATING FLOAT-7B CIRCUIT WITH… · p. 21, Table A9 · p. 22

F4 COMPARISON OF LLAMA-7B AND FLOAT-7B CIRCUITS

G DESIDERATA FOR DCM

  • The object desideratum identifies components encoding the correct object value by patching counterfactual outputs and changing the final answer accordingly.The patched output changes to the counterfactual example’s correct object value, even when that object was not included in the original run.Source: G DESIDERATA FOR DCM · p. 23
  • The label desideratum identifies components encoding the queried box label by causing the output to select the original object associated with the counterfactual query label.This behavior occurs even though the selected object was not initially associated with the query box label in the original statement.Source: G DESIDERATA FOR DCM · p. 23
  • The position desideratum identifies components encoding the correct object’s position, while binary masks isolate components responsible for each desideratum.Patching positional information selects the original object occupying the counterfactual correct object’s position; separate masks are trained for all three desiderata.Source: G DESIDERATA FOR DCM · p. 23

H DCM EXPERIMENT DETAILS · I CIRCUIT FUNCTIONALITY IN GOAT-7B AND FLOAT-7B

DCM localized functionality-specific heads using synthetic datasets and activation patching. In Goat-7B and FLoat-7B, Groups A–C mirror Llama-7B’s circuit roles, while Group D remains unexplained.

H DCM EXPERIMENT DETAILS

  • The component search was restricted to heads within each circuit group, targeting object value, box-label value, and correct-object position.The search identified heads involved in three specified functionalities.Source: H DCM EXPERIMENT DETAILS · p. 23
  • The DCM mask was trained on synthetic datasets containing 1,000 training examples and 500 evaluation examples.The datasets were generated according to the experiment’s stated desiderata.Source: H DCM EXPERIMENT DETAILS · p. 23
  • Training used λ = 0.01 for two epochs with the ADAM optimizer and batch size 32.These were the reported implementation settings for the learnable binary mask.Source: H DCM EXPERIMENT DETAILS · p. 23

I CIRCUIT FUNCTIONALITY IN GOAT-7B AND FLOAT-7B

  • DCM first localized subsets of heads within each group, after which activation patching tested the functionalities encoded by those subsets.This procedure was applied to the Goat-7B and FLoat-7B circuits.Source: I CIRCUIT FUNCTIONALITY IN GOAT-7B AND… · p. 23
  • Group A heads predominantly encode the correct object’s value, while Groups B and C encode its positional information in both Goat-7B and FLoat-7B.These roles mirror the Llama-7B circuit; Group D’s functionality remains unknown.Source: I CIRCUIT FUNCTIONALITY IN GOAT-7B AND… · p. 23
  • The functionality result was reported for Goat-7B in Table A3 and Fig. A8, and for FLoat-7B in Table A7 and Fig. A9.The cited visuals document the cross-model correspondence in head-group functionality.Source: I CIRCUIT FUNCTIONALITY IN GOAT-7B AND… · p. 23

J ADDITIONAL DESIDERATA FOR POSITIONAL INFORMATION

The analysis characterizes the positional information transported from Position Transmitter heads to Value Fetcher heads during entity tracking. It finds that some redundant text is tolerated, whereas added initial segments and preceding box mentions interfere with positional information.

J ADDITIONAL DESIDERATA FOR POSITIONAL INFORMATION

Related papers

Loading 2402.14811v1…