SuperMimicは、ヒューマノイドの全身動作を学習するために、従来のPPOを中心とするモデルフリー強化学習ではなく、学習した微分可能な世界モデルを通じた誤差逆伝播法（BPTT）によって方策を直接最適化する。その結果、従来128基のGPUで7日間かけて学習したモデルの性能を、わずか1基のGPU、3時間の学習で上回った。 問題設定は、あらかじめ正解となる人間の参照動作データが存在し、それをロボットに模倣させるタスクである。正解となる動作の位置、姿勢、並進速度、角速度がロボットの実際の動作に近ければ近いほど、高い報酬を与える。 学習は以下の3段階を繰り返す。 ・現在の方策をMuJoCo Warp上で実行し、状態・行動・次状態をリプレイバッファに保存する。 ・8ステップの自己回帰予測を使い、ロボットのダイナミクスをニューラルネットワークで近似する世界モデルを学習する。 ・学習した世界モデル上で32ステップの仮想ロールアウトを実行し、報酬から行動、方策へと勾配を逆伝播（BPTT）して方策を最適化する。 これらを同一の学習ループで交互に繰り返し、世界モデルと方策を並行して学習する。 世界モデル自体は3層・各2048ユニットのMLPであり、巨大なモデルではない。これだけ小さなモデルでも機能するのは、ロボットの運動が視覚や言語とは異なり、比較的低次元の状態空間と共通の物理法則によって記述できるためではないかと考えられる。 高速に学習できる理由は、大きく2つある。1つ目は世界モデルによるシミュレーションが高速であることと、2つ目は誤差逆伝播法によって方策を効率的に最適化できることである。 世界モデルはMuJoCo Warpによる物理シミュレーションより高速に状態遷移を計算でき、実際に使われる並列環境数（2K〜16K）では10倍以上、最大64Kの並列環境では62倍高速に実行できる。 さらに、従来のPPOのようにサンプリングした行動から方策勾配を推定するのではなく、微分可能な世界モデルを通じて報酬の勾配を直接計算するため、効率的に方策を改善できる。 実験では、同じ世界モデル上でPPOを学習させた場合に比べ、約64分の1の仮想ステップ数で同等の報酬を達成している。 今回、世界モデルは8ステップの自己回帰予測で学習するが、方策の最適化では32ステップ先まで展開して勾配を計算する。なお、32ステップは50Hzの制御周期で0.64秒に相当する短い時間である。 また、32ステップだけでは長期的な安定性や将来の報酬を十分に評価できないため、Criticによる価値推定を組み合わせる。これにより、短いロールアウトで勾配を計算しながら、それ以降の長期的な報酬も考慮できる。 世界モデルは摩擦係数、重心位置、キャリブレーション誤差などのシミュレーションコンテキストを入力として使う。一方、実機に配置する方策はこれらを直接利用せず、関節位置・速度やIMUなどの観測情報を使う。いわゆるPrivileged Informationを学習時に利用する構成になっている。 コメント === 微分可能な世界モデルを使って効率的に学習するアイデアは古くから数多く提案されてきたが、今回の成功には、長期のダイナミクスを直接ロールアウトするのではなく、短い時間の予測と勾配計算に限定したことが大きく寄与していると考えられる。 具体的には、世界モデルを8ステップ（0.16秒）の自己回帰予測で学習し、方策のBPTTは32ステップ（0.64秒）に限定する。それ以降の長期的な価値についてはCriticを利用することで、予測誤差の蓄積を抑えながら方策を改善できる。 また、並列化可能で高速なMuJoCo Warpをさらに、学習された世界モデルで高速化し、1秒あたり合計数万秒分のロボット動作を仮想的に計算できることになる。 さらに、世界モデルを現在の方策が訪れる状態分布に合わせて継続的に学習することも重要だと考えられる。方策の改善によって訪れる状態や選択する行動が変化しても、新しいシミュレーションデータを収集して世界モデルを更新するため、方策が世界モデルの学習範囲から大きく逸脱する問題を抑えられる。 Appendix D.1に書かれている、方策が失敗しやすい区間を重点的にサンプリングする工夫も興味深い。動作データを1秒単位に分割し、Returnの低い区間を優先して学習するようにしている。例えば、宙返りの着地が難しい場合、動作の最初から繰り返すのではなく、着地直前の状態から重点的に学習できる。 また、今回のタスクは、あらかじめ正解となる動作が存在し、位置や姿勢の誤差から各ステップで滑らかな報酬を計算できるという特徴がある。こうした密な報酬も、微分可能な世界モデルを通じた方策最適化が成功した重要な要因だと考えられる。 そのため、正解動作が存在しないタスクや、物体との複雑な接触を伴う操作タスクにおいても同じような学習効率を達成できるかは、今後の重要な研究課題となる。

