# Multi-Stage Balanced Distillation: Addressing Long-Tail Challenges in Sequence-Level Knowledge Distillation

arXiv: https://arxiv.org/abs/2406.13114 · HTML: https://mlsanity.com/paper/2406.13114/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 51 of its 80 references that are in the map.
- First posted on arXiv 19 Jun 2024.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 30 of 51)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Improved Baselines with Visual Instruction Tuning](https://mlsanity.com/paper/2310.03744/) (2023)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [FrugalGPT: How to Use Large Language Models While Reducing Cost and Improving Performance](https://mlsanity.com/paper/2305.05176/) (2023)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [Aligning Large Multimodal Models with Factually Augmented RLHF](https://mlsanity.com/paper/2309.14525/) (2023)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [VILA: On Pre-training for Visual Language Models](https://mlsanity.com/paper/2312.07533/) (2023)
- [It's Not Just Size That Matters: Small Language Models Are Also Few-Shot Learners](https://mlsanity.com/paper/2009.07118/) (2020)
- [Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds](https://mlsanity.com/paper/1906.03671/) (2019)
- [Large Language Models Struggle to Learn Long-Tail Knowledge](https://mlsanity.com/paper/2211.08411/) (2022)
- [Large Language Models Are Reasoning Teachers](https://mlsanity.com/paper/2212.10071/) (2022)
- [From Quantity to Quality: Boosting LLM Performance with Self-Guided Data Selection for Instruction Tuning](https://mlsanity.com/paper/2308.12032/) (2023)
- [Parametric Contrastive Learning](https://mlsanity.com/paper/2107.12028/) (2021)
- [Distilling Reasoning Capabilities into Smaller Language Models](https://mlsanity.com/paper/2212.00193/) (2022)
- [Active Learning by Acquiring Contrastive Examples](https://mlsanity.com/paper/2109.03764/) (2021)
- [Superfiltering: Weak-to-Strong Data Filtering for Fast Instruction-Tuning](https://mlsanity.com/paper/2402.00530/) (2024)
- [Holistic Analysis of Hallucination in GPT-4V(ision): Bias and Interference Challenges](https://mlsanity.com/paper/2311.03287/) (2023)
- [Cold-start Active Learning through Self-supervised Language Modeling](https://mlsanity.com/paper/2010.09535/) (2020)
- [Mementos: A Comprehensive Benchmark for Multimodal Large Language Model Reasoning over Image Sequences](https://mlsanity.com/paper/2401.10529/) (2024)
- [LLMaAA: Making Large Language Models as Active Annotators](https://mlsanity.com/paper/2310.19596/) (2023)
- [Enhancing Visual-Language Modality Alignment in Large Vision Language Models via Self-Improvement](https://mlsanity.com/paper/2405.15973/) (2024)
- [Distilling Virtual Examples for Long-tailed Recognition](https://mlsanity.com/paper/2103.15042/) (2021)
- [Explore Spurious Correlations at the Concept Level in Language Models for Text Classification](https://mlsanity.com/paper/2311.08648/) (2023)
- [Sampling Bias in Deep Active Classification: An Empirical Study](https://mlsanity.com/paper/1909.09389/) (2019)

## Built on by (most cited first, 5 of 5)

- [MergeME: Model Merging Techniques for Homogeneous and Heterogeneous MoEs](https://mlsanity.com/paper/2502.00997/) (2025)
- [DISCO Balances the Scales: Adaptive Domain- and Difficulty-Aware Reinforcement Learning on Imbalanced Data](https://mlsanity.com/paper/2505.15074/) (2025)
- [Feature Alignment and Representation Transfer in Knowledge Distillation for Large Language Models](https://mlsanity.com/paper/2504.13825/) (2025)
- [CSRec: Rethinking Sequential Recommendation from A Causal Perspective](https://mlsanity.com/paper/2409.05872/) (2024)
- [Memorization Inheritance in Sequence-Level Knowledge Distillation for Neural Machine Translation](https://mlsanity.com/paper/2502.01491/) (2025)

## Abstract (arXiv)

Large language models (LLMs) have significantly advanced various natural language processing tasks, but deploying them remains computationally expensive. Knowledge distillation (KD) is a promising solution, enabling the transfer of capabilities from larger teacher LLMs to more compact student models. Particularly, sequence-level KD, which distills rationale-based reasoning processes instead of merely final outcomes, shows great potential in enhancing students' reasoning capabilities. However, current methods struggle with sequence level KD under long-tailed data distributions, adversely affecting generalization on sparsely represented domains. We introduce the Multi-Stage Balanced Distillation (BalDistill) framework, which iteratively balances training data within a fixed computational budget. By dynamically selecting representative head domain examples and synthesizing tail domain examples, BalDistill achieves state-of-the-art performance across diverse long-tailed datasets, enhancing both the efficiency and efficacy of the distilled models.