WISE、再帰型モデルのアテンション対象範囲を再利用し、アテンション処理で1.36倍の高速化を計測

バージニア大学の研究者らは9月23日、再帰型言語モデルがグローバルアテンションを繰り返し計算するコストの削減を目指すWISEを公開した。この手法は再学習を必要とせず、初期の計算で重要なコンテキストのブロックを特定し、後続のステップでその選択範囲を再利用する。モデルの内部表現は引き続き更新される。[論文の概要](https://arxiv.org/abs/2609.27373)

公開実装では、デフォルトで32回の再帰処理を実行し、最初の12回はグローバルアテンションを維持する。システムは32トークンを1ブロックとし、第9〜12ステップのそれぞれで、アテンション重みの累積が95%に達するまでブロックを選択する。その後、選択したブロックを統合し、後続ステップでアクセス可能な集合とする。第13ステップ以降に固定されるのは対象ブロックの範囲であり、Q、K、Vとブロック内のアテンション重みは引き続き再計算されるため、後半のステップで表現をさらに精緻化する能力が保たれる。[リポジトリの説明](https://github.com/tbn5pj/WISE_code)

この設計の工学的な価値は、不規則なアテンションの疎性を、GPUで実行可能な処理スケジュールに変換する点にある。著者らが提供するTritonカーネルは、選択されていないkeyブロックをスキップし、スコア計算、因果マスク、softmax、valueの集約を融合する。これにより、マスクを適用するだけで依然としてすべてのブロックを走査してしまう処理を回避する。[実装ドキュメント](https://github.com/tbn5pj/WISE_code/blob/main/docs/REPRODUCING.md)

著者らはRTX A6000上で、Huginnモデルの30組のアテンションワークロードを測定した。4Kコンテキストの後半20ステップでは、標準のFlashAttentionに対して約1.76倍の高速化を達成した。最初の12ステップの探索を含めると、全32ステップのアテンション処理では約1.36倍となる。これらの数値には対応するスケジュールの構築コストが含まれるが、モデル全体の推論は対象に含まれていない。[論文の速度測定表](https://arxiv.org/html/2609.27373v1#S7.SS4)

品質面の代償も見過ごせない。論文の対応のある検定では、2Kまでのコンテキストで平均F1スコアの低下は検出されなかったが、4Kでは約3.3ポイント低下した。つまり、最大の高速化が得られた条件は、同時に品質とのトレードオフを検討すべき条件でもあり、あらゆる長文タスクで従来の性能を維持できると一般化することはできない。[品質と効率の評価](https://arxiv.org/html/2609.27373v1#S7.SS5)

また、論文ではコンテキストの品質テストを100例と説明している一方、公開された再現用ドキュメントには、各コンテキスト長につき30件、合計120件のプロンプトと記載されている。両者のサンプルの対応関係は、なお明確にする必要がある。保存済みの結果表から図を再作成しても、元の評価を再実行したことにはならない。[再現用ドキュメント](https://github.com/tbn5pj/WISE_code/blob/main/docs/REPRODUCING.md)

再現用資料にも不足がある。リポジトリにはカーネル、設定、データ識別子、プロンプトのハッシュが含まれているが、速度測定に用いたテンソル一式は添付されていない。また、ドキュメントには、公開版をまとめる際に大規模モデルの評価とGPUでの速度測定を再実行していないと記されている。導入を検討する側の次の一歩は、照合可能なテスト入力を入手し、同じハードウェアとソフトウェアの環境でエンドツーエンドのレイテンシを再測定するとともに、中国語、マルチターン対話、さらに長いコンテキストでの品質を別途検証することだ。[公開版の整合性確認に関する説明](https://github.com/tbn5pj/WISE_code/blob/main/PAPER_CODE_CONSISTENCY.md)