WASD: Wasserstein-based Knowledge Distillation for Large Language Models

• Common distillation divergences compare probability values at vocabulary indices without explicitly using token semantics, so they do not distinguish similar-meaning tokens such as sofa and couch.LLMs also incur substantial inference computation and memory costs, motivating the goal of preserving large-model performance with smaller models.Source: §1 · pp. 1–2
• WASD defines transport costs from token embeddings and uses Sinkhorn divergence to align teacher and student distributions semantically.Its loss assigns lower cost to shifts between similar-meaning tokens than to shifts toward unrelated tokens.Source: §1 · p. 2
• WASD consistently improves distillation across multiple LLM families and tasks, including instruction following, mathematical reasoning, and code generation.The paper identifies token-space semantic information as an important ingredient for effective knowledge transfer.Source: §1 · p. 2
• WASD uses dual potentials computed through matrix scaling, avoiding additional neural networks and backpropagation through the transport solver.This procedure addresses the computational demands of optimal transport.Source: §1 · p. 2

2.1 Knowledge distillation for large language models

• LLM knowledge distillation aligns a smaller student's next-token distribution with a larger teacher's, using a divergence as the central design choice.The paper assumes teacher and student share a vocabulary; the dataset may use predefined or model-generated input-output pairs.Source: §2.1 · pp. 2–3
• Common KL, reverse KL, f-, Jensen-Shannon, adaptive KL, and α-β divergences compare probability values at corresponding vocabulary indices rather than explicitly using token semantics.Alternatives addressing numerical instability still use index-wise comparisons without explicitly leveraging token-space semantics.Source: §2.1 · p. 3

2.2 Wasserstein distance and Sinkhorn divergence

• The Wasserstein distance measures the cost of transporting probability mass between distributions using valid transport plans and a pairwise cost matrix.The cost matrix can reflect information encoded in the distributions' support.Source: §2.2 · p. 3
• Entropy regularization enables efficient matrix-scaling computation but biases transport; Sinkhorn divergence corrects this bias, motivating its use for LLM vocabulary relationships.The paper introduces a Sinkhorn-based distillation objective for autoregressive LLMs.Source: §2.2 · pp. 3–4

3.1 From Wasserstein distance to Sinkhorn divergence for LLM distillation

• The transport cost matrix is built from distances between teacher token embeddings, allowing token-space information to shape distribution alignment.The cost matrix is computed once and kept fixed during training.Source: §3.1 · p. 4
• Sinkhorn divergence addresses entropy-regularization bias: under suitable cost conditions, it is nonnegative and zero iff student and teacher distributions match, making the teacher the unique minimizer.Entropy regularization enables matrix-scaling computation, but its objective may not be minimized at the teacher distribution.Source: §3.1 · pp. 4–5

3.2 WASD: Wasserstein-based knowledge distillation for LLMs

• WASD forms a loss from stop-gradient dual-potential weights whose gradient matches the Sinkhorn objective, avoiding backpropagation through the transport solver.The optimal potentials can be computed with matrix-scaling and fixed-point iterations.Source: §3.2 · pp. 5–6
• Unlike forward KL's teacher-probability weighting at vocabulary indices, WASD's student-probability gradient weights depend on dual potentials shaped by token-embedding costs.This lets token geometry influence the update through semantic relationships between tokens.Source: §3.2 · pp. 6–7

4 Experiments

• Experiments cover instruction following and task-specific translation, summarization, arithmetic reasoning, and code generation, using task-matched datasets, model pairs, and evaluation metrics.General instruction-following tests include GPT-2, OpenLLaMA2, and Qwen2.5; task-specific tests use Gemma and Qwen2.5-Coder models.Source: §4 · p. 7, §4.1 · p. 7, §4.2 · p. 10

4.1 General instruction-following experiments

• WASD achieves the best average performance among the compared divergence objectives, without pretraining loss, student SFT initialization, or on-policy techniques.It also forms a favorable quality-diversity frontier as decoding temperature varies, with token-embedding costs offered as the explanation for retaining generation quality as diversity increases.Source: §4.1 · pp. 7–8
• WASD consistently achieves the best average performance when combined with recent distillation frameworks across GPT-2 student scales and a larger OpenLLaMA2 teacher.The comparison incorporates AMiD's assistant distribution and attributes additional improvements to WASD's token-embedding cost structure.Source: §4.1 · p. 8
• WASD uses approximately twice the training time and 20% more training GPU memory than baselines, but outperforms them at equal wall-clock time and adds no inference-time overhead.The paper attributes this result to faster convergence offsetting the extra per-epoch cost; Figure 3 presents training behavior on Dolly Eval.Source: §4.1 · pp. 8–9, Figure 3 · p. 8
• GPT-4 feedback scores favor WASD over ABKD, CSD, and AMiD on all three benchmarks, while GPT-4o-mini win rates favor WASD over AMiD across all three Qwen2.5 benchmarks.The GPT-4 scores are percentages of reference-answer scores; the Qwen2.5 comparison replaces AMiD's forward-KL term with WASD's Sinkhorn divergence while holding the other training configuration fixed.Source: §4.1 · p. 9, Table 4 · p. 9, Table 5 · p. 9
• Semantic transport cost gives the highest average ROUGE-L; permuted token relations approach AMiD, while uniform cost falls below the baseline.The controlled ablation uses GPT-2 XL to Base, holds other settings fixed, and averages results over three training seeds and five instruction-following benchmarks.Source: §4.1 · p. 9, Table 6 · p. 9
• Entropy-regularized Wasserstein distance beats AMiD but underperforms WASD, consistent with debiasing helping use semantic transport costs.As entropy regularization increases, performance declines; the paper relates this to the transport kernel becoming less sensitive to cost-matrix differences.Source: §4.1 · p. 10

