深度递归Transformer中,每个词元独立收敛,可提前停止以节省计算。
Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers

- 通过测量每词元输出变化,发现其在6~16轮间趋于稳定
- 平均仅需4.94轮即达深度8性能,比均匀深度降低38%计算量
- 无需训练即可根据收敛状态动态停顿,适合高效推理场景
深度递归Transformer通过权重共享的模块重复执行,先前研究显示随机递归次数训练可获得跨深度可用的检查点。本文直接测量每个词元的实际计算过程:在135M参数、基于FineWeb-Edu训练的模型中,词元输出的连续差异(均值后继输出KL)从第2轮的3.9e-1降至第16轮的8.5e-6,状态变化同步衰减。关键发现是收敛不均——中位词元在第6轮稳定,约10%在平均训练深度8仍持续更新,且收敛深度依词性排列(空格最浅,内容词最深)。该词元级差异是核心对象。我们证明其可直接读取,并优于学习预测:一个无需训练的规则,在输出稳定时停止每词元计算,实现深度8的统一质量,平均仅需4.94轮(减少38%),而使用相同数据训练的线性路由器需接近完整深度且无加速效果。弹性机制表现为验证损失从1轮的3.80单调下降至8轮的3.20并保持稳定至32轮。平均深度作为FLOP代理,实验采用三点实测时间区间,未声称与等量FLOP对比,结果基于单一规模与种子。全实验仅用一块RTX 4090,耗时约100 GPU小时。
原文摘要 · Abstract (English)
A depth-recurrent transformer applies a weight-tied core a variable number of times, and prior work has shown that training with a randomized recursion count yields one checkpoint usable across a range of inference depths. We ask what such a model actually computes per token, and measure it directly. On a 135M-class model trained on FineWeb-Edu, the recurrent state converges to a per-token fixed point: mean successive-output KL divergence falls from 3.9e-1 at the second loop to 8.5e-6 by the sixteenth, and per-token state change decays in step. Crucially, this convergence is not uniform across tokens. The median token converges by loop six, while approximately 10 percent of tokens continue to update at the training-mean depth of eight, and mean convergence depth is ordered by token type (whitespace shallowest, content words deepest). This per-token variation is the central object of the paper. We show it is directly readable and that reading it outperforms learning to predict it: a training-free rule that halts each token once its output stabilizes attains uniform depth-8 quality at 4.94 average loops (a 38 percent reduction in average depth) and matches uniform depth across the average-depth range, whereas a linear router trained on convergence labels harvested from the same model requires nearly full depth and yields no reduction. The elasticity that makes this possible reproduces here as background (validation loss decreases monotonically from 3.80 at one loop to 3.20 at eight and remains stable to 32 loops). We report average depth as a FLOP proxy with a three-point wall-clock bracket rather than a realized speedup, make no FLOP-matched parity claim, and note that the allocation results are established at a single scale and seed. The complete study runs on a single RTX 4090 in approximately 100 GPU-hours.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。