机器之心报道
编辑:佳琪、泽南
不需要特殊训练数据,能处理语言难以形容的任务,3.5B 参数能实现 50B 的性能。
这是一种全新的语言模型架构,能够通过使用循环语言模型在潜在空间中隐式推理,显著提升模型的计算效率,尤其是在需要复杂推理的任务上。
近日,马里兰大学的一篇论文在 AI 研究社区中引发了关注,其提出的语言模型通过迭代循环块来工作,能在测试时展开到任意深度。这与当前通过生成更多 token 来扩展计算的主流推理模型形成了鲜明的对比。
仅在上个月,Hugging Face 上的下载量就达到了 4500+。
深度循环(Recurrent Depth)方法无需生成大量「思考」token 即可获得高性能。与基于思维链的方法不同,马里兰大学的方法不需要任何专门的训练数据,可以使用小的上下文窗口,并且可以捕获不易用文字表示的推理类型。
该工作构建的概念验证模型达到 35 亿参数和 8000 亿 token,实验验证了新方法可以提高其在推理基准上的性能,尤其是在需要复杂推理的数学和编程问题上,最高相当于 500 亿参数的计算负载。
- 论文链接:https://arxiv.org/abs/2502.05171
- 模型下载: https://huggingface.co/tomg-group-umd/huginn-0125
- 代码链接: https://github.com/seal-rg/recurrent-pretraining
人类在解决某些问题时自然会花费更多的脑力。虽然人类能够通过用语言表达中间结果并写下来进行长时间的思考,但在说出答案的第一个字之前,大量的思考是通过大脑中复杂、反复的放电模式进行的。
而在 AI 一侧,早期提升语言模型能力的尝试聚焦于扩大模型规模,这种方法需要海量数据和算力支撑。近年来,研究者开始探索通过扩展测试时计算量来提升模型推理能力。主流方法是对长链条的思维链示例进行后训练,以开发模型在上下文窗口中语言化中间计算过程的能力,从而实现思维外显。
然而,将昂贵的内部推理过程始终压缩为单一的语言化下一个 token 的约束显然是低效的 —— 如果模型能原生地在连续的潜在空间中「思考」,其能力很可能得到显著提升。释放这一未被开发的计算维度的一种方法是为模型添加循环单元。该单元以循环方式运行,持续处理和更新隐藏状态,使得计算得以无限延续。尽管这并非当前主流范式,但这一思想贯穿机器学习发展史,每隔十年便以新形式重现:从循环神经网络(RNN)到扩散模型(Diffusion model),再到通用 Transformer 或循环 Transformer 架构。
在新工作中,研究人员展示了深度循环语言模型(depth-recurrent language models)可以有效学习、高效训练,并在测试时计算扩展的情况下展示出明显的性能改进。作者提出的 Transformer 架构建立在潜在深度循环块之上,在训练期间运行随机采样的迭代次数。作者展示了这种范式可以扩展到数十亿个参数和超过 5 万亿个预训练数据标记。在测试时,该模型可以通过潜在空间中的循环推理来提高其性能,使其能够与其他受益于更多参数和训练数据的开源模型竞争。
此外,作者展示了循环深度模型在推理时自然支持许多功能,这些功能在非循环模型中需要大量的调整和研究工作,例如每个 token 的自适应计算、(自)推测解码和 KV 缓存共享。通过跟踪潜在空间中的 token 轨迹来表明,许多有趣的计算行为会随着规模的扩大而出现,例如模型在潜在空间中旋转形状以进行数值计算。
为什么要使用深度循环来训练模型?
循环层使得 Transformer 模型能够在生成一个 token 之前执行任意数量的计算。从原理上讲,循环机制为测试时计算量的扩展提供了一种简洁的解决方案。作者认为,与更传统的长上下文推理方法(如 OpenAI 的 o1、DeepSeek-AI 的 R1)相比,潜在循环思考具备多项优势:
- 无需定制化训练数据:链式思维推理要求模型在特定领域的长示例上进行训练。而潜在推理模型则能在标准训练数据上训练,无需专门示例,且可根据计算预算灵活调整,在测试时通过额外计算资源增强能力;
- 更低的内存需求:潜在推理模型在训练和推理时所需内存少于链式思维推理模型。后者因需极长上下文窗口,可能需采用如 token 并行化等专门训练方法;
- 更高的计算效率:循环深度网络相较于标准 Transformer,每参数执行的浮点运算(FLOPs)更多,显著降低了大规模训练时 AI 计算卡之间的通信成本,尤其在低速互连条件下更能提升设备利用率;
- 促进「思考」而非记忆:通过构建计算密集但参数规模较小的架构,该工作期望强化模型构建「思考」(即学习元策略、逻辑与抽象)而非单纯记忆来解决问题的先验倾向。此前,循环先验在学习复杂算法方面的优势已在「深度思考」相关文献中得到验证。
从哲学视角看,我们肯定希望 AI 的潜在推理能捕捉人类推理中难以言表的部分,如空间思维、物理直觉或(运动)规划。通过循环过程的多次迭代,在高维向量空间中进行推理将允许同时深度探索多个方向,而非线性思维,从而催生能够展现新颖且复杂推理行为的系统。
以这种方式扩展计算能力并不排斥通过延长(语言化)推理或增加预训练参数数量来提升模型性能。因此,它可能构建出模型性能扩展的第三维度。
可扩展的循环架构
该研究提出的具有潜在循环深度的 Transformer 架构,模型主体结构基于仅解码器(decoder-only)的 Transformer 模块。然而,这些模块被组织成三个功能组:前奏(Prelude)P,通过多层 Transformer 将输入数据嵌入潜在空间;核心循环块(Core Recurrent Block)R,作为循环计算的核心单元,负责修改状态 s ∈ R n×h;以及尾声(Coda)C,通过若干层从潜在空间解嵌入,并包含模型的预测头。
核心块置于前奏与尾声之间,通过循环核心块,我们能在「歌曲」中插入无限数量的「诗节」。
模型架构如下图所示:
给定一定数量的循环迭代 r 和一系列输入标记 x ∈ V n,这些组以以下方式使用来产生输出概率 p ∈ R n×|V |:
其中 σ 是初始化随机状态的某个标准差。该过程如图 2 所示。给定一个初始随机状态 s0,模型重复应用核心块 R,它接受潜状态 si−1 和嵌入输入 e 并输出一个新的潜在状态 si 。完成所有迭代后,coda 块处理最后一个状态并产生下一个 token 的概率。该架构基于深度思考文献,其中表明在每一步中注入潜在输入 e 并用随机状态初始化潜在向量可以稳定递归并促进收敛到与初始化无关的稳定状态,即路径独立性。
在每个组中,模型大致都遵循标准的 Transformer 层设计。每个块包含多个层,每个层包含一个标准的因果自注意力块,使用 RoPE,基数为 50000,以及一个门控 SiLU MLP。作者使用 RMNSorm 作为规范化函数。为了稳定递归,模型按以下「三明治」格式对所有层进行排序,使用范数层 ni:
为了在训练时保持较低的计算量和内存,研究人员在训练时随机采样迭代次数,并通过截断反向传播来减少计算和内存开销。模型只对最后几次迭代进行反向传播,在保持训练效率的同时能够适应不同深度的循环计算。
实验结果
由于这个模型很难与其他基于 Transformer 架构大语言模型直接比较。它仅包含 3.5B 参数,在预训练阶段仅需少量互连带宽。然而,其实际算力(FLOPs)消耗要接近 32B 参数 Transformer 的水平。在测试中,该模型能够无限制地提升算力,最终达到与 50BTransformer 相当的性能水平。
该模型的训练有一些瑕疵:只训练了 47000 步,学习率从未降低,数据集也比较小,只有 800B tokens。这和现在动辄上万亿 tokens 的工业模型有很大差距。
尽管如此,它在 ARC、HellaSwag、MMLU、OpenBookQA、PiQA、SciQ 和 WinoGrande 的成绩上已经可以和 OLMo-7B 媲美,在多项任务上超过了老牌的 Pythia 模型。
作为首个在这个量级上训练的「深度循环」模型,这个结果已经可喜可贺了。这也证明了「潜在循环」是一个可行的思路,值得探索。在推理时动态地增加算力,也有希望达到上亿 token 级别的水平。
数学和代码测试
该团队在 GSM8k、MATH、Minerva 和 MathQA 上测试了模型的数学能力;在 MBPP 和 HumanEval 上测试了编程水平。
在数学推理方面,该模型大大超过了除 OLMo-2 模型以外的所有模型。在编程基准测试中,该模型击败了所有其他通用开源模型。不过没有 StarCoder2 等「编程专家模型」。
研究团队还注意到,如下图所示,虽然语言建模的进步会放缓,在这个训练规模下是正常的,但编程和数学推理在整个训练过程中都处于稳步上升的状态。
递归在哪里起效了?
递归到底在哪方面带来了提升?模型表现的提升到底是递归结构的功劳,还是数据集、分词方式、模型架构等其他因素在起作用?
为了找到答案,研究团队做了对比实验:他们用完全相同的条件训练了一个非递归模型,让两个模型都学习了 1800 亿个 token。
结果显示,递归模型的整体表现更胜一筹,尤其是在处理难度较大的任务时尤为明显,比如 ARC 挑战测试集。不过,像 SciQ 这样主要考验科学知识记忆的基础任务上,两个模型的差距就不那么大了。而在 GSM8k 数学推理任务上,才训练到 1800 亿 token,递归模型的成绩就已经比基准模型好了整整 5 倍!
如果限制递归模型只能递归一次,它从 1800 亿到 8000 亿 token 期间的表现就几乎没有进步。这说明模型的提升完全来自于递归模块的多次运算,而非前后的非递归层。
此外,研究团队还做了一个测试,看看在不同任务上增加计算量会带来什么效果。结果显示,简单任务很快就能达到性能上限,而复杂任务则能从更多的计算中持续获益。
递归次数和上下文长度的关系
下图展示了模型在 ARC-C 测试中的表现如何随着递归次数和少样本示例数量的变化而变化,当上下文中没有少样本示例时,模型仅需 8 到 12 次递归即可达到性能上限。
但是,当给模型提供更多上下文信息时,它就像得到了更大的「思考空间」。比如提供 1 个示例时,模型需要 20 次递归才能发挥最佳水平;如果给到 25-50 个示例,则需要 32 次递归。
OBQA 测试也出现了同样的情况。当研究团队为每道题提供相关背景知识时,递归模型的表现突飞猛进,几乎追平了 OLMo-2 的水平。这个结果其实很好理解:递归模型虽然不太擅长记忆大量事实,但在推理和利用上下文信息方面却相当出色。
递归深度如何简化 LLM?
测试阶段的零样本自适应计算
该团队发现,其模型能够根据不同的任务灵活调整计算量。
对于简单问题,模型只需进行较少的计算;而对于复杂问题,则会增加计算轮次。这不正是我们期望的效果吗?
为了量化观察,他们采用了一个简单的判断标准:观察连续两步计算结果之间的 KL 散度。如果这个差异小于 0.00005,便判断为可以停止计算了。
如图 10 所示,在处理 MMLU 测试的前 50 个问题时,不同类型的题目所需的计算步数差异很大。例如,高中数学题相对简单,计算步数较少;面对需要道德判断题时,平均需要多算 3.5 步。而在 MTBench 测试中,这种自适应计算机制几乎不会影响模型的对话表现。
零样本 KV 缓存共享
为了提高效率,研究团队还尝试减少了 KV 缓存的内存占用。他们给每个 token 设定了固定的 KV 缓存预算,比如 16 步。这样,第 17 步的计算就会覆盖第 1 步的缓存,以此类推。在 MTBench 测试中,即使把缓存预算降到 4,模型表现依然稳定。
零样本连续思维链
为了做到这一点,研究团队设计了一种方法:在生成新 token 时,不是每次都从随机状态开始,而是利用上一个 token 的最终状态进行「热启动」。图 10 显示,这种方法可以减少 1 到 2 步的计算量。特别是在处理哲学问题时,通过重复利用之前的计算结果,模型更容易提前完成任务。
零样本自推测解码
递归模型还有一个天然优势:不需要额外的「草稿模型」就能实现推测解码。换句话说,模型可以用较少的迭代次数先生成后面 N 个 token 的草稿,然后用更多的迭代次数 M(M>N)来验证。
这个过程可以分多个阶段进行,「草稿模型」还可以使用自适应计算。最棒的是,草稿阶段的计算结果可以在验证阶段重复使用,不会浪费计算资源。
模型在递归过程中到底在做什么?
通过从潜在空间沉浸式观察模型的「思考」过程 ,研究团队发现了和我们通常理解的「思维链」截然不同的现象。
研究团队分析了每个递归状态 s_i 与极限点 s^∗(通过 128 次迭代计算得到)之间的距离变化。结果显示,模型的收敛行为与上下文密切相关。
在处理问题的关键部分以及开始生成回答时,模型会更多「深思熟虑」。即使是完全相同的符号(例如省略号中有两个重复的...),模型也会因上下文不同而表现出不同的处理方式。
他们对所有 token 的轨迹进行了 PCA 分解,并将个别轨迹投射到前六个 PCA 方向上。结果显示:
- 简单标记往往直接收敛到一个固定点。
- 在处理复杂问题时,token 的状态会形成程式。
- 一些关键词(如「makes」和「thinks」)经常出现程式,这些词往往决定了回答的结构。
- 某些 token(如「wrong」)会「滑动」,其轨迹会朝着特定方向漂移。
这表示模型并没有采用传统的线性推理方式(如思维链),而是自主学会了利用高维空间来实现更丰富的几何模式,包括程式、收敛路径和漂移等。
尽管模型内部呈现出复杂的动态特征,但如下图所示,研究团队证实了它仍然保持着路径独立性。无论从哪个起点 s_0 开始,模型都会形成相似的轨迹,展现出一致的行为模式 —— 同样的轨迹、固定点或方向漂移。
更多研究细节,请参看论文原文。
参考内容:
https://x.com/iScienceLuvr/status/1888792081382137966
https://arxiv.org/abs/2502.05171