シャオミが先にMiMo-V3の新アーキテクチャを公開：100万トークンのプリフィル計算量が80%削減

動察 Beating AI 速報、Xiaomi MiMo チームは次世代 MiMo-V3 が採用する HySparse2 アーキテクチャを公開した。これは主に Agent が実行するほどコストが高くなる問題を解決する。Agent は毎回非常に短い指令を出すだけだが、Webページ、ターミナル、ツールから大量の内容を受け取る可能性がある。モデルはまずこれらの新しい内容を読み切る必要があり、コンテキストが長くなるほど、このステップの計算コストが増大し、KV Cache もより多くの VRAM を占有する。

HySparse2 はまずモデルを前後2段に分ける。前半段は入力の処理を担当し、後半段は推論と生成を続ける。後半段が必要とするコンテキスト情報は、前半段ですでに計算済みの結果から直接生成できるため、prefill で後半段を完全に再実行する必要がない。この発想は Microsoft Research が2024年に提唱した YOCO に由来し、その名前は「You Only Cache Once」を意味し、核心は後半段に前半で計算済みの情報を共有させ、キャッシュを少なく保存し、重複計算も減らすことにある。

同時にスパースアテンションも再設計した。通常のフルアテンションは毎回完全なコンテキストを見る必要があるが、HySparse2 では少数の層だけがそうする。それらはまず長いコンテキストから最も関連性の高い 1024 個の token を選び出し、後続の層はこの結果をそのまま利用し、同時に直近 128 個の token を固定で保持する。前世代の HySparse は 64 token ごとにグループ化してから選択し、グループ内に重要な情報が1つでもあればグループ全体を残す必要があった。今回は token 単位の選択に変更され、同じ計算量でより多くの真に関連する内容に割り当てられる。この変更単独のテストで、2つの長文検索スコアがそれぞれ 6.57 点と 8.14 点向上した。

論文では総パラメータ 800 億、毎回のアクティブパラメータ約 30 億のモデルを対照として用いている。49 層モデルは prefill 段階で前半 25 層のみを実行すればよい。入力が 100 万 token に達すると、MiMo-V2 シリーズが採用するハイブリッドスライディングウィンドウアテンションと比較して、prefill 計算量は約 1/5 に低下し、KV Cache は 12.09GB から 2.69GB に減少した。