论文《Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA》的核心是通过参数共享将现有的大语言模型压缩为更小、更高效的递归模型同时尽量保持原模型的性能。以下是全面总结1. 研究动机与问题LLM部署成本高内存大、计算重。参数共享层间复用权重可以降低模型尺寸和推理成本但在现代LLM中效果有限。已有工作多在小模型上从头训练缺乏对现有预训练大模型的有效压缩方案。2. 核心方法2.1 递归TransformerRecursive Transformer将原模型的 LL 层压缩为一个只有 K 层的块然后在推理时循环重复使用该块多次CYCLE策略。参数量大幅减少例如 Gemma 2B → 递归 Gemma 1B参数减半。关键点从预训练模型初始化共享层而不是随机初始化。三种初始化策略方法说明Stepwise步进法等间隔选取原模型的层保留首尾层 → 递归模型效果最好Average平均法将绑定层的权重取平均Lower下层法直接取原模型的前 KK 层2.2 放松递归TransformerRelaxed Recursive Transformer在共享层上附加层级专属的LoRA模块为每一层引入少量可训练的低秩增量参数。这样既保留了参数共享的紧凑性又给每层一定的灵活性。LoRA初始化用截断SVD分解原模型权重与共享权重的残差矩阵使共享权重LoRA能近似原权重。秩 r控制放松程度r0 退化为纯递归模型r 足够大则逼近原全尺寸模型。3. 推理加速连续深度方向批处理Continuous Depth-wise Batching递归模型中不同样本可以在不同的循环深度迭代次数退出。结合提前退出Early-Exiting如果中间循环的预测已经足够好与最终输出一致该样本可以提前结束计算。这使得在批处理中新请求可以立即填充已完成样本的空位提高GPU利用率。理论分析显示相比普通Transformer可获得2~3倍的吞吐量提升。4. 实验与主要结果实验设置模型Gemma 2B、TinyLlama 1.1B、Pythia 1B上升训练Uptraining在SlimPajama上继续训练15B或60B token评估7个少样本任务LAMBADA、HellaSwag、PIQA、WinoGrande、ARC等主要结果结论说明递归模型优于同等参数量的从头训练模型递归Gemma 1B比非递归Gemma 1B准确率提升13.5%放松模型可逼近全尺寸模型秩512的放松Gemma达到58.4%接近原Gemma 2B的58.6%初始化策略关键Stepwise对递归模型最优Average对放松模型最优延长训练蒸馏进一步提升60B token 知识蒸馏性能显著接近全尺寸教师模型提前退出训练有效中间循环输出与最终输出准确率差距仅约4.6%有利于提前退出理论吞吐量提升显著递归模型相比普通模型可达2~3倍加速部分场景接近4倍5. 主要贡献总结提出递归Transformer框架将现有LLM压缩为参数共享的循环模型。提出放松递归Transformer通过LoRA在共享基础上增加灵活性并设计SVD初始化方法。提出连续深度方向批处理结合提前退出实现推理吞吐量的大幅提升。在大规模实验上验证了方法在性能、模型尺寸和推理速度上的优越性。6. 局限与未来方向目前主要在1B~2B规模上验证扩展到7B需要更多上升训练成本。提前退出基于“神谕”假设实际需开发基于置信度的退出策略。LoRA多模块推理效率需进一步优化如专用CUDA内核。可与稀疏化、量化、MoE等技术结合进一步提升效率。这里是自己的论文阅读记录感兴趣的话可以参考一下如果需要阅读原文的话可以看这里如下所示大型语言模型LLM的部署成本高昂。参数共享提供了一条可能减小模型尺寸和成本的途径但其在现代LLM中的有效性仍然相当有限。在这项工作中我们重新审视了Transformer中作为参数共享形式的“层绑定”并介绍了将现有LLM转换为更小、更高效的“递归Transformer”的新方法这些模型在各层之间共享参数且性能损失极小。我们的递归Transformer能够从标准预训练Transformer高效初始化但只使用一个独特层的块然后在循环中多次重复该块。我们通过引入放松的递归Transformer进一步提升了性能该模型通过深度方向低秩适应LoRA模块为层绑定约束增加了灵活性同时仍然保持了整体模型的紧凑性。我们展示了我们的递归模型例如递归Gemma 1B在性能上优于类似规模的普通预训练模型如 TinyLlama 1.1B 和 Pythia 1B和知识蒸馏基线甚至可以恢复原始“全尺寸”模型例如无参数共享的 Gemma 2B的大部分性能。最后我们提出了连续深度方向批处理这是一种由递归Transformer与提前退出相结合所实现的有前景的新推理范式。在理论分析中我们展示了这有可能带来显著的2-3倍推理吞吐量提升。1. 引言大型语言模型的高效部署需要在性能和资源之间取得平衡Leviathan 等人2023Raposo 等人2024Riviere 等人2024Wan 等人2024Zhou 等人2024。虽然具有更多参数的大型模型持续展现出优越的性能Hoffmann 等人2022Rae 等人2021Rosenfeld 等人2020但它们巨大的内存和计算需求是昂贵的Pope 等人2023。参数共享方法例如 Dehghani 等人2019Lan 等人2020Takase 和 Kiyono2023Xia 等人2019通过在模型层之间重用权重可以减少内存占用从而降低这些成本并允许使用更少或等级更低的加速器或使用更大的批处理大小以获得更好的吞吐量。虽然参数共享在先前的工作中已显示出令人鼓舞的能力Giannou 等人2023Lan 等人2020但其在现代LLM中的应用取得的成功报道有限。在这项工作中我们重新审视了LLM的参数共享并提出了新颖的方法将现有的非共享模型转换为更小、更高效的递归Transformer。这些模型使用一个单一的独特层块在多个循环中递归重用但相对于其缩减的尺寸仍然实现了令人印象深刻的性能。为了缓解与参数共享相关的潜在性能下降我们首先基于原始模型的预训练参数初始化共享层块然后对得到的递归模型进行有限步数的“上升训练”。重要的是我们展示了我们的初始化策略使我们能够以最少的训练时间实现强大的性能。这与以下观察结果一致模型压缩技术如层跳过Elhoushi 等人2024Fan 等人2020Zeng 等人2023Zhang 等人2024a、剪枝Frankle 和 Carbin2019Ramanujan 等人2020或嵌套Devvrit 等人2023可以保留出奇高的性能——这进一步激励了我们将模型压缩为更紧凑但性能更强的架构此处为具有低秩适配器的重复层的方法。图 1 |从普通 N 层Transformer到具有 N/K 个 K 层共享块的递归Transformer的转换概述。递归Transformer通过将单一的 K 层块多次重复获得从而形成一个循环架构。递归Transformer也可以通过添加层级特定的LoRA模块转换为放松的递归Transformer。这保留了许多权重共享的优势同时也允许更好的性能。如图1所示我们进一步提出了放松的递归Transformer它是递归Transformer的一种扩展其中通过引入多个层级特定的低秩适应LoRA模块Hu 等人2022略微放宽了跨重复层块的权重绑定。尽管简单但这种策略提供了几个非平凡的优势。首先它允许在共享层之间存在低秩增量同时只增加最小的开销。其次LoRA矩阵的秩可以调整以控制放松程度这直接影响模型容量。此外由于放松后的模型与原始Transformer具有相同的整体形状我们可以通过截断奇异值分解Hansen1987在原始层权重和共享层权重之间的残差矩阵上高效地初始化LoRA模块。因此秩值作为一个关键的超参数使得放松的递归Transformer能够在普通Transformer和递归Transformer架构的两个极端之间无缝过渡。虽然本文的主要焦点在于如何构建和训练递归Transformer但我们也强调了它们通过一种新的批处理推理范式——连续深度方向批处理——实现显著吞吐量提升的潜力这种范式由其递归性质所促成。先前的工作引入了连续序列方向批处理Kwon 等人2023Yu 等人2022它利用了这样一个事实无论序列中的token位置如何计算一个新token所执行的计算在功能上是相同的并且使用相同的模型参数。这允许当批处理中的槽位可用时可以持续调度新的请求。例如当一个响应完成时下一个要形成的响应的开始可以立即占据已完成响应在批处理中的位置而无需等待可能更长的其余批处理响应。在我们的递归Transformer中参数共享不仅发生在不同的时间步之间还发生在不同的深度循环迭代之间。这启用了一个额外的动态分组维度在同一批次内为各个响应联合计算循环层块的不同迭代。我们的主要贡献如下我们引入了一个用于初始化和训练放松的递归Transformer的框架并展示了与类似规模的非递归模型相比的强大性能。例如当我们从预训练的 Gemma 2BTeam 等人2024转换而来的递归 Gemma 1B 模型进行上升训练时我们观察到在少样本任务上与非递归的 Gemma 1B 模型从头开始预训练相比绝对准确率提高了多达 13.5 个百分点错误率降低 22%。此外我们展示了通过结合知识蒸馏Hinton 等人2015Kim 和 Rush2016我们的递归 Gemma 模型在 600 亿个token上进行上升训练后其性能与在 3 万亿token语料库上训练的全尺寸 Gemma 模型相当详见 §3.3。基于我们的放松的递归Transformer我们还评估了一个关键用例即结合提前退出的连续深度方向批处理Bae 等人2023Elbayad 等人2020Graves2016aSchuster 等人2022该方法在早期阶段有机会性地对高置信度的样本进行预测。从我们的模拟来看提前退出显示与具有相同架构的普通Transformer相比吞吐量显著提高了2-3 倍。值得注意的是性能优于普通 Pythia 模型的递归 Gemma 模型理论上可以实现近4 倍的吞吐量提升详见 §3.8。2. 利用递归模式进行有效的模型压缩在本节中我们将介绍将普通Transformer模型转换为参数共享模型的方法细节该模型在性能上优于同等规模的模型。我们首先简要概述Transformer架构§2.1。然后我们介绍递归Transformer并提出利用原始预训练模型权重来初始化其循环层的有效技术§2.2。在 §2.3 中我们放宽了模型设计中的参数共享约束并添加了一组有限的层级特定参数以在保持紧凑表示的同时进一步提高模型的准确性。最后我们展示了除了减少内存之外递归Transformer如何通过一种新颖的推理范式§2.4天然地支持进一步的吞吐量优化。2.1. 基础Transformer架构2.2. 递归Transformer循环层绑定在这项工作中我们在LLM的背景下重新审视参数共享并提出了递归Transformer架构。在各种循环策略中参见附录 B我们特别为递归Transformer采用了CYCLE策略Takase 和 Kiyono2023其中单一独特层块被递归重用。这种固有的设计与提前退出机制无缝对齐可能提供显著的加速。模型的隐藏状态计算如下图 2 |左图一个具有 6 层的非共享全尺寸模型示例。中图在递归Transformer中初始化循环层的三种提议方法。每个层编号指示用于初始化的全尺寸模型中的源层。右图通过 SVD 方法初始化的放松递归Transformer示例。此处循环层使用平均法进行初始化。循环层的初始化技术为了缓解参数共享模型中容量减少带来的潜在性能下降我们提出了几种新颖的初始化方法以促进从非共享的预训练模型到递归Transformer的有效知识迁移。图2展示了三种这样的技术。步进法在固定间隔选择中间层同时保持第一层和最后一层固定。这受到先前工作的启发Fan 等人2020Liu 等人2023Zeng 等人2023Zhang 等人2024a这些工作表明在LLM中跳过几层对生成质量影响很小。平均法通过对其权重矩阵进行平均来初始化绑定层之间的共享权重而下层法则直接使用非共享模型的前 K 层的权重。我们对这些初始化模型进行了 150 亿个token的简短上升训练以研究性能恢复的程度§3.4发现步进法对递归Transformer表现最好。然而我们发现平均法对接下来讨论的放松的递归Transformer表现最好。2.3. 放松的递归Transformer多层LoRA虽然完全的层绑定在压缩模型大小的同时保持了强大的能力但它有两个明显的局限性1可能的模型大小仅限于缩放层数以及2每个模型层最终不得不承担与模型不同深度相关的多个角色。为了解决这个问题我们引入了放松的递归Transformer其中我们为每一层加入了独立的适配器模块Houlsby 等人2019Hu 等人2022放宽了严格的参数共享。虽然我们尝试了各种方法如层级特定的前缀Liu 等人2021参见附录 H但我们发现低秩适应LoRA模块Hu 等人2022能有效捕获绑定层之间的细微差异。具体来说我们将公式 2 修改为图 3 |结合提前退出的连续深度方向批处理策略的说明性示例。我们假设最大批处理大小为 32三个模型“阶段”例如层块以及一个按时间顺序到达的批处理输入流。在 (a) 中所有三个阶段必须完成第一批非最大16 个样本的处理然后才能开始下一批到达的 32 个样本。然而在 (b) 中第二批 32 个样本中的一半可以与仍在完成的第一批 16 个样本共享计算。最后(c) 展示了每个批次中的一些样本可以在阶段 2 后提前退出的情况它们批次中的空槽位随后被立即填充。2.4. 连续深度方向批处理与提前退出在实际部署中用户请求是顺序且异步到达的。最近的研究引入了连续序列方向批处理Kwon 等人2023Yu 等人2022这是一种允许新请求立即替换批处理中已完成终止序列的服务策略。表 1 |三个模型的关键参数和预训练细节。每个模型的尺寸指的是嵌入参数嵌入矩阵和分类头的数量以及所有其他非嵌入参数。Gemma 和 TinyLlama 利用多查询Shazeer2019和分组查询Ainslie 等人2023注意力机制这导致键值头数量减少。\* 我们使用 TinyLlama 的一个早期检查点来研究在 SlimPajama 上欠训练模型基础上的递归转换。经过更长预训练的普通性能报告在表 D.1 中。这种方法利用了这样一个事实为新token执行的计算在功能上是相同的并且使用相同的模型参数。通过以这种方式持续调度请求模型可以以其最大批处理容量运行从而提高服务效率。递归Transformer的重复结构允许将相同的函数不仅应用于序列还应用于深度循环迭代。这为连续批处理引入了一个新维度我们称之为连续深度方向批处理。该技术使得能够为不同样本同时计算循环层块的不同迭代图3 展示了一个单次前向传播的示例这可以轻松扩展到每个请求的多次解码迭代。对于最大批处理大小为 32 的情况标准Transformer必须等待所有模型阶段完成后才能处理新请求。相比之下我们的递归Transformer由于在所有阶段共享层函数可以在时间步 2 立即调度新到达的请求从而最大化批处理大小利用率。这种策略可以通过更快的调度在生成中产生显著的加速并减少首个token的时间Fu 等人2024Miao 等人2023。当与提前退出相结合时Bae 等人2023Elbayad 等人2020Schuster 等人2022深度方向批处理带来的吞吐量提升会进一步放大。如图 3c 所示一旦一些样本在特定的循环迭代后退出排队的请求就可以立即被调度。虽然递归Transformer利用了提前退出的加速但它们也固有地解决了提前退出方法在批处理推理中的一个关键挑战当服务大批次时的同步问题因为提前退出的token可能需要等待其他token完成整个模型的处理。我们证明了配备这种在不同深度进行动态样本调度的递归Transformer在评估的LLM上理论上可以实现高达 2−3 倍的加速。3. 实验3.1. 实验设置我们在三个流行的预训练LLM上评估我们的方法Gemma 2BTeam 等人2024TinyLlama 1.1BZhang 等人2024b和 Pythia 1BBiderman 等人2023。表1总结了每个模型的架构和预训练方案其少样本性能总结在附录 D。除非另有说明所有实验的循环块数 (B) 设置为 2。Gemma 在 B3 时的结果在附录中提供。转换为递归Transformer后我们在 SlimPajama 数据集Soboleva 等人2023上对模型进行上升训练。我们使用 Language Model Evaluation Harness 框架Gao 等人2023在七个少样本任务上评估准确率并对其取平均值以进行性能比较。详细的实验设置可以在附录 E 中找到。表 2 |在与预训练数据集质量或分布显著不同的数据集上进行上升训练可能导致性能下降。我们在 SlimPajama 数据集上对模型进行了上升训练后的评估。我们测量了在 SlimPajama RedPajama 和 PG19 测试集上的困惑度以及在 LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC-challenge 和 OpenBookQA 基准上的少样本准确率。3.2. 非递归模型基线鉴于我们利用了预训练模型权重进行初始化随后对模型进行了上升训练为我们的参数共享模型定义清晰的性能目标变得至关重要。全尺寸模型我们的最终目标是让递归Transformer在不需要太多上升训练的情况下达到与原始全尺寸预训练模型相当的性能。然而我们观察到预训练和上升训练数据集之间的分布差异可能会阻碍达到预期的性能。特别是在新数据集特别是那些质量相对较低的数据集上进行上升训练有时会导致在某些基准上的性能下降。表2总结了基于上升训练token数的全尺寸模型的评估结果。例如对于 Gemma其预训练数据集未公开但可能经过精心策划Team 等人2024在 SlimPajama 数据集上进行上升训练后所有少样本性能指标逐渐下降。这表明使用 SlimPajama 数据集可达到的性能上限可能远低于原始模型性能。因此我们将 Gemma 和 Pythia 模型的目标性能设定为对全尺寸预训练模型进行相同数量token上升训练后达到的性能。由于 TinyLlama 已经在 SlimPajama与我们用于上升训练的数据集相同上进行了预训练消除了任何分布偏移且训练时间略长于我们的运行我们使用原始检查点的性能作为参考。缩减尺寸模型为了展示递归Transformer与具有相同参数数量的模型相比的性能优势我们引入了另一个基线缩减尺寸模型。这些模型的参数数量是其全尺寸模型的一半或三分之一与我们的递归模型的参数数量相匹配。然而这些缩减模型是使用相同的训练方案训练token数和来自全尺寸模型的蒸馏从头开始预训练的但没有预训练权重和循环机制带来的好处。这种比较有助于突出我们的初始化技术和递归函数本身在实现强大性能方面的有效性即使在模型尺寸受限的情况下也是如此。3.3. 主要结果图4展示了具有两个块及其放松变体的递归Transformer的少样本性能。即使在未放松的情况下递归Transformer也表现出显著的高性能尽管其参数仅为全尺寸模型的一半。与同样使用蒸馏损失在 600 亿个token上训练的缩减尺寸模型相比Gemma 模型获得了 10% 的性能提升。值得注意的是递归 TinyLlama 模型甚至超过了普通模型的性能尽管后者是在更大的 1050 亿个token语料库上预训练的。我们的初始化技术在实现这一优越结果方面非常有效同时上升训练数据集SlimPajama与其预训练数据集相同也带来了好处。放松模型有效地在全尺寸模型和递归Transformer之间进行插值具体取决于 LoRA 秩。随着 LoRA 模块增大导致模型尺寸增加SVD 初始化方法允许更精确地逼近全秩矩阵从而带来性能提升。值得注意的是秩为 512 的放松 Gemma 模型在参数更少且仅在 600 亿个token上进行上升训练的情况下达到了与在 3 万亿个token上预训练的原始模型相当的性能58.4% 对比 58.6%。这种权衡为在各种部署场景中选择最佳配置提供了灵活性。我们相信额外的上升训练和更高质量的数据集可以在更精简的模型上带来更好的性能。在接下来的章节中我们将全面概述在达到最终性能之前进行的广泛消融研究。在 §3.4 中我们深入分析了递归Transformer的各种初始化方法。关于放松模型的见解详见 §3.5。最后我们探讨了增强的训练策略如知识蒸馏§3.6。图 4 |递归和放松的递归Transformer达到了与全尺寸模型相当的性能并显著优于缩减尺寸模型。递归模型使用步进法初始化而放松模型使用平均法和 SVD 方法初始化循环层和 LoRA 模块。我们展示了四个不同秩值64 128 256 和 512的性能。递归和缩减尺寸模型均使用知识蒸馏目标在 600 亿个token上进行了上升训练递归模型或从头开始预训练缩减尺寸模型。图 5 |(a) 使用不同初始化方法的递归 Gemma 模型的训练损失曲线具有 2 个块。(b) 使用不同初始化方法的递归 Gemma 模型具有 2 个块的少样本准确率以及 (c) Gemma 模型具有 3 个块和 (d) TinyLlama 模型具有 2 个块的少样本准确率。图 6 |使用不同方法初始化循环层的递归 Gemma 模型的最终循环输出18 层和中间循环输出9 层的性能比较。在平均准确率方面中间输出与最终输出之间的差异在 4 到 5 个百分点之间。表现出与原始全尺寸模型相当的性能。这一观察结果强调了即使在面对分布偏移时利用预训练权重初始化也能实现强大的性能从而证明了循环层初始化的有效性。在本节的其余部分我们将提供更深入的分析以增强对这些初始化方法的理解。我们在附录 F.1 中提供了额外的比较。平均法是放松递归Transformer的最佳初始点当涉及使用 LoRA 模块进行放松时我们发现平均法比步进法表现更好见图 4。与标准零初始化相比我们提出的 SVD 初始化方法显著提高了这些性能提升参见附录 G。结果平均法在松弛模型中成为最有效的初始化技术。3.5. 放松对递归Transformer的影响图 7a 描绘了在将 LoRA 模块整合到递归 Gemma 模型中时的性能变化。随着秩的增加放松模型在准确率上持续提升逐渐缩小与全尺寸模型之间的差距并最终在秩为 512 时达到与全尺寸模型相当的性能。相比之下在 TinyLlama 和 Pythia 模型中没有观察到明显的性能提升。值得注意的是虽然大型模型如 Gemma仅需要少量额外的参数即可实现显著的性能提升但小型模型似乎无法从这种放松中受益。我们推测这是由于较小的模型容量限制了它们利用额外参数的能力或者是因为 TinyLlama 在其预训练阶段没有完全收敛。我们注意到即使在秩为 2048 时也没有进一步改进这表明初始化的全尺寸模型权重已经接近最优或由于教师模型表现不佳而受到限制。此外我们研究了针对不同 Transformer 组件如查询、键值、输出投影和 FFN 权重分配不同秩的影响结果总结在附录 G 的表 G.1 中。我们的分析显示即使仅在查询权重上使用全秩 LoRA也能实现与在所有组件上使用全秩 LoRA 相当的改进。这表明注意力的 query 投影矩阵在捕获不同层深度之间的差异中起着最关键的作用。3.6. 延长上升训练与知识蒸馏为了达到接近其原始全尺寸模型的性能我们纳入了两项关键的训练技术延长上升训练和知识蒸馏。我们在 SlimPajama 数据集上将上升训练的token数从 150 亿扩展到 600 亿。此外我们采用知识蒸馏KD损失Hinton 等人2015Kim 和 Rush2016利用每个架构的全尺寸模型已在 150 亿个 SlimPajama token上进行了上升训练作为教师。在如此大量的上升训练token下我们从教师模型中在线提取 logits。在评估了各种损失函数包括前向 KLFKLKim 和 Rush (2016)反向 KLRKLGu 等人 (2024)Jensen-Shannon 散度JSDAgarwal 等人 (2024)和总变差距离TVDWen 等人 (2023)之后我们选择了前向 KL因为其简单性和优越的性能。如图 7b 所示延长上升训练和知识蒸馏都带来了显著的性能提升总共提升超过 3%使我们的递归模型更接近全尺寸教师模型的性能。图 7 |(a) 增加 LoRA 秩通常会导致放松 Gemma 模型的性能提升这归功于使用了 SVD 初始化。(b) 延长上升训练和知识蒸馏为 Gemma 模型带来了显著的准确率提升。请注意全尺寸模型是进一步在 600 亿个token上进行上升训练的预训练模型。(c) 递归和放松的递归Transformer在模型尺寸和性能方面实现了令人信服的帕累托前沿。递归和放松模型分别使用步进法和平均法初始化循环层。3.7. 为提前退出训练递归Transformer我们将这种用于提前退出的训练后策略应用到了我们最终的上升训练模型在 §3.3 中展示所有实验结果详见附录 J。激进系数策略结合自蒸馏在所有模型中持续为中间输出取得了最佳性能同时保持了最终循环输出的强劲性能。然而由于从非放松模型得出的最优策略直接应用于放松模型为放松的递归Transformer量身定制的训练方法可能会进一步提升其中间循环输出的性能。3.8. 通过连续深度方向批处理实现的理论生成加速我们如何理论近似实际吞吐量由于开发实用的提前退出算法超出了本工作的范围我们基于一种神谕退出方法Bae 等人2023Schuster 等人2022呈现了假设性的吞吐量提升。这假设token在其预测与最终循环预测一致的最早循环块处退出。我们模拟了语言建模数据集的生成过程就好像它们是由我们的模型生成的一样以获得每个token的退出轨迹。然后我们在特定约束下例如不同的内存限制或上下文长度测量了平均每token生成时间。利用这些测量值和退出轨迹数据我们进行了模拟以估计理论吞吐量。详细解释和局限性在附录 K 中讨论。图 8 |带有提前退出的连续深度方向批处理CDB使递归Transformer理论上能够实现显著的吞吐量提升。吞吐量token/秒是在 SlimPajama RedPajama 和 PG19 上取平均值然后相对于普通 Pythia 模型的吞吐量进行归一化。附带的表格给出了 Gemma 的详细吞吐量和性能测量值。ΔV​ 衡量相对于普通 Gemma 模型的吞吐量而 ΔSeq​ 衡量相对于具有连续序列方向批处理CSB的普通 Gemma 模型的吞吐量。与普通模型相比生成速度提升了2.66 倍。此外递归 Gemma 模型显著优于普通预训练的 Pythia 模型吞吐量提升了近4 倍。放松的递归模型在可实现的少样本性能和吞吐量之间显示出明确的权衡这种权衡通过 LoRA 秩的放松程度来调节。这种特性使得可以根据特定的部署场景灵活选择模型。完整结果见表 K.2 和 K.4。连续深度方向批处理的要点我们通过连续深度方向批处理分析了递归Transformer吞吐量提升的潜力这是一种新颖的推理范式。理论上我们发现与普通Transformer相比可以实现高达 2−3 倍的加速。这甚至超过了现有普通模型中的连续序列方向批处理方法所带来的吞吐量增益。4. 相关工作跨层参数共享已被证明是在深度学习模型中实现参数效率的有效方法如 RNNGraves2016bSherstinsky2018CNNEigen 等人2014Guo 等人2019Savarese 和 Maire2019Shen 等人2022和流行的Transformer架构。通用TransformerDehghani 等人2019一种循环自注意力模型展示了以显著更少的参数获得优于非递归对应物的性能。这种跨层参数共享方法随后在各种任务中得到了探索包括语言理解Lan 等人2020语言建模Bai 等人2019Csordás 等人2024Glorioso 等人2024Liu 等人2024bMohtashami 等人2023和机器翻译Dabre 和 Fujita2019Ge 等人2022Milbauer 等人2023Takase 和 Kiyono2023Xia 等人2019。这些方法通常声称以更紧凑的模型和更高的计算速度实现相当的性能同时也为有效的自适应计算解决方案奠定了基础Dehghani 等人2019Graves2016bSchuster 等人2021。与此同时利用循环架构进行算法或逻辑推理任务的兴趣日益增长Saunshi 等人2024。先前的研究McLeish 和 Tran-Thanh2022Schwarzschild 等人2021表明循环网络可以将从简单问题中学到的推理策略通过在推理过程中进行额外的循环外推到更困难、更大的问题。循环Transformer结构也被用于模拟程序执行的基本计算块Giannou 等人2023为数据拟合问题学习迭代算法Yang 等人2024在算法任务中实现长度泛化Fan 等人2024并在少样本学习方面展现出有前景的理论潜力Gatmiry 等人2024。然而先前的工作主要集中在相对较小的Transformer模型上这些模型从头开始训练没有利用预训练模型权重。我们的工作通过在LLM背景下研究参数共享并提出了利用现有LLM内部知识的有效初始化策略从而与先前的工作区分开来。据我们所知我们是第一个提出参数共享模型通用框架的通过层级特定模块放松了权重绑定约束。在本文中我们还讨论了递归Transformer如何非常适合提前退出技术以加速LLM中的解码。固有的递归结构天然地使得在大型服务批次中为单个响应实现提前退出成为可能而这通常是此类技术在实际中的一个限制。普通Transformer在提前退出时遇到同步问题即即使批次中只有一个token需要完整处理模型也必须前向传播所有层已退出的token必须等待它们。几种方法试图通过为已退出的token计算后续层中缺失的KV缓存来利用这段空闲时间这对于后续序列生成至关重要。这些技术包括状态传播Elbayad 等人2020Schuster 等人2022SkipDecodeDel Corro 等人2023和并行解码可与推测解码结合Bae 等人2023Chen 等人2024bElhoushi 等人2024Liu 等人2024aTang 等人2024。尽管如此跨不同模型深度的异构参数仍然阻碍了已退出token高效地进入后续序列。相比之下我们的递归Transformer能够为不同深度和序列的token在连续深度方向批处理范式中实现并行计算——也允许在内存受限的解码阶段以最小的开销并行计算缺失的KV缓存。5. 结论与未来工作在这项工作中我们介绍了递归Transformer通过跨递归循环层块进行参数共享来压缩LLM。此外我们提出了一种新颖的放松策略通过将层级特定的LoRA模块集成到完全绑定的结构中允许在共享层之间存在低秩增量。通过针对循环层和LoRA模块的新颖初始化技术我们实现了显著的性能提升非常接近原始的预训练模型。最后通过利用递归模式和提前退出方法我们提出了一种为递归Transformer高效服务系统量身定制的连续深度方向批处理范式。我们从理论上证明一种神谕退出策略可以产生显著的吞吐量增益最高可达 2−3 倍加速。这项工作激励了对现代LLM中递归模式的进一步研究例如与稀疏设计的兼容性基于稀疏性的方法如剪枝Han 等人2015量化Jacob 等人2018或层跳过机制Raposo 等人2024最近也给出了有前景的模型压缩结果。事实上这些技术中的许多与我们的方法是互补的例如我们可以无缝地拥有一个递归的、稀疏的架构。在这项工作中我们选择专注于递归密集设计一个仍相对未被探索的领域这些设计也具有非常有前景的实际性能特征例如允许连续深度方向批处理以实现更快的吞吐量。尽管如此虽然在这项工作中我们迈出了研究具有密集Transformer层的放松递归Transformer的第一步但我们确实相信在循环块中整合专家混合Fedus 等人2022激活跳过Liu 等人2023和SSM组件Glorioso 等人2024是未来研究的有希望的方向。通过递归深度进行潜在推理除了通过递归模式缩小物化参数以实现效率提升之外另一个研究方向在于扩展递归深度以促进潜在推理。具体来说递归计算可以通过处理每个深度的内部隐藏状态来纵向地体现“思考”。一个有前景的方法涉及利用思考tokenGoyal 等人2024Pfau 等人2024或潜在连续空间表示Cheng 和 Van Durme2024Hao 等人2024来增强数学和代码生成任务中的推理。另一个有价值的方向侧重于提高递归扩展深度的方法的效率和训练稳定性建立在深度思考的概念之上Geiping 等人2025Schwarzschild 等人2021。扩展递归Transformer将我们的方法扩展到更大的LLM70亿参数及以上是未来研究的一个有前景的方向。虽然我们的方法预计仍然有效但实现相当的性能可能需要显著更高的上升训练成本。模型尺寸的增加为通过递归模式减少内存占用提供了潜力然而考虑到隐藏维度的相应增加这是否能转化为更大的批处理大小尚不清楚。尽管如此我们的连续深度方向批处理将在服务效率方面产生可观的收益。超越理论生成加速我们的神谕退出方法假设任何与最终输出匹配的中间预测都可以退出。然而准确的吞吐量测量需要基于置信度的提前退出算法Bae 等人2023Schuster 等人2022。此外实际部署需要解决解码瓶颈问题例如为已退出token在剩余循环中计算键值缓存。尽管如此存在潜在的解决方案例如缺失的KV缓存计算可以通过利用连续深度方向批处理来解决允许在后续循环中为已退出位置并行计算KV缓存同时进行下一个序列样本的计算。此外我们可以在未来的工作中探索键值缓存共享策略Brandon 等人2024Sun 等人2024。多LoRA层的高效服务放松模型需要在批处理推理期间计算不同的LoRA模块类似于多任务学习Feng 等人2024Wang 等人2023这阻碍了并行计算。我们将LoRA权重连接成一个单一的权重以提高顺序计算的效率但这引入了冗余。为了缓解这个问题我们可以探索用于LoRA服务的优化CUDA内核Chen 等人2024aSheng 等人2023以及跨加速器的并行化这受到专家混合分布式训练的启发Fedus 等人2022Gale 等人2023。