UniSD: Towards a Unified Self-Distillation Framework for Large Language Models

Yiqiao Jin; Yiyang Wang; Lucheng Fu; Yijia Xiao; Yinyi Luo; Haoxin Liu; B. Aditya Prakash; Josiah Hester; Jindong Wang; Srijan Kumar

Source-linked AI summary

UniSD: Towards a Unified Self-Distillation Framework for Large Language Models

Yiqiao Jin, Yiyang Wang, Lucheng Fu, Yijia Xiao, Yinyi Luo, Haoxin Liu, B. Aditya Prakash, Josiah Hester, Jindong Wang, Srijan Kumar

arXiv:2605.06597v2May 7, 2026cs.CLcs.AIcs.LG

TL;DR

Self-distillation for autoregressive LLMs lacks systematic evidence on supervision reliability, component roles, and interactions. UniSD unifies mechanisms across these dimensions, and its integrated variant achieves the strongest overall performance, improving the overall score by +5.4 points over the base model and +2.8 over the strongest baseline.

Problem
Self-distillation in autoregressive LLMs lacks systematic understanding of how to assess unstable self-supervision and how design choices contribute across tasks.
Method
UniSD systematically studies self-distillation through complementary mechanisms targeting supervision reliability, representation alignment, and training stability.
Results
+5.4 overall-score points over the base model and +2.8 over the strongest baseline, with UniSD* achieving the strongest performance across six benchmarks and six models.
Takeaways & Limitations
UniSD supports self-distillation as an efficient, controllable approach for LLM adaptation using self-derived supervision rather than stronger external teachers.
Takeaways & Limitations
The study mainly focuses on single-turn scenarios, leaving broader settings for future work.

Abstract

from arXiv · show

Self-distillation (SD) offers a promising path for adapting large language models (LLMs) without relying on stronger external teachers. However, SD in autoregressive LLMs remains challenging because self-generated trajectories are free-form, correctness is task-dependent, and plausible rationales can still provide unstable or unreliable supervision. Existing methods mainly examine isolated design choices, leaving their effectiveness, roles, and interactions unclear. In this paper, we propose UniSD, a unified framework to systematically study self-distillation. UniSD integrates complementary mechanisms that address supervision reliability, representation alignment, and training stability, including multi-teacher agreement, EMA teacher stabilization, token-level contrastive learning, feature matching, and divergence clipping. Across six benchmarks and six models from three model families, UniSD reveals when self-distillation improves over static imitation, which components drive the gains, and how these components interact across tasks. Guided by these insights, we construct UniSDfull, an integrated pipeline that combines complementary components and achieves the strongest overall performance, improving over the base model by +5.4 points and the strongest baseline by +2.8 points. Extensive evaluation highlights self-distillation as a practical and steerable approach for efficient LLM adaptation without stronger external teachers.

1 Introduction

UniSD presents a unified, reliability-aware framework for self-distillation in autoregressive LLMs, addressing the challenges of free-form generation and noisy self-supervision. It systematically analyzes complementary mechanisms across supervision reliability, representation alignment, and training stability, then combines them into an integrated variant for stronger in-domain and OOD adaptation without external teachers.

Challenges: Self-distillation is difficult because autoregressive LLMs generate free-form trajectories with task-dependent correctness and unreliable, unstable self-derived supervision.Outputs may be partially correct, stylistically different, or locally misleading, while on-policy trajectories expose models to their own errors.
This Work: UniSD is a unified and extensible framework that studies self-distillation through supervision reliability, representation alignment, and training stability.It frames self-distillation as reliability-aware self-correction over on-policy trajectories using comparison and supervision across multiple teacher views.
Contributions: UniSD evaluates six benchmarks and six models from three model families to identify component contributions and interactions affecting robustness, transfer, and retention.These findings guide UniSD∗, an integrated variant combining complementary components for improved in-domain and OOD performance without external teachers.
Framework: UniSD integrates Multi-Teacher Agreement, EMA Teacher, Token-Level Contrastive Learning, Feature Matching, and Divergence Clipping as modular objectives for controlled analysis.The modular design is extensible to additional strategies.

2 Method

