Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA

📄 arXiv: 2410.20672v3 📥 PDF

作者: Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, Tal Schuster

分类: cs.CL, cs.LG

发布日期: 2024-10-28 (更新: 2025-02-28)

备注: ICLR 2025; 49 pages, 17 figures, 19 tables

💡 一句话要点

提出Relaxed Recursive Transformers，通过层间LoRA实现高效参数共享，缩小LLM体积。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 参数共享 递归Transformer 低秩适应 模型压缩 语言模型 高效推理 层绑定

📋 核心要点

  1. 现有大型语言模型部署成本高昂，参数共享是降低成本的有效途径，但传统层绑定的效果有限。
  2. 提出Relaxed Recursive Transformers，通过循环重复单个层块并引入深度方向LoRA，在参数共享的同时保持性能。
  3. 实验表明，递归模型在性能上优于同等规模的预训练模型和知识蒸馏模型，并能恢复原始模型的性能。

📝 摘要（中文）

大型语言模型（LLM）的部署成本很高。参数共享提供了一种降低模型大小和成本的可能途径，但其在现代LLM中的有效性仍然相当有限。本文重新审视了Transformer中的“层绑定”作为参数共享的一种形式，并引入了将现有LLM转换为更小的“递归Transformer”的新方法，这些Transformer在层之间共享参数，同时最大限度地减少性能损失。我们的递归Transformer有效地从标准预训练Transformer初始化，但只使用一个独特的层块，然后在循环中重复多次。我们通过引入Relaxed Recursive Transformers来进一步提高性能，该方法通过深度方向的低秩适应（LoRA）模块为层绑定约束增加了灵活性，同时仍然保持了整体模型的紧凑性。我们表明，我们的递归模型（例如，递归Gemma 1B）优于类似大小的vanilla预训练模型（例如，TinyLlama 1.1B和Pythia 1B）和知识蒸馏基线——甚至可以恢复原始“全尺寸”模型（例如，没有共享参数的Gemma 2B）的大部分性能。最后，我们提出了连续深度方向批处理，这是一种有前景的新推理范式，当递归Transformer与提前退出结合使用时，可以实现这种范式。在理论分析中，我们表明这有可能带来显著（2-3倍）的推理吞吐量提升。

🔬 方法详解

问题定义：论文旨在解决大型语言模型部署成本高昂的问题，特别是模型体积带来的挑战。现有参数共享方法，如简单的层绑定，在现代LLM中效果有限，会导致显著的性能下降。因此，需要一种更有效的参数共享策略，能够在减小模型体积的同时，尽可能地保持模型的性能水平。

核心思路：论文的核心思路是利用递归结构和低秩适应（LoRA）来改进层绑定。通过循环重复使用单个层块，可以显著减少模型参数量。同时，为了弥补层绑定带来的性能损失，引入深度方向的LoRA模块，为每一层提供一定的灵活性，从而在参数共享和模型性能之间取得平衡。

技术框架：整体框架包括以下几个主要步骤：1）从预训练的Transformer模型出发；2）提取或设计一个核心的层块；3）将该层块循环重复多次，构建递归Transformer；4）在每一层或每几层之间插入LoRA模块，形成Relaxed Recursive Transformer。在推理阶段，可以结合early exiting策略，进一步提高效率。

关键创新：论文的关键创新在于将递归结构与LoRA相结合，提出了Relaxed Recursive Transformers。这种方法既能实现高效的参数共享，又能通过LoRA模块保持模型的性能。此外，提出的Continuous Depth-wise Batching推理范式，能够充分利用递归结构的特点，进一步提高推理吞吐量。

关键设计：关键设计包括：1）LoRA模块的秩（rank）的选择，需要在参数量和性能之间进行权衡；2）循环重复的层块的设计，需要保证其能够捕捉到模型的核心特征；3）Continuous Depth-wise Batching的具体实现，需要考虑如何有效地利用硬件资源。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，递归Gemma 1B模型优于类似大小的TinyLlama 1.1B和Pythia 1B模型，以及知识蒸馏基线。更重要的是，递归模型能够恢复原始Gemma 2B模型的大部分性能，而参数量显著减少。理论分析表明，结合Continuous Depth-wise Batching，推理吞吐量有望提升2-3倍。

🎯 应用场景

该研究成果可应用于资源受限环境下的LLM部署，例如移动设备、边缘计算等。通过减小模型体积，降低计算和存储成本，使得LLM能够在更多场景下应用。此外，该方法还可以用于模型压缩和加速推理，提高LLM的实用性。

📄 摘要（原文）

Large language models (LLMs) are expensive to deploy. Parameter sharing offers a possible path towards reducing their size and cost, but its effectiveness in modern LLMs remains fairly limited. In this work, we revisit "layer tying" as form of parameter sharing in Transformers, and introduce novel methods for converting existing LLMs into smaller "Recursive Transformers" that share parameters across layers, with minimal loss of performance. Here, our Recursive Transformers are efficiently initialized from standard pretrained Transformers, but only use a single block of unique layers that is then repeated multiple times in a loop. We further improve performance by introducing Relaxed Recursive Transformers that add flexibility to the layer tying constraint via depth-wise low-rank adaptation (LoRA) modules, yet still preserve the compactness of the overall model. We show that our recursive models (e.g., recursive Gemma 1B) outperform both similar-sized vanilla pretrained models (such as TinyLlama 1.1B and Pythia 1B) and knowledge distillation baselines -- and can even recover most of the performance of the original "full-size" model (e.g., Gemma 2B with no shared parameters). Finally, we propose Continuous Depth-wise Batching, a promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting. In a theoretical analysis, we show that this has the potential to lead to significant (2-3x) gains in inference throughput.