← 返回 2026-10-28

截断与否?截断反向传播下循环Transformer的缩放研究 Truncate or Not? A Scaling Study of Looped Transformers with Truncated Backpropagation

Xuyuan Xiong, Siyu Chen, Zihan Shao, Yue Su, Weiyang Jin, Ping Luo, Zhuoran Yang, Xihui Liu 📅 2026-10-20 👍 8 2026-11-11 18:30
参数效率 循环Transformer 截断BPTT 缩放律 计算最优训练

统一缩放律表明循环Transformer全BPTT最优,两次完整循环以27%更少参数媲美非循环模型

前置知识

循环Transformer(looped transformer)

一种深度参数共享架构:模型由 prelude(含输入嵌入)、共享循环块、coda(含输出头)三段组成,循环块 $R$ 与注入矩阵 $W_{inj}$ 只有一份参数,但在一次前向中被重复执行 $T$ 次。每步更新为 $s_t=\mathrm{RMSNorm}(R(W_{inj}[s_0;s_{t-1}]))$,即把初始状态 $s_0$ 注入每一步(input injection)再通过共享块。这样计算的深度随 $T$ 增长,可训练参数却不变,代表工作有 Universal Transformer、ALBERT 和 Huginn 的 prelude–recurrent–coda 架构。

本文的研究对象就是这种架构:缩放律里只有循环块参数 $N_{rec}$ 会被循环"放大",prelude/coda 参数 $N_{once}$ 只算一次,这一结构划分直接决定了有效参数 $N_{eff}$ 的构造方式。

截断BPTT(truncated backpropagation through time)

循环模型训练时梯度需沿 $T$ 个循环反向传播(BPTT),开销大。截断 BPTT 设置为 $(T,B)$:前向跑满 $T$ 个循环,但只对最后 $B$ 个循环回传梯度,在第 $T-B$ 步的状态 $s_{T-B}$ 上施加 stop gradient,前面的循环仅前向执行、不存激活。$B=T$ 即全 BPTT,$B<T$ 为截断,$(1,1)$ 就是普通 Transformer。一次反向循环的 FLOPs 约是前向的两倍,所以截断省算力,但梯度有偏。

论文标题"截断与否"问的正是这个选择:截断的算力节省与梯度偏差之间的权衡,是全文缩放律要定量刻画的核心变量 $B$。

神经缩放律与计算最优训练

缩放律用幂律刻画验证损失与参数量 $N$、训练token数 $D$ 的关系,如 Chinchilla 的 $L=E+A N^{-\alpha}+H D^{-\beta}$;本文采用 Videau 等的耦合形式 $L=E+[A(N/N_0)^{-\alpha}+H(D/D_0)^{-\beta}]^{\kappa}$。在固定算力 $C\approx 6ND$ 下最小化损失可得计算最优的 $N,D$ 配比(IsoFLOP 分析),这是现代大模型"该训多大、训多久"的标准决策工具。

本文把循环数引入这套框架:用有效参数替换 $N$,再从驻点条件解析导出循环数、token、模型大小随算力的最优增长指数,是典型的缩放律建模工作。

有效参数与计算等价参数

论文分析循环设置 $(T,B)$ 时使用两个"等效的非循环模型大小":按损失等价的有效参数 $N_{eff}=N_{once}+g(T,B)\cdot N_{rec}$(它达到相当于多大的非循环模型的损失),以及按FLOP等价的计算参数 $N_{comp}=N_{once}+\frac{T+2B}{3}N_{rec}$(它花费相当于多大的非循环模型的算力,反向按2倍前向计)。二者之比 $\chi=N_{comp}/N_{eff}$ 表示"每一份有效参数付出的算力"。

整篇论文的决策逻辑都建立在这对概念上:固定 $T$ 时该不该截断,归结为枚举整数 $B$ 最小化 $\chi$;理解了 $N_{eff}$ 与 $N_{comp}$ 的分离,才能看懂全部结论。

研究动机