DeepSeekは、自社の推論フレームワークDeepInferを公開するとともに、DeepSeek V4.1 Flashの大規模推論システムに関する技術報告も公開した。 この報告によれば、V4.1 Flashは単一クラスタで1日あたり155兆トークンを処理し、GPUあたり毎時38ドルのAPI売上換算値を達成している。また、Decodeのベンチマークでは、GPUあたり毎秒14,466トークンという出力スループットを報告している。 この報告では、PrefillとDecodeの分離、Expert Parallelism（EP）、KVキャッシュの長期保存、投機的デコード、負荷分散などについて、実際の運用を踏まえた設計が詳しく説明されている * PrefillとDecodeの並列化 PrefillではExpert Parallelism（EP）を2〜8、Data Parallelism（DP）を2〜8とし、モデルの重みがHBMに収まる範囲で、できるだけ小さいEPを利用する。 DecodeではEP64〜128、DP64〜128という大規模な並列化（Wide EP）を採用している。 V4.1 Flashは各MoE層に256個のRouted Expertを持つためEP64であれば、単純に均等配置した場合、GPUあたり平均4個のExpertが配置される。人気のあるExpertを複製し、各GPUの負荷に応じて再配置するExpert Parallel Load Balancing（EPLB）も導入している。 このようにEPを大きくすることで、各GPUが読み込むExpert重みの量を減らし、同時に多数のトークンを集約してExpertのGEMMを効率化できる。一方で、GPU間通信の負担は増えるため、通信側の負担は大きくなっている。 * KVキャッシュの長期保存と再利用 本番トラフィックでは、入力長の平均が77Kトークンで、そのうち約96%がキャッシュヒットしていると報告されている。 KVキャッシュは分散ファイルシステム3FS上に保存され、どのPrefillノードからでもアクセスできるように設計されている。これにより、リクエストを特定のGPUに固定する必要がなく、柔軟な負荷分散が可能になる。 また、Sliding Window Attention（SWA）はキャッシュ再利用時に必要な状態を厳密に再構築するには、層数分だけ、より長い範囲の再計算が必要になるが、そこを直近128トークンだけの再計算に置き換えている。これは近似だが、Prefill全体の計算量をほぼ半減できると報告している。 * GPUあたりの処理性能と売上 DeepSeekは、GPUあたり毎時38ドルのAPI売上換算値を報告している。これを単純に年間換算すると、GPUあたり約33万ドル（1ドル150円換算で約5,000万円）となる。 主要な処理性能は次の通りである。 - Prefill処理量（キャッシュ込み）：約150万トークン/秒/GPU - Prefill処理量（新規計算）：約6万トークン/秒/GPU - Decode処理量：約14,000トークン/秒/GPU - Decode同時実行数：約60系列/GPU - 1ユーザーあたりのDecode速度：約230トークン/秒 - PrefillとDecodeのGPU数：約1対1 これらの値とAPI料金を組み合わせると、GPUあたり毎時38ドルという売上換算値と概ね整合する。 1日155兆トークンという処理量からGPU規模を逆算すると、上記のスループットを継続的に達成できるという仮定のもとで、Prefill約1,200基、Decode約1,200基、合計約2,400基という規模が推定される。 この構成が年間を通じて同じ稼働率と売上単価を維持した場合、年間売上換算値は約8億ドルとなる。 例えば、GPU原価を毎時5ドル、その他のインフラ原価を毎時3ドルと仮定すると、GPUあたり毎時38ドルの売上に対して粗利益率は約79%となる。これは原価を仮定した試算であり、実際の利益率を示すものではないが十分採算がとれていることを示している。 * GPUの演算効率 次に、Decode性能をB200やGB200のピーク演算性能と比較してみる。 V4.1 FlashのDecode時のActive Parametersを16Bとすると、1トークンの主要な行列積演算量は約32 GFLOPsとなる。 GPUあたり14,466トークン/秒で生成する場合、出力トークンに対応する演算量は約463 TFLOPS/GPUである。 さらに投機的デコードにおける検証トークンの増加を考慮すると、主要モデルの演算量は約516 TFLOPS/GPUと推定される。 これはB200やGB200のFP8・FP4ピーク演算性能に対して約5〜12%に相当する。 また、GPUあたり60系列を同時に処理し、投機的デコードによって1ステップで平均3.21トークンを生成するとすると、1ステップの処理時間は約13.3ミリ秒と推定される。 * HBM帯域とExpert重みの読み込み メモリ帯域の観点からも考えてみると、仮に552BパラメータをすべてFP4で格納し、EP64で均等に分割すると、GPUあたり約4.31 GBの重みを保持することになる。 これを1ステップごとにすべて読み込むと仮定すると、13.3ミリ秒のステップ時間で平均約324 GB/sの読み込み帯域が必要になる。 B200のHBM帯域は約8 TB/sなので、ステップ全体で平均すればピーク帯域の約4%に相当するので十分間に合う上、重みの読み込みとGEMMは一部をオーバーラップできる。 ただし、13.3ミリ秒の間にはAttention、Expert間通信、GEMMなどの処理も必要であり、重みの読み込みに13.3ミリ秒すべてを使えるわけではなく、例えば重み読み込みに使える時間が1ミリ秒しかないなら、必要帯域は4.31 TB/sとなり、B200のピーク帯域の約54%を使用する計算になる。 * SSDを利用したKVキャッシュ保存 今回の推論インフラでSSDは中心的な役割をはしている。 推論中のKVキャッシュと状態はGPUのHBMに保持される。一方、SWAのKVキャッシュはCPU RAMに、Paged Block KV Cacheは3FS上のSSDに保存される。 SSDはDecode中のKVキャッシュをHBMから退避させるためではなく、主に数日間にわたるPrefixの再利用に使われる。 本番システムでは、KVキャッシュ用のSSDストレージを約100 PiB確保しており、GPUあたりの論理ストレージ容量は約14.4 TBである。 キャッシュ保持期間は、Agentトラフィックで3日間、Chatbotトラフィックで7日間に設定されている。 また、キャッシュヒット率は約96%である一方、API料金ではキャッシュヒットした入力トークンと新規計算が必要な入力トークンの価格差が50倍に設定されており、Prefill売上においてキャッシュヒットした場合とそうでない場合がちょうどバランスされている。 コメント === 今回の報告は、現在の大規模LLM推論インフラについて、実際の利用状況と、それを支えるシステムの構成や設計上の工夫を具体的に示している点で参考になる。 KVキャッシュをSSD上に長期間保存してPrefillの再計算を削減し、Decode側でWide EPを採用してExpert重みの読み込み負担を減す上、単に推論カーネルを高速化するのではなく、ストレージ、通信、負荷分散、投機的デコードまで含めて最適化することで、システム全体のスループットと経済性を改善している。 また、推論フレームワークDeepInferが公開されたことで、今後コミュニティによる性能検証や改善が進むことも期待できる（ただし、本番で使われているKVキャッシュ管理や負荷分散システムなどは公開されていないようである）

