WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

現代のテキスト生成を行う人工知能は、トランスフォーマーと呼ばれる構造に依存しており、これは一連の積み重なった層を通じて情報を処理します。情報を加工するステーション1、次にステーション2、というように、原材料が各ステーションを通過していく工場の組立ラインを想像してみてください。各ステーションは製品を洗練させていきます。標準的なセットアップでは、システムが次のテキストに移動するとき、各ステーションは直上のステーションが前のステップで行った作業しか見ることができません。これは、スタックの上部で形成されたより深く複雑な洞察が、新しい情報を処理している際の下層や初期のステーションにはアクセスできないことを意味します。システムは豊かな理解の履歴を作り上げているものの、その全容を効率的に活用することに苦慮しており、それが長期的な文脈の追跡や複雑な問題の解決能力を制限しています。

南カリフォルニア大学の研究者たちは、このボトルネックを解決するために、「WhiteMatter」と呼ばれる新しいアーキテクチャを提案しました。WhiteMatterは、各層を自身の深さに限定する代わりに、システム全体が過去のあらゆる層からの混合された要約情報にアクセスすることを可能にします。チームは、ダイナミック・ルーターとして機能するメカニズムを構築しました。これは、前の単語のすべての層からの隠れ状態を取り込み、それらを少数の共有チャネルへと混合するものです。現在の単語の各層は、これらのチャネルから一つを選択して読み取ります。この設計は、脳の遠隔領域が白質線維を通じて接続される仕組みを模倣しており、これにより、ネットワークの浅い部分が過去の深い処理済み情報を受け取れる一方で、深い部分も依然として生の即時的な文脈にアクセスできるようになります。

研究者たちは、80億個のテキストトークンからなる大規模なデータセットを用いて、このシステムをゼロから学習させることでテストを行いました。彼らは、この新しいWhiteMatterモデルを、同サイズの標準的なモデル、およびより大規模な標準モデルと比較しました。その結果、16層のWhiteMatterモデルは、24層の標準モデルよりも優れた性能を示し、次の単語の予測におけるエラー率を大幅に低下させました。この向上は、システムに必要なメモリ保存容量を半分に圧縮した場合でも達成され、モデルが高い性能を維持しながら、より少ないメモリを使用できることを証明しました。また、この研究は、コンピュータが一度に多くの単語をまとめて処理できる特定の反復手法を用いることで、このシステムを効率的に学習できることを示しました。これは、通常このような深い接続のために必要とされる「一つずつ順番に処理する」という制約を回避するものです。

核心となる革新は、システムが過去の単語のメモリをどのように扱うかにあります。典型的なセットアップでは、単語のメモリは層ごとに保存され、次の単語は同じ深さにあるメモリにしかアクセスできません。WhiteMatterは、単語の内容に基づいて更新されるメモリチャネルのプールを作成することで、このルールを打破します。ルーターは、過去の単語のすべての層の隠れ状態を分析し、それらを少数のチャネルへと混合します。これらのチャネルの数は調整可能であり、研究者たちは、チャネルの数を総層数よりも少なく設定することで、複雑なテキストの理解能力を損なうことなく、メモリのフットプリントを削減できることを見出しました。この柔軟性により、システムはすべての層の履歴を詳細に保持する必要はなく、最も関連性の高い混合物のみを保持すればよいため、より効率的になります。

これを学習フェーズで実現するために、研究者たちは循環依存を克服しなければなりませんでした。なぜなら、単語のメモリは単語の処理に依存し、かつ処理もメモリに依存するため、システムは単純に一度のパスですべてを計算することができないからです。チームは、テキストをグループごとに処理する「サイクリック・ガウス＝ザイデル反復法（cyclic Gauss–Seidel iteration）」と呼ばれる手法を開発しました。これにより、システムは段階的に理解を更新でき、シーケンス内の後続のグループが、同一パス内であっても先行するグループから更新された情報を即座に利用できるようになります。このアプローチは、同じレベルの精度に達するために何度もパスを必要とする従来のメソッドよりも大幅に高速であると判明し、このような深い相互接続を持つモデルの学習を標準的なハードウェア上で実現可能にしました。

実験の結果、このアーキテクチャは生成されるテキストの質を向上させるだけでなく、さまざまな推論タスクにおけるモデルの性能も高めることが確認されました。言語理解に関する標準的なベンチマークでテストした際、WhiteMatterモデルは、同等の深さを持つ標準的なモデルを一貫して上回り、さらに大型の標準モデルをも凌駕しました。また、16層のモデルに対して8つのチャネルのみを使用する「ハーフメモリ構成」は、フルキャッシュを使用する場合よりも大幅に少ないメモリを使用しながら、性能向上の大部分を維持しました。このことは、過去の相互作用のメモリを動的に混合する能力が、単にメモリストレージのサイズを大きくすることよりも価値があることを示唆しています。

研究では、システムが異なる学習条件下でどのように振る舞うかについても調査されました。研究者たちは、学習中にデータを反復処理する方法が、最終的なパフォーマンスに大きな影響を与えることを発見しました。最終的なステップバイステップのデコーディング・プロセスを密接に模倣したスケジュールで学習されたモデルは、より優れた性能を示し、より安定していました。しかし、より単純な学習スケジュールを用いた場合でも、WhiteMatterモデルは標準的なモデルを上回っており、これはアーキテクチャ自体の変更が改善の主要な要因であることを示しています。研究者たちは、学習プロセスには標準的なモデルよりも多くの計算能力が必要となるものの、デコーディング（実際のテキスト生成）プロセスはほぼ同等の速度であり、実世界のアプリケーションにおいて実用的であると指摘しました。

人工知能のより広い文脈において、この研究は、ディープラーニングモデルが情報を処理する際の根本的な限界に対処するものです。すべての層がネットワーク全体の履歴の合成されたビューにアクセスできるようにすることで、WhiteMatterは、深いアーキテクチャを典型的に悩ませる「孤立」を克服します。今回の知見は、システムの効率性が単に規模を大きくしたり深くしたりすることから来るのではなく、内部構造間でどのように情報を接続し、共有するかから来るものであることを示唆しています。過去の相互作用のメモリを、データの豊かさを失うことなく圧縮できる能力は、より有能で効率的な言語モデルを構築するための有望な道筋を提供します。

研究者たちは、彼らのアプローチが「深層から浅層へのフィードバック」をうまく統合し、システムが表現の全深度を活用できることを成功裏に示したと結論付けました。彼らは、これが減少したメモリ・フットプリントで実現可能であることを証明し、「優れた性能にはより多くのストレージが必要である」という仮説に異を唱えました。学習プロセスには多少の複雑さが伴うものの、モデルの使用時における精度と効率性の向上は多大です。この研究は、ニューラルネットワークの異なる部分間の接続を再考することが、いかに性能の劇的な向上につながるかを示す具体的な例であり、次世代の人工知能システムの発展に向けた新しい方向性を提示しています。