Source-linked AI summary

Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models

Gongbo Zhang, Wen Wang, Ye Tian, Li Yuan

arXiv:2604.26951v1cs.CLcs.AIcs.LG

TL;DR

Existing dLLM distillation methods compress inference steps within one architecture, leaving transfer across differing architectures, attention mechanisms, and tokenizers unresolved. TIDE addresses this gap with three modular components for scheduling, contextual enrichment, and cross-tokenizer alignment, and improves performance when distilling large teachers into a 0.6B student. The best configuration gains 1.53 points on the eight-benchmark average, while HumanEval reaches 48.78 versus 32.3 for the same-size AR model.

  • Problem

    Existing dLLM distillation methods focus on step compression within a single architecture, leaving cross-architecture transfer across differing architectures, attention mechanisms, and tokenizers unexplored.

  • Method

    TIDE distills heterogeneous dLLM teachers into a smaller student using TIDAL for timestep-aware scheduling, COMPDEMO for complementary contextual enrichment, and Reverse CALM for cross-tokenizer alignment.

  • Results

    +1.53 on the eight-benchmark average is achieved over the non-distilled baseline, with HumanEval reaching 48.78 versus 32.3 for the same-size AR model.

  • Takeaways & Limitations

    Cross-architecture dLLM distillation is effective across distinct tokenizers and attention mechanisms, with the two pipelines favoring different component configurations.

  • Takeaways & Limitations

    COMPDEMO increases overall training time by approximately 50% because it doubles the frozen teacher’s forward passes.

Abstract

from arXiv · show

Diffusion large language models (dLLMs) offer parallel decoding and bidirectional context, but state-of-the-art dLLMs require billions of parameters for competitive performance. While existing distillation methods for dLLMs reduce inference steps within a single architecture, none address cross-architecture knowledge transfer, in which the teacher and student differ in architecture, attention mechanism, and tokenizer. We present TIDE, the first framework for cross-architecture dLLM distillation, comprising three modular components: (1) TIDAL, which jointly modulates distillation strength across training progress and diffusion timestep to account for the teacher's noise-dependent reliability; (2) CompDemo, which enriches the teacher's context via complementary mask splitting to improve predictions under heavy masking; and (3) Reverse CALM, a cross-tokenizer objective that inverts chunk-level likelihood matching, yielding bounded gradients and dual-end noise filtering. Distilling 8B dense and 16B MoE teachers into a 0.6B student via two heterogeneous pipelines outperforms the baseline by an average of 1.53 points across eight benchmarks, yielding notable gains in code generation, where HumanEval scores reach 48.78 compared to 32.3 for the AR baseline.

1 Introduction

TIDE addresses the unexplored problem of distilling dLLMs across architectures, attention mechanisms, and tokenizers by integrating three modular components. Across heterogeneous teacher–student pipelines, it improves benchmark performance while transferring large teachers into a much smaller student.

  • Research gap: Existing dLLM distillation methods focus on inference-step compression within a single architecture, leaving cross-architecture transfer unexplored.Cross-architecture transfer must handle heterogeneous architectures, attention mechanisms, and tokenizers.
  • Framework: TIDE integrates TIDAL, COMPDEMO, and Reverse CALM to address temporal, contextual, and vocabulary barriers in heterogeneous dLLM distillation.TIDAL schedules distillation, COMPDEMO enriches teacher context, and Reverse CALM aligns distinct vocabularies.
  • Framework: TIDAL modulates distillation strength across training progress and diffusion timestep, selectively using teacher signals when they are reliable.The framework distinguishes timestep-dependent teacher reliability from the evolving ability of the student to absorb knowledge.
  • Evaluation: TIDE transfers knowledge from 16B MoE and 8B dense teachers into a 0.6B student across two heterogeneous pipelines.The framework includes both cross-tokenizer and shared-tokenizer distillation settings.
  • Results: +1.53 points on the eight-benchmark average is achieved over the non-distilled baseline, with HumanEval reaching 48.78 versus 32.3 for the same-size AR model.The best configuration scores 34.20 versus 32.67 on the eight-benchmark average.

2 Method

TIDE distills across heterogeneous dLLM architectures using timestep- and training-aware scheduling, enriched teacher context, and tokenizer-compatible objectives. Its components address unreliable teacher signals, context scarcity under masking, and vocabulary misalignment.

  • Framework: TIDE distills a large teacher dLLM into a smaller student even when their architectures, attention mechanisms, and tokenizer vocabularies differ.The student predicts clean tokens at masked positions using ground-truth labels and teacher distributions.
  • TIDAL: TIDAL jointly schedules distillation strength over diffusion timestep and training progress.It reduces teacher influence at high noise and increases it later in training to avoid unreliable supervision and representation collapse.
  • TIDAL: At high noise, TIDAL makes the target student-dominated, whereas at low noise it permits stronger teacher supervision.The timestep coefficient approaches 0 as t ≈1 and approaches λtrain as t ≈0.
  • COMPDEMO: COMPDEMO randomly splits masked positions into complementary subsets and uses two teacher passes so each subset provides context for predicting the other.The merged logits assign each subset predictions from the pass in which its complement was revealed.
  • COMPDEMO: COMPDEMO doubles teacher forward passes and increases overall training time by approximately 50%.The teacher is frozen, so no gradient computation is required for these additional passes.
  • Distillation Objectives: TIDE uses token-level KL divergence for shared tokenizers and chunk-level CALM for cross-tokenizer alignment.Byte-level alignment forms chunks from complete tokens in both vocabularies, while Reverse CALM reverses BCE to produce bounded gradients and filter alignment noise.