SCALLOPは、Flow Matchingによる高速なサンプル生成に加えて、生成したサンプルの確率密度（尤度）も数ステップで推定できる手法である。 従来の尤度推定で必要だったヤコビアンの計算や、分散の大きなHutchinson推定を、Flow Matchingと同様の条件付き期待値の学習に置き換えることで回避している。 サンプルの確率密度（尤度）を評価することは、さまざまなアプリケーションにおいて重要である。 例えば、エネルギーU(x)を持つ分子系のボルツマン分布は、 π(x) = exp(-U(x)) / Z で表される。ここではエネルギーを無次元化し、Zは正規化項（分配関数）である。 生成モデルの確率密度をp(x)としたとき、これによって生成したサンプルはボルツマン分布とは異なるが、重要度重み w(x) = exp(-U(x)) / p(x) を使って重要度サンプリングを行うことで、生成モデルによるサンプリング結果を目標のボルツマン分布に補正できる。 ただし、この補正が正しく機能するためには、生成したサンプルの確率密度p(x)を正確に評価できる必要がある。 従来、フローマッチングを含むContinuous Normalizing Flowなどで生成したサンプルの確率密度を求めるには、速度場の発散を積分する必要があった。 d log p_t(x_t)/dt = -∇·v_t(x_t) dx_t/dt = v_t(x_t) ここで∇·v_tは、速度場v_tのヤコビアンのトレース（対角成分の和）である。 この計算は高次元になるほど高コストになるため、HutchinsonのTrace Estimatorを使うことが多かった。しかし、確率的な推定による分散が大きく、学習や密度推定のボトルネックになっていた。 SCALLOPは、この発散を直接計算するのではなく、ニューラルネットワークで学習する。 まず、Flow Matchingで一般的に利用される線形補間を考える。 x_t = (1-t)x_0 + tx_1 ここでx_0は標準ガウス分布からのサンプル、x_1はデータ分布からのサンプル、0 < t < 1は時刻である。 各サンプル対(x_0, x_1)を結ぶ速度は、 u_t = x_1 - x_0 である。 一方、周辺分布p_tを輸送する速度場は、この速度の条件付き期待値になる。 v_t(x) = E[x_1 - x_0 | x_t = x] つまり、時刻tに位置xを通るような、x_0を始点、x_1を終点とするさまざまな輸送経路の速度の平均である。 Flow Matchingは、この関係を利用して速度場を学習する。SCALLOPは同様の考え方を、発散の学習に適用する。 まず、対数確率密度の全微分を考える。 d log p_t(x_t)/dt = ∂_t log p_t(x_t) + ∇log p_t(x_t) · v_t(x_t) 一方、確率密度の連続の方程式から、 d log p_t(x_t)/dt = -∇·v_t(x_t) が成立する。これらを組み合わせると、 -∇·v_t(x_t) = ∂_t log p_t(x_t) + ∇log p_t(x_t) · v_t(x_t) となる。 つまり、負の発散は、時間方向のスコアと空間方向のスコア、および速度場の組み合わせとして表現できる。 次に、データx_1を固定して考える。このとき中間分布はガウス分布になる。 p_t(x | x_1) = N(tx_1, (1-t)^2 I) この性質を利用すると、x_1で条件付けたスコアは解析的に求められる。 空間スコアは、 ∇log p_t(x_t | x_1) = -x_0/(1-t) 時間スコアは、 ∂_t log p_t(x_t | x_1) = [d + x_0·(x_1-x_0)]/(1-t) となる。ここでdはデータの次元数であり、·は内積を表す。 さらに、周辺スコアは条件付きスコアの期待値として表現できる。 ∇log p_t(x) = E[∇log p_t(x | x_1) | x_t = x] ∂_t log p_t(x) = E[∂_t log p_t(x | x_1) | x_t = x] これらを組み合わせると、負の発散を次の条件付き期待値で表現できる。 -∇·v_t(x) = E[[d - x_0·(v_t(x)-(x_1-x_0))]/(1-t) | x_t=x] このように、右辺には速度場v_t(x)は登場するものの、そのヤコビアンは一切登場しない。 SCALLOPでは、まず事前学習したFlow Matchingモデルを使ってv_t(x)を近似する。次に、その出力とx_0、x_1から条件付き期待値の内側は求められるのでそれを教師信号として使い、これをニューラルネットワークを二乗誤差最小化で学習し回帰する。 MSE回帰の最適解は教師信号の条件付き期待値になるため、理想的な速度場が得られている場合には、負の発散を出力できるように学習できる。 さらに、論文では発散をスカラーではなく、ベクトルとして学習する方法も提案している。 発散はヤコビアンの対角成分の和である。 通常はこの合計値を一つのスカラーとして学習するが、SCALLOPではヤコビアンの対角成分を個別に予測し、最後に合計する。 これにより各次元に個別の学習信号を与えられ、異なる次元の誤差が合計時に相殺されることを避けられる。 この場合の教師信号は、 g_t = [1 - x_0 ⊙ (v_t(x_t) - (x_1-x_0))]/(1-t) となる。 ここで⊙は要素ごとの積、1は全成分が1のベクトルである。 その条件付き期待値は、 E[g_t | x_t=x] = -diag(∇v_t(x)) となり、ヤコビアンの対角成分に負号を付けた値を学習できる。 論文では、この結果を対数密度の全微分からだけでなく、条件付き期待値として表された速度場のヤコビアンを直接導出する方法からも示している。 結果として学習したFlow Mapでは、ネットワークのforward計算だけでサンプル生成と確率密度推定を同時に行える。 実験では、アラニンの2〜6残基からなる分子系を用いて評価している。 その結果、正規化フローベースのSBGと比較して約10倍、従来手法FALCONと比較して約100倍の推論速度を報告している。 また、Hutchinson推定を利用するF2D2と比較して学習の安定性が向上し、Importance Samplingの効率を表すEffective Sample Size（ESS）も改善している。 ただし、サンプリング精度ですべての既存手法を上回ったわけではなく、特に大きな分子系では改善の余地がある。 コメント === Flow Matchingにおける尤度推定は、基本的かつ重要な問題でありながら、効率的に実現することが難しかった。 通常の確率モデルでは、尤度を計算できる場合、尤度を最大化するようにパラメータを推定する最尤推定が利用できる。 一方、Flow Matchingは、尤度を直接最大化するのではなく、ガウス分布などの簡単な分布から目標分布へ変化する速度場を学習する。 原理的には、学習した速度場の発散を積分することで尤度を計算できるが、その計算は高コストであり、特に高次元のモデルでは扱いづらい。 そのため、Likelihoodベースのモデルとの比較や、確率密度を必要とする下流タスクへの利用には課題があった。 尤度はモデルの確率分布を評価する基本的な指標であり、分布推定や重要度サンプリング、自由エネルギー計算など、多くの問題で重要になる。 今回の研究は、Flow Matchingで速度場を条件付き期待値の回帰によって学習するのと同様に、尤度推定に必要な発散も条件付き期待値の回帰に帰着させたという点で、非常に興味深い。

