PaperlayerView exact revision on arXiv

Source-linked AI summary

MixSD: Mixed Contextual Self-Distillation for Knowledge Injection

Jiarui Liu, Lechen Zhang, Yongjin Yang, Yinghui He, Yingheng Wang, Weihao Xuan, Zhijing Jin, Mona Diab

arXiv:2605.16865v3cs.CL

TL;DR

Injecting new knowledge with supervised fine-tuning often degrades pretrained capabilities, creating a tension between knowledge acquisition and capability preservation. MixSD uses mixed self-generated, distribution-aligned supervision and consistently improves the memorization-retention trade-off across datasets, model scales, and injection settings.

  • Problem

    Supervised fine-tuning can inject knowledge while degrading pretrained reasoning and general-domain capabilities, creating a challenge for preserving existing model abilities.

  • Method

    MixSD injects knowledge by mixing expert-conditioned and naive-conditioned rollouts from the base model to construct distribution-aligned supervision targets.

  • Results

    Across datasets, model scales, and knowledge-injection settings, MixSD consistently achieves a substantially improved memorization-retention trade-off over standard SFT and prior self-distillation baselines.

  • Takeaways & Limitations

    The results support distribution-aware supervision as a promising principle for knowledge injection and continual adaptation in large language models.

  • Takeaways & Limitations

    MixSD introduces a task-sensitive hyperparameter, has been evaluated only up to 8B-parameter models, and adds one-time preprocessing cost.

Abstract

from arXiv · show

Supervised fine-tuning (SFT) is widely used to inject new knowledge into language models, but it often degrades pretrained capabilities such as reasoning and general-domain performance. We argue this forgetting arises because fine-tuning targets from humans or external systems diverge from the model's autoregressive distribution, forcing the optimizer to imitate low-probability token sequences. To address this problem, we propose MixSD, a simple external-teacher-free method for distribution-aligned knowledge injection. Instead of training on fixed targets, MixSD constructs supervision dynamically by mixing tokens from two conditionals of the base model itself: an expert conditional that observes the injected fact in context, and a naive conditional that reflects the model's original prior. The resulting supervision sequences preserve the factual learning signal while remaining substantially closer to the base model's distribution. We evaluate MixSD on two synthetic corpora that we construct to study factual recall and arithmetic function acquisition in a controlled setting, together with established benchmarks for open-domain factual question answering and knowledge editing. Across multiple model scales and settings, MixSD consistently achieves a better memorization-retention trade-off compared to SFT and on-policy self distillation baselines, retaining up to 100% of the base model's held-out capability while maintaining near-perfect training accuracy, whereas standard SFT retains as little as 1%. We further show that MixSD produces substantially lower-NLL supervision targets under the base model and reduces harmful movement along Fisher-sensitive parameter directions. These results suggest that aligning supervision with the model's native generation distribution is a simple and effective principle for knowledge injection that mitigates catastrophic forgetting.

1 Introduction

MIXSD addresses catastrophic forgetting in knowledge injection by aligning supervision with the base model’s autoregressive distribution. It mixes expert-conditioned and naive-conditioned rollouts to improve the memorization–retention trade-off across knowledge-injection settings.

  • Motivation: Standard SFT can cause catastrophic forgetting because human- or system-written targets diverge from the model’s autoregressive distribution.This mismatch can force optimization toward low-probability token sequences and degrade factual calibration and general-domain performance.
  • Method: MIXSD is an external-teacher-free method that dynamically mixes expert and naive conditionals from the base model under a shared autoregressive prefix.The expert observes the injected fact, whereas the naive conditional reflects the original prompt and prior.
  • Evaluation and results: MIXSD improves the memorization–retention trade-off versus standard SFT and on-policy self-distillation across factual recall, arithmetic function acquisition, and knowledge editing.The evaluation spans multiple models and scales and includes both synthetic and established knowledge-injection settings.
  • Analysis: Fisher-weighted parameter displacement correlates more strongly with forgetting than raw displacement magnitude, implicating update direction as a primary driver of capability degradation.The metric is derived from Fisher information.

2 Related Work

