PaperlayerView exact revision on arXiv

Source-linked AI summary

SelFusion: Self-distillation for Diffusion Language Models

Hyeongsoo Lim, Jinyoung Kim, Eunseo Seo, Minho Jang, Jiwon Yoon

arXiv:2608.22898v1cs.CL

TL;DR

DLMs offer lower-latency generation than autoregressive LLMs but typically produce lower-quality text, while conventional distillation provides limited or degraded gains. SelFusion uses bidirectional self-distillation between differently masked modes and consistently outperforms external-teacher KD, sometimes surpassing the teachers.

  • Problem

    DLMs typically underperform LLMs in generation quality, while conventional KD for DLMs yields marginal gains or performance degradation.

  • Method

    SelFusion performs bidirectional knowledge distillation between easy and hard modes created by different masking ratios within one DLM.

  • Results

    SelFusion consistently outperforms conventional KD baselines across configurations, improving over the strongest baseline by 2 to 4 points and surpassing LLM teachers on multiple benchmarks.

  • Takeaways & Limitations

    SelFusion provides a practical DLM distillation framework that transfers knowledge within one model without external LLM or DLM teachers.

  • Takeaways & Limitations

    Validation is limited by the availability of DLM backbones and focuses only on English instruction tuning, leaving broader domains and multilingual settings untested.

Abstract

from arXiv · show

Diffusion language models (DLMs) alleviate the inherent latency bottleneck of autoregressive (AR) large language models (LLMs), but their degraded generation quality limits practical applicability. Although knowledge distillation (KD) can be a promising direction for improving performance, we empirically find that naively applying conventional KD yields only marginal gains, or even degrades generation quality. Based on these observations, we propose a novel self-distillation framework for DLMs, namely SelFusion. To enable effective KD without an external teacher model, SelFusion performs two forward passes with different masking levels, defining the hard mode with a larger masking probability and the easy mode with a smaller masking probability. However, the easy mode is not always more accurate than the hard mode and can be overconfident on incorrect tokens. Thus, we introduce bidirectional KD between the two modes, which can dynamically determine the distillation direction based on token-level correctness. Experimental results on instruction-following tasks show that the proposed self-distillation substantially outperforms other KD methods with external LLM and DLM teachers. In many configurations, the student trained with SelFusion even surpasses the performance of the LLM teacher, providing a practical path toward improving DLM generation quality. Source code can be found at https://github.com/scai-research/SelFusion_official

1 Introduction

DLMs offer faster inference than autoregressive LLMs but generally lag in generation quality, while conventional knowledge distillation provides limited or degraded gains. SelFusion addresses this gap through bidirectional self-distillation between easy and hard masking modes, consistently outperforming external-teacher baselines and sometimes surpassing the teachers themselves.

  • DLMs use parallel decoding for faster inference but typically underperform autoregressive LLMs in generation quality because of their non-autoregressive nature.This tradeoff motivates methods that improve DLM quality without sacrificing their latency advantages.
  • Distilling DLM students from LLM or DLM teachers yields only marginal gains or performance degradation, with autoregressive-to-non-autoregressive distribution mismatch limiting logit-level transfer.Sequence-level KD improves over logit-level supervision but remains limited in performance gains.
  • SelFusion uses one DLM with two masking-based forward modes: an easy mode with lower masking and a hard mode with higher masking.The easy mode is expected to produce more accurate predictions, but it is not always more accurate and can be overconfident on incorrect tokens.
  • SelFusion introduces bidirectional knowledge distillation that dynamically selects the distillation direction according to token-level correctness.This mechanism transfers knowledge between the easy and hard modes within a single model rather than requiring an external teacher.
  • SelFusion consistently outperforms conventional KD baselines across configurations and can surpass both DLM and LLM teacher models on instruction-following benchmarks.On the Dolly dataset, existing KD methods underperform the SFT baseline, whereas SelFusion significantly outperforms all same-size baselines.

2 Related Work

Related work characterizes DLMs by continuous or discrete generation, with masked diffusion emerging as the dominant discrete approach. Although DLMs support fast parallel inference, they still trail AR models in generation quality, motivating knowledge distillation methods that transfer teacher knowledge through logits or sequences.

  • DLMs: DLMs comprise continuous methods that embed tokens in continuous space and discrete methods that operate directly in token space.Masked diffusion is predominantly adopted among discrete approaches, progressively masking tokens before reverse-process prediction.
  • DLMs: Increasing denoising steps generally improves DLM generation quality, while LLaDA demonstrates 8B-parameter masked diffusion with parallel generation for fast inference.These advances establish the potential of DLMs for efficient text generation.
  • DLMs: 10-32%: recent DLMs underperform AR baselines by 10-32% in perplexity, highlighting the need for further improvement.The passage attributes this gap to generation quality despite DLM advances.
  • Knowledge Distillation: Knowledge distillation transfers knowledge from a teacher model to improve model performance, including logit-level and sequence-level approaches.Logit-level KD matches output logits, whereas sequence-level KD trains on teacher-generated sequences.