循环Transformer通过重复使用同一组共享层(如Huginn的prelude–recurrent–coda架构)在不增加可训练参数的前提下提升计算深度,参数与显存效率出色。但训练它需要沿循环深度做时间反向传播(BPTT),代价高昂;为省算力,近年的循环语言模型普遍采用截断BPTT——只对最后 $B$ 个循环回传梯度,前面 $T-B$ 个循环仅前向执行。截断节省反向计算和激活显存,却引入有偏梯度,可能削弱循环带来的收益(Schwethelm等2026年已观察到这种现象)。于是出现一个关键的资源分配问题:固定训练算力下,算力应如何在 $T$、$B$、模型大小 $N$ 和训练token数 $D$ 之间划分?现有缩放研究都没把 $T$ 和 $B$ 当独立变量联合建模:Schwethelm等对"全BPTT"和"截断一半"分别各拟一条缩放律,无法回答"该不该截断";Parcae在小模型上变动过 $B$,但在缩放实验中固定 $B=\lceil T/2\rceil$,把联合优化留作未来工作。

本文的目标是本文要回答一个明确问题:"模型大小、训练token数、前向与反向循环数如何共同决定损失?在固定训练算力预算下到底该不该截断反向传播?"为此作者训练了近600个模型(共592次运行),缩放套件覆盖58M–919M非嵌入参数的五个尺寸、每参数1.1–1291个token(三个数量级)、最多16个前向循环共15种 $(T,B)$ 设置,独立变动前向和反向循环数。在此基础上拟合一统一缩放律,希望它做到三件事:(1)用"有效参数量"统一刻画循环收益与截断折扣;(2)解释截断在什么条件下划算、什么条件下不划算,并解释已有文献表面矛盾的结论;(3)分别给出模型大小固定(受显存约束)和模型大小自由两种情形下的计算最优分配规则,并给出可落地的参数高效训练配方。

与已有工作不同的是,本文的独特切入是把反向循环数 $B$ 作为独立变量写进缩放律,并用一个有界饱和的循环增益函数同时统一"循环收益"与"截断折扣"。此前对循环语言模型的缩放律要么对每种架构单独拟合曲线(SMELT、Lee等),要么只建模单一循环数 $T$ 而把 $B$ 固定或绑定于 $T$:Schwethelm等用无界幂律增益 $T^{\phi}$,Parcae用线性增益,Ouro用加性损失项。本文则把耦合的规模–数据缩放律(Videau等)与有界循环增益(Chen等的饱和观察)结合,提出 $g(T,B)=\exp[c\,(1-T^{-p}B^{-q})]$,同时满足归一化 $g(1,1)=1$、单调性、零反向极限 $g(T,0)=0$ 和饱和上界 $e^c$ 四条建模要求。正是这一形式让"固定 $T$ 时截断有利"与"联合优化时全BPTT最优"这两组看似矛盾的实验现象在同一律下自洽地被解释。

核心方法

直觉上,循环Transformer里只有循环块 $R$ 和注入矩阵 $W_{inj}$ 被每个循环重复执行,循环只会"放大"这部分参数的价值,因此作者把非嵌入参数拆成只跑一次的 $N_{once}$(prelude与coda)和每循环复用的 $N_{rec}$,用有效参数量替换损失律中的 $N$:$L=E+[A(N_{eff}/N_0)^{-\alpha}+H(D/D_0)^{-\beta}]^{\kappa}$,其中 $N_{eff}=N_{once}+g(T,B)\cdot N_{rec}$,循环系数 $g(T,B)=\exp[c\,(1-T^{-p}B^{-q})]$ 随 $T$、$B$ 增大单调升向有限上界 $e^c$,截断则打折扣 $g(T,B)/g(T,T)\le 1$。实验遵循Schwethelm等的prelude–recurrent–coda架构与训练配方:在100B token的FineWeb-Edu语料上训练五个尺寸(60M/120M/240M/480M/1B,非嵌入参数57.8M–918.7M,宽度768–2176,prelude/循环/coda层数从2/4/2到4/8/4),矩阵参数用MuonH、嵌入/输出/归一化用AdamW,报告去重held-out分片上的验证交叉熵(nats/token)。全网格592次运行、347个配置,$T\in\{1,2,4,8,16\}$,230个配置用于拟合、117个留作独立检验。