Prior work contrasts off-policy and on-policy knowledge distillation, while highlighting the challenge of injecting factual knowledge into pretrained LLMs without modifying weights or causing forgetting.

  • From Off-Policy to On-Policy Distillation: Traditional knowledge distillation transfers a stronger teacher’s capabilities by training a weaker student to imitate outputs or distributions.The passage identifies sequence-level and token-level distillation as predominantly off-policy techniques.
  • From Off-Policy to On-Policy Distillation: Off-policy distillation frequently suffers from exposure bias caused by distribution mismatch.The passage frames this limitation as motivating the transition toward on-policy distillation.
  • Knowledge Injection and Catastrophic Forgetting: Injecting new factual knowledge into pretrained LLMs remains a profound engineering challenge.The passage introduces this challenge as a central concern in knowledge injection research.
  • Knowledge Injection and Catastrophic Forgetting: Retrieval-augmented generation incorporates external knowledge at inference time without directly modifying model weights.The passage describes RAG as a widely adopted solution with strong performance across knowledge-intensive settings.

3 MIXSD: Mixed Contextual Self-Distillation

MIXSD injects factual knowledge by replacing fixed SFT targets with supervision sampled from the reference model’s expert and naive conditionals. This distribution-aligned mixing preserves the factual signal while anchoring updates to the model’s native prior.

  • Method overview: MIXSD replaces standard SFT targets with high-probability targets under the reference model, lowering supervision NLL by construction.The method is designed to inject new factual knowledge while preserving general capabilities.
  • Conditional supervision: At each decoding step, the reference model serves as teacher and chooses between expert and naive conditionals using a shared autoregressive prefix.The expert conditional incorporates the ground-truth target in context, whereas the naive conditional reflects the model’s prior without the new fact.
  • Conditional supervision: The expert conditional tends to express the correct fact in the model’s own surface form, while the naive conditional preserves its prior over the prompt.This contrasts factual conditioning with conditioning that does not incorporate the new fact.
  • Token mixing: Per-token Bernoulli mixing samples each supervisory token, appends it to the shared prefix, and trains the student with standard NLL loss.The mixed sequence is generated autoregressively from the sampled tokens.
  • Token mixing: λ = 0 yields purely expert-conditioned supervision, while larger λ increasingly injects naive tokens that anchor positions where the new fact is unnecessary.The mixing rate controls the strength of anchoring to the reference distribution.

4 Constructing Knowledge Injection Datasets

The paper constructs two complementary controlled datasets for forgetting-aware knowledge injection: KGFACT tests factual recall, while KGFUNC tests arithmetic function acquisition. Each includes evaluation splits designed to probe retention beyond direct training performance.

  • Dataset Overview: The datasets target two distinct knowledge forms: factual knowledge from a simulated world graph and arithmetic function learning through novel digit-level operations.KGFACT and KGFUNC provide complementary controlled settings for studying knowledge injection and forgetting.
  • KGFACT (Factual Recall): KGFACT uses novel entities unseen during pretraining, organized across semantic domains with randomly assigned directed relations and unique query answers.The graph contains D domains, each with E entities, and relational edges defined for ordered domain pairs.
  • KGFACT (Factual Recall): KGFACT-RETRIEVAL prepends ground-truth statements and same-graph distractors to evaluate retrieval while separating parametric retention failures from reasoning failures.The split complements direct recall testing of trained atomic facts.
  • KGFUNC (Arithmetic Function Acquisition): KGFUNC consists of opaque-labeled deterministic digit-level operations mapping inputs and outputs in [0, 99999], with each training example providing 10-shot input-output pairs.Operations are compositions of digit-level primitives, and opaque labels prevent reliance on surface cues.
  • KGFUNC (Arithmetic Function Acquisition): KGFUNC-UNSEEN holds out simple operations, such as digit-sum and reverse-number, to test whether fine-tuning on novel operations degrades pre-existing arithmetic capabilities.The held-out operations are unseen during training but inferable from few-shot examples.

5 Experimental Setup

The experiments study factual retrieval and function acquisition using controlled synthetic datasets, SimpleQA, and general-domain benchmarks across three Qwen3 model scales. They compare Base, SFT, OPSD, and MIXSD under specified rollout and mixing settings.

  • Datasets: KGFACT-SMALL spans 5 domains with 10 entities each, while KGFACT-LARGE spans 7 domains with 25 entities each and tests context-based fact retrieval.Each retrieval instance includes 50 additional atomic facts involving the query entities and has a paired test query targeting the same fact.
  • Datasets: KGFUNC covers 7 operations with 1,600 training and 175 test instances per operation, plus 20 unseen operations for generalization.Examples provide 10-shot input-output pairs; the unseen split contains 500 total instances.
  • Datasets: The evaluation additionally uses 4,326 SimpleQA questions and general-domain tests spanning math, code generation, and knowledge understanding.The benchmarks include AIME2024, MATH500, GSM8K, HumanEval, and MMLU.
  • Models: Experiments benchmark Qwen3-1.7B, Qwen3-4B-Instruct-2507, and Qwen3-8B to examine performance across model scales.The models provide different sizes while controlling for architectural and training differences.
  • Methods: Methods compare Base, SFT, OPSD, and MIXSD; OPSD uses 8 rollouts per query, while MIXSD trains on Bernoulli-mixed rollouts with λ ∈{0, 0.3, 0.5, 0.7}.SFT applies NLL to canonical targets, whereas OPSD provides token-level KL supervision from a context-aware teacher.

