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

Problem Statement

While inference-time scaling for foundation models has largely focused on increasing model depth (e.g., layer recurrence) or employing reasoning strategies like Chain-of-Thought, a promising alternative involves width-based scaling. This approach inserts special <PAUSE> tokens into the input stream, granting the model additional forward passes to process context before generating a response.

However, existing methods for utilizing pause tokens, specifically the "Think Before You Speak" (TBYS) framework, suffer from a critical limitation: the number of compute steps is determined externally by the data loader. In TBYS, a fixed number of <PAUSE> tokens are inserted, and the model's output is ignored until the final pause step. This forces the model to act greedily at every intermediate step, unable to anticipate future pauses or regulate its own processing. Consequently, the model cannot signal readiness to respond or adaptively allocate compute resources based on the complexity of the current token, treating additional steps as a static barrier rather than a dynamic resource.

Methodology: Catch Your Breath (CYB)

The authors propose Catch Your Breath (CYB), a supervised loss function framed as a sequential-decision problem that enables models to autonomously determine the number of compute steps required for each input token.

Core Mechanism

Instead of a fixed pause sequence, CYB allows the model to decide at each step $i$ whether to:

1. Respond: Generate the next token in the sequence.
2. Abstain: Emit a special <DON'T KNOW> token, requesting an additional <PAUSE> step to continue processing.

This creates a variable-latency system where the model can abstain multiple times consecutively, effectively scaling its computation depth per token based on internal uncertainty.

The Loss Function

The CYB loss is derived from the expected discounted accuracy of the model's decision process. Let $d_i$ be the probability of the model abstaining at step $i$, and $t_i$ be the probability of selecting the correct target token at step $i$. The model generates an output at step $S$ (either self-induced via abstention or world-induced via a forced stop).

The loss is defined as the negative log-likelihood of the expected discounted accuracy:
$$\ell_{CYB}(t, d; \omega, \gamma) = -\log \mathbb{E}_{S \sim \text{Pr}(S|d, \omega)} [\gamma_S t_S]$$
Where:

• $\omega$ represents the distribution of world-induced stop times (allowing for forced termination).
• $\gamma$ represents step-wise discount factors, allowing the model to be penalized for delaying responses unless accuracy improves.
• The expectation is computed by summing over all possible stopping steps, weighted by the probability of stopping at that step.

Training Implementation

To ensure efficiency, the authors propose a parallel training procedure. During training, $W_{max} - 1$ pause tokens are inserted after every real input token. The transformer processes these steps in parallel with causal masking. The loss is computed such that if the model chooses to respond early (low $d_i$) or if the world forces a stop, subsequent pause steps do not contribute to the loss for that specific token. This approach maintains the same computational and memory footprint as standard cross-entropy training with fixed pauses.

Key Contributions

1. Principled Loss Function: Introduction of CYB, a loss that trains transformers to self-determine compute steps, contrasting with the static, externally imposed delays of TBYS.
2. Efficient Parallel Training: A training procedure that matches the compute and memory efficiency of standard cross-entropy while enabling adaptive computation.
3. Hypothesis Validation: The paper validates four hypotheses:
   • H1: CYB outperforms TBYS because it self-determines readout timing.
   • H2: CYB yields a well-calibrated model where <DON'T KNOW> probability correlates with the potential for accuracy improvement via delay.
   • H3: Removing model control over stop times (via forced early stopping) degrades performance, confirming the value of autonomy.
   • H4: The speed-accuracy trade-off can be manipulated via discount factors ($\gamma_0$), allowing the model to respond faster at the cost of some accuracy.
4. Empirical Superiority: Demonstration that CYB reduces perplexity and enhances downstream accuracy in both pretraining and fine-tuning scenarios without additional computational cost compared to TBYS.

Results

Experiments were conducted using Gemma1-2B and Gemma3-4B models on the C4 dataset and various downstream benchmarks (HellaSwag, MMLU, GSM8K).

• Perplexity Reduction: CYB consistently outperformed TBYS. Notably, a CYB model with only one pause token achieved lower perplexity than a TBYS model with three pause tokens. In pretraining, CYB with 3 pauses showed a 3.6% relative perplexity improvement over TBYS, compared to 2.7% in fine-tuning.
• Data Efficiency: CYB required only 33% of the data to match the baseline performance and less than 66% of the data to beat the 3-pause TBYS model, indicating genuine improvements in learning efficiency rather than data saturation artifacts.
• Downstream Accuracy: CYB-trained models achieved higher top-1 accuracy on HellaSwag, MMLU, and GSM8K compared to both the baseline and TBYS.
• Calibration and Behavior: Analysis showed a statistically significant positive correlation between the <DON'T KNOW> probability and the potential accuracy gain from an additional step. Qualitative analysis revealed that pause durations are context-sensitive; for instance, the model pauses longer on complex tokens (e.g., "traditional") but not on predictable sequences (e.g., digits).

Significance and Claims

The paper claims that autonomy in compute allocation is the key differentiator between CYB and existing width-based scaling methods. By allowing the model to decide when it is ready to respond, CYB transforms pause tokens from a static barrier into a dynamic resource for integrating information.

The authors emphasize that CYB is not primarily motivated by raw inference-time compute savings, but rather by providing the architecture with flexible, context-dependent control over its computation. This mirrors human reading behavior, where fixation duration varies based on processing demands. The paper positions CYB as a structurally superior alternative to depth-based adaptive methods (like looped transformers) due to its parallelizability, simplicity, and formal expressivity, offering a new paradigm for inference-time scaling that improves model performance without increasing computational overhead.