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

📄 arXiv: 2601.00894v1 📥 PDF

作者: Gihyeon Sim

分类: cs.LG, cs.CL

发布日期: 2025-12-31

备注: 14 pages, 1 figure, 14 tables, code available at https://github.com/deveworld/ponderTTT

💡 一句话要点

提出PonderTTT，通过测试时训练自适应分配代码生成计算资源。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 代码生成 自适应计算 测试时训练 语言模型 门控策略

📋 核心要点

  1. 现有大型语言模型对所有输入采用统一计算，忽略了不同代码生成的难度差异。
  2. PonderTTT利用TTT层的自监督重构损失，自适应地决定何时进行测试时训练更新。
  3. 实验表明，PonderTTT在代码语言建模任务上显著优于随机跳过基线，且无需额外训练。

📝 摘要（中文）

大型语言模型对所有输入都采用统一的计算量，而忽略了输入难度的差异。本文提出PonderTTT，一种门控策略，利用TTT层的自监督重构损失来选择性地触发测试时训练(TTT)更新。该门控决策本身是免训练的——不需要学习的分类器或辅助网络；只需要一个标量阈值，最初在无标签数据上校准，并通过EMA持续调整以维持目标更新率。在代码语言建模（The Stack v2，教师强制困惑度）上使用GPT-2模型（124M到1.5B）的实验表明，这种信号与推理兼容，不需要ground-truth标签。我们的重构门控实现了82-89%的Oracle Recovery，同时完全免训练，显著优于随机跳过基线（在OOD语言上损失最多降低16%）。

🔬 方法详解

问题定义：现有的大型语言模型在代码生成任务中，对所有输入都采用相同的计算量，没有考虑到不同代码片段的复杂度和难度差异。这种统一的计算分配方式效率低下，对于简单的代码片段可能造成资源浪费，而对于复杂的代码片段则可能计算不足，影响生成质量。因此，如何根据输入代码的难度自适应地分配计算资源是一个关键问题。

核心思路：PonderTTT的核心思路是利用测试时训练(TTT)层的自监督重构损失作为门控信号，来判断是否需要对当前输入进行额外的计算（即测试时训练更新）。重构损失可以反映模型对当前输入的理解程度，损失越高，说明模型对该输入的理解越差，越需要进行更新。通过这种方式，PonderTTT可以自适应地为不同的输入分配不同的计算资源。

技术框架：PonderTTT主要包含以下几个模块：1) TTT层：用于进行测试时训练更新。2) 重构损失计算模块：计算TTT层的自监督重构损失。3) 门控模块：根据重构损失和一个自适应阈值来决定是否触发TTT更新。阈值通过EMA（指数移动平均）进行更新，以维持目标更新率。整体流程是：输入代码片段 -> TTT层 -> 计算重构损失 -> 门控模块 -> 决定是否进行TTT更新 -> 输出代码片段。

关键创新：PonderTTT的关键创新在于提出了一种完全免训练的门控策略，该策略不需要学习任何额外的分类器或辅助网络，只需要一个标量阈值。该阈值最初在无标签数据上校准，并通过EMA持续调整以维持目标更新率。这种方法简单有效，并且与推理兼容，不需要ground-truth标签。与现有方法相比，PonderTTT更加轻量级，易于部署，并且不需要额外的训练数据。

关键设计：PonderTTT的关键设计包括：1) 使用TTT层的自监督重构损失作为门控信号。2) 使用EMA来更新阈值，以维持目标更新率。3) 阈值的初始校准在无标签数据上进行。4) 目标更新率是一个可调节的超参数，可以根据具体任务进行调整。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，PonderTTT在代码语言建模任务上取得了显著的性能提升。在The Stack v2数据集上，PonderTTT实现了82-89%的Oracle Recovery，同时完全免训练。与随机跳过基线相比，PonderTTT在OOD语言上的损失最多降低了16%。这些结果表明，PonderTTT能够有效地自适应地分配计算资源，提高代码生成的效率和质量。

🎯 应用场景

PonderTTT可应用于各种代码生成场景，例如软件开发、自动化测试、代码补全等。通过自适应地分配计算资源，可以提高代码生成的效率和质量，降低计算成本。此外，该方法还可以扩展到其他自然语言处理任务中，例如机器翻译、文本摘要等，具有广泛的应用前景。

📄 摘要（原文）

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).