MAI-Base-2は、2.8Tパラメータ、80B activeのMoEモデルを、16000基のGB200を用いて38日間、50Tトークンで事前学習したと報告している。 前回のMAI-Baseに続き、公開情報が限られている大規模事前学習について、技術的な詳細や実験結果、開発手法がまとめられており、参考になる。 今回も、改善ループが回せるよう、小規模モデルの実験結果から大規模モデルの性能を正確に予測するために、Source Scalingという手法を導入している。 具体的には、あるデータソースについて本番モデルで8 epoch繰り返す場合、小規模実験でも同じ8 epochを繰り返すように、使用するデータソースのサイズを調整する。 通常、小規模実験では本番学習よりも学習トークン数が少ないため、同じデータソースを使うと繰り返し回数が少なくなる。しかし、それではデータの反復による過学習などの影響を再現できず、小規模実験で得られた最適な設定が本番モデルでも最適になるとは限らない。 Source Scalingでは、学習トークン数だけでなく、ユニークなデータ量やデータの反復率も調整することで、大規模学習の挙動を小規模実験から予測しやすくしている。 もう一つの重要な結果は、学習率（LR）とWeight Decay（WD）を独立に探索するのではなく、その積に着目することで、最適なハイパーパラメータについて次の経験則を得ていることである。 LR × WD = 17.04 / S また、最適な学習率は次の関係を満たす。 log LR = -1.51 + 0.175 log(TPP) - 0.573 log S ここで、Sは学習ステップ数、TPPは学習トークン数をActive Parameter数で割った値である。 なお、LRとWDの積に着目する考え方は、既存研究でもいくつか見られ、初出ではない。 Weight Decayによる累積的な重みの縮小効果を考えると、学習率が一定であると単純化した場合、 W_{t+1} ≒ (1 - LR × WD) W_t となる。これをSステップ繰り返すと、 W_S ≒ exp(-S × LR × WD) W_0 となる。 今回の経験則では、S × LR × WD = 17.04なので、 W_S ≒ exp(-17.04) W_0 となり、学習全体を通じたWeight Decayの累積的な強さは、設定によらず一定に保つという意味を持つ。 なお、実際には学習率のスケジュールや勾配による重みの更新が存在するため、重みのスケールがこのようになるわけではない。 また、これらの係数は今回の実験設定で得られた経験値であり、異なるOptimizerやモデル構造でもそのまま成り立つとは限らない。 このほか、Activation Gradient Clippingによって学習の安定性を大きく改善したことも興味深い。 通常のGradient Clippingでは、重みに対する勾配のノルムを制限する。しかし、巨大なバッチのうち1サンプルだけが異常に大きな勾配を発生させた場合でも、バッチ全体の勾配が縮小されてしまう。 そこで、ネットワーク内の特定箇所におけるActivationの勾配をクリップする方法を採用している。 Activationに対する勾配は、Backward Passではバッチやトークンの次元が保持されているため、重み勾配として合算される前に、個別の勾配を制限できる。 これにより、異常なサンプルの影響を局所的に抑えつつ、正常なサンプルの学習信号を残すことができる。 さらに、異常に大きな勾配を発生させたサンプルを調査したところ、データの前処理のバグにより、想定しない位置でトークンが出現していることが判明し、データ前処理の修正にも役にたった。 計算効率についても注目に値する。MAI-Base-2は、総パラメータ数2.8T、活性化率約2.9%という非常に大規模かつ疎なMoEモデルであるにもかかわらず、モデル全体でMFU 21%を達成している。 GB200のピークBF16性能を2.5 PFLOPSとすると、これは1GPUあたり525 TFLOPSの有効学習演算性能に相当する。特に、MoE層単体ではMFUを24%から35%まで改善しており、全層MoEを採用できるようになったことも重要である。 コメント === 前回のMAI-Baseに引き続き、大規模モデルの事前学習について詳細な技術情報を提供している。 前回は、様々な後続タスクに関連するテキストの予測誤差を指標として、同じ損失に到達するための投入計算資源をどこまで減らせるかという、学習効率の最適化に重点が置かれていた。今回は、こうした開発手法をさらに大規模なモデルへと適用し、実際に学習を完了させた成果といえる。 特に、総パラメータ数が約3T、活性化率が3%を下回るモデルを実際に学習している点は重要であり、モデル規模としてはフロンティアモデルに近い。 今回の報告において独自の評価指標ではDeepSeek-V4-Proの事前学習などを上回る結果が得られたと報告している。今後、Post-trainingによって推論能力やエージェント性能がどの程度向上するのかが注目される。 もう一つ注目したいのは、これほど大規模な事前学習を38日間で完了させたことである。 従来、大規模モデルの事前学習には3〜4か月程度をかける例も多いが、今回は16,000基のGB200を投入し、高い疎率にもかかわらずMFU 21%という学習効率を実現していることで、1ヶ月強で終了させている。 このような短期間で終わるようになれば大規模事前学習でも、アーキテクチャ、データ、Optimizer、学習条件などを改善して再学習するサイクルをすすめらえっるようになる。