6 Main Results

Across four corpora and three model scales, MixSD improves the trade-off between memorizing injected knowledge and retaining pre-existing capabilities. It consistently outperforms SFT and OPSD on held-out capability benchmarks while maintaining strong training performance, with larger models and higher mixing rates favoring retention.

  • SFT baseline: SFT achieves near-perfect training performance but reduces average held-out capability scores by 30-40% on KGFACT-SMALL and nearly collapses KGFUNC-UNSEEN generalization.SFT performs well on in-domain objectives while substantially degrading held-out capabilities and generalization to unseen operations.
  • OPSD baseline: 5.1 is OPSD’s average held-out capability score versus 14.3 for SFT on KGFACT-SMALL with Qwen3-1.7B.OPSD can preserve more capability than SFT in some settings, but its performance varies across datasets and model scales; the default uses n=8 rollouts per prompt.
  • MixSD results: MixSD maintains strong training performance while preserving substantially more existing capability across datasets and model scales, consistently outperforming SFT and OPSD on held-out benchmarks.On KGFACT-SMALL, MixSD traces a better Pareto frontier between injected-knowledge memorization and held-out capability retention; SimpleQA results show the same ordering.
  • Scaling: Larger models exhibit substantially less forgetting than smaller ones, although MixSD improves over SFT at every model size.The paper hypothesizes that mixed supervision requires sufficient existing capability to integrate tokens from different conditional distributions without destabilizing generation.
  • Mixing rate: The mixing rate λ balances memorization and capability retention: smaller values emphasize expert-conditioned supervision, whereas larger values add naive tokens that anchor the model to its prior.This trade-off is reported as consistent across datasets, with λ providing a direct control over the injection-retention balance.

7 Discussion

The discussion attributes forgetting more to directional movement into Fisher-sensitive parameter regions than to update magnitude alone, and shows that MixSD’s lower-NLL targets support retention across settings. It also examines qualitative failure modes, cross-family generality, and knowledge editing, where interference limits effectiveness.

  • Parameter-space mechanism: Fisher-sensitive directions, rather than displacement magnitude alone, are proposed as a better explanation of catastrophic forgetting.The diagonal empirical Fisher identifies parameters whose perturbations strongly affect base-model likelihood; R > 1 denotes updates concentrated in sensitive directions.
  • Parameter-space mechanism: Displacement magnitude alone weakly predicts forgetting, with correlations of +0.34, +0.02, and +0.10 for 1.7B, 4B, and 8B models.These correlations use approximately 40 checkpoints per model size, and SFT and MixSD often show similar parameter displacement despite different forgetting outcomes.
  • Token-level training signal: 27–42% of SFT targets contain base-model NLL above 8 on knowledge datasets, versus 4–8% for MixSD targets.Across model scales, increasing λ shifts MixSD targets toward lower base-model NLL because more tokens come from model rollouts.
  • Cross-family generality: On Llama-3.2-1B-Instruct, MixSD with λ=0.5 retains 78% of the base model’s average held-out score at 98% training accuracy, versus 21% retention for SFT.This evaluation uses KGFACT-SMALL and reports the same qualitative pattern as the Qwen experiments.
  • Knowledge editing: Knowledge editing remains challenging because modifying many related facts can cause destructive interference when entities are shared across multiple knowledge triples.The discussion evaluates MixSD on MQuAKE, which requires overwriting facts already stored in model parameters.

8 Conclusion

MIXSD is an external-teacher-free method for forgetting-aware knowledge injection that uses mixed self-generated, distribution-aligned targets. The results indicate that this approach improves the memorization-retention trade-off by addressing mismatch between external supervision and the model’s native autoregressive distribution.

  • 8 Conclusion: MIXSD provides a simple external-teacher-free approach to forgetting-aware knowledge injection in large language models.The method is presented as a knowledge-injection technique designed to account for forgetting.
  • 8 Conclusion: The paper attributes part of catastrophic forgetting to mismatch between externally authored targets and the model’s native autoregressive distribution.This mismatch motivates constructing supervision that better reflects the model’s own distribution.
  • 8 Conclusion: MIXSD constructs distribution-aligned supervision through mixed self-generated targets and consistently improves the memorization-retention trade-off across multiple datase.The supplied passage states that the improvement is consistent across multiple datase, though the text is truncated there.

