[論文レビュー] ReLoRA: High-Rank Training Through Low-Rank Updates

Vladislav Lialin, Namrata Shivagunde+2 more|arXiv (Cornell University)|Jul 11, 2023

Advanced Neural Network Applications被引用数 11

ひとこと要約

ReLoRAは再起動を伴う複数の低ランク更新を蓄積して高ランクのトランスフォーマーネットワークを訓練し、GPU RAMの使用量を削減しつつ、フルランク訓練と同等の性能を達成し訓練を高速化する。

ABSTRACT

Despite the dominance and effectiveness of scaling, resulting in large networks with hundreds of billions of parameters, the necessity to train overparameterized models remains poorly understood, while training costs grow exponentially. In this paper, we explore parameter-efficient training techniques as an approach to training large neural networks. We introduce a novel method called ReLoRA, which utilizes low-rank updates to train high-rank networks. We apply ReLoRA to training transformer language models with up to 1.3B parameters and demonstrate comparable performance to regular neural network training. ReLoRA saves up to 5.5Gb of RAM per GPU and improves training speed by 9-40% depending on the model size and hardware setup. Our findings show the potential of parameter-efficient techniques for large-scale pre-training.

研究の動機と目的

• 非常に大規模なトランスフォーマーモデルのパラメータ効率的な事前学習を動機づける。
• 連続する低ランク更新を通じて高ランク更新が実現できるかを調査する。
• 再起動、ギザギザした学習率スケジュール、部分的なオプティマイザリセットを備えたReLoRAを開発する。
• 最大1.3BパラメータのトランスフォーマーでReLoRAを実証し、LoRAおよびフルランク訓練と比較する。

提案手法

• ウォームスタートとなるフルランク訓練のベースラインから開始する。
• 線形層へランクr=128のLoRA風低ランク更新を適用する。
• 低ランク更新を基底ウェightsへ統合する複数の再起動を使用する（更新の総和）。
• 各統合と再初期化の後、ゼロウォームアップを伴うギザギザのコサイン学習率スケジュールを採用する。
• 更新を導く古い勾配モーメントを避けるため、振幅剪定による部分的なオプティマイザ状態リセットで、陳腐化した更新を回避する。
• 埋め込みと正規化はフルランクのまま維持し、線形層をReLoRAで更新する。

Figure 1: Training loss for 250M models. ReLoRA learns a high-rank network through a sequence of low-rank updates. It outperforms networks with the same trainable parameter count and achieves similar performance to training a full network at 100M+ scale. The efficiency of ReLoRA increases with the m

実験結果

リサーチクエスチョン

• RQ1 高ランクのネットワークは低ランク更新の連続によって効果的に訓練できるか？
• RQ2 モデルサイズ全体で、ReLoRAの性能と効率はLoRAおよびフルランク訓練とどう比較されるか？
• RQ3 ReLoRAの成功にとって不可欠な訓練技術は何か（再起動、オプティマイザリセット、ウォームスタート）？
• RQ4 より大きなトランスフォーマーモデル（最大1.3Bパラメータ）で、ReLoRAは効率と性能を拡張できるか？

主な発見

• ReLoRAはGPUあたり最大5.5GBのRAMを節約し、モデルサイズとハードウェアに応じて訓練を9–40%高速化する。
• ReLoRAはフルランク訓練に近いパープレキシティを達成し、LoRAを上回る。1.3Bモデルは最終時点でパープレキシティ17.27、フル訓練は16.83。
• 特異値分析は、ReLoRAの更新が高ランク／フル訓練に近い分布を示すのに対し、LoRAは主にゼロ/低ランクスペクトルであることを示している。
• 1.3Bモデルでは、ウォームスタートと再起動を組み合わせたReLoRAが訓練全体を通じてLoRAより良い性能を達成し、フルランク訓練とのギャップを縮める（最終パープレキシティ17.27対16.83）。
• 訓練速度の向上はハードウェアによって異なる。8x A100環境では約9%の実時計速アップ、安価なハードウェアではより大きな利得。
• オンラインReLoRA（非常に頻繁なリセット）は、本研究では標準のReLoRAと比べて結果を改善しなかった。

Figure 2: Jagged cosine scheduler used in ReLoRA. As a base for our scheduler we follow a standard cosine decay schedule as in Touvron et al. ( 2023 ) . On every optimizer reset, we set the learning rate to zero and perform a quick (50-100 steps) learning rate warm-up back to the cosine schedule.