Transformer models save computing time by exiting early when confident
One Size Does Not Fit All: Setting Inference Depth from the Questions a Deployment Actually Asks
Computation and LanguageMachine Learning
Summary
Transformer language models usually use the same amount of computing power to answer every question, even if some questions are simpler than others. The authors show that by adding a small extra part inside the model, it can decide to stop processing early when it's confident about its answer, which saves time. How much time is saved depends a lot on the type of questions the model is asked. They found that adjusting the stopping point based on the kind of questions the model gets makes a big difference. However, some ways to check if the model’s quick answer is actually good don’t always reflect the true accuracy.
What this means in practice
- •For mobile app developers: Reduce response time and power usage for language features on phones by customizing early exit thresholds for typical user queries.
- •For customer support teams: Speed up automated handling of delivery complaints by running fewer model layers without losing answer quality.
Authors
Jerry Kaplan
Abstract
A transformer language model is trained to respond to any prompt, but each deployment asks only a narrow range of questions: a support assistant sees delivery complaints, a coding tool sees Python. Every deployment nonetheless pays the same computation per token. This paper measures how much of that cost is avoidable when the range of prompts is known in advance. The mechanism examined is early exit: a small, trained component - called a readout - is attached to an intermediate layer and proposes a token, and a confidence test decides whether to emit it or to run the remaining layers. The models are frozen, and the only supervision used is the model's own output on ordinary traffic. Three findings are reported. First, achievable savings depend strongly on the kind of traffic: at half depth on a 1.5-billion-parameter model, 96 percent of tokens could be emitted early for arithmetic word problems and 8 percent for Chinese-language explanations, at matched token-level fidelity to the full model (a measure whose limits the third finding exposes). Second, of three ways a deployment might use knowledge of its traffic, only customizing the threshold for exiting early is worthwhile: calibrating it per deployment raised exit rates by up to 59 percentage points across three models, and by more than 10 points on most corpora tested. Third, token-level fidelity - the standard evaluation measure in the early-exit literature - fails in domains where tokens can be checked against ground truth: on arithmetic word problems, three models each answered sixty questions correctly when run in full, and between 10 and 28 correctly under early exit, in the configuration that scored highest on fidelity. The intended setting is small models on personal devices, where generation is limited by memory bandwidth rather than computation.