9 Limitations

MixSD has practical and scaling limitations: its performance depends on an additional mixing hyperparameter, its behavior on substantially larger models is unknown, and mixed-rollout generation adds preprocessing cost.

  • Hyperparameter sensitivity: λ = 0.3 performs consistently well, but the optimal value may differ across tasks.MixSD introduces λ as an additional hyperparameter.
  • Model scaling: The largest evaluated backbone is 8B parameters, leaving scaling to 70B-and-beyond models unresolved.The paper identifies substantially larger-model scaling as an open question.
  • Computational overhead: Mixed-rollout generation adds a one-time preprocessing cost over standard SFT, though it is substantially smaller than on-policy self-distillation’s per-step cost.The extra cost is incurred before training rather than at every optimization step.

A Experimental Results on KGFACT-LARGE and SimpleQA · B Additional Related Work: Evaluation Datasets for Knowledge Injection · C Broader Impacts

The paper reports MixSD evaluations on KGFACT-LARGE and SimpleQA, contextualizes dataset choices for knowledge-injection evaluation, and discusses potential accessibility and safety benefits. These sections emphasize controlled assessment of acquisition, retention, and knowledge editing.

  • A Experimental Results on KGFACT-LARGE and SimpleQA: The paper presents KGFACT-LARGE and SimpleQA results in Tables 3 and 4, respectively.These tables organize the experimental results for the two evaluation settings.
  • A Experimental Results on KGFACT-LARGE and SimpleQA: Table 3 reports KGFACT-LARGE results across three model backbones, with values expressed as percentages and best and second-best methods marked.Methods and column semantics follow Table 1; Base is excluded from best-method comparisons.
  • A Experimental Results on KGFACT-LARGE and SimpleQA: Table 4 evaluates SimpleQA across three model backbones, including training accuracy and five unrelated held-out capability benchmarks summarized by an unweighted average.The Train column uses the 4,326-question SimpleQA training split, and all values are percentages.
  • B Additional Related Work: Evaluation Datasets for Knowledge Injection: Evaluation should distinguish genuine acquisition from memorization, annotation artifacts, and noisy supervision; SimpleQA provides objectively verifiable answers for short information-seeking questions.SimpleQA is useful for testing factual recall after fine-tuning but does not provide controlled target knowledge to inject.
  • B Additional Related Work: Evaluation Datasets for Knowledge Injection: Knowledge-editing benchmarks test explicit factual updates: ZsRE and CounterFact assess individual-fact revision and locality, while MQuAKE and RippleEdits probe related or multi-hop propagation.These datasets may remain confounded by real-world facts, knowledge bases, or generated counterfactual variants.
  • C Broader Impacts: MixSD may lower data and compute costs for specialized-domain adaptation, broaden accessibility, and help preserve safety-relevant behaviors such as instruction following and refusal calibration.These are presented as potential broader impacts of improving the memorization-retention trade-off.

D Additional Experiment Setup Details … E.2 High-NLL Token Counts