4.2 Task-specific experiments

• WASD consistently improves translation, summarization, arithmetic-reasoning, and code-generation benchmarks across model series and task domains.The task-specific results are reported for Gemma models on translation, summarization, and arithmetic reasoning, and Qwen2.5-Coder models on code generation.Source: §4.2 · p. 10, Table 7 · p. 10, Table 8 · p. 10

5 Conclusion

• WASD incorporates token-level semantic information through a transport-based formulation for LLM knowledge distillation.The method adopts Sinkhorn divergence and an efficient objective based on optimal-transport dual variables, without additional networks or backpropagation through the transport solver.Source: §5 · p. 10
• Across multiple model families and tasks, WASD consistently improves distillation performance over existing methods, particularly in quality-diversity trade-offs.The conclusion identifies token-level information encoded in the vocabulary space as important for effective knowledge transfer.Source: §5 · p. 10
• The reported findings point to token-level information in the vocabulary space as relevant to effective knowledge transfer in LLM distillation.The authors also identify semantic-aware distribution alignment as a direction for future work.Source: §5 · p. 10

A.1 Proof of Lemma 1

• Lemma 1 concerns optimal dual variables for the entropic optimal-transport problem associated with the regularized distance.The passage identifies the variables as ϕ∗ and ψ∗ for distributions p and qθ in context (x, y<l).Source: §A.1 · p. 17
• Strong duality equates the entropy-regularized Wasserstein distance with the optimal value of a dual problem.The proof then substitutes the optimal dual variables into the dual objective.Source: §A.1 · p. 17
• The proof uses the optimality condition that the transport plan has total mass one, making the last term in Eq. (15) equal to −ϵ.Source: §A.1 · p. 17

A.2 Proof of Proposition 2

• Proposition 2 gives the gradient of the entropy-regularized Wasserstein distance in terms of optimal dual variables for the student marginal.Source: §A.2 · p. 17
• The proof assumes positive entropy regularization, a finite fixed cost matrix, a fixed teacher distribution, strictly positive distributions, and a student distribution differentiable in its parameters.It invokes strong duality to express the distance as the optimal value of the dual objective.Source: §A.2 · p. 17
• The envelope theorem permits differentiating the optimal value while treating the implicitly parameter-dependent dual variables as fixed; only the student marginal term depends explicitly on the parameters.The proof also notes that additive ambiguity in the dual potentials does not affect the gradient.Source: §A.2 · p. 18

A.3 Proof of Theorem 3

• The proof obtains the divergence expression from its definition and combines the resulting gradients using the coefficients in Eq. (22).The intermediate derivation applies the stop-gradient rule and uses that stop-gradient preserves its input value.Source: §A.3 · pp. 18–19
• Symmetry of the cost matrix lets the self-transport dual potentials be chosen equal when both marginals are the student distribution.The proof uses this equality when differentiating the self-transport term.Source: §A.3 · p. 18
• The teacher self-transport term is independent of the student parameters, while both marginals in the other dual objective depend on them.These dependencies determine which terms contribute to the gradient.Source: §A.3 · p. 18
• The stop-gradient convention keeps the dual-potential difference fixed during backpropagation, yielding the WASD gradient in Theorem 3.The proof differentiates the WASD loss and applies the stop-gradient rule to the potential difference.Source: §A.3 · pp. 18–19

B.1 Wasserstein distance and Sinkhorn divergence

• Unlike WPR’s rollout reward penalty, WASD uses debiased Sinkhorn divergence as the teacher–student training signal and a stop-gradient objective for its gradient.WPR regularizes a policy against a reference model; WASD corrects entropic bias with self-transport terms.Source: §B.1 · pp. 19–20

B.2 Wasserstein distance in knowledge distillation

• Prior OT distillation includes representation-, feature-, label-, sample-, and class-level transfer, while WASD targets semantic geometry among next-token vocabulary items.SinKD models sample-level geometry, and WKD targets class- and feature-level transfer in vision models.Source: §B.2 · pp. 20–21
• Cross-tokenizer methods address vocabulary or support mismatch, whereas WASD assumes a shared vocabulary and uses its embedding geometry as transport cost.The two settings are described as complementary.Source: §B.2 · pp. 20–21

C.1 Baselines

