[論文レビュー] Retentive Network: A Successor to Transformer for Large Language Models

Yutao Sun, Li Dong+6 more|arXiv (Cornell University)|Jul 17, 2023

Topic Modeling被引用数 107

ひとこと要約

リテンションネットワーク（RetNet）は、アテンションに代わるマルチスケールリテンション機構を導入し、並列トレーニング、O(1)推論、長大系列メモリの線形性を実現し、Transformersと競合する性能を持ちます。

ABSTRACT

In this work, we propose Retentive Network (RetNet) as a foundation architecture for large language models, simultaneously achieving training parallelism, low-cost inference, and good performance. We theoretically derive the connection between recurrence and attention. Then we propose the retention mechanism for sequence modeling, which supports three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent. Specifically, the parallel representation allows for training parallelism. The recurrent representation enables low-cost $O(1)$ inference, which improves decoding throughput, latency, and GPU memory without sacrificing performance. The chunkwise recurrent representation facilitates efficient long-sequence modeling with linear complexity, where each chunk is encoded parallelly while recurrently summarizing the chunks. Experimental results on language modeling show that RetNet achieves favorable scaling results, parallel training, low-cost deployment, and efficient inference. The intriguing properties make RetNet a strong successor to Transformer for large language models. Code will be available at https://aka.ms/retnet.

研究の動機と目的

• 性能を犠牲にすることなく、LLMの導入時に推論コストとメモリ使用量を削減する動機を示す。
• Trainingの並列性を維持しつつ、Transformerの後継アーキテクチャを開発する。
• トレーニングと推論を最適化するため、並列・再帰・チャンク単位再帰表現をサポートするリテンション機構を導入する。

提案手法

• マルチヘッドアテンションの代替として、マルチスケールリテンション（MSR）モジュールを提案する。
• リテンションの二重表現を導出する：並列表現（トレーニングに適する）と再帰表現（推論に適する）。
• 長大な系列のため、3つの計算パラダイムを実装する：並列リテンション、再帰リテンション、チャンク単位再帰リテンション。
• 表現力と学習安定性を高めるため、ゲーティング（swish）とマルチヘッドデケイ（gamma）を組み込む。
• マルチスケールヘッドによるヘッドごとの分散を扱うため、GroupNormを用いる。
• MSR + FFNブロックと訓練/推論戦略を備えた、エンドツーエンドのRetNetアーキテクチャを提供する。
• スケーリング、トレーニングコスト、推論指標の観点で、RetNetをTransformerおよび効率的なTransformer系と比較する。

実験結果

リサーチクエスチョン

• RQ1 RetNetはTransformerと比較して推論効率を維持または改善しつつ、トレーニングの並列性を達成できるか？
• RQ2 リテンション機構は、トレーニングとデプロイメント中のメモリ、待機時間、および計算コストを削減しつつ、言語モデリング性能をTransformerと同等に提供できるか？
• RQ3 並列・再帰・チャンク単位再帰表現は、長大系列のモデリングとスケーラビリティにどう影響するか？
• RQ4 RetNetを用いた大規模モデルと長文コンテキストで、どの程度のメモリ・スループット・待機時間の向上が得られるか？
• RQ5 Transformerと比較して、ゼロショット・少数ショットの下流タスクでRetNetはどのように性能を発揮するか？

主な発見

• RetNetはTransformerと比較して、好適なスケーリング、並列学習、低コスト展開、効率的な推論を実現する。
• 7Bモデルで8kコンテキストの場合、RetNetはトランスフォーマーのキー・バリューキャッシュを使用した場合と比べてデコードが8.4倍速く、メモリを70%節約します。
• トレーニング時、RetNetは25–50%のメモリ節約を実現し、標準のTransformerより7倍速く、FlashAttentionと良い競合をする。
• RetNetの推論待機時間は長さに依存せず、バッチサイズにも影響を受けにくいため、デコード時のスループットを向上させる。
• RetNetは言語モデリングの困惑度をTransformerに匹敵する水準に達し、いくつかのタスクでゼロショット/少数ショット学習の有利な性能を示す。
• アブレーションにより、ゲーティング、 GroupNorm、およびマルチスケールデケイが性能向上に寄与することが示され、より大きなヘッド次元は結果を改善する。