[RetNet] Retentive Network: A Successor to Transformer for Large Language Models

Transcript

1. (RetNet) Retentive Network: A Successor to Transformer for Large Language Models 2023/08/04 飯田啄巳
2. 背景：RetNetの目指すところ Transformer 言語モデルのデファクトになってるけど、 👍高パフォーマンス 👎🏽O(N) Complexity 👎🏽Key-Valueストアでメモリ効率悪い 👎🏽シーケンス長いと、レイテンシも悪い Linearized Attention アテンションスコアのexp(𝑞 ⋅ k)を𝜙 𝑞 ⋅ 𝜙(𝑘)で書き換え、自己回帰推論可能にした 👎🏽性能悪い Recurrent Network 👎🏽学習の並列性ない Others: Attentionを書き換える（S4など） 👎🏽性能悪い RetNet 並列表現 → 👍並列学習 リカレント表現（実装も楽） → 👍メモリと計算の両面でO(1)推論 チャンクごとのリカレント表現（chunkwise recurrent） → 👍長いシーケンスに対応
3. 手法：Retentive Networks 全体の流れ 1. 入力 ：𝑥 = 𝑥1 , … 𝑥 𝑥 のシーケンス 2. 各トークンの次元を𝑑𝑚𝑜𝑑𝑒𝑙 にする ：𝑋0 = 𝑥1 , … , 𝑥 𝑥 ∈ ℝ 𝑥 ×𝑑𝑚𝑜𝑑𝑒𝑙 3. 自己回帰的に次の状態を推定 ：𝑋𝑙 = 𝑅𝑒𝑡𝑁𝑒𝑡𝑙 𝑋𝑙−1 , 𝑙 ∈ 1, 𝐿 こんなイメージ？ 𝑑𝑚𝑜𝑑𝑒𝑙 𝑑𝑚𝑜𝑑𝑒𝑙 𝑥1 𝑥2 𝑥3 𝑥 𝑥 MSR (Multi-Scale Retention) FFN (Feed-Forward Network) RetNet
4. 手法：Retentionモジュールの仕組み 状態𝒔𝑛 を介して、𝒐𝑛 を出力 Aの対角化の式 を使うと xPosという相対位置埋め込みの表現形式 𝛾をスカラ化 共役転置 Transformerの式 （並列化可能） リカレントモデルの式 RNNとTransformerの式の関係性を考えてみる ハイブリッド表現（Chunkwise Recurrent Representation） 長いシーケンスの学習効率化 要素数Bのチャンクを作る チャンク内では並列化 チャンク外では再帰 𝑖番目のチャンク 三段階正規化（スケーリング） 未来情報使わないように 近傍の重み強めに
5. 手法：Retentionモジュールの仕上げ（ゲート化・マルチスケール化） マルチヘッド化 マルチスケール化 ヘッドごとに異なる𝛾を使う ヘッドの数 ゲート化 swishを使う 正規化層とかもちゃんと書くと… GroupNormは各ヘッドの出力を正規化（SubLNという方法に基づくらしい） ヘッドごとに異なる𝛾を使うとヘッドごとに分布が変わってくるので、ヘッドごとに正規化
6. 全体まとめ 𝑑𝑚𝑜𝑑𝑒𝑙 𝑑𝑚𝑜𝑑𝑒𝑙 𝑥1 𝑥2 𝑥3 𝑥 𝑥 MSR (Multi-Scale Retention) FFN (Feed-Forward Network) RetNet さっきのMSR (Multi-Scale Retention)をTransformerブロックみたいに積んで完成！ 学習時 parallel（シーケンス内並列化）or chunkwise recurrent（チャンク内並列化） parallel chunkwise recurrent 推論時 recurrentを使う → 自己回帰推論 = O(1)
7. 実装的には 未来情報使わないように 近傍の重み強めに
8. 実装的には 要素数Bのチャンクを作る チャンク内では並列化 チャンク外では再帰
9. 実験：モデルサイズとその性能 パラメータ数 Transformer RetNet 𝑠𝑒𝑙𝑓 𝑎𝑡𝑡𝑒𝑛𝑡𝑖𝑜𝑛 𝑊𝑄 , 𝑊𝐾 , 𝑊𝑉 , 𝑊𝑂 = 4𝑑2 𝐹𝐹𝑁 = 8𝑑2 𝑖𝑛𝑡𝑒𝑟𝑚𝑒𝑑𝑖𝑎𝑡𝑒 𝑑𝑖𝑚𝑒𝑛𝑠𝑖𝑜𝑛 = 4𝑑 𝑅𝑒𝑡𝑒𝑛𝑡𝑖𝑜𝑛 𝑚𝑜𝑑𝑢𝑙𝑒 𝑊𝑄 , 𝑊𝐾 ∈ ℝ𝑑×𝑑, 𝑊𝐺 , 𝑊𝑉 ∈ ℝ𝑑×2𝑑, 𝑊𝑂 ∈ ℝ2𝑑×𝑑 = 8𝑑2 混乱度、低いほうが良いらしい。 （確率分布を比較する指標）
10. 実験：Zero-shot, Few-shotの性能 Transformerよりもいいです
11. 実験：Transformerとのメモリ＆スループットの比較 Kernel FusionとFlashAttentionは除外
12. 実験：推論コスト GPUメモリ • TransformerはKVキャッシュで線形に増加 • RetNetは長いシーケンスでも同じ スループット • Transformerは長くなると低下 • RetNetはずっと高いスループット レイテンシ • Transformerはバッチサイズ大 → レイテンシ遅 • RetNetはずっと速い
13. 実験：周辺技術との関係性と性能比較 Query, Keyが Content-unaware attention free& 位置埋め込みを指数減衰に置換 →再帰 関係性 性能比較
14. Ablation Study (𝛾 = 1)
16. 使い方や実際の実装 切り替えて使う感じ モデルのロード方法は書いてあるケド…運用方法は？