核心创新有三点。第一,用"一对尺寸"刻画任意循环设置 $(T,B)$:同一模型对应两个等效的非循环尺寸——按损失等价的有效尺寸 $N_{eff}$(它值多少参数)和按FLOP等价的计算尺寸 $N_{comp}=N_{once}+\frac{T+2B}{3}N_{rec}$(它花多少算力,前向每参数每token算2 FLOPs、反向算4),二者之比 $\chi=N_{comp}/N_{eff}$ 就是"每有效参数付出的算力",于是固定 $T$ 的截断决策化成与总算力无关的整数优化:选 $B$ 最小化 $\chi(T,B)$(Lemma D.1)。第二,凸优化论证:固定循环算力 $S=T+2B$ 时最大化 $g$ 等价于最小化 $T^{-p}B^{-q}$,这是标准凸问题,当 $q>2p$ 时唯一解为 $T=B=S/3$,即全BPTT(Lemma D.2);拟合值 $q\approx1.022$ 远大于 $2p\approx0.474$。第三,有界饱和增益 $g(T,T)=\exp[c(1-T^{-(p+q)})]$ 趋于 $e^c\approx3.04$,取代先前工作的无界幂律,使"截断在固定 $T$ 时有利"与"联合优化时全BPTT最优"在同一律下自洽。

方法步骤详情

流程分五步。(1)建实验网格:五个尺寸的 $N_{rec}/N$ 相近(0.505–0.564),每个尺寸在15种 $(T,B)$ 设置与多个token预算下训练,缩放套件347个配置跨1.1–1291 token/参数,共592次运行,同配置多次运行取平均验证损失。(2)联合拟合:对九个参数 $E,A,H,\alpha,\beta,\kappa,c,p,q$ 最小化平方误差,配置等权,参考尺度固定 $N_0=2.4\times10^8$、$D_0=2^{30}$,从64个随机初值出发取最优,全部收敛;结果 $E=0.299$、$A=6.487$、$H=7.633$、$\alpha=0.306$、$\beta=0.393$、$\kappa=0.325$、$c=1.112$、$p=0.237$、$q=1.022$。(3)留一验证:分别留出一个配置、一个模型尺寸、一个token预算、一个循环设置重新拟合再预测被留出的测量,RMSE为 $6.4$–$8.7\times10^{-3}$ nats,$R^2=0.9992$。(4)解释性检验:用PCHIP插值从非循环运行反推经验有效尺寸与经验循环系数,与仅含 $c,p,q$ 的 $g$ 预测面对比,吻合良好。(5)决策与分配:固定 $T$ 时枚举整数 $B$ 最小化 $\chi$;$T$ 自由时做凸优化;再从驻点条件解析导出固定尺寸与自由尺寸的算力分配幂律,并用专门实验验证。

技术新颖性

与已有工作的本质区别有四。(a)统一性:Schwethelm等需要为每种BPTT配方各拟一条律,结构上无法回答"该不该截断";本文的律把 $T,B$ 当变量,一个框架内同时预言"固定 $T$ 时截断有利"与"$T$ 可选时全BPTT最优"两种结论,并用分别拟合的律和等算力实验交叉验证。(b)增益形式:循环增益取有界指数饱和而非幂律(Schwethelm)、线性(Parcae)或加性(Ouro);作者用对held-out前向循环数的预测精度证明有界形式更符合数据。(c)决策还原:把截断决策还原为只依赖 $\chi(T,B)$、与总算力无关的整数优化问题,并给出严格引理(D.1、D.2),说明结论 $q>2p\Rightarrow$ 全BPTT的解析依据。(d)分配规则的来源:$T^*\propto C^{0.24}$、$D^*\propto C^{0.76}$ 等指数是从拟合律的渐近驻点条件推导出来的,而非对经验最优点的二次拟合,把Parcae只在固定尺寸、固定截断规则下的分析推广到联合选择 $T,B$。

实验结果

