ReplaySSMによるハイブリッドモデルの高速化
ポッドキャスト、The Signal Shift by A.I.R. Labs始まりました。
SpotifyやApple Podcastなどで配信中です。
日々刻々と変化するAIの最先端トレンドと大きな変化のシグナルを読み解き、
ビジネスや日常にどう影響するか、初心者にでも分かりやすくお届けする番組です。
本日の10月13日、朝の配信号では、
NVIDIAなどの研究によるハイブリッドモデルの高速化手法、リプレイSSM、
強化学習のデータ合成技術、Magic Beanstalk、
そしてニューラルネットのサブ空間モジュール性という3つのテーマを取り上げます。
それでは最初のシグナルから見ていきましょう。
最初の話題は、NVIDIAなどが発表したリプレイSSMですね。
ハイブリッドモデルの文章生成を早くする手法とのことですが、
具体的にどのようなことが起きたのでしょうか。
はい、NemotronやQuen3.5などのハイブリッドモデルは、
文脈全体を参照するアテンション層に固定サイズの状態で情報をまとめるSSM層、
つまり状態空間モデルの層を組み合わせています。
リプレイSSMは、このSSM層の状態の持ち方をチェックポイント、
つまりある時点の状態の保存と直近の入力をためたバッファの組み合わせに変えました。
これで出力の計算と状態の更新を切り離します。
共同研究者はNVIDIAとプリンストンのDAO AIラボで、
TensorRT LLMやVLLMなどにも統合されています。
SSM層は早いのが売りだったと思うのですが、なぜそこがボトルネックになるのでしょうか?
理由は2つあります。
まず、Nemotron 3 Superの例では、SSM層の数がアテンション層の5倍あります。
次に、1層あたりの状態が1つの系列で4MBあり、
これはアテンション層の約8000トークン分の記憶に相当します。
従来は1トークンごとに状態を読み込み、更新し書き戻していました。
計算は少ないのでメモリの読み書き速度が足を引っ張り、
バッチサイズ256、つまり同時に256件を処理する場面では、
約10万トークンまでSSM層が最大の処理時間を占めます。
他にも課題があったのですか?
あと2つあります。
1つは、統計的デコードです。
小さなモデルが数トークン先を予想し、本命のモデルが一括検証して、
外れた分を巻き戻す高速化手法です。
アテンションなら巻き戻しは簡単ですが、
SSMは更新を元に戻せないため、
予想トークンごとに状態の写しを保存する必要があり、
メモリを圧迫します。
もう1つは量子化、つまり数値を荒い精度で持ってメモリを節約する技術です。
SSMの状態は次の計算にそのまま使われるので、
丸め誤差が毎ステップ積み重なってしまいます。
なるほど、状態を毎回読み書きしていたのを、
バッファに入力だけ書き足す形に変えたわけですね。
実際にどのくらい早くなったのでしょうか?
4つのモデルで、バッチサイズ256の時、
状態を処理する部分が1.46倍から1.80倍、
1ステップ全体では1.20倍から1.47倍に早くなりました。
バッチサイズ1では効果はありませんでした。
つまり、聞くのは多数のリクエストを裁く大規模な提供環境です。
GLM5.3 FlashやKimi K3では既に本番で使われています。
本番で使われているのは心強いです。
ところで、陶器的デコードの方ではどんな成果が出たのでしょうか?
予想トークンの入力をバッファに足し、検証後に採用分だけポインタを進めるので、
状態の移しも復元も不要です。
GSM8系の問題で、バッチサイズ512の時、
VLLM標準の陶器的デコードの1.85倍と2.10倍のスループットでした。
従来方式は8トークンの検証が通常の1ステップのほぼ5倍かかりますが、
リプレイSSMは約1ステップ分です。
精度面でも高精度のチェックポイントを保ち、
読み取り専用に低精度のコピーを使うので誤差が積み重なりません。
QUAN3.54Bの128系トークン時点で、誤差の指標が0.136から0.0038に下がりました。
Magic BeanstalkとRLVRのデータ合成
続いて2つ目の話題は、強化学習のデータ構成に関する論文、Magic Beanstalkです。
強化学習というのは、AIが個体の整合をもとに学ぶ方法ですよね。
どのようなアプローチなのですか?
はい。正解を自動で確かめられる問題で学ぶ手法をRLVRと呼びます。
この研究は、問題の数が多くても学習の手がかりになるとは限らないと指摘し、
元の問題を分解して再起的に組み立て直し、
新しい訓練問題の家族を作る方法を提案しました。
ゼロアドバンテージの罠とは、どういう状態を指すのですか?
対象のモデルと試行回数を固定した時、同じ問題を何度解かせても、
毎回正解か毎回失敗だと差が生まれず、学びの手がかりが得られません。
広く使われるデータセットを集めた大規模な数学コーパスでは、約88%がこの状態でした。
それなら問題を絞り込めば解決しませんか?
フィルタリングは残りを濃くできますが、選ぶことしかできず、難しい問題は捨てられてしまいます。
かといって、先生役のモデルに難問を言い換えさせるだけでは不十分でした。
難易度の調整も必要で、成功する思考も失敗する思考もあり、
さらに正解が訓練の長さの上限内に収まる必要があります。
Magic Beanstalkは、元の問題の考え方の骨組みを保ったまま、この条件に合う問題を作ります。
ちょうどいい難しさの問題を自動で作るわけですね。実際の成績はどうだったのでしょうか?
8つの数学ベンチマークで、GPT OSS120Bは61.2から69.8へ8.6点、
DeepSeek V4 Flashは72.4から77.9へ5.4点上がりました。
元の78系コーパス全体で訓練しても改善は1.1点、最良の絞り込みでも5.8点でした。
金融では340件の回答から14000問長を作り、石球で6.9点、ファイナンスベンチで2.4点の改善です。
データの量より質が効くということですね。
でも、わずか340件の回答から14000問を作れるなら、他の分野にも応用できるのでしょうか?
論文が示したのは、数学と金融での結果です。他の分野は未検証なので断定はできません。
ただ、約88%が全成功か全失敗なら、データを増やしても学習の手掛かりは増えにくいわけです。
質を見るには、まずそのモデルと試行回数で正解率を測るのが出発点になります。
ニューラルネットのサブ空間モジュール性
そして3つ目の話題は、ニューラルネットワークのサブ空間モジュール性に関する研究ですね。
個別のニューロン単位ではなく、サブ空間単位でモジュール性を見るというのはどういうことでしょうか?
従来は、1つのニューロンが複数の機能に関わるため、ネットワークは機能ごとに分かれていないと考えられてきました。
本研究は、モジュール性を3段階の入れ子で定義しました。
ニューロン単位、回転すれば機能が独立した方向に分かれるか、そして線形の変換で分けられるかです。
それぞれで分解を学習し、特定の方向だけを取り除く実験で因果的に確かめました。
ニューロン単位の分析では、なぜ見落とされてしまうのでしょうか?
複数の機能が同じニューロンの集団を共有していると、個々のニューロンを見ても役割が混ざって見えるからです。
機能はニューロンではなく、神経細胞の活動が作る空間の特定の方向に分かれていることがあり、座標軸を回すように見方を変えないと見えません。
実際、画像や言語のモデルでは、内容と様式、構文と意味が重なったニューロン集団を使いながら方向を取り除く実験で別々に切り離せました。
見かけはごちゃ混ぜでも、見方を変えれば整理されているわけですね。
この結果の確かさはどう裏付けられているのですか?
脳のデータとの対応が一つです。内容と様式の方向は、それぞれ視覚やの高次と低次の活動を予測し、意味の方向は構文の方向より言語やの活動をよく予測しました。
ニューロン単位では分かれていないと言われてきたAIの内部の仕組みを理解する研究、メカニスティックインタプリタビリティにとって重要な結果です。
以前は構造がないと言われた小さなネットワークでも確かめられたのでしょうか?
はい。複数の認知課題を学習したRNNは、これまで分かれた構造がないとされていました。
しかし、この手法では課題ごとの鋭い二重乖離が見つかりました。
つまり、片方の方向を取り除くと、一方の課題だけが壊れ、もう片方は壊れないのです。
個々のニューロンに専門性がなくても、特定の計算だけを狙った介入ができます。
ただ、論文が述べているのは、この解釈研究や生物の理解への願意までです。
AI運用と訓練データ設計への示唆
今回の3つのシグナルも、推論の効率化からデータの質、そして内部構造の解明まで、AIの進化を支える重要な基盤に迫る内容でしたね。
それでは、本日のビジネスアクションを、貝からお願いします。
本日のビジネスアクションは2点です。
1点目は、大量のリクエストを処理するAIを運用する場合、モデルの構造と推論の仕組みの最新動向をインフラ選定の判断材料として確認することです。
2点目は、強化学習やデータ合成を検討する際、データの量より、モデルが解いて正解と失敗が混ざる難易度かどうかを確認することです。
番組のフォローや評価、そしてコメントやフィードバックも是非お待ちしています。
リスナーの皆さんからの温かいフィードバックは、毎週丁寧に目を通して今後のアップデートに反映させていきます。
それでは、また次回の配信でお会いしましょう。バイバイ!