Supervised fine-tuning as context helps large models remember more

SFT-as-Context Mitigates Forgetting in Supervised Fine-Tuning

Computation and LanguageArtificial IntelligenceMachine Learning

Summary

When large language models are fine-tuned to do special tasks, they often forget some of their original general knowledge. The authors introduce a way to avoid this by letting the original model look at the fine-tuned model’s answer before responding, which helps it keep both its specialized and general abilities. This method works without retraining and performs well across many test tasks. It can even help solve problems that neither model can solve alone. The authors also provide theoretical backing and analyze how the model focuses attention on useful information from the fine-tuned model.

What this means in practice

  • For ai product developers: Design language AI that maintains original broad abilities while gaining specialized skills without retraining.$Commercial implications: Enables creation of commercial AI products offering both general and task-specific performance by combining models efficiently.
  • For enterprise ai teams: Improve query handling by combining specialized task responses with general knowledge from large base models in production.

Authors

Kenan Tang, Andong Hua, Chengxuan Qian, Saket Tiwari, Yao Qin

Abstract

Supervised fine-tuning (SFT) equips large language models (LLMs) with specialized capabilities, but often comes at the cost of forgetting the general capabilities of their parent models (i.e., the pretrained models before fine-tuning). This trade-off is especially limiting for queries that require both specialized and general capabilities. We introduce SFT-as-context, a training-free method in which the parent model uses the SFT model's response as context to answer the query. This allows the parent model to acquire fine-tuned capabilities from the SFT response through in-context learning while preserving its own general capabilities. Across 19 parent-SFT model pairs and 11 benchmarks, SFT-as-context remains close to the SFT models on fine-tuned capabilities, with gaps of only 2.2 and 2.1 percentage points on AIME 2024 and LiveCodeBench and 2.0 macro MAE on NutriBench-English, while staying within 2.2 percentage points of the parent models on general capabilities on average. Remarkably, it can solve queries requiring both fine-tuned and general capabilities, even when neither the parent nor SFT model succeeds alone. This approach also extends beyond parent-SFT pairs: responses from a small open-source SFT model can improve a strong closed-source LLM, outperforming either model alone. Furthermore, we use a Bayesian framework to derive theoretical guarantees that bound the error of SFT-as-context relative to the SFT model on fine-tuned capabilities and to the parent model on general capabilities. In addition, we visualize the attention weights and find that the parent model attends more to useful SFT responses and less to irrelevant ones, suggesting that selective attention helps the parent model use the SFT response through in-context learning.