Research map / Large language models / Language, safety and interpretability
Adversarial Moment-Matching Distillation of Large Language Models
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 adversarial moment-matching distillation for LLMs, casting KD as imitation learning that minimizes on- and off-policy action-value moment-matching bounds via a min-max game trained with policy gradients.
Problem
Standard LLM knowledge distillation minimizes step-wise distribution distances (KL, RKL, JS, TV) between teacher and student token distributions, which is a behavior-cloning objective that may not capture the full teacher knowledge or the quality of whole outputs.
Approach
Derive on- and off-policy upper bounds on the imitation gap via the performance difference lemma using Q-value functions, and optimize them jointly through an adversarial min-max game in which Q-value networks are maximized and the student policy is updated with policy gradients, replacing explicit distribution-distance losses.
Key results
- On instruction following with OpenLLaMA-7B to 3B, the method reaches 59.8 GPT4 on DollyEval versus 59.2 for DistiLLM and 58.7 for MiniLLM, and is best on all five test sets.
- On task-specific T5-XL to T5-Large, the method scores 52.7 ROUGE-L on SAMSum, 36.0 BLEU on IWSLT'17, and 65.3 accuracy on StrategyQA, all best among compared methods.
- On GPT-2 XL to GPT-2 distillation, the method reports 22.7 GPT4 on SelfInst versus 21.7 for DistiLLM, best across five datasets.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use joint on-policy and off-policy adversarial moment matching instead of the compared KD baselines | OpenLLaMA-3B student · Instruction-following datasets | GPT4 feedback and ROUGE-L | Helped |
| Use joint on-policy and off-policy adversarial moment matching instead of the compared KD baselines | T5 student models · SAMSum, IWSLT'17 (en-de), and StrategyQA | ROUGE-L, BLEU, and accuracy | Helped |
| Use moment-matching distance instead of KL, RKL, JS, or TV distance | Knowledge distillation · Four tasks | task performance | Helped |
| Combine on-policy and off-policy moment-matching objectives rather than use either alone | Knowledge distillation · Instruction-following and task-specific datasets | task performance | Helped |
Surprising findings
- Plain KD and SeqKD (SFT on teacher outputs) hardly improve over the SFT student, and several are worse than SFT on many datasets.
Limitations
- Inner-loop gradient ascent for Q-value functions adds computational time cost.
- Auxiliary Q-value networks add memory overhead.
- Evaluation relies largely on GPT-4 judging and ROUGE-L on a limited set of models and tasks.
Experimental setup
Datasets: databricks-dolly-15k OpenWebText SAMSum IWSLT'17 (en-de) StrategyQA
Benchmarks: DollyEval SelfInst VicunaEval S-NI UnNI SAMSum IWSLT'17 (en-de) StrategyQA
Metrics: GPT-4 feedback score ROUGE-L BLEU accuracy
Baselines: SFT KD SeqKD ImitKD MiniLLM GKD DistiLLM KL RKL JS TV distance
Builds on: Swamy et al. moment-matching imitation learning MiniLLM GKD DistiLLM ImitKD
Closest papers by meaning
- DistiLLM: Towards Streamlined Distillation for Large Language Models (ICML 2024 · similarity 0.88)
- TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models (ICLR 2025 · similarity 0.88)
- Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models (NeurIPS 2026 · similarity 0.88)
- On Teacher Hacking in Language Model Distillation (ICML 2025 · similarity 0.88)
- DA-KD: Difficulty-Aware Knowledge Distillation for Efficient Large Language Models (ICML 2025 · similarity 0.88)
- Impartial Adversarial Distillation: Addressing Biased Data-Free Knowledge Distillation via Adaptive Constrained Optimization (AAAI 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.