六组发现均有数字支撑。(1)律拟合极好:230配置拟合集 $R^2=0.9992$,RMSE $6.14\times10^{-3}$ nats/token;留出模型尺寸或token预算的LOO误差仅7.48和 $6.42\times10^{-3}$,留出循环设置最难($8.67\times10^{-3}$)。(2)循环增加有效参数、截断打折:$N_{eff}/N$ 从 $(1,1)$ 的1.00升至 $(2,2)$ 的1.46、$(4,4)$ 的1.76、$(8,8)$ 的1.91,系数趋向上界 $e^c\approx3.04$;且 $q\approx1.02>p\approx0.24$,同起点下 $B$ 翻倍比 $T$ 翻倍更值钱。(3)固定 $T$ 时截断有利:$T=8$ 时律选 $B^*=2$,$\chi(8,2)=1.580$ 远小于 $\chi(8,8)=2.376$,预测等损失省34%算力($N^*$、$D^*$ 各增1.51、1.20倍);单独拟合的律证实 $T=4$ 省5–11%、$T=8$ 省33–35%。(4)$T$ 自由时全BPTT最优:等算力实验($D=16D_0\approx172$亿token)中 $(4,4)$ 在60M–1B所有尺寸上均低于 $(8,2)$,$\Delta L$ 从 $2.78\times10^{-3}$ 降至 $0.87\times10^{-3}$ nats;等损失下省8–14%算力。(5)固定尺寸最优分配 $T^*\propto C^{0.24}$、$D^*\propto C^{0.76}$:算力×10对应循环×约1.7、token×约5.8,五个尺寸一致。(6)尺寸自由时不循环最优、两次完整循环最佳:$(2,2)$ 与计算最优非循环模型差距≤0.008 nats却少27–32%参数;$2.1\times10^{19}$ FLOPs处多训8%后,$(2,2)$(344M)损失2.36反超非循环(471M)的2.37。

任务指标本文基线提升
缩放律对验证损失的预测(含留一验证) RMSE(nats/token,×10⁻³) 拟合集内 6.14;留出一个配置 6.40;留出一个模型尺寸 7.48;留出一个token预算 6.42;留出一个循环设置 8.67 此前工作为各架构/各配方分别拟合,无统一 $T,B$ 联合律可比 单条统一律覆盖15种循环设置与5个尺寸,$R^2=0.9992$,对未见过的尺寸和预算误差仅约0.006–0.007 nats
固定前向循环数 $T=8$ 时截断与全BPTT的等损失计算 等损失所需训练计算节省 律预测 $B^*=2$ 时节省34%($\chi(8,2)=1.580$ vs $\chi(8,8)=2.376$);对 $T=8$ 单独拟合的律证实33–35% 全BPTT $B=T=8$ 截断在 $T$ 被规定时确有优势,且节省随 $T$ 增大($T=4$ 时为5–11%)
等算力匹配实验:$(4,4)$ vs 截断 $(8,2)$($D=16D_0$≈172亿token) 验证损失差 $\Delta L$(nats/token,×10⁻³) $(4,4)$ 全BPTT在所有尺寸更低:60M 2.78、120M 1.47、240M 1.36、480M 1.07、1B 0.87 $(8,2)$ 截断BPTT(同为 $T+2B=12$ 循环算力) 等损失下 $(4,4)$ 少用8–14%计算,验证 $T$ 可选时全BPTT最优
参数效率:$(2,2)$ 两次完整循环 vs 计算最优非循环模型 验证损失(nats/token)与非嵌入参数量 等算力下差距≤0.008 nats;$2.1\times10^{19}$ FLOPs处344M参数、+8%算力后损失2.36 非循环模型471M参数、损失2.37 少27–32%非嵌入参数达到近乎相同损失;多花8%计算即反超0.007 nats

局限与改进