Priming Modelsは、拡散モデルやFlow Matchingのような複雑な生成過程を学習しなくても、単純な復元学習による事前学習と、生成分布を目標分布に合わせる事後学習を組み合わせることで、高品質な1ステップ生成モデルを実現する。 ImageNet 256×256の1ステップ画像生成ではFID 1.19を達成し、同程度の品質を持つ従来手法と比較して学習計算量を大幅に削減している。さらに、事後学習にFD-lossを使用した別の実験では、FID 0.94を達成している。 この手法は画像生成だけでなく、行動を条件として次時刻の画像を予測する世界モデルにおいても有効であることが示されている。 現在の画像生成では、複数ステップの生成過程を用いる拡散モデルやFlow Matchingが広く使われている。また、これらの生成過程を1ステップに短縮するMeanFlowやConsistency Models、あるいは分布を直接学習するDrifting Modelsなども登場している。 1ステップ生成で実現したいことは、ガウスノイズzを入力したとき、生成器f(z)の出力分布qが実データ分布pと一致することである。これを最も直接的に実現するには、適切な分布間距離D(q,p)を定義し、それを最小化すればよい。 しかし、ニューラルネットワークをランダムに初期化した状態から、勾配降下法によって直接生成分布を実データ分布に一致させることは不可能なくらい難しい。 そこでPriming Modelsは、この問題を2段階に分けて解決する。 第1段階：ノイズから画像を復元する事前学習 まず、実画像にノイズを混ぜた画像を作り、そこから元の画像を復元するように学習する。クラスラベルやテキストなどの条件がある場合には、それらも入力する。 通常の拡散モデルでは、複数のノイズ強度を使って復元を学習する。しかし、Priming Modelsでは、固定したノイズ強度での復元学習だけでも、その後の生成学習に有効な初期化が得られることが示された。 興味深いのは、元画像の情報を完全に消すほどノイズをいれると、最終的な生成性能が大幅に悪化することである。つまり、第1段階から純粋なノイズを入力して元画像を予測するように学習する方法は、うまく機能しなかった。 ImageNetの実験では、ノイズ強度をs=0〜1としたとき、s=0.8程度が最も良い結果となった。 ノイズからの復元が可能な程度に元画像の情報を残しつつ、最終的な生成時に用いる純粋なガウスノイズに近い入力で学習することが重要だと考えられる。 なお、ImageNetの最高性能モデルでは、単純なピクセル単位の復元損失だけでなく、LPIPSやConvNeXtの特徴量を用いた知覚損失も併用している。 第2段階：生成画像の分布を実画像の分布に合わせる事後学習 復元学習を終えたモデルに純粋なガウスノイズを入力すると、1ステップで何らかの画像を生成できる。しかし、この時点では生成分布は実データ分布と一致しておらず、生成画像の品質も低い。 そこで、生成画像と実画像を学習済みの特徴抽出器に通し、特徴空間で両者の分布が一致するよう追加学習する。 この際の主な学習目標は、最適輸送（Optimal Transport）を利用したSinkhorn divergenceである。 最適輸送は、ある分布を別の分布に変換する際、確率質量をどのように移動させれば全体の輸送コストが最小になるかを求める方法である。 エントロピー正則化付き最適輸送の目的関数をD(q,p)としたとき、Sinkhorn divergenceは次のように表される。 S(q,p) = D(q,p) - 1/2 D(q,q) - 1/2 D(p,p) 第1項は生成分布と実データ分布の間の輸送コストであり、両者を近づける。 第2項は生成分布自身の輸送コストを補正するもので、生成サンプルが過度に集中することを防ぐ方向に働く。第3項は実データ分布自身に関する補正項であるが、今回、実データ分布は固定されているため、モデルのパラメータ更新には影響しない。 この目的関数をそのまま微分することもできるが、本研究では同じパラメータ勾配を与える次の損失関数を使用する。 L = D(q,p) - D(q,sg[q]) ここでsgはstop-gradientを意味し、その引数を通じた誤差逆伝播を止める操作である。 元の目的関数ではD(q,q)の前に1/2が付いているが、D(q,q)を微分すると、2つの引数の両方から勾配が流れるため、打ち消される。 さらに、最適輸送の勾配計算には包絡定理（Envelope Theorem）を利用する。 最適輸送では、生成サンプルiと実サンプルjの間の輸送コストC_{ij}を定義し、それぞれの間でどれだけの確率質量を輸送するかを輸送計画π_{ij}として求める。 エントロピー正則化を含む目的関数を、輸送計画πについて最小化することで、最適な輸送計画π*が得られる。 ここで生成画像が変化すると、輸送コストCだけでなく、最適輸送計画πも変化するため、勾配を計算するにはπの変化まで微分する必要があるように思える。 しかし、π*はすでに輸送計画の最適解である。したがって、輸送計画が制約を満たしながら微小に変化しても、その変化が最適化済みの目的関数に与える一次の寄与は、最適性条件により0になる。 このため、最適輸送計画が正確に求められ、輸送計画の制約が固定されている場合には、π*を固定して輸送コストCのみを微分すればよい。 具体的には、 ∇θ D(q,p) = Σij π*_{ij} ∇θ C_{ij} となる。ここでDはエントロピー正則化項を含む最適輸送の最適値を表す。 つまり、Sinkhornアルゴリズムによって最適輸送計画を計算した後、その輸送計画をstop-gradientで固定し、生成画像と実画像の特徴量間の距離を変化させる勾配だけを逆伝播すればよい。 実際の有限ミニバッチと有限回のSinkhorn反復では近似となるが、この方法によって複雑な輸送計画の微分を計算せずに、分布を一致させるための学習を実行できる。 ここで、単純な復元事前学習によって分布学習が容易になるのかについて、著者らは、復元事前学習によって、分布学習に適したJacobianが形成されるのではないかと考え、実験的に分析している。 モデルの出力特徴量をhとし、分布損失Lの特徴量に関する勾配をvとする。 v = ∇h L モデルのパラメータθに関する出力特徴量のJacobianを J = ∂h / ∂θ とする。 連鎖律より、損失のパラメータに関する勾配は ∇θ L = J^Tv となる。 これは、特徴量に関する損失勾配vを、Jacobianの転置によってパラメータ空間の勾配に変換していることを意味する。 学習率ηで勾配降下した場合、パラメータの更新量は Δθ = -ηJ^Tv となる。 一方、パラメータが微小に変化したとき、出力特徴量の変化は一次近似によって Δh ≈ J Δθ と表される。これは、通常の微分における「出力の変化 ≈ 微分係数 × 入力の変化」を多変数に拡張したものである。 ここにパラメータの更新式を代入すると、 Δh ≈ -ηJJ^Tv となる。 つまり、分布損失が要求する修正方向がvであっても、実際の出力変化はJJ^Tvによって決まる。 JJ^Tは出力特徴量に対するNeural Tangent Kernel（NTK）に相当する。この行列の固有ベクトル方向のうち、固有値が大きい方向には、同じ勾配に対して出力を大きく変化させやすい。 逆に、固有値が小さい方向では、分布損失が適切な修正方向を示していても、パラメータ更新によって出力を十分に動かせない。 著者らは、復元事前学習を行ったモデルについて、分布損失が要求する主要な修正方向に対するJacobianの応答を測定した。その結果、復元事前学習が長いほど、一般にその方向へのJacobianの相対的な応答が大きくなり、その後の分布学習でもFIDが速く改善する傾向が確認された。 つまり、復元事前学習によって、分布を修正するために必要な方向へ出力を変化させやすいモデルが形成されていると考えられる。 ただし、現時点ではJacobianの応答と学習効率の相関を示したものであり、それが性能改善の直接的な原因であることまで証明したわけではない。 コメント === 生成モデルの出力分布を目標分布に直接一致させる方法は、GANやMMDを利用した生成モデルなど、以前から研究されてきた。しかし、ランダム初期化からの直接的な分布学習には安定性や学習効率の問題があり、高品質な生成モデルを簡単に学習することは難しかった。 今回の研究では、単純な復元タスクによる事前学習によって、直接的な分布学習が効率的に進むモデルを構築できることが示された。 構成は非常に単純だが、従来手法を大幅に上回る学習計算効率を達成しており、重要な結果だと思う。 また、事前学習が単に画像の特徴や構造を獲得するだけでなく、その後の学習に適したパラメータ空間の幾何学的構造を形成している可能性である。これは他の事前学習に対しても同様の示唆を与える。 また、ノイズ強度として0.8程度が有効だった点について、画像の大局的な構造や意味的な特徴が形成される重要な領域に対応する。他のデータや問題はどう設定するかは未解決である。