3 Experiments

Experiments evaluate TIDE across heterogeneous teachers, a shared 0.6B student, eight benchmarks, component ablations, and controlled inference-efficiency settings. Cross-architecture distillation improves benchmark performance while retaining substantial deployment efficiency.

  • Experimental Setup: The evaluation uses a 0.6B BD3LM student distilled from 8B dense and 16B MoE teachers across two pipelines and eight benchmarks.The benchmarks span reasoning, knowledge, commonsense, and code generation tasks.
  • Main Results: 34.20 is the highest eight-benchmark average, exceeding the undistilled BD3LM baseline of 32.67; the shared-tokenizer pipeline reaches 33.55.The cross-tokenizer pipeline uses the native TIDE-Cross strategy, while the shared-tokenizer pipeline uses TIDE-Shared.
  • Main Results: TIDE-Cross exceeds its swapped strategy by 0.37 on average, while TIDE-Shared exceeds TIDE-Cross by 2.76 in the shared-tokenizer pipeline.The results support matching Reverse CALM to cross-tokenizer transfer and TIDAL plus COMPDEMO to exact token alignment.
  • Main Results: 48.78 on HumanEval is achieved by TIDE-Shared, compared with 32.30 for the equivalent-sized autoregressive model.On MBPP, the best distilled model scores 38.60 versus 36.60 for the autoregressive baseline.
  • Ablation Studies: 33.14 is the highest average ablation score, and removing the timestep axis causes the largest average drop of 0.26.The timestep-axis removal also reduces HumanEval by 3.05, while removing COMPDEMO reduces the average by 0.17.
  • Inference Efficiency: The distilled student uses 1.4 GB of peak memory and takes 6.25 s for 256 tokens, versus 31.3 GB and 32.55 s for LLaDA2.Throughput decreases only 2.6% relative to undistilled BD3LM, from 42.1 to 41.0 tokens/s, with identical memory footprint.

4 Conclusion

TIDE is presented as a cross-architecture distillation framework for heterogeneous diffusion language models. Across two pipelines and eight benchmarks, it improves the baseline on average, with pipeline-specific strategies and strong HumanEval performance.

  • Conclusion: TIDE is a cross-architecture distillation framework for heterogeneous diffusion language models.Its modular components address timestep reliability, complementary masking, and cross-tokenizer alignment.
  • Conclusion: +1.53 is the average improvement across eight benchmarks, while distilled dLLMs exceed the same-size autoregressive model by +16.48 on HumanEval.Reverse CALM is favored in the cross-tokenizer pipeline, whereas TIDAL plus COMPDEMO is favored in the shared-tokenizer pipeline.

Ethics Statement

The work improves diffusion language-model efficiency through knowledge distillation using public datasets and pretrained models, with moderate computational requirements.

  • The study focuses on improving diffusion language-model efficiency through knowledge distillation.
  • The experiments use publicly available datasets and pretrained models.
  • The authors characterize the computational requirements as moderate and foresee no direct negative societal impacts beyond general language-model deployment concerns.

LLM Usage

The paper situates diffusion language models and their distillation methods within prior work, contrasting established autoregressive approaches with dLLM-specific step compression and TIDE's cross-architecture focus.

  • Diffusion Language Models: Diffusion language models developed from discrete-diffusion foundations into practical architectures including full bidirectional attention and block diffusion.
  • Knowledge Distillation of Large Language Models: Knowledge distillation transfers knowledge from large teachers to smaller students, with autoregressive methods using diverse objectives and feedback mechanisms.
  • Distillation for Diffusion Language Models: Existing dLLM distillation methods focus on reducing inference steps while keeping teacher and student architectures and tokenizers aligned.
  • Distillation for Diffusion Language Models: TIDE addresses cross-architecture distillation, where teacher and student differ in architecture, attention mechanism, and tokenizer.

B Training, Inference, and Evaluation Details

The paper specifies training, inference, and evaluation protocols for component ablations and benchmark comparisons across two pipelines, including controlled sampling settings and documented hyperparameters.

  • Training: Training configurations and hyperparameters are documented for both distillation pipelines, with ablation studies trained for 3 epochs.
  • Training: The ablations remove TIDAL's timestep axis, replace dual-axis scheduling with a timestep-only baseline, or remove complementary demonstrations.
  • Evaluation: Evaluations use block size 32, classifier-free guidance scale 0.0, and task-dependent sampling steps ranging from 3 to 256.
  • Evaluation: Benchmark evaluation averages results over 50 randomly sampled examples per benchmark.

