跳过层还是循环利用?预训练大语言模型的测试时深度自适应

发表
Tianyi ZhouTianyi Zhou 提交
作者: Ziyue Li, Yang Li, Tianyi ZhouTianyi Zhou

摘要

预训练神经网络能否在不进行任何微调的情况下,根据不同的输入调整其架构?对于简单任务,我们是否需要所有层?对于挑战性任务,所有层是否都足够胜任?我们发现,预训练大型语言模型(LLM)的层可以作为独立的模块进行操作,从而为每个测试样本构建一个更好甚至更浅的定制模型。具体而言,预训练模型中的每一层都可以被跳过/剪枝,或者作为循环神经网络(RNN)重复多次,并以任意顺序与其他层堆叠,从而为每个样本生成一个层链(CoLa)。这种组合空间大大扩展了现有关于循环/递归预训练模块、层剪枝或提前退出网络等工作的范围。我们开发了一种蒙特卡洛树搜索(MCTS)协议,用于探索并识别数学和常识推理基准测试中每个样本的最佳CoLa。与固定深度的静态模型相比,CoLa允许快捷路径(快速思考)、相同层(或多层)的递归(慢速思考),并能结合两者,为不同的输入提供更灵活、动态的架构。我们对经MCTS优化的CoLa进行了广泛分析,这带来了两个关键发现:(1) 对于超过75%的原始LLM预测正确的样本,我们能找到更短的CoLa,这表明在提高推理效率方面有很大的空间;(2) 对于超过60%的原始预测不正确的样本,我们能识别出实现正确预测的CoLa,这表明在性能提升方面有很大的空间。我们的结果凸显了使用预训练LLM的固定架构对不同样本进行推理的不足,并为释放测试时深度自适应的泛化能力铺平了道路。
查看 arXiv 页面查看 PDF

评论

Tianyi ZhouTianyi Zhou
论文作者
论文提交者

我们发现,预训练大型语言模型(LLM)的层可以被当作独立的模块来操作,以构建针对每个测试样本定制的、更好甚至更浅的模型。具体来说,来自预训练LLM的每一层可以被跳过或作为循环神经网络(RNN)重复多次,并以任意顺序与其他层堆叠,从而为每个样本生成一个层链(CoLa)。这种组合空间显著扩展了现有关于循环或递归预训练模块、层剪枝或早期退出网络等工作的范围。

Screenshot 2025-07-11 at 1.14.03 AM.png

我们开发了一种蒙特卡洛树搜索(MCTS)协议,以探索并识别适用于数学和常识推理基准测试中每个样本的最优CoLa。与固定深度的静态模型相比,CoLa允许捷径(快速思考)、同一层(或多层)的重复(慢速思考),并结合两者,为不同的输入提供更灵活、动态的架构。具体来说,

  • 我们引入了一个新的泛化维度,它将静态预训练LLM转变为自适应深度的动态架构,而无需训练任何参数:对于不同的测试样本/任务,预训练的层可以被跳过、重复和组装,以创建更好(更准确和/或更浅)的CoLa模型,而无需进一步训练。

  • 我们开发了一种MCTS协议,用于为每个样本高效地搜索具有自适应深度的CoLa架构。对所获得的CoLa模型中模式的深入分析揭示了关于不同大小的预训练/微调模型在不同深度处层的重要性及冗余性的关键见解,这些见解也因不同难度级别的任务而异。

我们对MCTS优化的CoLa进行了广泛分析,得出了两个关键发现:

(1) 对于原始LLM预测正确的超过75%的样本,我们可以找到更短的CoLa,这表明在提高推理效率方面有很大的空间;

(2) 对于原始LLM预测错误的超过60%的样本,我们可以识别出能够实现正确预测的CoLa,这表明在性能提升方面有很大的空间。

Screenshot 2025-07-11 at 1.21.28 AM.png

Screenshot 2025-07-11 at 1.21.44 AM.png

我们的结果突出了使用固定架构的预训练LLM对不同样本进行推理的缺点,并为释放测试时深度自适应的泛化能力铺平了道路。

Michael YMichael Y

你好 @zhoutianyi, 感谢分享你的论文——发现非常有趣!我对算法1有一个小问题。模拟步骤提到它将“评估在保留输入上的路径准确性”。算法是如何实际选择这些保留输入的?我猜测它们需要与当前应用CoLA的输入(紧密?)相关,但我找不到细节,或者我可能理解错了什么。

谢谢!

Tianyi ZhouTianyi Zhou
论文作者
论文提交者

你好 @myeesw,很高兴听到你觉得我们的发现很有趣。这里的“保留输入”指的是来自保留测试集,且未用于模型训练的输入。我们会在后续版本中使其更清晰。在这项工作中,我们没有探讨将CoLA从一个输入推广到其他(相似)输入的可能性。但你提到的正是我们正在尝试的方向,我们会在后续的预印本中分享结果。很好的观点!谢谢!

Michael YMichael Y

感谢你的澄清!并期待看到该领域的未来工作。

Eni GrandEni Grand

你好 @zhoutianyi。感谢分享你的论文。

关于表1有一个问题:你们是针对不同的数据集搜索了不同的层推理策略,还是只使用了一个策略?

Tianyi ZhouTianyi Zhou
论文作者
论文提交者

嗨,@Enigrand,我们针对每个样本搜索/优化CoLA所用的MCTS策略在所有数据集中都是相同的,但搜索到的架构在不同样本之间有所不同。

Eni GrandEni Grand

感谢 @zhoutianyi 的回复。你的意思是,在你进行的所有测试中,你总共只对一项全球策略执行了一次搜索吗?

Mostafa ElhoushiMostafa Elhoushi

令人印象深刻的成果!恭喜!👏

Tianyi ZhouTianyi Zhou
论文作者
论文提交者

谢谢!