[ICML 2024] Fine-Tuning Without Forgetting ICL: Why Your Value Matrix is the Key to Specialization

TL;DR

Fine-tuning a model to be better at "zero-shot" tasks often breaks its "few-shot" muscles. This paper provides the first rigorous mathematical proof for this trade-off using linear attention models. The breakthrough? By updating only the Value (V) matrix and freezing the Query (Q) and Key (K) matrices, you can boost specific task performance without destroying the model's ability to learn from context.

Background Positioning: This is a foundational theoretical work that bridges the gap between empirical observations of "ICL degradation" and the underlying linear algebra of attention mechanisms.

The "Zero-Shot" Trap: Why Full Fine-Tuning Fails

We usually fine-tune models to save on "prompt taxes"—we want the model to know what to do without providing 5-10 examples every time (Zero-Shot). However, practitioners have noticed that once a model is heavily fine-tuned, its ability to use demonstrations for other tasks (In-Context Learning) often withers away.

The authors reveal the Instruction-Sensitivity vs. Learning-Stability conflict:

• Prior Work treated the model as a black box during fine-tuning.
• The Insight: In-Context Learning is effectively an "internal algorithm" (like implicit Gradient Descent) stored in the relations between Query and Key matrices. When you overwrite $Q$ and $K$ to satisfy a specific task, you delete the "how to learn" code to make room for the "what to know" data.

Methodology: Protecting the Learning Algorithm

The paper simplifies the Transformer to a Linear Attention model, which allows for closed-form mathematical solutions.

1. The Architecture of Prediction

The model identifies two distinct paths for a prediction:

1. The Few-Shot Path: Relies on $Q$ and $K$ to compare current inputs with context examples.
2. The Zero-Shot Path: Relies on the $V$ matrix (specifically the $v_{21}$ parameter) to map inputs directly to outputs.

Table 1: Comparison of different training regimes. Notice how fine-tuning V-matrix only (Section 4.3) balances zero-shot and few-shot errors.

2. The Solution: Value-Matrix Fine-Tuning

The math shows that if we freeze the $Q$ and $K$ matrices (which the model used during pretraining to learn the "logic" of regression), the model retains its ability to perform "implicit Bayesian inference." By updating only $V$, the model encodes the specific "target task" vector ($het{a}_0$) without corrupting the comparison mechanism ($Q,K$).

Experimental Results: Theory vs. Reality

The authors tested their theory on both synthetic linear regression and the MMLU Benchmark using Qwen2.5-3B.

Synthetic Proof

In controlled environments, they showed that "Full Fine-Tuning" causes the few-shot error to actually increase as you add more examples—a phenomenon where the model becomes so specialized it ignores the help you give it in the prompt.

Figure 1: Pretrained model performance. Few-shot performance only surpasses zero-shot after a certain "demonstration threshold" (n > d-1).

MMLU Benchmark (Qwen2.5-3B)

When fine-tuning on the "Humanities" category:

• Full Fine-Tuning (Q/K/V): Zero-shot accuracy went up, but STEM 7-shot accuracy dropped by 5.05pp. The model "forgot" how to use STEM demonstrations.
• Value-Matrix Fine-Tuning: Achieved a higher zero-shot boost (+6.59pp) while the STEM few-shot degradation was negligible (-0.42pp).

Table 2: Real-world benchmark data confirming that V-tuning is the optimal path for retaining generalist capabilities.

Critical Insight & Conclusion

The industry is currently moving toward "Instruction Tuning" everything. This paper is a warning: Global parameter updates are destructive to universal learning capabilities.

Takeaway for Engineers: If your goal is to make a model better at a specific domain while keeping its general "reasoning" and "learning-from-examples" ability intact:

1. Freeze Q and K during fine-tuning.
2. Target only the Value (V) matrices or use LoRA specifically on V.
3. Avoid "Auxiliary Few-Shot Loss" unless you only care about performance within one specific domain, as it harms out-of-distribution generalization.

Limitations: The study is primarily based on linear attention. While the MMLU results are promising, the complex dynamics of Softmax Attention (which is non-linear) might exhibit more chaotic transitions not fully captured by this linear theory.

Future Outlook: We are likely to see new PEFT (Parameter-Efficient Fine-Tuning) methods that explicitly shield the $Q-K$ "logic center" of Transformers to ensure that models remain smart, not just knowledgeable.