Research map / Large language models / Language, safety and interpretability

WASD: Wasserstein-based Knowledge Distillation for Large Language Models

NeurIPS 2026 · Large language models › Language, safety and interpretability · approach: LLM unlearning and benchmarks

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes WASD, a Sinkhorn-divergence-based distillation objective for autoregressive LLMs that uses teacher token-embedding costs and a stop-gradient dual-potential loss, improving instruction-following, summarization, translation, math and code distillation.

Problem

Existing LLM knowledge distillation divergences compare probabilities index-wise and ignore token-level semantic similarity; entropy-regularized Wasserstein objectives are biased away from the teacher distribution.

Approach

Use the debiased Sinkhorn divergence with a cost matrix from teacher token embeddings, and show its gradient equals that of a simple stop-gradient loss weighting student next-token log-probabilities by the difference of dual potentials (cross-transport minus self-transport), avoiding extra networks or backprop through the solver.

Key results

• On GPT-2 XL to GPT-2 Base distillation, WASD reaches 24.02 average ROUGE-L across five instruction benchmarks vs. 23.46 for the strongest baseline AMiD (GPT-2 Medium setting: 25.51 vs 24.74).
• On OpenLLaMA2-7B to 3B, WASD achieves 29.61 average ROUGE-L vs. 29.30 for AMiD.
• WASD achieves the best average among compared divergences in the GPT-2 XL to Base divergence comparison (22.47 vs 21.46 for CSD).

Surprising findings

• Temperature scaling of the student during training (about 2) was needed; without it WASD produces more diverse outputs and lower ROUGE-L at default decoding temperature.
• A permuted token-cost ablation performs comparably to AMiD and a uniform cost performs below the baseline, indicating meaningful token geometry matters.

Limitations

• Requires about twice the training time and 20% more training GPU memory than baselines.
• Requires a shared vocabulary between teacher and student; cross-tokenizer settings are not covered.
• Performance degrades as entropy regularization epsilon increases.

Experimental setup

Datasets: databricks-dolly-15k OpenWebText Flores-200 DialogSum GSM8k WizardCoder

Benchmarks: Dolly Eval Self-Instruct Vicuna Super NI UnNI AlpacaEval Evol-Instruct UltraFeedback Flores-200 DialogSum GSM8k HumanEval

Metrics: ROUGE-L Self-BLEU win rate COMET accuracy pass@1

Baselines: KL RKL TV GJS SKL SRKL alpha-beta divergence CSD GKD TAID DistiLLM ABKD

Builds on: Wasserstein distance Sinkhorn divergence AMiD CSD WPR

Closest papers by meaning

• DistiLLM: Towards Streamlined Distillation for Large Language Models (ICML 2024 · similarity 0.92)
• Multi-Level Optimal Transport for Universal Cross-Tokenizer Knowledge Distillation on Language Models (AAAI 2025 · similarity 0.92)
• DA-KD: Difficulty-Aware Knowledge Distillation for Efficient Large Language Models (ICML 2025 · similarity 0.91)
• TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models (ICLR 2025 · similarity 0.91)
• DDK: Distilling Domain Knowledge for Efficient Large Language Models (NeurIPS 2024 · similarity 0.91)
• Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models (NeurIPS 2026 · similarity 0.91)