Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs

一个预训练的神经网络能否在不进行微调的情况下，根据不同的输入自适应地调整其架构？对于简单的任务，我们是否需要用到所有网络层？而对于具有挑战性的任务，这些网络层又是否足够？我们发现，预训练大语言模型（LLM）中的每一层都可以被当作独立模块进行操作，从而为每个测试样本构建出更优甚至更浅的定制化模型。具体而言，预训练模型中的每一层可以像循环神经网络（RNN）一样被跳过/剪枝或多次重复使用，并且可以以任意顺序与其他层堆叠组合，从而为每个样本生成一种“层链”结构（chain-of-layers，简称CoLa）。这种组合空间大大拓展了现有研究的范围，包括循环/递归式预训练模块、层剪枝以及提前退出网络等方面的工作。 我们开发了一种蒙特卡洛树搜索（MCTS）协议，用于探索并识别来自数学和常识推理基准测试中每个样本的最佳CoLa配置。与采用固定深度的静态模型相比，CoLa允许存在捷径路径（快速思考）、某些层的重复使用（慢速思考），或将两者结合，从而为不同输入提供更加灵活、动态的架构选择。我们对MCTS优化后的CoLa进行了广泛分析，得出了两个关键发现：（1）对于原LLM模型预测正确的大于75%的样本，我们能够找到更短的CoLa配置，表明推理效率还有很大的提升空间；（2）对于原模型预测错误的大于60%的样本，我们也能够找到能实现正确预测的CoLa配置，说明性能提升同样存在巨大潜力。我们的结果突显了在处理不同样本时使用固定架构的预训练LLM在推理上的局限性，并为解锁测试时深度自适应所带来的更强泛化能力提供了新思路。

论文试图解决预训练大型语言模型（LLM）在推理过程中使用固定架构所带来的效率和性能瓶颈问题，验证了是否可以通过动态调整网络结构（如跳过、重复某些层或重新排列层的顺序）来提升推理效率和准确性。

论文的核心思路是将预训练LLM中的每一层视为可组合的模块，通过蒙特卡洛树搜索（MCTS）为每个测试样本寻找最优的“层链”（Chain-of-Layers, CoLa），从而实现无需微调的架构自适应。相比传统方法中模型结构固定的做法，该方法引入了测试时深度与结构的灵活性，结合了“快速思考”（浅层路径）和“慢速思考”（层重复）机制。

1. 提出了一种新颖的测试时动态架构调整方法CoLa，突破了传统的静态深度限制 2. 实验显示，在超过75%的原本正确预测样本上可以找到更短的CoLa链以提升推理效率 3. 在超过60%的原本错误预测样本上，CoLa能够找到更优的层组合来提升准确率 4. 实验基于数学和常识推理基准数据集进行，具有实际意义 5. 尚未提及代码开源情况，但方法本身具有较强的可复现性

1. Looped and recurrent use of pretrained modules (如Retro-reading等机制) 2. Layer pruning in transformers (如《What Does BERT Compress?》系列研究) 3. Early-exit networks for adaptive inference (如《Deeply Informed Neural Networks》) 4. Neural architecture search (NAS) at test time 5. Dynamic computation graphs and conditional execution in deep learning