UpLoop:零额外算力训练循环Transformer,性能显著提升

UpLoop通过渐进式训练调度,在总FLOPs不变下训练循环Transformer,提升推理能力。
UpLoop是一种训练循环Transformer的预训练方案,旨在解决循环模型训练算力随循环次数线性增长的问题。它采用渐进式深度增长策略:训练初期使用浅层模型以减少FLOPs,后期通过复制层并引入循环迭代,将早期节省的算力用于重复执行共享中间块。最终在参数数量、训练token数和总FLOPs上与标准密集基线完全对齐。在1.3B参数规模下,22个下游任务平均得分从22.99提升至27.52,且经过监督微调后数学推理任务同样受益。该方法为资源有限的研究团队提供了不增加算力预算即可获得更强推理能力的训练范式。
循环语言模型通过共享权重重复计算来增强推理能力,在不增加独立参数的情况下提升执行深度。但每多一次前向传播,训练浮点运算量(FLOPs)就会增加,导致这些增益的获取成本更高。为此,研究者提出UpLoop——一种利用upcycling训练循环Transformer的预训练方案,在最终参数数量、总训练token数和训练FLOPs上与标准密集基线完全匹配。

循环Transformer的算力困境
循环Transformer的核心思想是让模型多次经过同一组共享权重的层,以此模拟更深网络的计算效果。这种设计在推理任务上表现出色,因为重复计算相当于给了模型更多"思考"时间。然而,训练阶段的算力开销不容忽视:每次额外的循环都意味着一次完整的前向和反向传播,FLOPs随循环次数线性增长。
这形成了一个矛盾:想要更强的推理能力,就需要更多循环;但更多循环又推高了训练成本。现有方案往往在参数效率和训练效率之间做取舍,难以两全。UpLoop的出发点正是打破这种取舍——在不增加总训练FLOPs的前提下,获得循环结构带来的性能收益。
循环Transformer中的“循环”指的是让输入序列多次通过同一组Transformer层(通常是一个中间块)。每次通过称为一次“循环迭代”或“前向传播”。在标准Transformer中,每一层有独立的参数;而在循环Transformer中,这些层共享权重,因此模型的总参数量不随循环次数增加。这种设计借鉴了深度递归网络的思想,通过增加计算深度(而非参数深度)来提升模型的表达能力。然而,训练时需要为每次循环迭代计算梯度,导致反向传播的链式法则需要沿循环展开,从而显著增加训练时的内存和计算开销。UpLoop的核心挑战就是在不增加总FLOPs的情况下,让模型学会利用循环计算。
UpLoop如何做到"零额外算力"
据论文描述,UpLoop采用了一种渐进式的训练调度策略。训练初期,模型是一个较浅的Transformer,计算负担较轻。随着训练推进,通过复制已学习的层来逐步增加模型深度。关键在于,早期浅层模型节省下来的算力被"存储"起来,用于训练后期对共享中间块进行重复执行。
换句话说,UpLoop把训练过程重新分配:前期少算,后期多循环。最终在参数数量、训练token数和训练FLOPs三个维度上,与标准密集基线完全对齐。这种设计不需要额外的预训练算力,却能让模型在训练结束时具备循环执行的能力。
渐进式训练调度(progressive training schedule)是一种在训练过程中动态调整模型结构或计算图的策略。常见做法包括逐步增加网络深度、宽度或输入分辨率。UpLoop的具体做法是:训练开始时使用较浅的模型(例如只有标准深度的一部分),此时每个训练步骤的FLOPs较低;随着训练进行,通过复制已训练好的层来增加深度,并逐步引入循环迭代。这样,早期节省的FLOPs被累积起来,用于后期训练循环块。最终总FLOPs与标准密集基线持平。这种调度要求精心设计复制时机和循环引入点,以避免训练不稳定。
实验效果:下游任务平均分从22.99提升至27.52
实验结果覆盖了多个模型规模。在1.3B参数规模下,22个下游任务的平均得分从基线的22.99提升到27.52,增幅明显。这一提升不仅体现在语言建模指标上,也反映在下游任务的实际表现中。
更值得关注的是,经过监督微调后,UpLoop带来的增益还能延伸到数学推理任务。这意味着循环结构所强化的计算能力具有可迁移性,并非只对预训练阶段的损失函数优化有效。
监督微调(Supervised Fine-Tuning, SFT)是指在预训练模型基础上,使用标注数据对模型进行额外训练,使其适应特定任务。在UpLoop的实验中,经过SFT后,模型在数学推理任务上的表现得到提升,说明循环结构带来的计算能力增强可以迁移到需要多步推理的任务上。数学推理任务通常需要模型进行链式思考(chain-of-thought),而循环计算提供的额外“思考时间”恰好有助于此类任务。需要注意的是,SFT的数据和任务选择会影响迁移效果,UpLoop的增益是否广泛适用于其他推理任务仍需验证。
对循环模型训练范式的启示
UpLoop的价值在于提供了一种新的训练思路:循环模型的高算力成本并非不可规避,通过训练调度的重新设计,可以在总预算不变的情况下获得额外收益。这对于资源有限的研究团队尤其有意义——无需增加GPU集群规模,就能训练出推理能力更强的模型。
当然,该方法目前主要验证了语言建模和数学推理场景,在更广泛任务上的泛化能力还有待进一步观察。但至少从现有数据看,UpLoop为循环Transformer的实用化迈出了扎实的一步。
相关推荐

Claude Sonnet 6疑泄露,Anthropic内部模型加速惊人
Anthropic疑似泄露Sonnet 6,Meter报告称其AI研发加速1.9倍;OpenAI发布GPT-6 Cyber,GPT-6.2 Astra或将至;Google Gemini Omni 1.2 Flash登顶视频生成榜;Meta stealth模型与实时头像亮相;市场波动下Google仍投170亿美元建数据中心。

科幻小说对AI的预测有多准?
科幻小说对AI的预测有多准?本文结合Hacker News讨论,分析科幻作品在AI发展中的远见与盲区,探讨其真正价值在于启发思考而非精准预测。

SwarmEvolve:潜空间智能体集群加速进化搜索
SwarmEvolve框架让数百个低成本智能体在潜空间中集群推理,通过KV缓存共享信息,在数学发现和科学计算问题上实现快速进化搜索,中位数仅3-8分钟达到其他框架60分钟的最佳分数。