作者明确承认四点局限:只研究稠密(dense)模型,未涉及循环MoE架构(而循环MoE常采用两次循环),截断与算力分配结论能否迁移未知;只测量验证损失,未验证观察到的收益能否转化为下游任务表现或推理能力;实验止步于 $T\le16$(Schwethelm等到8循环、Parcae平均12循环),增益上界 $e^c$ 与更深循环处的切换点均属外推,而更深循环还会加剧残差爆炸与loss spike等不稳定问题;每个尺寸的prelude/coda与循环块参数比例固定,未探索不同比例下的循环收益与最优循环数。我的补充观察:FLOP计数采用leading-order近似(每参数每token前向2、反向4,忽略注意力项、输出投影与优化器更新),长上下文或大词表场景下 $\chi$ 的排序可能改变;§6.3的关键差距只有0.007–0.008 nats、每模型仅3个种子,已接近种子间噪声量级;所有系数在单一语料FineWeb-Edu(教育网页)上拟合,对代码、数学等分布的普适性存疑;饱和假设基于信息存储论证,属经验建模而非理论必然。

独立分析的弱点

弱点一:计算模型粗糙。$C\simeq 6D\cdot N_{comp}$ 忽略注意力的序列长度项、输出投影与优化器更新,在实验的短上下文小模型上尚可,但注意力占比随上下文上升,$\frac{T+2B}{3}$ 的等价关系会失真,可能动摇 $B^*=2$ 与全BPTT结论;改进方向是把注意力与优化器FLOPs显式纳入重新拟合并重解分配。弱点二:结论完全依赖拟合值 $q\approx1.022>2p\approx0.474$,论文未报告 $c,p,q$ 的置信区间,也没有数据重采样下的稳健性检验;若 $q$ 逼近 $2p$,"全BPTT最优"将失去解析依据,建议bootstrap给出参数区间。弱点三:只看验证损失——循环模型常被宣传对推理/思维链有优势,仅凭损失选 $T,B$ 可能与下游能力错位。弱点四:所有结论绑定单一架构选择(RMSNorm、输入注入、ReLU² FFN、固定 $N_{rec}/N$),残差/超连接、MoE路由等改良结构下 $g$ 的形状可能不同。弱点五:整数循环数的台阶效应使固定尺寸分配律只是渐近结果,小算力区间偏差可能不小,实践应结合逐点IsoFLOP扫描。

未来方向

作者提出的方向:把统一律扩展到循环MoE模型;验证损失的收益能否转化为下游任务与推理能力;突破 $T=16$ 并发展稳定训练配方以应对残差爆炸与loss spike;变化prelude/coda与循环块的参数占比,研究参数分配对循环收益的影响。基于成果可延伸的方向:把截断BPTT与不动点梯度/截断Neumann级数的既有理论联系做定量刻画,从理论上解释为何 $q\gg p$(反向深度远比前向深度重要);训练–推理协同设计——循环次数可变的自适应深度与早退机制下,最优训练深度与推理预算可能不同,律可以给出先验;在代码、数学等不同数据分布上重拟合 $c,p,q$,检验饱和上界 $e^c$ 的普适性;把"两次完整循环"的参数效率结论放大到7B级实用规模,并结合蒸馏、深度共享微调等手段落地;以及把该律用作实际预训练的循环配置决策工具(类似Chinchilla在非循环模型中的角色)。

复现评估

复现条件总体较好:代码开源于 https://github.com/xjxyys/Truncate-or-Not,数据为公开的FineWeb-Edu(100B token),prelude–recurrent–coda架构细节(输入注入 $W_{inj}$ 初始化为 $[I\ 0]$、RMSNorm、stop gradient位置)与训练超参(MuonH用于矩阵、AdamW用于嵌入/输出/归一化,学习率与batch规则)在附录B.3有完整描述,五个模型配置(宽度、层数、$N_{rec}/N$)在表2给出。但完整复现缩放套件需592次训练、最大919M非嵌入参数、单次运行算力最高约 $10^{20}$ FLOPs量级,估计要数十万GPU小时,普通实验室难以承受。务实的做法是部分复现:在中尺寸(如240M)上做几组等算力匹配对比($(4,4)$ vs $(8,2)$,或 $(2,2)$ vs 非循环),以及在1–2个尺寸上重拟合九参数看量级是否一致,即可低成本检验核心主张。难度评估:部分复现中等,全量复现困难。

相似论文推荐