3 Methodology

SelFusion addresses ineffective external-teacher distillation by constructing easy and hard modes within one DLM and adaptively distilling between them at the token level. Its training objective jointly combines both modes’ diffusion losses with bidirectional knowledge distillation, using reliability-based direction selection and calibration to mitigate overconfidence.

  • Motivation: External-teacher KD is limited by distribution mismatch and weak DLM teachers, producing marginal gains despite higher training costs.AR models assign near 100% probability to their top-1 token, whereas DLMs assign approximately 60% at a 50% masking ratio; sequence-level DLM distillation also remains limited.
  • Two Modes with Different Masking: SelFusion replaces external teachers with self-distillation between easy and hard modes created by applying different masking levels within the same model.The hard mode follows the original random masking scheme, while the easy mode uses a lower masking ratio and more visible context.
  • Two Modes with Different Masking: The easy mode generally provides relatively more accurate knowledge because it predicts with less masking, but it is not always correct.This limitation motivates adaptively selecting the distillation target rather than always distilling from easy to hard.
  • Bidirectional KD: Bidirectional KD selects the token-level distillation direction from correctness and confidence across three cases: both correct, both wrong, or only one correct.When both are correct, the higher-probability prediction is selected; when both are wrong, the mode assigning higher probability to the ground truth is selected; when one is correct, that mode is selected.
  • RMSNorm-based Logit Calibration: Lower masking produces more peaked distributions, and easy-mode overconfidence can affect incorrect tokens as well as correct ones.Both modes predict incorrectly in approximately 40% of cases, so easy-mode confidence may hinder effective distillation; SelFusion introduces RMSNorm-based logit calibration to address this issue.
  • Training Objective: The final objective jointly optimizes hard-mode diffusion loss, easy-mode diffusion loss, and token-wise bidirectional distillation loss.Because both modes share parameters, the combined loss updates them simultaneously in a single backward pass.

4 Experiments

Experiments show that conventional teacher-based distillation often harms DLM performance, whereas SelFusion improves benchmark quality without an external teacher while reducing total training computation. Ablations and efficiency analyses support bidirectional, confidence-aware distillation as a key component of the method.

  • Distillation Baselines: Conventional LLM-to-DLM logit-level KD reduced Dolly performance to 15.11 from 18.77 for the DLM SFT baseline and was ineffective across benchmarks.Sequence-level KD also failed to surpass the DLM SFT baseline on most benchmarks.
  • Distillation Baselines: DLM-to-DLM distillation also degraded performance, with logit-level KD reaching 19.87 versus 23.84 for the DLM SFT model.The results indicated that the DLM teacher’s generation quality was insufficient for beneficial sequence- or logit-level distillation.
  • Main Results: SelFusion improved performance by 1.5 to 3 points over the DLM SFT baseline across all five benchmarks and by 2 to 4 points over the strongest competing baseline.The improvement over the competing baseline corresponded to approximately 16% relative improvement, and SelFusion surpassed the LLM teacher on multiple benchmarks.
  • Training Cost: SelFusion used about 2× more per-iteration computation than SFT but reduced overall training cost by roughly 2× by removing teacher training.The comparison used the same training configuration and 485 iterations for all methods.
  • Ablation Studies: Removing bidirectional KD or RMSNorm consistently degraded performance, and neither easy→hard nor hard→easy distillation matched the bidirectional setting.The method dynamically switches token-level distillation direction according to accuracy and confidence; hard→easy outperformed easy→hard on several benchmarks.
  • Inference Efficiency: At 32 diffusion steps, SelFusion generally outperformed baseline distillation methods and achieved the best average benchmark performance, while DLMs remained approximately 1.5× faster than LLMs.The latency advantage derives from DLM parallel token generation versus sequential LLM generation.

5 Conclusions

SelFusion is a self-distillation framework for effective logit-level knowledge distillation in diffusion language models. It uses simultaneous forward passes with different masking ratios to train two modes within a single model, outperforming conventional KD without external teachers.

  • SelFusion enables effective logit-level knowledge distillation for diffusion language models through a novel self-distillation framework.
  • Different masking ratios and simultaneous forward passes decompose the model into two modes with a more suitable learning distribution.
  • SelFusion outperforms conventional knowledge-distillation methods without relying on external teacher models.

Limitations

