Catch Your Breath: Adaptive Computation for Self-Paced Sequence Production

TL;DR

The paper introduces Catch Your Breath (CYB), a framework that enables adaptive, token-level computation in transformers by allowing the model to emit a <don't know> token to request additional compute and by inserting <pause> tokens when delays are granted. It formalizes three loss variants—CYB-AP (anytime prediction), CYB-VA (variational stop-time), and CYB-DP (distributional penalty)—to train the model to calibrate when extra compute is beneficial, balancing speed and accuracy. Empirical results show CYB-AP delivering the best data efficiency and perplexity improvements on a C4 fine-tuning setup, with adaptive pause patterns that correlate with token difficulty and context, and with the model learning to self-calibrate its use of <don't know> to improve outputs. The approach advances efficient inference for large language models by coupling adaptive computation with principled stopping-time control, offering practical benefits for faster, more data-efficient generation and insights into token-level processing dynamics. Future work may refine pause conditioning in attention, explore selective pause strategies, and further analyze epistemic uncertainty signaling through <don't know> in complex contexts.

Abstract

We explore a class of supervised training objectives that allow a language model to dynamically and autonomously scale the number of compute steps used for each input token. For any token, the model can request additional compute steps by emitting a <don't know> output. If the model is granted a delay, a specialized <pause> token is inserted at the next input step, providing the model with additional compute resources to generate an output. The model can request multiple pauses. To train the model to use <don't know> outputs judiciously and to calibrate its uncertainty, we frame the selection of each output token as a sequential-decision problem with a time cost. We refer to the class of methods as $\textit{Catch Your Breath}$ losses and we study three methods in this class: CYB-AP frames the model's task as anytime prediction, where an output may be required at any step and accuracy is discounted over time; CYB-VA is a variational approach that aims to maximize prediction accuracy subject to a specified distribution over stopping times; and CYB-DP imposes a penalty based on a computational budget. Through fine-tuning experiments, we identify the best performing loss variant. The CYB model needs only one third as much training data as the baseline (no pause) model needs to achieve the same performance, and half as much data as a model with pauses and a cross-entropy loss. We find that the CYB model requests additional steps when doing so improves accuracy, and the model adapts its processing time to token-level complexity and context. For example, it often pauses after plural nouns like $\textit{patients}$ and $\textit{challenges}$ but never pauses after the first token of contracted words like $\textit{wasn}$ and $\textit{didn}$, and it shows high variability for ambiguous tokens like $\textit{won}$, which could function as either a verb or part of a contraction.

Figures (13)

• Figure 1: Sequence processing in a transformer with <don't know> outputs and <pause> inputs.
• Figure 2: Gaze fixations of a human reader. Indices indicate order of gaze within a sentence along with fixation duration (in msec). Reprinted from just1980 (just1980, Figure 1).
• Figure 3: Catch Your Breath: A sequential decision task to obtain a single output token. The blue boxes represent model <don't know> decisions, the purple boxes model choices of output tokens. The orange boxes are choices made by the environment about whether the model is granted a delay in responding. The green and red boxes represent end-state rewards.
• Figure 4: C4 evaluation set perplexity (lower is better) for a pretrained Gemma1 2B model (black bar) and various models fine tuned on the C4 train set. Model variants include a baseline condition, where no modification is made to training objective or token sequence (brown bar), Think Before You Speak (TBYS) with 1, 2, and 3 pause tokens after each input token (green bars), and the three proposed variants of the CYB loss. CYB-AP obtains the best performance among the methods.
• Figure 5: C4 evaluation set perplexity for various hyperparameter settings of the three CYB variants. The black dashed line indicates the perplexity of the best performing variant. Hyperparameters $\omega$ and $\rho$ refer to the distribution of stopping times and are indicated by ratios for steps 0-3 (i.e., the initial input token and the following 3 pause tokens). Hyperparameters $\gamma$ and $\alpha$ are penalty terms associated with CYB-AP and CYB-DP, respectively. See Tables \ref{['table:notation']} and \ref{['table:cyb']} for notation.