UniSD is a unified framework for self-distillation in autoregressive LLMs that addresses unreliable self-derived supervision and temporal instability. It combines mechanisms for signal reliability, representation alignment, and training stability on on-policy student trajectories.

Framework: UniSD systematically studies LLM self-distillation, where the model learns from supervision derived from its own behavior rather than stronger external teachers.The framework targets open-ended generations whose self-distillation signals can be unstable and task-dependent.
Supervision reliability: Multi-teacher agreement evaluates whether self-derived supervision is stable across task-preserving teacher views at token and sequence granularities.Context variation supplies alternative perspectives using retrieved or randomly sampled few-shot examples or induced high-level instructions, without extra teacher copies.
Training stability: EMA Teacher smooths teacher targets across optimization steps, complementing agreement, which controls which signals are trusted within the current step.EMA reduces temporal drift that can propagate transient errors or overconfident predictions into later updates.
Supervision reliability: Token-level contrastive learning separates positive supervision from plausible but incorrect alternatives by encouraging student trajectories to favor correct signals.The margin-based objective is particularly motivated by examples such as code solutions that share surface structure but differ in key implementation details.
Representation alignment and training stability: Feature matching aligns selected internal representations, while divergence clipping limits rare high-divergence token losses without changing teacher construction or agreement estimation.The implementation matches final-layer hidden states on completion tokens, and clipping recovers the unclipped objective when no threshold is specified.

3 Evaluation

UniSD is evaluated across six benchmarks and six language models spanning three families, with analyses of component effectiveness, agreement design, cross-family generalization, and distributional retention. The integrated UniSD* pipeline achieves the strongest overall results while improving task performance and preserving alignment with the base model.

Evaluation Setup: Evaluation spans six benchmarks across four task categories and six LLMs from three model families, centered on Qwen2.5-7B-Instruct with scale and cross-family tests.Qwen2.5-0.5/1.5/3B-Instruct, Llama-3.1-8B-Instruct, and gemma-3-4b-it extend the evaluation beyond the primary model.
Main Results: Static imitation improves format learning but is unreliable overall, raising ToolAlpaca by +4.4 while degrading ScienceQA, GPQA, MBPP, and HumanEval.SFT yields only +0.7 on CoS-E, indicating that off-policy maximum-likelihood training can learn output conventions without consistently improving reasoning or coding.
Main Results: Token-level agreement achieves the strongest ScienceQA result (85.2) and ranks best or second-best on four of six datasets by down-weighting high-disagreement supervision.Sequence-level agreement is more conservative but more stable, reflecting a reliability-versus-adaptivity trade-off across agreement granularity.
Main Results: UniSD* raises the overall score from 67.9 to 73.3 (+5.4) and outperforms the strongest baseline GKD by +2.8.It is best or tied-best on MBPP, ToolAlpaca, GPQA, and HumanEval, and second-best on ScienceQA and CoS-E.
Sensitivity Analysis: Agreement performance varies non-monotonically with context count K, task, granularity, and penalty γ rather than improving monotonically with more auxiliary contexts.For example, sequence-level agreement peaks at K = 3 on ScienceQA with γ = 0.01 (85.2), whereas token-level agreement peaks at K = 7 on ScienceQA with γ = 0.01 (84.4).
Generalization and Retention: UniSD* generalizes across Qwen2.5, Llama-3.1, and Gemma-3, improving over base models by +5.4, +3.1, and +2.2, respectively, while reducing distributional drift versus SFT.On Qwen2.5-7B, UniSD* improves accuracy from 80.8 to 85.0 and lowers mean token-level JSD from 0.054 for SFT to 0.041.

4 Related Work

Related work situates UniSD at the intersection of continual or on-policy learning and knowledge distillation for LLMs. Prior approaches address adaptation through model-induced trajectories or by transferring predictions, representations, rationales, and reasoning traces from teachers.

Continual Learning and On-Policy Learning: Continual learning adapts models to new knowledge and skills while preserving existing capabilities, with catastrophic forgetting as a central challenge.In LLM post-training, standard supervised fine-tuning is off-policy because it uses fixed expert demonstrations rather than trajectories induced by the current policy.
Knowledge Distillation and Self-Distillation for LLMs: Knowledge distillation transfers teacher knowledge by matching predictions, logits, hidden states, generated outputs, or reasoning traces.Prior LLM work distills token-level distributions, attention patterns, intermediate representations, rationales, and step-by-step reasoning traces, while recent variants include VLA-OPD, SCOPE, and StableOPD.

