Source-linked AI summary
Confident Adaptive Language Modeling
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, Donald Metzler
TL;DR
Large Transformer-based LLMs make autoregressive inference costly even though different inputs and generation steps require different amounts of computation, while naive early exiting can harm sequence quality. CALM calibrates local token-level exits against global textual or risk constraints and addresses confidence and missing-computation challenges. Across diverse generation tasks, it reports substantial efficiency gains, including about ×3 acceleration, while maintaining high performance under statistical guarantees.
Problem
Autoregressive LLM decoding is computationally costly, while naive early-exit decisions can cause inefficient computation and unpredictable performance degradation in dependent sequence generation.
Method
CALM calibrates local per-token early-exit decisions to satisfy global sequence-level textual or risk constraints using confidence measures, threshold selection, and distribution-free statistical control.
Results
About ×3 inference acceleration was demonstrated across summarization, machine translation, and question answering while reliably controlling high performance.
Takeaways & Limitations
CALM shows that adaptive token-level computation can reduce average LLM generation complexity while preserving user-specified sequence-quality guarantees.
Takeaways & Limitations
The rigorous performance guarantees can be conservative, and relaxing constraints or using larger calibration sets may improve empirical gains.
Abstract
from arXiv · showhide
Recent advances in Transformer-based large language models (LLMs) have led to significant performance improvements across many tasks. These gains come with a drastic increase in the models' size, potentially leading to slow and costly use at inference time. In practice, however, the series of generations made by LLMs is composed of varying levels of difficulty. While certain predictions truly benefit from the models' full capacity, other continuations are more trivial and can be solved with reduced compute. In this work, we introduce Confident Adaptive Language Modeling (CALM), a framework for dynamically allocating different amounts of compute per input and generation timestep. Early exit decoding involves several challenges that we address here, such as: (1) what confidence measure to use; (2) connecting sequence-level constraints to local per-token exit decisions; and (3) attending back to missing hidden representations due to early exits in previous tokens. Through theoretical analysis and empirical experiments on three diverse text generation tasks, we demonstrate the efficacy of our framework in reducing compute -- potential speedup of up to $\times 3$ -- while provably maintaining high performance.
1 Introduction
CALM addresses the computational cost of autoregressive LLM decoding by allocating computation adaptively while connecting local early-exit decisions to global sequence-level quality constraints. The framework combines confidence-based calibration with empirical validation showing substantial efficiency gains while maintaining high performance.
- Autoregressive decoding can impose massive computational loads because the full model stack is used repeatedly during generation.
- Early exiting dynamically selects intermediate representations for token generation, but naive decisions can waste computation or unpredictably degrade autoregressive performance.
- CALM calibrates local per-token exit decisions to provably satisfy global sequence-level textual or risk constraints with arbitrarily high probability.The framework can target metrics such as ROUGE or BLEURT and support guarantees such as 95%.
- The method uses calibration data and confidence measures to find generations that exit as early as possible while preserving desired performance guarantees.Its objectives include textual consistency with full-model outputs or bounded relative increases in risk against references.
- About ×3 inference acceleration was demonstrated across text summarization, machine translation, and question answering while reliably controlling performance.
- CALM contributes a framework, an analysis of token-wise early exiting and confidence functions, and empirical efficiency results on three diverse generation datasets.
2 Related Work
Prior work studies inference efficiency through compression, conditional computation, and early exiting, but much of the existing early-exit literature focuses on encoder-only models or local predictions. CALM extends this direction to sequence generation by deriving local exits from global quality constraints and adding model and calibration improvements.
- Inference-efficiency research includes knowledge distillation, quantization, pruning, vector dropping, and conditional computation that activates only part of a larger network.
- Adaptive compute dynamically allocates computational power per example to reduce overall complexity while maintaining high performance, and early exiting is complementary to other efficiency methods.
- Earlier language-model studies identified saturation events in which top-ranked predictions remain unchanged after some layer, but local prediction analyses do not address sequence-generation challenges.
- CALM introduces theoretically grounded calibration that derives local early exits from desirable global sequence constraints, alongside a decaying threshold, improved classifier training, and further empirical analyses.
3 Early Exiting for Adaptive Language Modeling
CALM enables Transformer language models to exit early at selected layers during decoding, while addressing state propagation and the relationship between local decisions and sequence-level quality. Experiments analyze these effects and show substantial compute savings with limited performance loss under suitable confidence and threshold mechanisms.
- 3.2 Decoding with early exiting: Early exiting selects an intermediate decoder layer for confident token predictions instead of always using the final layer.A local confidence score determines whether decoding exits or proceeds to the next representation.
- 3.3.1 State propagation: Self-attention makes prior tokens’ hidden states necessary, so the method copies an exited token’s state into subsequent skipped layers as an approximation.The paper explicitly notes that copied states introduce error and analyzes their impact on performance.
- 3.3.1 State propagation: 38.24 ROUGE-L versus 38.32 for the full model was achieved with an oracle using 1.53 layers per token.This experiment isolates state-copying effects because the oracle exits at the earliest layer matching the final layer’s top prediction.
- 3.3.1 State propagation: Up to ×5.2 compute savings were observed while preserving performance when a good confidence measure was available.An alternative copying strategy for projected key-value states instead reduced performance to 23.02 ROUGE-L.
- 3.3.2 Sensitivity to local errors: Earlier token perturbations reduce sequence-level scores more strongly, motivating thresholds that become more permissive as decoding continues.The decaying threshold offers finer performance-efficiency control, whereas a single threshold caused a textual-similarity drop of more than 10 points when efficiency increased.
- 3.5 Local confidence measures: The method trains local consistency because perfect local consistency implies global consistency, then relaxes the local requirement during inference.This preserves the efficient teacher-forcing training strategy while targeting quality of the complete output sequence.
4 Calibrating Local Early Exits from Global Constraints
CALM calibrates a shared exit threshold from calibration data so local early exits satisfy global textual or risk consistency constraints with controlled error. It selects the lowest statistically valid threshold, retaining the full model as a fallback.
- The procedure defines textual consistency against the full model and risk consistency against task references as alternative global constraints.Textual consistency uses a bounded dissimilarity function, whereas risk consistency limits the early model’s expected risk relative to the full model.
- CALM frames threshold selection as multiple hypothesis testing over candidate exit thresholds, addressing non-monotonic performance effects and finite-sample uncertainty.Naive selection from average calibration performance can be statistically invalid in the finite-sample, distribution-free setting.
- LTT tests each candidate threshold with valid p-values and selects the minimum threshold whose null hypothesis is rejected under family-wise error control.The selected threshold is valid with probability at least 1 − ϵ when the p-values are super-uniform and the testing procedure controls FWER.
- The empirical consistency statistic is bounded using a conservative risk-difference construction that does not reward cases where early-exit risk is lower.This restriction is introduced to keep the random variable within [0, 1].
- The efficient search evaluates descending thresholds and stops at the first failed rejection, returning the preceding threshold or 1 as fallback.Thresholds can use a coarse grid, such as increments of 0.05; λ = 1 reproduces the full model by construction.
5 Experimental Setting
The experiments evaluate CALM on three text-generation tasks using an eight-layer T5 model and task-specific standard metrics. They measure efficiency by decoder-layer use, FLOPs, and estimated end-to-end speedup, with static and oracle baselines.
- The evaluation covers CNN/DM summarization, WMT15 EN-FR translation, and Open-book SQuAD 1.1 question answering.These datasets vary in target generation length and how extractive the output is relative to the input.
- CALM is implemented on an eight-layer T5 1.1 encoder-decoder model with shared output embeddings and an early-exit classifier across decoder layers.The implementation uses the T5X framework.
- ROUGE-L, BLEU, and Token-F1 are used for CNN/DM, WMT, and SQuAD, respectively, with BLEURT replacing BLEU for expectation-based risk or distance control.Distance and risk are computed as 1 − m for the relevant metric m.
- Efficiency is measured primarily by average decoder layers per output token, supplemented by decoder FLOPs reduction and estimated full-sequence encoder-decoder speedup.The layer metric is intended to measure complexity reduction without implementation-specific confounds.
- The comparisons include static baselines using the same layer count for every token and a local oracle estimating the upper-bound performance-efficiency tradeoff.The framework is presented as compatible with different autoregressive multilayer LMs and confidence measures.
6 Experimental Results
CALM achieves substantial efficiency gains by selecting confidence measures and exit thresholds that reduce decoder-layer use while preserving calibrated performance. Softmax confidence generally performs best, although its vocabulary-wide computation can offset layer savings.
- Softmax confidence allows minor performance loss while reducing more than half of the layers across all three tasks.The early-exit classifier is also effective, especially when targeting high performance, while state saturation often falls below the static baseline.
- The dynamic oracle uses 1.5, 1.3, and 1.2 layers on average for summarization, WMT, and QA, respectively, without losing performance.These results illustrate CALM’s potential and the value of inference-time state propagation for skipped layers.
- CALM’s calibrated guarantees remain valid while efficiency gains increase with larger tolerance values.The reported calibration process guarantees risk at most δ at least 95% of the time.
- Softmax generally produces the greatest decoder-layer reduction and enables speedups of up to about three times over using all layers.Gains are larger for longer summarization and translation outputs than for short QA outputs.
- Rigorous performance guarantees can be conservative, leaving potential for larger empirical gains through relaxed constraints or larger calibration sets.The authors identify both calibration conservatism and confidence-interval tightening as avenues for improvement.
- Softmax confidence can increase total FLOPs because computing the full vocabulary softmax is expensive, despite using fewer decoder layers.State-based and early-exit classifier measures offer lower-FLOP alternatives.
- In example generations, most tokens exit after one or a few layers, while only a few tokens use the model’s full capacity.The examples compare different confidence thresholds and report textual and risk consistency alongside efficiency gains.
- Lowering the confidence threshold generally increases textual distance from the full-model output, while risk consistency also depends on the reference output.When references are available, risk consistency can permit more aggressive early exiting for imperfect models.
7 Conclusion
The conclusion presents CALM as an adaptive-compute framework that allocates computation per generated token while controlling full-output quality within explicit tolerances. It combines modeling and analysis to support efficient generation with rigorous guarantees.
- CALM dynamically allocates different amounts of compute per generated token under explicitly defined tolerances on the full output.The framework requires minimal modifications to the underlying language model.
- CALM combines local early exiting, a decaying threshold function, and strategies for restoring missing computations needed by future tokens.These components address performance-efficiency control and dependencies created by skipped layers.
- The framework enables efficiency gains while satisfying rigorous quality guarantees, with oracle and runtime analyses indicating further room for efficiency improvements.
A.1 Proof of Proposition 1
The proof establishes that the selected exit threshold is well-defined and risk-controlling with probability at least 1−ϵ. It relies on the valid threshold set together with the full-model threshold λ = 1.
- Every λ in the valid threshold set is risk-controlling with probability at least 1−ϵ when that set is non-empty.
- The threshold λ = 1 is risk-controlling because it produces the full-model output by definition.
- Choosing the minimum of the valid thresholds and 1 yields a well-defined threshold guaranteed to be risk-controlling with probability at least 1−ϵ.
- The appendix supplements the proof with calibration plots, example early-exit visualizations, and results for a larger 12-layer model.
B.1 Calibration results for the full tolerance range
Calibration experiments show that textual and risk consistency remain within specified tolerances while efficiency gains grow with larger tolerances. The examples illustrate uneven compute allocation, with fewer layers for easy tokens and more for hard tokens.
- Empirical risk values generally remain no greater than the specified δ, often falling below the target because the risk measure is conservative.The appendix reports this pattern for textual and risk consistency objectives.
- Efficiency gains increase with larger tolerance values even under rigorous and conservative theoretical guarantees.
- The calibration figures plot consistency in the top row and average decoder layers in the bottom row across choices of δ, with shaded standard deviations.Figure B.1 covers textual consistency and Figure B.2 covers risk consistency, both with ϵ = 0.05.
- Example outputs remain relatively similar to the full model even when decoding uses very few layers, although textual distance generally increases with acceleration.
- CALM uses very few layers for easy tokens and additional computation for hard tokens, often near sentence beginnings or when generating verbs.
B.3 T5-base results
The 12-layer T5-base model shows trends broadly similar to the 8-layer model, with adaptive early exiting often outperforming static baselines and calibration retaining similar behavior across consistency objectives.
- The 12-layer T5-base model has 12 attention heads and attention-head, embedding, and MLP dimensions of 64, 768, and 2048, respectively.
- The empirical performance-efficiency tradeoffs across three tasks are broadly similar to those observed with the 8-layer model.
- On SQUAD, static baselines using only one or two decoder layers completely fail, suggesting predictions begin forming from the third layer.
- Adaptive methods often outperform static baselines, while the local oracle matches full-model performance on two datasets.
- Textual- and risk-consistency calibration show largely similar validity and efficiency behavior on the 12-layer model, supporting generality across backbone configurations.
C Implementation Details
The implementation builds CALM on T5 encoder-decoder models, trains early-exit components separately, and exits when calibrated confidence exceeds a threshold while propagating skipped hidden states when needed.
- CALM is implemented with T5X using T5 1.1 models and adds early-exit heads to all layers.
- The main experiments use an 8-layer encoder-decoder with 6 attention heads, dimensions 64, 512, and 1024, and a 32,128-token vocabulary.
- The models are fine-tuned for up to 500K steps, with checkpoints selected by validation performance using full-model predictions.
- During inference, CALM exits when per-token confidence exceeds a threshold calibrated from user-defined δ and ϵ tolerances for textual or risk consistency.
- Speedup is measured with batch size one using full generation time, including encoding and all decoding steps until completion.
D Training Details of the Early Exit Classifier
The early-exit classifier is trained against full-model oracle decisions, with independent per-layer objectives outperforming a geometric-like alternative across evaluated layers and tasks.
- The early-exit classifier predicts whether an intermediate layer’s top-ranked token matches the top layer’s prediction.
- Backbone parameters are frozen during classifier training, so only newly added, layer-shared classifier parameters are optimized.
- The geometric-like alternative models the first correct exit as a Bernoulli event and maximizes summed log probabilities for numerical stability.
- The per-layer independent objective outperforms the geometric-like objective across all evaluated layers and tasks, especially at higher layers.
- Calibration uses FST-based LTT to select a threshold providing global textual or risk consistency within δ and ϵ tolerances.