Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

When to Ponder: Adaptive Compute Allocation for Code Generation via Test-Time Training

Gihyeon Sim

TL;DR

PonderTTT introduces Reconstruction Gating, a training-free mechanism that gates Test-Time Training updates in a TTT-augmented transformer based on the TTT layer's self-supervised reconstruction loss. By adaptively tuning a threshold via EMA to maintain a target update rate, the method achieves 82–89% Oracle Recovery across GPT-2 scales with about 2.0x TTt-layer FLOPs, significantly outperforming Random Skip while remaining inference-compatible with no ground-truth labels. The approach demonstrates robust in-distribution and out-of-distribution performance on code-language tasks, offering determinism, explainability, and practical compute savings. Overall, it provides a principled, scalable path to adaptive compute during inference for language models applied to code generation and similar domains.

Abstract

Large language models apply uniform computation to all inputs, regardless of difficulty. We propose PonderTTT, a gating strategy using the TTT layer's self-supervised reconstruction loss to selectively trigger Test-Time Training (TTT) updates. The gating decision itself is training-free--requiring no learned classifier or auxiliary networks; only a single scalar threshold is initially calibrated on unlabeled data and continuously adapted via EMA to maintain target update rates. Our experiments with GPT-2 models (124M to 1.5B) on code language modeling (The Stack v2, teacher-forced perplexity) demonstrate that this signal is inference-compatible, requiring no ground-truth labels. Our Reconstruction Gating achieves 82-89% Oracle Recovery while being fully training-free, significantly outperforming Random Skip baselines (up to 16% lower loss on OOD languages).

When to Ponder: Adaptive Compute Allocation for Code Generation via Test-Time Training

TL;DR

PonderTTT introduces Reconstruction Gating, a training-free mechanism that gates Test-Time Training updates in a TTT-augmented transformer based on the TTT layer's self-supervised reconstruction loss. By adaptively tuning a threshold via EMA to maintain a target update rate, the method achieves 82–89% Oracle Recovery across GPT-2 scales with about 2.0x TTt-layer FLOPs, significantly outperforming Random Skip while remaining inference-compatible with no ground-truth labels. The approach demonstrates robust in-distribution and out-of-distribution performance on code-language tasks, offering determinism, explainability, and practical compute savings. Overall, it provides a principled, scalable path to adaptive compute during inference for language models applied to code generation and similar domains.

Abstract

Large language models apply uniform computation to all inputs, regardless of difficulty. We propose PonderTTT, a gating strategy using the TTT layer's self-supervised reconstruction loss to selectively trigger Test-Time Training (TTT) updates. The gating decision itself is training-free--requiring no learned classifier or auxiliary networks; only a single scalar threshold is initially calibrated on unlabeled data and continuously adapted via EMA to maintain target update rates. Our experiments with GPT-2 models (124M to 1.5B) on code language modeling (The Stack v2, teacher-forced perplexity) demonstrate that this signal is inference-compatible, requiring no ground-truth labels. Our Reconstruction Gating achieves 82-89% Oracle Recovery while being fully training-free, significantly outperforming Random Skip baselines (up to 16% lower loss on OOD languages).
Paper Structure (35 sections, 9 equations, 1 figure, 14 tables)

This paper contains 35 sections, 9 equations, 1 figure, 14 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Preliminaries: TTT-Linear Update
  5. Reconstruction Gating
  6. Experiments
  7. Setup
  8. Main Results: Efficiency and Performance
  9. Out-of-Distribution Generalization
  10. Computational Cost
  11. Analysis of Gating Behavior
  12. Discussion
  13. Limitations
  14. Future Work
  15. Conclusion
  16. ...and 20 more sections

Figures (1)

  • Figure 1: PonderTTT Architecture. Each chunk undergoes an initial TTT forward pass. The reconstruction loss $\mathcal{L}_{\text{rec}}$ determines whether to UPDATE (adapt weights and re-forward) or SKIP (reuse initial forward result). The gating decision requires no learned classifier.