5 Conclusion · A Algorithm Details of UniSD

UniSD is presented as a unified framework for self-distillation in LLMs without stronger external teachers, identifying component contributions and interactions across tasks. Its integrated pipeline achieves the strongest overall performance, while Algorithm 1 provides the detailed procedure.

5 Conclusion: UniSD studies self-distillation in LLMs without relying on stronger external teachers.
5 Conclusion: Across six benchmarks and six models from three families, UniSD identifies which components drive self-distillation gains.
5 Conclusion: UniSD also reveals how self-distillation components interact across tasks.
5 Conclusion: These findings motivate UniSD∗, an integrated pipeline with the strongest overall performance.
5 Conclusion: The authors position UniSD as a foundation for future work on efficient, controllable self-distillation of LLMs.
A Algorithm Details of UniSD: Algorithm 1 presents the detailed procedure of UniSD.

B Additional Experiments · B.1 Training Time

UniSD’s training cost is driven primarily by teacher-conditioned scoring passes, making agreement-based variants substantially more expensive than standard SFT. Lightweight stabilization methods reduce this overhead, while EMA, contrastive learning, and feature matching occupy an intermediate cost regime.

B.1 Training Time: Agreement-based UniSD variants incur substantially higher training costs because each on-policy completion requires multiple teacher-conditioned scoring passes.Standard SFT is the cheapest baseline, and the distillation loss itself is not the main cost driver.
B.1 Training Time: Clipping and feature matching add much smaller overhead than full multi-context agreement.These methods apply lightweight stabilization on top of a single teacher signal.
B.1 Training Time: EMA, contrastive learning, and joint matching require additional teacher or auxiliary forward passes but avoid agreement’s aggressive multiplication of context-conditioned scoring.They therefore lie between lightweight stabilization and full agreement in training overhead.
B.1 Training Time: UniSD samples x ∼ D and rolls out an on-policy trajectory ˆy from the student policy πθ(· | x) during each training iteration.The algorithm repeats this process while training has not converged.
B.1 Training Time: When enabled, the EMA teacher is incorporated into the UniSD training loop as an additional teacher mechanism.The algorithm explicitly checks whether the EMA teacher is enabled during training.
B.1 Training Time: UniSD uses a primary teacher under the primary condition c∗ and evaluates teacher-conditioned token probabilities across auxiliary conditions.These auxiliary scoring operations explain why agreement-based variants require more computation.
B.1 Training Time: Token-level contrastive learning is an optional component in the UniSD algorithm.Its activation introduces an auxiliary computation branch within training.
B.1 Training Time: Feature matching is another optional UniSD component that provides a representation-level auxiliary signal.The algorithm checks whether feature matching is enabled before applying this signal.

B.2 Resource Consumption

UniSD’s resource analysis uses runtime-derived energy and emissions estimates under fixed hardware, utilization, facility-overhead, and carbon-intensity assumptions. Single-teacher variants are most efficient, while agreement-based variants add modest throughput and memory overhead.

Estimation methodology: Resource consumption is estimated from runtime, GPU count, assumed utilization, PUE, and carbon intensity rather than facility-level power measurements.The estimates support relative comparison under fixed assumptions.
Token-normalized cost: 0.08 kWh per million tokens: Match (Joint) is the most efficient variant, compared with 0.10–0.11 kWh per million tokens for Contrast, EMA, and Match (Repr.).The passage identifies single-teacher stabilization methods as more efficient overall.
Throughput and overhead: 2.32–3.22M tokens/GPU-hour: variants preserve high throughput despite representation, contrastive, or temporal stabilization.The added mechanisms incur only modest overhead according to the resource analysis.
Throughput and overhead: 13–17GB (+21–28%): agreement-based variants increase peak memory over single-teacher variants because reliability estimation re-scores completions under multiple auxiliary contexts.Agreement-based variants consume 0.16–0.18 kWh per million tokens.