SelFusion’s validation is constrained by limited DLM backbone availability and an English-only instruction-tuning evaluation, while its teacher-free design remains practically advantageous.

  • Limited availability of DLM backbones constrains validation across a broader range of architectures.
  • Evaluation is restricted to English instruction tuning, leaving broader domains and multilingual settings for future work.
  • SelFusion does not rely on external LLM or DLM teacher models, providing a practical distillation advantage despite these limitations.

A Appendix · A.1 Training Configuration

The appendix presents detailed training configurations for the experiments. Learning rates were selected through grid search while accounting for model scale and architectural differences.

  • A.1 Training Configuration: The appendix reports detailed training configurations used in the experiments.
  • A.1 Training Configuration: Learning rates were determined through grid search.
  • A.1 Training Configuration: The search accounted for both model scale and architectural differences.
  • A.1 Training Configuration: For 1.4B teacher models, the explored learning rates were {1e-5, 5e-5, 1e-4, 2e-4}.This range was used for both AR and DLM architectures.
  • A.1 Training Configuration: The 1.4B teacher-model search covered both AR and DLM architectures.The passage attributes the wider range to their distinct architectural characteristics.
  • A.1 Training Configuration: For 0.472B DLM target models, the grid-search learning rates were {5e-5, 1e-4, 2e-4}.

A.2 Prompt Formatting for Instruction Tuning

Instruction-tuning examples are standardized into Dolly-style prompts, with separate templates depending on whether an input field is present. The input-present format includes contextual task wording, while the input-absent format omits that clause.

  • With input field: Examples containing an input field are converted into a Dolly-style prompt template.The standardized format begins with task instructions and further context from the input field.
  • With input field: The input-present template states: “Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.”It then uses the heading “### Instruction: {instruction}”.
  • Without input field: Examples without an input field use a shorter template stating: “Below is an instruction that describes a task. Write a response that appropriately completes the request.”The corresponding heading remains “### Instruction: {instruction}”.
  • Data availability: All experimental datasets are publicly available through Microsoft’s MiniLLM data release.The cited release is hosted in the Microsoft LMOps repository.

A.3 Masking Strategy for Easy and Hard Modes

SelFusion constructs paired easy and hard masking modes using the standard DLM noising procedure. The easy mode has lower masking and more visible context, while the hard mode has higher masking and reduced visibility, with independently sampled mask positions.

  • A.3 Masking Strategy for Easy and Hard Modes: Each example samples a noise level t ∼ U(0, 1), converts it to a masking probability, and independently masks tokens with that probability.Masked tokens are replaced with a dedicated mask token using vocabulary index V.
  • A.3 Masking Strategy for Easy and Hard Modes: The hard mode samples t_hard ∼ U(0, 1), then conditions the easy mode noise level so that t_easy ≤ t_hard.This coupling is applied within a single training step to construct paired modes.
  • A.3 Masking Strategy for Easy and Hard Modes: Because p_mask(t_easy) ≤ p_mask(t_hard), the easy mode uses lower masking and the hard mode uses higher masking.The easy mode therefore observes more visible context, whereas the hard mode operates under reduced visibility.
  • A.3 Masking Strategy for Easy and Hard Modes: Despite the conditioned masking ratios, masked token positions are sampled independently for the easy and hard modes.Thus, the two modes share an ordered masking level but not identical corruption locations.

A.4 Token level Accuracy Comparison Details · A.5 Generalization Across Tasks

The appendix examines token-level correctness patterns underlying bidirectional KD and evaluates SelFusion’s generalization on SAMSum. SelFusion outperforms SFT and existing KD methods across diffusion-step settings, with larger gains under the tighter 8-step budget.

  • A.4 Token level Accuracy Comparison Details: On hard-mode masked positions, only the hard mode predicts correctly in approximately 6% of cases, while the easy mode alone is correct in about 37%.These measurements assess cases where only one mode correctly predicts the masked token.
  • A.4 Token level Accuracy Comparison Details: Figure 8 measures, for each mode, the probability that only one mode correctly predicts a masked token.The analysis tests whether bidirectional KD is activated during training.
  • A.5 Generalization Across Tasks: SelFusion is evaluated beyond instruction-following on the SAMSum summarization benchmark.The benchmark is attributed to Gliwa et al. (2019).
  • A.5 Generalization Across Tasks: On SAMSum, SelFusion outperforms SFT and existing KD methods in both the 8-step and 16-step settings.Results are measured using ROUGE-L across diffusion steps.
  • A.5 Generalization Across Tasks: SelFusion consistently achieves the best SAMSum ROUGE-L results across the evaluated diffusion-step settings.Table 7 marks the best result in bold.
  • A.5 Generalization Across Tasks: The improvement is more pronounced with 8 diffusion steps, indicating effectiveness under tighter inference budgets.The 8-step setting is described as more efficient than the 16-step setting.
Loading 2608.22898v1…