# Cascade-Aware Training of Language Models

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

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 51 of its 65 references that are in the map.
- First posted on arXiv 29 May 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems](https://mlsanity.com/paper/1905.00537/) (2019)
- [PaLM 2 Technical Report](https://mlsanity.com/paper/2305.10403/) (2023)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [Learning to Reweight Examples for Robust Deep Learning](https://mlsanity.com/paper/1803.09050/) (2018)
- [MentorNet: Learning Data-Driven Curriculum for Very Deep Neural Networks on Corrupted Labels](https://mlsanity.com/paper/1712.05055/) (2017)
- [BranchyNet: Fast Inference via Early Exiting from Deep Neural Networks](https://mlsanity.com/paper/1709.01686/) (2017)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing](https://mlsanity.com/paper/2404.14618/) (2024)
- [DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference](https://mlsanity.com/paper/2004.12993/) (2020)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [Matryoshka Representation Learning](https://mlsanity.com/paper/2205.13147/) (2022)
- [Consistent Estimators for Learning to Defer to an Expert](https://mlsanity.com/paper/2006.01862/) (2020)
- [RouterBench: A Benchmark for Multi-LLM Routing System](https://mlsanity.com/paper/2403.12031/) (2024)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [Predict Responsibly: Improving Fairness and Accuracy by Learning to Defer](https://mlsanity.com/paper/1711.06664/) (2017)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [Routing to the Expert: Efficient Reward-guided Ensemble of Large Language Models](https://mlsanity.com/paper/2311.08692/) (2023)
- [Large Language Model Cascades with Mixture of Thoughts Representations for Cost-efficient Reasoning](https://mlsanity.com/paper/2310.03094/) (2023)

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

- [Universal Model Routing for Efficient LLM Inference](https://mlsanity.com/paper/2502.08773/) (2025)
- [What is the Role of Small Models in the LLM Era: A Survey](https://mlsanity.com/paper/2409.06857/) (2024)
- [Active Data Curation Effectively Distills Large-Scale Multimodal Models](https://mlsanity.com/paper/2411.18674/) (2024)
- [Rational Tuning of LLM Cascades via Probabilistic Modeling](https://mlsanity.com/paper/2501.09345/) (2025)
- [Gatekeeper: Improving Model Cascades Through Confidence Tuning](https://mlsanity.com/paper/2502.19335/) (2025)
- [Privacy-preserved LLM Cascade via CoT-enhanced Policy Learning](https://mlsanity.com/paper/2410.08014/) (2024)
- [Economic Evaluation of LLMs](https://mlsanity.com/paper/2507.03834/) (2025)
- [Honesty over Accuracy: Trustworthy Language Models through Reinforced Hesitation](https://mlsanity.com/paper/2511.11500/) (2025)
- [Learning to Partially Defer for Sequences](https://mlsanity.com/paper/2502.01459/) (2025)
- [Data-Centric Lessons To Improve Speech-Language Pretraining](https://mlsanity.com/paper/2510.20860/) (2025)
- [Streaming Model Cascades for Semantic SQL](https://mlsanity.com/paper/2604.00660/) (2026)
- [Bi-directional Model Cascading with Proxy Confidence](https://mlsanity.com/paper/2504.19391/) (2025)
- [Stable Curves, Unstable Items: Item-Level Scaling Heterogeneity in Video LLMs](https://mlsanity.com/paper/2608.07014/) (2026)
- [LaCy: What Small Language Models Can and Should Learn is Not Just a Question of Loss](https://mlsanity.com/paper/2602.12005/) (2026)
- [AutoRelAnnotator: Calibrated Model Cascades for Cost-Efficient Relevance Evaluation in Sponsored Search](https://mlsanity.com/paper/2606.25871/) (2026)

## Abstract (arXiv)

Reducing serving cost and latency is a fundamental concern for the deployment of language models (LMs) in business applications. To address this, cascades of LMs offer an effective solution that conditionally employ smaller models for simpler queries. Cascaded systems are typically built with independently trained models, neglecting the advantages of considering inference-time interactions of the cascaded LMs during training. In this paper, we present cascade-aware training(CAT), an approach to optimizing the overall quality-cost performance tradeoff of a cascade of LMs. We achieve inference-time benefits by training the small LM with awareness of its place in a cascade and downstream capabilities. We demonstrate the value of the proposed method with over 60 LM tasks of the SuperGLUE, WMT22, and FLAN2021 datasets.