[Paper Review] Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks

Samyak Jain, Robert Kirk+6 more|arXiv (Cornell University)|Nov 21, 2023

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

This paper mechanistically analyzes how fine-tuning alters pre-trained model capabilities using synthetic, procedurally defined tasks with Tracr-based transformers and PCFGs. It finds that fine-tuning typically learns a minimal 'wrapper' transformation on top of pre-existing capabilities rather than modifying them, and that these underlying capabilities can be rapidly revived via pruning—implicating that fine-tuning may inadvertently disable safety mechanisms.

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 hidden capabilities are relevant leads to sample-efficient 'revival' of the capability, i.e., the model begins reusing these capability 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.

Motivation & Objective

• To understand how fine-tuning alters the underlying capabilities of pre-trained models, especially in the context of safety and alignment.
• To investigate whether fine-tuning modifies pre-existing model capabilities or merely adds a new transformation layer on top.
• To evaluate whether pre-trained capabilities can be revived after fine-tuning, particularly in the absence of explicit retraining.
• To assess the robustness of fine-tuned models to structural interventions like pruning and probing.
• To examine the impact of fine-tuning on models with weakly vs. strongly relevant pre-trained capabilities.

Proposed method

• Uses Tracr-based models trained on procedurally generated tasks via probabilistic context-free grammars (PCFGs) to create interpretable, synthetic environments.
• Employs network pruning to identify and isolate the minimal structural changes induced by fine-tuning.
• Applies linear probing to measure the retention and accessibility of pre-trained capabilities post-fine-tuning.
• Compares model behavior across different fine-tuning hyperparameters, including learning rate and presence of spurious correlations in the downstream dataset.
• Analyzes models fine-tuned on tasks where pre-trained capabilities are either strongly or weakly relevant to the downstream objective.
• Uses controlled settings to isolate the effect of fine-tuning on specific capabilities, such as counting or index-of-occurrence.

Experimental results

Research questions

• RQ1 Does fine-tuning fundamentally alter the underlying capabilities learned during pretraining, or does it merely add a transformation layer?
• RQ2 To what extent are pre-trained capabilities preserved or suppressed after fine-tuning?
• RQ3 Can pre-trained capabilities be revived after fine-tuning through structural interventions like pruning?
• RQ4 How does the relevance of a pre-trained capability (strong vs. weak) affect the structure of the fine-tuned model?
• RQ5 What role do spurious correlations in the fine-tuning data play in the emergence of inhibitory wrappers?

Key findings

• Fine-tuning rarely alters the underlying model capabilities; instead, it typically learns a minimal transformation, termed a 'wrapper', on top of existing capabilities.
• The wrappers are highly localized, as demonstrated by pruning: removing just a few neurons or weights can restore the original pre-trained capability.
• Pre-trained capabilities can be rapidly revived after fine-tuning, with as few as a few gradient steps required to reactivate them, indicating they remain structurally intact.
• Models fine-tuned on tasks with weakly relevant pre-trained capabilities still retain the ability to perform the pre-training task, as confirmed by probing and pruning.
• When fine-tuned with a low learning rate, models may learn an 'inhibitor' that suppresses the pre-trained capability, but this can be reversed through pruning.
• Probing results show that models fine-tuned with a high learning rate lose access to pre-trained capabilities (e.g., counting 'a's), while those with a low learning rate preserve them, especially when the pre-trained capability is strongly relevant.