The appendix specifies training, rollout, checkpoint, and Fisher-estimation procedures, then analyzes high-NLL tokens as memorization targets and measures their overlap across SFT and MIXSD targets. These details support comparisons of capability retention, parameter sensitivity, and distribution-aligned factual supervision.

  • D Additional Experiment Setup Details: SFT and MIXSD use one rollout, whereas OPSD uses eight, a fixed base-model teacher, and a forward KL objective unless otherwise specified.Additional OPSD variants vary the loss, rollout count, and generation temperature.
  • D Additional Experiment Setup Details: MIXSD uses up to 10 retries with a rule-based verifier; failed examples are discarded, and this budget suffices for correct outputs on KGFACT and KGFUNC.SimpleQA examples that remain incorrect after all retries are also discarded.
  • D Additional Experiment Setup Details: Rollouts use VLLM V1 with up to 8192 new tokens and total sequence length 10,000; temperature T = 0 makes Bernoulli mixing the sole token-proposal randomness source.Experiments use 4×H100 GPUs and approximately 2000 GPU hours.
  • E.1 Fisher Information: Estimation and Diagnostics: The Fisher estimate uses N=979 reference examples from AIME-2024, MATH-500, GSM8K, HumanEval, and MMLU, with per-sample gradients accumulated and averaged.Batch size 1 provides per-sample gradients, which are cast to fp32 before squaring to avoid underflow.
  • E.1 Fisher Information: Estimation and Diagnostics: QF measures movement along sensitive base-model directions, while R compares update-direction curvature with isotropic curvature; R>1 indicates concentration and R<1 avoidance.Table 5 reports displacement and R across methods, model sizes, and task cohorts, while Table 6 correlates these metrics with mean forgetting.
  • E.2 High-NLL Token Counts: At NLL τ = 5, predicted probability is approximately 0.7%, whereas at τ = 8 it is approximately 0.034%, marking increasing uncertainty and near-complete prediction failure.Counts at both thresholds distinguish very uncertain tokens from tokens the base model essentially fails to predict.
  • E.2 High-NLL Token Counts: High-NLL memorization targets are factual tokens such as named entities, numerical values, dates, and domain-specific terms likely absent from the base model’s pretraining distribution.The procedure extracts these tokens from SFT ground-truth targets before comparing them with MIXSD targets.
  • E.2 High-NLL Token Counts: 81–98% of high-NLL token types in SFT ground-truth sequences also appear in corresponding MIXSD targets, regardless of absolute position.The overlap indicates that factual content is largely preserved while surrounding phrasing changes across training configurations.

E.3 Error Analysis: Per-Method Breakdown … E.6 KL-based Distillation Ablations on KGFACT-SMALL

Across error analysis, generalization, knowledge editing, and distillation ablations, MIXSD preserves base-model behavior while maintaining strong knowledge acquisition. Its main trade-offs are reduced collapse and leakage, improved retention, and stronger NLL-based acquisition than KL-based variants.

  • E.3 Error Analysis: Per-Method Breakdown: The error taxonomy assigns incorrect responses in priority order format > leakage > collapse > genuine across five benchmarks and multiple methods, backbones, and training corpora.Format denotes unparsable outputs, leakage denotes unrelated fictional KGFACT entities, collapse denotes parseable answers lacking reasoning, and genuine denotes ordinary reasoning errors.
  • E.3 Error Analysis: Per-Method Breakdown: 96.0% of SFT responses on AIME-2024 collapse, while 72.6% of HumanEval generations are boxed stubs without code.On Qwen3-1.7B / MMLU, collapse reaches 59.9% and leakage 4.3%, including fictional-entity emissions unrelated to the prompt.
  • E.3 Error Analysis: Per-Method Breakdown: OPSD largely eliminates collapse but raises format errors, reaching 4.3% unparsable responses on Qwen3-1.7B / KGFACT-SMALL MMLU.The analysis attributes this pattern to on-policy KL training occasionally destabilizing instruction-following behavior.
  • E.3 Error Analysis: Per-Method Breakdown: MIXSD preserves the base model’s genuine reasoning errors while keeping collapse and leakage near zero across benchmarks and backbones.At λ=0.7, format errors increase slightly because long chain-of-thought generations may fail to terminate with a valid \boxed{} answer.
  • E.4 Llama-3.2-1B-Instruct on KGFACT-SMALL: 98% training accuracy for SFT on Llama-3.2-1B-Instruct accompanies a held-out average decline from 6.8 to 1.4, whereas MIXSD reaches 4.8 and 5.3 at λ=0.3 and λ=0.5.OPSD achieves only 31% training accuracy in this setting.
  • E.5 Knowledge Editing on MQuAKE: 93–99% editing accuracy for MIXSD at λ=0.3 preserves held-out averages of 17.6, 76.0, and 70.0, unlike SFT’s 45–86% capability reduction.These results span the three Qwen3 scales and compare knowledge-editing accuracy with held-out capability retention.
  • E.5 Knowledge Editing on MQuAKE: 84–95% retrieval-augmented chain-of-thought performance for MIXSD exceeds SFT’s 41–58% across all model scales.MEMIT instead preserves held-out capabilities within roughly 1–2 points of the base model but achieves only 53–70% editing accuracy.
  • E.6 KL-based Distillation Ablations on KGFACT-SMALL: 96–100% training accuracy with a single greedy rollout makes NLL the default, while MIXSD-KL reaches only 34–60% despite comparable or slightly stronger held-out preservation.The KL variants use top-64 forward-KL objectives under different student rollout settings, whereas MIXSD uses the NLL counterpart at the same λ.
Loading 2605.16865v3…