リラックスした再帰トランスフォーマー：層ごとのLoRAを用いた効果的なパラメータ共有

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

October 28, 2024

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

cs.AI

要旨

大規模言語モデル（LLMs）の展開には高コストがかかります。パラメータ共有は、そのサイズとコストを削減する可能性のある方法ですが、現代のLLMsにおける効果はかなり限定されています。本研究では、Transformersにおけるパラメータ共有の形態である「レイヤー結合」を再考し、既存のLLMsをより小さな「再帰Transformers」に変換するための新しい手法を導入しました。ここでは、当社の再帰Transformersは、標準の事前学習済Transformersから効率的に初期化されますが、一意のブロックのレイヤーを1つだけ使用し、それをループ内で複数回繰り返します。さらに、深さ方向の低ランク適応（LoRA）モジュールを介してレイヤー結合制約に柔軟性を追加するRelaxed Recursive Transformersを導入することで、全体モデルのコンパクト性を維持しつつ性能を向上させます。当社の再帰モデル（例：再帰Gemma 1B）は、同様のサイズのバニラ事前学習済モデル（TinyLlama 1.1BやPythia 1Bなど）および知識蒸留のベースラインを上回り、元の「フルサイズ」モデル（共有パラメータのないGemma 2Bなど）の性能のほとんどを回復できることを示しています。最後に、再帰Transformerと早期終了を組み合わせた場合に可能となる有望な新しい推論パラダイムである連続深さ方向バッチングを提案します。理論的な分析では、これが推論スループットの大幅な（2〜3倍）向上につながる可能性があることを示しています。

English

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.