MSAR（Multi-Stream Attention Residuals）は、Transformerの残差接続を改良し、過去の層の情報を複数の異なる視点から選択的に取り出す手法である。特に、従来のAttnResをマルチヘッド化し、異なる局所文脈に基づいて過去の情報を読み出すことで、同じ学習損失に到達するための計算量を、通常の残差接続と比べて約37%削減できると報告している。 通常のTransformerでは、各層の出力を残差接続によって足し合わせる。 h_{l+1} = h_l + f_l(h_l) ここで、h_lは第lサブレイヤーへの入力、f_lはAttentionまたはMLPによる変換である。 各サブレイヤーが新しく計算した出力をv_l = f_l(h_l)とすると、この式は h_l = h_1 + Σ_{i=1}^{l-1} v_i と表せる。つまり、通常のTransformerでは、過去のすべてのサブレイヤーの出力を係数1で加算している。 この仕組みは勾配を安定して伝え、各層が既存の表現を徐々に更新する形で計算を分担できる一方、初期の層で生成された情報を後続の層が選択的に取り出すことは難しい。 この問題に対して、HC、mHC、AttnResなど、残差接続を改良する研究が進んでいる。 特にAttnResでは、 h_l = Σ_{i=0}^{l-1} α_{i,l} v_i とし、過去の各層の出力に対する読み出し重みα_{i,l}をAttentionによって決める。ここでv_0はEmbeddingであり、v_iは先程と同様に各層の計算結果である。 このAttentionでは、各層に学習可能なQueryを用意し、過去の各層の出力をKeyとValueとして利用する。通常のSelf-Attentionが系列方向に情報を選択的に取り出すのに対し、AttnResは深さ方向に情報を選択的に取り出す。 MSARは、このAttnResをマルチヘッド化する。通常のMulti-Head Attentionでは、特徴ベクトルを複数のHeadに分割し、それぞれ独立したAttentionを計算する。 しかし、Depth Attentionでは、各層の出力を単純に特徴次元方向に分割すると、性能が悪化することが先行研究で報告されている。 そこでMSARでは、系列方向のCausal Depthwise Convolutionを利用し、異なる局所文脈を取り込んだ複数の表現を生成する。 標準構成では4つのStreamを使用する。 ・Stream 0：元の表現をそのまま使用 ・Stream 1~3：直近4, 8, 12トークンに対するCausal Depthwise Convolution 各チャネルで独立した畳み込みを行うため、追加パラメータ数とFLOPsを小さく抑えられる。 こうして得られた4つのStreamは、いずれも元のHidden Dimensionと同じ次元数を持ち、それぞれに独立したDepth Attentionを適用する。 最後に、得られた4つの表現を適応的な重みで統合する。この重みは、4つのStreamの出力を連結したベクトルにRMSNorm、線形写像、Sigmoidを適用し、4つのスカラーゲートを計算することで得られる。 また、全サブレイヤーの出力を個別にAttentionの参照対象とすると、必要なメモリ容量やPipeline Parallelの通信量が増大する。 そこで、大規模実験ではBlock MSARを採用する。これはTransformer全体を約8ブロックに分割し、各ブロックに含まれるサブレイヤーの出力を加算した一つの状態を保持して、これをAttentionの参照対象とする。 また、4つのStreamを常時保持するのではなく、保存した単一の状態から必要なときだけStreamを生成することで、メモリ使用量と通信量を抑えている。 実験では18B-A1.7Bと28B-A2.7BのMoEモデルを用いて評価を行い、学習損失の改善に加えて、下流タスクでも大きな性能改善が得られた。 28B-A2.7Bモデルでは、MSARの追加学習FLOPsは約0.5%である一方、12種類のベンチマークすべてで通常の残差接続、mHC、AttnResを上回った。 さらにScaling Lawの分析では、同じ学習損失に到達するための計算効率が通常の残差接続に対して1.58倍となり、mHCの1.30倍、AttnResの1.35倍を上回った。ただし、これは実際の学習速度ではなく、学習損失と計算量の関係から推定した値である。 Ablationでは、単にDepth AttentionのHead数を増やすだけでは不十分であり、各Headに異なる局所文脈を反映した表現を与えることが重要であると示唆されている。 また、2 Streamでも4 Stream構成の改善の大部分が得られており、ブロック化による性能低下も小さい。 コメント === Transformerの残差接続の改良については、近年さまざまな手法が提案されている。 今回の論文ではMulti-Head Depth Attentionの有効性が主張されているが、単純にCausal Convolutionを導入した効果ではないかという点は、十分にAblationされていないため、追加検証が必要だと思う。 実際、他の研究でもCausal ConvolutionがInduction Headに相当する振る舞いの獲得を容易にする可能性が指摘されている。また、MambaなどのSSM系アーキテクチャでもCausal Convolutionが利用されており、局所的な情報混合が学習を容易にするうえで重要な役割を持つ可能性がある。 もちろん、Causal Convolution自体に効果があり、さらにMulti-Head Depth Attentionと組み合わせることで追加の改善が得られる可能性も十分にある。切り分けのためには、Single-HeadのAttnResにCausal Convolutionだけを追加した構成との比較が必要だろう。 今回のように、Attentionを適用する対象をどこまで拡張できるかが今後注目される。 通常のSelf-Attentionでは、同じ層の過去のトークンを参照する。一方、AttnResやMSARでは、同じトークン位置における過去の層の出力を参照する。さらにMSARでは、Causal Convolutionによって、過去のトークン位置にある下位層の表現も間接的に利用している。 これを一般化すると、系列方向と深さ方向を合わせた二次元の計算履歴の中から、必要な中間表現を選択的に取り出すという方向性が考えられる。 例えば、現在の層から過去のトークン位置にある上位層、特に最終層の表現を参照する構成も考えられる。こうした過去時刻の上位層の表現を再利用する場合は、入力からの計算ステップ数を増やせる（しかし学習時に全時刻を並列で学習できなくなるという大きな問題がある） 一方、計算、メモリ、並列計算の容易性から、中間表現をすべて保存することは現実的ではないため、過去の情報のうち何を捨て、何を集約し、何を将来の参照のために選択的に残すべきかという問題も、今後の重要な研究テーマになると思う。