• The listed baselines span divergence-based and assistant-distribution approaches, including GJS, skew KL, α-β-divergence, score matching, and generalized assistant frameworks.The passage also lists contrastive distillation motivated by direct preference optimization.Source: §C.1 · p. 21

C.2 Training and evaluation details

• Most experiments use cosine distance between teacher-token embeddings, nearest-k sparse costs, k = 8, ϵ = 0.001, and 10 Sinkhorn and fixed-point iterations.Qwen2.5 uses k = 2; L2 distance is used for task-specific distillation experiments.Source: §C.2 · p. 21
• GPT-2 and OpenLLaMA2 instruction-following runs use 20 epochs and select the best validation ROUGE-L checkpoint; results average five evaluation seeds.Qwen2.5 instruction-following distills 7B-Instruct into 1.5B-Instruct and reports judge-assessed win rates on three benchmarks.Source: §C.2 · p. 21
• Figure 5 extends the loss-contour visualization to a toy next-token task with different teacher distributions.The caption identifies the probability simplex as the plotted domain.Source: Figure 5 · p. 22
• The controlled cost-matrix ablation compares embedding-based semantic costs with permuted and uniform costs over three training seeds.The passage distinguishes these training seeds from the evaluation seeds used elsewhere.Source: §C.2 · pp. 21–22
• Gemma task-specific distillation uses about 1K samples per task, while Qwen2.5 code generation follows a one-epoch setup with batch size 32.Teacher fine-tuning lasts 3 epochs for summarization and arithmetic reasoning and 10 for translation; KD uses the corresponding task durations.Source: §C.2 · p. 22

Loss contours of divergence variants

• WASD’s contours reflect token semantics, unlike the other compared divergences, which primarily follow probability assignments to vocabulary indices.The toy example distinguishes semantically similar tokens such as sofa and couch; Figure 5 extends the visualization across teacher distributions.Source: D Additional experimental results · pp. 22–23

Effect of entropy regularization in Sinkhorn divergence

• As entropy regularization ϵ increases, Sinkhorn loss contours flatten, and excessively large regularization is associated with degraded performance.Figure 6 visualizes the contour change across regularization settings.Source: Figure 6 · p. 23, D Additional experimental results · p. 23

Sensitivity analysis of the neighborhood size k

• For GPT-2 1.5B-to-0.1B, mean ROUGE-L is highest at k = 8; larger neighborhoods offer no consistent improvement, while smaller ones remain competitive.The sensitivity analysis uses three training seeds and tests k values from 2 to 64.Source: Figure 8 · p. 23, D Additional experimental results · pp. 23–24

Quality-diversity trade-off with recent baselines

• Across five instruction-following benchmarks, WASD forms a favorable ROUGE-L versus diversity frontier as decoding temperature varies.The figure reports higher ROUGE-L across diversity regimes than existing baselines.Source: Figure 7 · p. 23, D Additional experimental results · p. 23

Ablation study on WASD

• L2 transport costs remain comparable to cosine costs and outperform CSD, while increasing truncation from k = 8 to k = 16 yields nearly identical results.These ablations support sparse transport with relatively few nearest neighbors.Source: Table 9 · p. 24, D Additional experimental results · p. 24

Repeated training results

• WASD achieves the highest mean performance on translation, summarization, and arithmetic reasoning across three training seeds.For translation, it beats AMiD in all three seeds with a paired t-test p = 0.048; gains on the other tasks are small.Source: Table 10 · p. 24, D Additional experimental results · p. 24

Cross-model-family distillation

• WASD-S improves average ROUGE-L from 20.47 to 20.95, and WASD-S+T raises it to 21.31 while outperforming WASD-S on all five datasets.WASD-S+T improves over the baseline on four of five datasets; the setup uses greedy decoding after three training epochs.Source: D Additional experimental results · p. 24, Table 11 · p. 25

Semantic coherence and token-level analysis

• At comparable Distinct-2, WASD has higher semantic coherence than AMiD at every tested decoding temperature, with paired-test differences p < 0.01.The measured differences generally increase as sampling diversity grows.Source: Table 12 · p. 25, D Additional experimental results · p. 25

E Limitations

• WASD takes approximately twice the training time and 20% more GPU memory than conventional baselines, motivating faster or more memory-efficient Sinkhorn optimization.The authors describe the approach as still practically feasible for large-scale LLM distillation.Source: E Limitations · p. 25
• The teacher-embedding distance used for transport costs offers generic semantic similarity but may miss task- or domain-specific token information, motivating more expressive or adaptive matrices.The authors report that this simple design already provides consistent improvements.Source: E Limitations · p. 25

F Broader impacts

• Distillation may lower deployment costs and broaden LLM access, but could also make misinformation, bias, harmful content, and misuse more accessible, warranting safety alignment and monitoring.The authors encourage responsible deployment alongside advances in efficient LLM compression.Source: F Broader impacts · p. 26

10. Broader impacts

• The authors use GPT-4 and GPT-4o-mini as LLM judges to evaluate generated responses, with protocols and results described in Section 4.1 and Appendix C.2.Source: §10 · p. 33