今回の論文
今回取り上げるのは、Yutao Sunらによる「Retentive Network: A Successor to Transformer for Large Language Models」です。2023年にarXivで公開されたMicrosoft Researchと清華大学の研究で、研究分野は大規模言語モデルのモデル構造・高速推論です。論文URLは https://arxiv.org/abs/2307.08621 です。
選んだ理由は、長い会話や文書を扱うLLMで避けにくいKVキャッシュの増大を、後段の圧縮ではなくモデルの状態表現から見直しているためです。RetNetは「学習ではGPUを並列に使いたい」「生成では過去トークンを抱え続けたくない」という相反しがちな要求を、一つのRetention機構の計算表現を切り替えることで扱います。
どんな技術か
RetNet（Retentive Network）は、Transformerの自己注意をMulti-Scale Retentionへ置き換えた自己回帰型の言語モデルです。過去全トークンのKey/ValueをKVキャッシュとして保存して参照する代わりに、過去の情報を固定サイズの状態行列へ逐次的に畳み込みます。
重要なのは、同じRetentionを三通りに計算できる点です。学習時はトークン列を並列計算し、生成時は状態を一回更新して次トークンを出し、長い列の学習時はチャンク内を並列・チャンク間を再帰的に処理します。つまり、アルゴリズムを別々に切り替えるのではなく、等価な数式表現を用途に合わせて使い分けます。
課題
Transformerは強力ですが、自己回帰生成では新しいトークンごとに過去の全KVキャッシュを参照します。文脈が長くなるほど、キャッシュ容量、メモリ帯域、1トークン当たりの待ち時間が大きくなります。
- RNNは状態が固定サイズなので生成は軽い一方、時系列順にしか学習できずGPUの並列性を活かしにくいです。
- Linear AttentionはAttentionを近似して定数時間推論を狙えますが、位置や内容の表現力が十分でない場合がありました。
- KVキャッシュ量子化・選別は既存Transformerを延命できますが、長文・多セッションのサービングではキャッシュの運用自体が依然として主要コストです。
- チャット、長文RAG、エージェントの実行ログなどでは入力履歴が伸びやすく、同時接続数を上げるほどGPUメモリが先に制約になります。
RetNetは、学習の並列性、Transformer級の表現力、低コストな生成を同時に満たすという「impossible triangle」を課題に据えています。
用語解説
- KVキャッシュ
- 自己回帰Transformerが過去トークンのKeyとValueを保存する仕組みです。再計算を避けられる反面、文脈長と並行リクエスト数に応じてGPUメモリを消費します。RetNetはこれをトークン列ではなく固定形状の状態に置き換えます。
- 自己注意（self-attention）
- 現在位置のQueryが過去位置のKey/Valueを重み付けして参照する仕組みです。RetNetはその全対全参照をRetentionに置き換え、内容依存の参照と位置に応じた減衰を保とうとします。
- 再帰表現（recurrent representation）
- 前時刻の状態に現在の入力を加えて次の状態を作る計算です。RetNetでは生成時にこの表現を使うため、過去全トークンを走査せず固定サイズの状態だけを更新できます。
- 減衰率（decay rate, γ）
- 古い情報をどの程度残すかを決める係数です。γが1に近いほど長期記憶になり、小さいほど直近を重視します。RetNetはヘッドごとに異なるγを持たせ、短期・長期の時間尺度を並行して扱います。
- チャンク化
- 長い列を一定長のブロックに分ける方法です。本論文ではチャンク内を並列に計算し、過去チャンクの要約状態だけを受け渡します。学習時の並列性と長文時のメモリ効率を両立する鍵です。
技術の仕組み
RetNetの各ブロックは、Transformerと同様に残差接続・事前LayerNorm・FFNを持ちます。ただしMulti-Head Attentionの位置には、複数時間尺度を持つMulti-Scale Retention（MSR）が入ります。
Retentionは「内容に応じた参照」と「距離による忘却」を掛け合わせる
入力Xから線形変換でQ、K、Vを作り、位置回転を施します。位置nの出力は、過去位置mのValueを、内容の相性Q_n K_m^Tと距離減衰γ^(n-m)で重み付けして集約します。未来の位置はマスクするため、自己回帰生成を保てます。
直感的には、Attentionの「何を読むか」という内容依存性に、「どれくらい昔か」という指数的な忘却を明示的に足したものです。softmaxを使わない形へ整理することで、次の二つの等価な計算形に変形できます。
学習は並列、生成は状態更新
並列表現では、全位置のQK^Tと因果・減衰マスクをまとめて計算します。GPUはトークン方向の行列演算を並列に処理できるため、通常のTransformer学習に近い使い方ができます。
再帰表現では、各ヘッドの状態行列をS_n = γS_(n-1) + K_n^T V_nで更新し、出力をQ_n S_nとして取り出します。状態の形はヘッド次元で決まり、生成済みトークン数では増えません。したがって、各生成ステップの追加計算・追加メモリは系列長に対してO(1)です。これは「KVキャッシュを小さくする」のではなく、「成長するキャッシュを持たない」設計です。
長文の学習にはチャンク再帰を使う
全列を一度に並列化すると巨大な行列が必要になります。そこで長い列を例えば512トークンのチャンクへ分け、チャンク内の相互作用は並列に、過去チャンクからの情報は状態行列で受け取ります。チャンクの終了時に状態を更新して次へ渡すため、履歴全体のKVを保存する必要がありません。
多時間尺度・ゲート・正規化で表現力を補う
各ヘッドには異なる固定減衰率を割り当てます。速く忘れるヘッドは局所的な文法や直近指示を、遅く忘れるヘッドは長期の話題を担うことが期待されます。さらにSwishゲートで非線形性を加え、ヘッドごとに分散が異なる問題にはGroupNormを使います。論文のアブレーションでは、ゲート、GroupNorm、減衰、多尺度化のいずれを外してもパープレキシティが悪化しました。
実験と結果
論文は1.3B、2.7B、6.7BのRetNetとTransformerを、The Pile、C4、The Stackを混ぜたコーパスで各100Bトークン学習しました。言語モデリング、ゼロ・few-shot評価、学習時のメモリとスループット、推論時のメモリ・レイテンシを比較しています。
言語モデリングと下流タスク
検証パープレキシティではRetNetはTransformerと同等以上で、論文では約2Bパラメータを超えるとRetNetがTransformerを上回る傾向が報告されました。6.7Bモデルの7タスク平均精度は、zero-shotでTransformerの66.07に対しRetNetは69.51、4-shotで66.44に対し69.76でした。
200M規模でRWKV、H3、Hyena、Linear Transformerとも比較し、in-domainおよびPG22、QMSum、GovReport、SummScreenでRetNetが最も低いパープレキシティを示しました。たとえばQMSumはRetNet 21.33に対してRWKV 28.17、Hyena 28.18でした。少なくともこの条件では、一定コストの再帰化が単純な表現力の犠牲にはなっていません。
学習コスト
系列長8192、A100 80GB×8枚での測定では、6.7B RetNetのメモリ使用量は48.0GB、FlashAttention付きTransformerは51.4GBでした。スループットはRetNet 17,458.6 words/s、FlashAttention付きTransformer 16,230.1 words/sです。RetNetは専用カーネル融合を前提にしない素のPyTorch実装での比較であり、著者はさらに最適化の余地があると述べています。
推論コスト
6.7B・入力長8kでは、著者はRetNetがKVキャッシュ付きTransformerよりデコードを8.4倍高速化し、メモリを70%節約できたと報告しています。Transformerのメモリとスループットは生成長に伴って悪化しますが、RetNetの追加メモリはほぼ一定で、論文ではモデル重み以外は約3%とされています。
ただし、これは同論文のモデル・実装・GPUでの測定結果です。今日の最適化済みTransformerサービング基盤と比較する際は、実際のバッチサイズ、コンテキスト長、カーネル実装、品質要件で再測定する必要があります。
何に使える？
RetNetそのものの採用には基盤モデルの用意が必要ですが、「成長する履歴を固定状態へまとめる」という発想は複数のAIシステムに応用できます。
長時間の対話・音声ストリーミング
会議支援、音声対話、常駐コパイロットでは、過去ログを毎回KVキャッシュへ積み上げると、同時接続数が制約になります。Retention型のバックボーンは生成状態が長さに依存しないため、長いセッションでのメモリ予算を読みやすくできます。全文の厳密な再参照が必要な箇所はRAGに退避するハイブリッド構成が現実的です。
エッジ端末のオンデバイス生成
モバイル端末や組み込み機器では、モデル重みだけでなくKVキャッシュが実行可能性を決めます。固定状態で逐次生成できる設計は、少量メモリで長く動くローカル要約、キーボード補完、音声UIの候補になります。ただし端末向けの量子化・カーネル最適化は別途必要です。
長文RAGの文脈処理
RAGで取得した大量の根拠を一つの長いプロンプトに詰めると、生成側のキャッシュが増えます。RetNetの考え方は、局所的に重要な根拠を細かく処理しつつ、前方の要約状態を受け渡す設計に向きます。引用が必要な回答では、状態に圧縮した情報だけでなく原典チャンクIDを別経路で維持してください。
開発や事業へのヒント
この論文の実務的な価値は、「モデルを軽くする」ことより、状態の成長則をプロダクト設計の早い段階で意識させる点にあります。
コスト表をトークン長で作る
LLM機能の見積もりでは、平均入力長だけでなくp95・p99の会話長、ツール実行ログ、再試行を含めてください。TransformerならKVキャッシュがどのように増えるか、固定状態型なら品質劣化がどこで出るかを並べると、モデル選定と履歴保存方針を切り分けられます。
短期状態と外部メモリを分業させる
RetNetの多尺度減衰は、すべてを同じ保持期限で記憶しない設計です。アプリでも、直近の会話・実行中タスクはモデル文脈、長期の事実はベクトルDBや構造化DB、監査対象は原文ログというように役割を分けると、コストと再現性を両立しやすくなります。
置換前に品質の種類を評価する
固定状態は長い履歴を要約するため、逐語的な再現、離れた箇所の数字照合、例外条件の検索で弱くなる可能性があります。既存Transformerから移す場合は、perplexityだけでなく、長距離の根拠引用、長い指示の保持、ツール引数の正確さを個別にテストしてください。
限界
RetNetの状態は固定サイズであり、理論上は任意長の履歴を個別トークンとして完全に保持できません。指数減衰も古い情報を意図的に弱めるため、遠い過去の固有表現・数値・例外規則を正確に再利用する仕事には外部検索や明示的なメモリが必要です。
また論文の大規模比較は最大6.7Bの言語モデリングが中心で、現代の超大規模な指示追従、コード生成、マルチモーダル、エージェントワークロードでの優位性を直接保証しません。高速性も実装・GPU・バッチ構成に依存します。Attention向けに成熟したFlashAttentionやサービング最適化と比べ、RetNet系を運用する際にはカーネル、量子化、エコシステムの整備度を検証する必要があります。
さらに、元論文の減衰率は固定です。タスクに応じていつ忘れるべきかを動的に最適化すること、圧縮された状態の中身を説明・監査することは難しいままです。高い正確性や説明責任が求められる業務では、原文根拠を取り出せるRAGやログ設計を併用してください。
よくある質問
Q. RetNetはTransformerを完全に置き換えますか？
A. いいえ。論文は言語モデリング条件で有望な性能・効率を示しましたが、全タスク・全規模での置換を証明したわけではありません。成熟したTransformerの学習済みモデルやツール群も多いため、長文・低メモリという要件が強い箇所から評価導入するのが現実的です。
Q. RetNetではKVキャッシュが不要ですか？
A. Retention層の過去情報は、トークンごとのKV列ではなく固定サイズの状態行列に保持します。その意味で系列長とともに増えるKVキャッシュは不要です。ただし、モデル重み、現在の入出力、バッチ処理用のメモリは必要であり、アプリ側で原文や引用を保存する仕組みも別に必要です。
Q. なぜ学習は並列なのに推論は再帰的にできますか？
A. Retentionは、減衰付きの過去情報の和として一括計算する形と、状態へ逐次的に足し込む形が数式上等価になるよう設計されています。前者を学習時に、後者を自己回帰デコード時に使うことで、それぞれの計算環境に合う実装を選べます。
Q. 長文RAGならRetNetだけで検索は不要ですか？
A. 不要にはなりません。RetNetは履歴を効率良く状態にまとめますが、圧縮後に特定段落を原文のまま参照・引用する機能はありません。大量資料から根拠を見つけるRAGと、取得後の長いコンテキストを効率よく処理するモデルは補完関係です。
Q. 小規模チームがすぐ活かせる点は何ですか？
A. まず、会話長ごとのGPUメモリ、遅延、回答品質を測ることです。そのうえで、短期文脈、外部検索、永続ログの保持責務を分けてください。RetNetを自前学習しなくても、「状態を無制限に伸ばさず、必要な情報は外部メモリへ出す」という設計思想はすぐ適用できます。
今日の学び
この論文は、Transformerが長い生成履歴のKVキャッシュによって推論コストを増やす課題を扱いました。RetNetは、内容依存の参照と多尺度の減衰を持つRetentionを、並列・再帰・チャンク再帰の等価な計算形で実装して解こうとします。開発では、全履歴をモデル内に抱え続ける前提を見直し、短期状態と検索可能な外部メモリを分業させることが、長文AI機能のコストと品質を両立するヒントになります。