画像モデルと言語モデルを完全に独立に学習していても、埋め込み空間の幾何構造のみから、画像と文章の意味的な対応関係を推定できることが以下の論文で示された。 "Shared Geometry as a Rosetta Stone: Cross-Modal Alignment Without Paired Data" これはPlatonic Representation Hypothesis、つまり異なるモデルが共通の世界構造に対応する表現へと収束するという仮説を、実際の異種モダリティ間の接続を通じて部分的に裏付けた研究といえる。 従来、別々に学習された画像モデルと言語モデルの表現を接続するには、画像と文章の対応ペアを利用することが一般的だった。 しかし、例えば画像側に犬、猫、道路、自動車、人に対応する表現があり、文章側にもそれらの意味に対応する表現がある場合、「犬と猫は近く、自動車と道路も近い」といった表現間の相対的な位置関係が似ていれば、対応ペアを知らなくても両者を重ね合わせられる可能性がある。 この実現において中心となるのが、Wasserstein Procrustesという枠組みである。 \min_{W,T } \sum_{i,j}T_{ij} |x_iW-y_j|^2 これは、両者を対応付ける変換行列Wと、サンプル同士の対応関係Tを同時に推定する最適化問題である。ここでWは半直交行列、Tは輸送計画であり、変換は距離や角度を保存するように制約されている。 この最適化問題は非凸であり、適切な初期値がなければ解くことが難しい。そこで、本論文では次の3段階の方法を採用する。 1. 粗い幾何構造の整合 それぞれの埋め込み空間を30クラスタに分け、クラスタ中心間の類似度構造が可能な限り一致するように対応付ける。このために、CKA（Centered Kernel Alignment）を最大化する二次割当問題を解く。 2. 疑似ペアの生成 クラスタ間の対応関係から、個々のサンプル間の対応関係を推定する。具体的には、線形割当アルゴリズムを使って画像と言語の疑似ペアを構築する。 3. Procrustesによる反復改善 疑似ペアから最適な直交変換Wを求め、変換後の類似度から新しい対応関係を推定し直す。これを反復することで、対応付けの精度を改善する。 実験では、30クラスタ、30回の初期化、バッチサイズ10,000、100回の改善反復を使用している。 これにより、ペアなしの画像・文章対応付けにおいて従来手法を大きく上回る性能を達成している。さらに、少数の教師ペアを追加するだけで対応付けの精度が大幅に改善し、従来手法よりも高いデータ効率を実現していることも示している。 また、画像と言語だけでなく、RNA発現とクロマチンアクセシビリティ、異なる被験者のfMRI、質量スペクトルと分子構造といった異なるモダリティ間の対応付けにも適用し、一定の成功を収めている。 ただし、この手法が有効に機能するかどうかは、独立に学習された表現がどの程度似た幾何構造を持っているかに大きく左右される。 コメント === 異なるモダリティの表現をペアデータなしで対応付けることは、以前から重要かつ困難な研究課題として知られていた。 古くは異なる言語間の埋め込み空間を対応付ける研究があり、その後、画像と言語のような異種モダリティ間の対応付けへと研究が広がってきた。 今回、異なるモダリティ間でも粗い幾何構造が共通しており、細かな構造までは一致していなくても、比較的単純な手法によって十分に良い初期対応を得られることを示した点は興味深い。 この成果には、事前学習によって得られる表現の質が向上し、異なるモデル間でも共通する意味構造が現れやすくなっていることが大きく貢献していると考えられる。 一方で、なぜ独立に学習された異なるモダリティのモデルが、このような共通の幾何構造を獲得するのかは、まだ十分に理解されていない。さらに、どのような情報や構造がモダリティを超えて普遍的に現れ、どのような情報がモダリティ固有のものとして残るのかもまだ分かっていない。