C Gradient Analysis

The gradient analysis explains why Reverse CALM is more stable than forward CALM under cross-tokenizer mismatch and why combining Reverse CALM with TIDAL is counterproductive.

  • Forward CALM: Forward CALM can suffer gradient explosion when imperfect chunk alignment gives the student a low initial probability.
  • Reverse CALM: Reverse CALM's gradient coefficient depends only on fixed teacher probabilities, directing updates toward the teacher's high-probability modes.
  • Reverse CALM: Reverse CALM filters noise at both ends: poorly aligned chunks suppress the teacher-end signal, while low student probabilities suppress the student-end signal.
  • Reverse CALM: Reverse CALM is equivalent to minimizing Bernoulli KL divergence up to a parameter-independent additive constant, giving it an information-theoretic mode-seeking interpretation.
  • TIDAL and reverse direction: Applying TIDAL to Reverse CALM is counterproductive because late-stage weighting suppresses its gradient and destroys its self-selection mechanism.

D Limitations and Future Work

The study’s evidence is limited to a 0.6B block-diffusion student with staircase attention and a 512-token context, while broader architectures and longer sequences remain untested. Training efficiency and module interactions also require refinement.

  • Scope: The empirical evaluation is limited to a 0.6B-parameter student using block diffusion with staircase attention.The authors propose scaling to 1.3B or 3B parameters to test whether a narrower capacity gap improves distillation efficiency.
  • Scope: Alternative structures, including continuous-state diffusion and encoder-style dLLMs, still require empirical validation.Applying the losses to continuous densities is identified as a necessary step for broader applicability.
  • Sequence length: The 512-token training context leaves cross-tokenizer alignment and COMPDEMO on extended sequences unexplored.Future work should examine how increasing alignment chunks changes the components’ relative contributions.
  • Future objectives: The cross-tokenizer and shared-tokenizer pipelines are currently trained independently rather than with a unified multi-teacher objective.A unified objective is proposed as a possible route to complementary knowledge transfer.
  • Efficiency and optimization: COMPDEMO requires two frozen-teacher forward passes per step, increasing training duration by approximately 50%.The authors also report that combining Reverse CALM and TIDAL can be counterproductive because of gradient-suppression mechanisms.

E Case Study

The case studies examine whether distillation transfers teacher-like predictions and changes qualitative reasoning or code errors. TIDE reduces student–teacher KL divergence and resolves distinct baseline failures across mathematical reasoning and code generation.

  • Dark Knowledge Transfer: 46% lower KL divergence on GSM8K indicates closer student predictions to the WeDLM teacher than the non-distilled baseline.The distilled student reaches 6.69 versus 12.44 for the baseline.
  • Dark Knowledge Transfer: The cross-tokenizer KL comparison is omitted because the LLaDA2 and student vocabularies are misaligned.The shared-tokenizer analysis therefore provides the reported distribution-matching diagnostic.
  • HumanEval: On HumanEval #50, TIDE-Shared reverses the shift by subtracting 5, whereas the baseline calls encode shift again.The case study contrasts the incorrect +5 operation with the successful inverse operation.
  • GSM8K: On GSM8K #39, TIDE-Cross correctly allocates two running hours and two walking hours, obtaining the ground-truth total of 18 miles.The baseline instead uses incorrect distance calculations and answers 16.
  • Case-study design: Table 7 compares baseline and distilled reasoning chains or code, while Table 8 reports distinct error resolutions across distillation variants.The summary states that the baseline fails on all four cases.

TIDE-Shared (WeDLM)

The WeDLM-based shared-tokenizer pipeline shows mixed qualitative behavior across the examined cases. It succeeds on some arithmetic and decoding tasks but retains or introduces errors on others.

  • GSM8K: TIDE-Shared incorrectly solves GSM8K #39 by using incorrect distance calculations despite allocating two hours to each activity.The reported answer differs from the correct 18-mile result.
  • GSM8K: TIDE-Shared correctly solves GSM8K #77 by computing David’s 100 pounds, Raymond’s 200 pounds, and their 100-pound difference.The baseline instead produces an incorrect 400-pound difference.
  • Overall pattern: The qualitative cases show that TIDE-Shared’s improvements are task-specific rather than uniformly resolving the examined errors.Its successful GSM8K #77 and HumanEval #50 cases coexist with failures on GSM8K #39 and HumanEval #11.
  • HumanEval: On HumanEval #11, TIDE-Shared fails with the same erroneous code as the baseline, while TIDE-Cross passes the string-XOR case.The shared pipeline appends a longer incorrect explanation to the erroneous implementation.
  • HumanEval: On HumanEval #50, TIDE-Shared passes by subtracting 5 to reverse the character shift.This contrasts with the baseline and TIDE-Cross, which fail the decoding task.
Loading 2604.26951v1…