Retrofitting Linear Attention into Diffusion Language Models

人工知能の世界を、ある非常に賢いロボット司書が物語を書こうとしている、巨大で賑やかな図書館だと想像してみてください。長年、この司書の標準的な働き方は「自己回帰的（オートレグレッシブ）」なものでした。それは、一つの単語を書き、止まり、自分が書いたすべてについて考え、次の単語を書き、また止まって考える、というやり方でした。それはまるで、レンガを一つずつ積み上げ、次のレンガを置く前に接着剤が乾くのを待つようなものでした。信頼性は高いものの、この方法は長い物語を書くには非常に時間がかかりました。

最近、「拡散（ディフュージョン）」と呼ばれる新しいスタイルが登場しました。レンガを一つずつ積み上げる代わりに、拡散型の司書は、疑問符でいっぱいの白紙のページからスタートし、物語全体を一度に推測しようとします。そして、言葉が意味を成すようになるまで、その推測を何度も何度も洗練させていきます。これは、チームのアーティストたちが同時に壁画を描いているようなものです。前のレンガが乾くのを待つ必要がないため、彼らははるかに速く作業できます。しかし、この速い手法にも落とし穴があります。壁画が大きくなるにつれ、アーティストたちは新しい言葉が既存の言葉と適合していることを確認するために、すでに決定したすべての単語を絶えず振り返って確認しなければなりません。それは、101ページ目を書いている時に、10,000ページの書物を記憶しようとするようなものであり、最終的には全員の動きを鈍らせ、図書館のメモリを使い果たしてしまいます。

これが、ある研究チームが取り組んだパズルです。彼らはこう問いかけました。「過去の記憶の持ち方を変えることで、この高速な拡散法をさらに高速化できるだろうか？」彼らの答えは、「ブロック・ハイブリッド・アテンション（Block-Hybrid Attention）」と呼ばれる巧妙なトリックでした。彼らは、司書が現在の段落を完璧に理解する必要はあるものの（文章の流れを維持するため）、過去の章のすべてをハイビジョン級の解像度で読み直す必要はないことに気づきました。代わりに、古い章を小さな固定サイズの「要約シート」へと要約することができるのです。これにより、司書は膨大な歴史の重みに押しつぶされることなく、電光石火のスピードで物語の新しい部分を書き進めることができるようになります。

ビッグアイデア：二段階のメモリシステム

この論文は、既存の強力な拡散型言語モデル（具体的には160億パラメータを持つLLaDA 2.1）を、ゼロから再学習させることなく、この「要約シート」戦略を使えるように「レトロフィット（アップグレード）」する方法を紹介しています。

モデルの脳を、20層の思考レイヤーを持つものと考えてください。元のモデルでは、すべてのレイヤーが超注意力高い司書として機能し、現在の単語を理解するために、これまでのすべての単語を読み返します。これは正確ですが、非常に重たい処理です。研究者たちの革新である「ブロック・ハイブリッド・アテンション」は、作業を2つのモードに分割します。

1. 「今」モード（完全アテンション）： モデルが現在生成している単語のブロックに取り組んでいるときは、「超注意力」モードを維持します。文章を完璧に成立させるために、標準的な高負荷のメモリを使用して、現在の段落内のすべての単語を確認します。
2. 「過去」モード（線形アテンション）： モデルが（すでに完成した）前のブロックを思い出す必要があるときは、「線形アテンション」に切り替えます。本全体を読み直す代わりに、固定サイズの要約状態を参照します。これは、テキスト全体を読み直すのではなく、本の索引や1ページの要約を見るようなものです。この要約は、本がどれほど長くなってもサイズは変わりません。

実装方法：二段階のアップグレード

研究者たちは、単に部品を入れ替えて上手くいくのを待ったわけではありません。アップグレード中にモデルの知能が失われないよう、慎重な二段階のプロセスを用いました。

• ステージ1：シャドウ・トレーニング（影の訓練）。 彼らは元の賢いモデル（「教師」）を取り出し、それを凍結させました。次に、20層あるアテンション・レイヤーのうち6層を、新しい「ブロック・ハイブリッド」バージョン（「生徒」）に置き換えました。生徒は、入力された「破損した（corrupted）」バージョンのテキストを使用して、教師の出力を正確に模倣するように訓練されました。これは、生徒がマスターシェフの影となり、シェフの正確な動きをコピーし、割り当てられた特定の料理に対して全く同じ味を体験しようとするようなものです。このステージには約24時間を要しました。
• ステージ2：ファインチューニング（微調整）。 生徒が教師を模倣できるようになったところで、研究者たちはモデル全体を再び連携させました。彼らはLoRA（Low-Rank Adaptation）という技術を用いて、モデルの接続部に極めて精密な微調整を加えました。これは、新しいパーツと古いパーツが対話する際に生じる小さな不具合を修正するための、アップグレードされたモデルへの素早い「磨き上げ」のようなものです。このステージには約6時間を要しました。

結果：より速く、より軽く、そして依然としてスマートに

チームは、アップグレードされたモデル（名称：LLaDA-HYBRID）をテストし、それが依然として文章作成や問題解決において優れているか、そして実際に高速化されているかを確認しました。

思考能力は維持されているか？
驚くべきことに、そうです。20層のうち6層をこの「要約シート」方式に置き換えたにもかかわらず、モデルのパフォーマンスは元のモデルと非常に近い数値を維持しました。

• コーディングテストであるHumanEvalにおいて、元のモデルは**75.6%を記録しましたが、ハイブリッドモデルは72.0%**でした。
• 数学テストであるCMATHでは、元が**88.3%に対し、ハイブリッドは86.7%**でした。
• 興味深いことに、別のコーディングテストであるMBPP+では、ハイブリッドモデルは57.7%から63.0%へと向上しました。
  論文は、モデルが最も困難な推論タスク（GPQA-Diamondと呼ばれる難しい科学クイズなど。ここでは38.9%から30.8%へ低下）においては、わずかに精度が落ちるものの、コード作成や数学の問題を解く能力は概ね維持していることを示唆しています。

実際に速くなったのか？
ここが魔法の起こる場所です。モデルは会話の履歴をすべて読み返す必要がなくなったため、より多くのリクエストを同時に処理できるようになりました。

• モデルが1秒間に生成できる単語数をテストしたところ、ハイブリッドモデルはピーク時の性能（128の同時リクエストを処理する場合）において、元のモデルよりも1.7倍高速でした。
• 元のモデルは、物語が長くなるにつれてメモリ使用量が増大し、最終的にユーザーをこれ以上受け入れられなくなる壁に突き当たりました。しかし、ハイブリッドモデルは固定サイズの「要約シート」のおかげで、メモリ不足に陥る前により多くの同時リクエストを処理することができました。

まとめ

この論文は、必ずしも新しい高速なAIをゼロから構築する必要はないということを示唆しています。代わりに、強力な既存のAIを取り出し、それに「ハイブリッド・メモリ」システムを与えることができます。現在進行中の瞬間に鋭い集中力を保ちつつ、過去をコンパクトな形式に要約することで、知能を大きく損なうことなく、これらのモデルを大幅に高速化し、効率化できるのです。研究者たちは、このアップグレードが公開データを用いた約60時間のトレーニングで可能であることを発見しました。これは、AIサービスを、その賢さを犠牲にすることなく、より速く、より安価に提供するための有望な道筋を示しています。ただし、著者らは、これはまだ始まりに過ぎないとも述べています。彼らは今回6つのレイヤーのみをアップグレードしましたが、将来的にはさらに多くのレイヤーをアップグレードしたり、より長い物語を扱ったりする方法が見つかるかもしれません。