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

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

该论文试图解决大型语言模型（LLM）部署成本高昂的问题，通过参数共享来减少模型的大小和成本。这是一个持续存在的问题，尤其是在资源有限的环境中。

论文的关键思路是重新审视Transformer中的“层绑定”（layer tying）作为参数共享的一种形式，并引入了新的方法将现有的LLM转换为更小的“递归Transformer”，这些模型在不同层之间共享参数，同时保持性能的最小损失。与传统的参数共享方法相比，这种递归方法更加有效，且能够通过深度低秩适应（LoRA）模块进一步提高性能。

论文提出了Relaxed Recursive Transformers，通过在层绑定的基础上添加LoRA模块，增加了灵活性，同时保持了模型的紧凑性。实验结果显示，递归模型（如递归Gemma 1B）在性能上超过了类似大小的预训练模型（如TinyLlama 1.1B和Pythia 1B），甚至接近原始“全尺寸”模型（如Gemma 2B）。此外，论文还提出了一种新的推理范式——连续深度批处理（Continuous Depth-wise Batching），结合早期退出策略，理论上可以显著提高推理吞吐量（2-3倍）。论文使用了多个数据集进行实验，并提供了开源代码，便于复现和进一步研究。

近期在这个领域中，相关的研究包括： 1. “Parameter-Efficient Transfer Learning for NLP” - 介绍了LoRA技术，用于在微调过程中高效地适应新任务。 2. “Layer-Wise Adaptive Scaling for Efficient Fine-Tuning of Pre-trained Models” - 提出了层自适应缩放技术，以提高预训练模型的微调效率。 3. “Efficiently Adapting Pre-trained Language Models via Low-Rank Updates” - 研究了如何通过低秩更新来高效地适应预训练语言模型。 4. “Dynamic Layer Sharing in Deep Neural Networks” - 探讨了动态层共享在深度神经网络中的应用，以减少模型大小和计算成本。