C Additional Experimental Details

The paper uses a standardized LoRA and AdamW training setup with specified optimization, precision, and sequence-length configurations. Evaluations use greedy vLLM decoding and task-specific metrics across UniSD and self-distillation baselines.

Training Configuration: All methods use LoRA with rank 64, alpha 128, dropout 0.05, and AdamW with β1 = 0.9 and β2 = 0.999.Training otherwise runs for 1 epoch.
Training Configuration: Training uses a 2e-5 learning rate, cosine decay, 10% warmup, four-step gradient accumulation, and bf16 mixed precision.Maximum prompt and completion lengths are 3072 and 1024 tokens, respectively.
Training Configuration: On-policy completions use colocated vLLM generation at temperature 0.7, with maximum prompt and completion lengths of 3072 and 1024 tokens.These settings define the completion-generation procedure and sequence limits.
Evaluation: Evaluation uses vLLM greedy decoding at temperature τ = 0.0 and compares UniSD with SFT, SDFT, GKD, SSD, and OPSD.The baselines include state-of-the-art self-distillation methods.
Evaluation: Code-generation results report pass@1 through sandboxed execution with a 10-second timeout, while multiple-choice results report accuracy with automatic answer extraction.The code tasks are MBPP and HumanEval; the multiple-choice tasks are ScienceQA, CoS-E, and GPQA.

D Broader Impact · E Ethical Considerations · F AI Assistants Usage

UniSD frames self-distillation as a potentially lower-cost, privacy-sensitive approach to LLM adaptation while acknowledging inherited safety risks. The paper also reports limited AI-assistant use focused on manuscript preparation and experimental workflows, with authors retaining responsibility for design and verification.

D Broader Impact: UniSD derives supervision from LLMs’ own behavior rather than stronger external teachers.This framing targets self-distillation without relying on stronger external supervision.
D Broader Impact: Self-distillation may lower post-training cost and access barriers for academic groups, smaller organizations, and resource-constrained settings.The passage specifically identifies these groups as potential beneficiaries.
D Broader Impact: UniSD may reduce the need to transmit in-domain data to external models, supporting privacy-sensitive or local adaptation.The passage presents reduced external data transmission as an appeal of the approach.
D Broader Impact: Adding more contexts does not consistently improve accuracy, according to the sensitivity analysis over context count k and agreement weight γ.The figure examines sensitivity to both k and γ.
E Ethical Considerations: UniSD inherits the base model’s potential factual errors, social biases, and unsafe behaviors.Reliability weighting, divergence clipping, and stabilization may reduce reinforcement of unreliable signals but are not substitutes for standard safety procedures.
E Ethical Considerations: Adapted models should undergo safety, bias, factuality, and domain-specific risk evaluation before deployment.The passage emphasizes this evaluation especially for human-centric applications.
F AI Assistants Usage: AI assistants supported language refinement, clarity, organization, and limited experimental workflows in manuscript preparation.The authors made the experimental design and methodological choices.
F AI Assistants Usage: The authors manually checked and verified all results, analyses, and final content.This responsibility remained with the authors despite auxiliary AI-assistant use.

G Limitations and Future Work

UniSD is evaluated primarily in controlled single-turn settings, leaving extensions to long-horizon agentic tasks, finer-grained trajectory evaluation, and broader self-supervision objectives as future work.

Scope: UniSD mainly studies single-turn scenarios, providing a controlled setting for isolating self-distillation effects and motivating future extensions.The authors frame this scope as a starting point for future directions.
Long-Horizon Agentic Settings: Long-horizon agentic tasks would test whether reliability-weighted self-correction can provide stable supervision under sparse, delayed feedback and interdependent decisions.These settings involve extended trajectories with multiple interdependent decisions.
Finer-Grained Trajectory Evaluation: Finer-grained evaluation could credit partially correct reasoning and useful intermediate steps beyond final-answer accuracy.Current benchmark protocols score final answers as correct or incorrect.
Broader Self-Supervision Objectives: Future self-supervision work could explore richer contrastive objectives, alternative disagreement measures, and prompt optimization to improve self-derived supervision.UniSD’s five complementary mechanisms make the framework naturally extensible.