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、質量スペクトルと分子構造といった異なるモダリティ間の対応付けにも適用し、一定の成功を収めている。 ただし、この手法が有効に機能するかどうかは、独立に学習された表現がどの程度似た幾何構造を持っているかに大きく左右される。 コメント === 異なるモダリティの表現をペアデータなしで対応付けることは、以前から重要かつ困難な研究課題として知られていた。 古くは異なる言語間の埋め込み空間を対応付ける研究があり、その後、画像と言語のような異種モダリティ間の対応付けへと研究が広がってきた。 今回、異なるモダリティ間でも粗い幾何構造が共通しており、細かな構造までは一致していなくても、比較的単純な手法によって十分に良い初期対応を得られることを示した点は興味深い。 この成果には、事前学習によって得られる表現の質が向上し、異なるモデル間でも共通する意味構造が現れやすくなっていることが大きく貢献していると考えられる。 一方で、なぜ独立に学習された異なるモダリティのモデルが、このような共通の幾何構造を獲得するのかは、まだ十分に理解されていない。さらに、どのような情報や構造がモダリティを超えて普遍的に現れ、どのような情報がモダリティ固有のものとして残るのかもまだ分かっていない。

MoEのモデルサイズやExpert活性化率を変えながら、最適な学習率とバッチサイズがどのように変化するかを体系的に調べた研究。 1,800回の事前学習実験を行い、Expert活性化率が1/32、追加検証では1/64まで調べている。最大12Bパラメータ、159Bトークンまでの実験を行った。 主な結果は以下の3点である。 ・総学習計算量が増えるほど、最適学習率は低くなる。 ・学習トークン数が増えるほど、最適バッチサイズは大きくなる。おおむねトークン数の平方根に比例するペースで増加する。 ・MoEを疎にすると、最適学習率は下がり、最適バッチサイズは大きくなる。 例えば同じ総学習計算量・学習トークン数の条件で、Denseモデルの最適学習率と最適バッチサイズをそれぞれ1とすると、 ・活性化率1/16：学習率0.69、バッチサイズ1.26 ・活性化率1/32：学習率0.62、バッチサイズ1.34 ・活性化率1/64：学習率0.57、バッチサイズ1.42 となる。 MoEでは、各Expertが受け取るトークン数が少なくなるため、勾配ノイズが大きくなる。しかし、活性化率を1から1/16に下げても、最適バッチサイズは1.26倍にしかならない。 これは、各Expertの勾配ノイズを抑えるためにバッチサイズを大きくするメリットがある一方、バッチサイズを大きくしすぎると、同じ学習トークン数でのOptimizerの更新回数が減ってしまうためである。結果として、Expertあたりのトークン数を一定に保つほどバッチサイズを増やすのではなく、更新回数とのバランスを取ることが最適になる。 なお、これらはMuonを用いた特定のモデル構成で得られた実験結果であり、数百Bパラメータ、数十Tトークン規模の学習でも同じ関係が成立するかは未検証である。 コメント === 最近のMoEでは、Expert活性化率が5%や3%を下回るような、非常に疎なモデルも登場している。 このようなモデルでは、各Expertが受け取るトークン数が減るため、勾配ノイズが大きくなり、学習率やバッチサイズをどのように設定すべきかが重要になる。 従来もMoEの最適ハイパーパラメータに関する研究はあったが、活性化率の違いによる影響は十分に整理されていなかった。今回の研究は、最適学習率を総学習計算量と活性化率、最適バッチサイズを学習トークン数と活性化率で整理できることを示した点で実用的な価値が高い。

ReplaySSMは、Mamba-2やGated DeltaNet（GDN）などのSSMにおける状態の読み書きのコストを削減し、モデル全体のデコードを1.20〜1.47倍高速化するとともに、投機的デコード（Speculative Decoding）のスループットも最大約2倍に改善する。 SSMは通常のAttentionと異なり、過去のトークン情報を固定サイズの状態行列に圧縮する。しかし、この状態の読み書きには大きなメモリ転送コストがかかる。 例えば、Nemotron-3-Super-120BではSSM層が40層あり、1層あたりの状態は1シーケンスにつき4MiBである。従来の実装では、各トークンのデコード時に状態を読み出して更新し、再び書き戻すため、全層で約320MiBのデータ転送が必要になる。 特にバッチサイズが256のように大きい場合、コンテキスト長が100Kトークン程度になるまで、SSMの状態更新がモデル全体の主要なレイテンシ要因になり得る。 ReplaySSMは、状態の書き戻しを毎トークンではなく一定間隔で行う。例えばL=16の場合、16トークンごとに状態をチェックポイントとして保存し、その間はチェックポイントの状態と、チェックポイント以降の各トークンの入力情報を明示的に保持する。 Mamba-2の場合、保持するのは減衰係数、キー、値であり、これらから必要な状態を再構成できる。毎回巨大な状態行列を書き戻す必要がなくなるため、メモリ転送量を従来の約58%まで削減できる。 特に有効なのが、投機的デコードへの応用である。 従来のvLLMなどのSSM実装では、複数の候補トークンをまとめて検証する際、各トークンの処理後の状態を保存しておき、受理されたトークン数に応じた状態を選択する必要があった。 ReplaySSMでは、候補トークンの入力情報をバッファに保存しておけばよく、受理されたトークン数に応じてバッファのポインタを進めるだけで状態を確定できる。棄却されたトークンの影響を取り除くために、巨大な状態行列を復元する必要がない。 さらに、各トークンの出力を計算する際には、完全な状態行列を再構成する必要すらない。キーとクエリの内積を先に計算することで、状態行列を明示的に構築することなく、出力を直接計算できる。 GDNでは、現在の状態を参照して書き込む差分を決めるため、Mamba-2より計算が複雑になるが、複数トークン間の依存関係を小さな三角行列の連立方程式として解くことで、同様の処理を実現している。 また、SSMの状態を低精度化すると、量子化誤差が次の状態更新に引き継がれ、長い生成では誤差が蓄積するという問題がある。 ReplaySSMでは、状態更新に用いるチェックポイントを高精度で保持し、通常の出力計算には別の低精度チェックポイントを用いる。これにより、低精度化による誤差が次の状態更新に伝播することを防ぎ、長い生成でも精度劣化を大幅に抑えられる。 コメント === SSMは、Attentionのように過去の全トークンのKVキャッシュを保持する必要がなく、固定サイズの状態だけを維持すればよいため、効率的なデコードが期待されていた。 しかし実際には、状態行列そのものが大きく、その読み書きがボトルネックになっていた。また、過去の情報を一つの状態に集約するため、KVキャッシュのようにポインタ操作だけで過去の状態に戻ることができず、投機的デコードとの相性も悪かった。 ReplaySSMは、直近の入力は状態に畳み込まずキャッシュとして保持することで、こうした制約を大幅に緩和する。SSMの固定サイズ状態という利点を維持しながら、KVキャッシュのような柔軟な状態管理も部分的に実現した点が重要である。 一方、Attention側でも、過去のKVキャッシュを圧縮・集約したり、一定区間ごとにまとめて管理したりする研究が進んでいる。さらにContext Language Modelのように、モデル自身が過去のコンテキストを管理する仕組みも提案されている。 記憶の表現と管理方法という観点では、両者の境界が徐々に曖昧になっているように思える。