Co-founder and CEO, Preferred Networks (PFN) CEO, Matlantis. Co-Director of Development, Noetra

Japan Tokyo
Joined January 2008
Preferred Networksでは、仲間を募集しています。 私たちは今、AIを社会に届けるための技術と事業を、半導体、計算基盤、基盤モデル、AI製品・ソリューションまで、一気通貫で強化しています。 これほどの規模で、これほど幅広い領域を通じてAIの社会実装に本気で挑める環境は、世界を見ても多くありません。 次の時代に必要な技術や基盤がなければ、自らつくる。そして、つくった技術を研究成果で終わらせず、実際に社会で使われるところまで持っていく。それが私たちの仕事です。 現時点ではまだ公表できない、重要かつ大規模なプロジェクトも数多く動いています。これから世の中に出ていくものを、まさに今、自分たちの手でつくっています。 AIはこれから、国や企業、そして一人ひとりの生活を支える社会の基盤になっていきます。そうした基盤を、自分たちの手でつくっていきたいと考えています。 世界で通用する新しい技術を生み出し、それを社会に届く形にする。その挑戦のど真ん中で、一緒に未来をつくってくれる方に来てほしいと思っています。 スレッドに採用ページを貼ります。興味のある方はぜひご覧ください。質問などあれば、気軽にコンタクトしてください。
5
66
6
386
79,726
OpenAIの内部研究モデルが、本来禁止されていた外部インターネットへの接続に成功した事件を受け、OpenAIはフロンティアモデルの学習環境におけるセキュリティ対策を強化した。 その過程で、最新のフロンティアモデルを利用し、数千億トークンを投入した大規模なセキュリティ検証(Red Teaming)を実施した。 このRed Teamingにおける攻撃対象はOpenAI自身の学習インフラであり、次の3つの課題を提示した。 ・microVMからの脱出 ・外部インターネットへの接続 ・隔離されたAIエージェント間の通信 これらの課題に対し、AIエージェントは複数の脆弱性やシステムの特性を組み合わせることで、いずれの課題についても成功例を示した。 特に印象に残った例として、3番目の隔離されたAIエージェント間の通信では、ネットワーク接続や共有ファイルが利用できない環境で、共有CPUキャッシュを通信路として利用した。 具体的には、送信側のAgent Aは、特定のデータをCPUキャッシュから追い出すことでキャッシュの状態を変化させる。一方、受信側のAgent Bは、同じデータへのアクセス時間を測定し、キャッシュに存在するかどうかを推定する。このアクセス時間の差を0と1に対応させることで、情報を伝達できる。 ただし、AIエージェントがこれらの手法をすべて一から発見したわけではなく、多くの実験では、人間が調査対象の通信路を指定するか、符号化・復号のプロトコルを事前に与えていた。したがって、何の手がかりもなく通信方式を発見し、独立したエージェント同士が協調してプロトコルを開発したことまで実証されたわけではない。 一方、脆弱性に関する広範な情報だけを与えたところ、AIエージェントが通信方式を考案し、情報交換に成功した事例も報告されている。 こうした結果を受け、OpenAIはフロンティアモデルのRL学習について、仮想的なWeb環境を構築して原則オフラインで実施するとともに、複数の防御機構を組み合わせた多層防御を強化する方針を示した。 コメント === 強力なAIエージェント、特にサイバーセキュリティに関する高度な能力を持つAIを学習・検証する際には、従来とは異なるレベルのセキュリティ対策が必要になりつつある。 仮にネットワーク接続を完全に遮断したとしても、今回の例のように、共有ハードウェアなどを利用して、想定外の通信経路を構築できる可能性が示された。 今後は、技術的な脆弱性だけでなく、人間を間接的に誘導する形で、外部との通信や権限の獲得を実現しようとする事例も増えてくるだろう。 AIの問題解決能力が向上することは、与えられた制約を迂回する方法を発見する能力も向上することを意味し、研究開発で壁にぶつかっても、粘り強く創意工夫をして、解決策を考える能力は攻撃においてもそのまま有効となる。 以前から懸念されてきた、人間を超える能力を持つAIを本当に制御できるのかという問題が、今日のシステム設計や運用に関わる課題になりつつあると思われる。
1
20
5
93
22,988
LAWは拡散言語モデルにおいて、生成途中で離散トークンをサンプリングせず、トークン埋め込みの加重平均を繰り返し更新することで生成品質を改善する手法である。従来は補助的に使われていた自己条件付け(Self-Conditioning)を、生成過程そのものに拡張した手法といえる。 “Loopholing All the Way: Language Modeling with Pure Self-Conditioning” 従来の離散拡散言語モデルでは、各ステップでトークンの確率分布を予測し、そこから離散トークンをサンプリングして次のステップに進む。さらに、Loopholingなどの手法では、前のステップの隠れ状態を自己条件付け(Self-Conditioning)として次のステップに渡す。しかし、離散トークンのサンプリングによって、予測分布が持っていた他の可能性や不確実性に関する情報は失われてしまう。 LAWでは、各ステップで予測したトークンの確率分布に基づき、トークン埋め込みの加重平均(Expected Embedding)を計算し、それを次のステップの入力とする。これにより、途中で離散トークンをサンプリングすることなく、連続的な状態を繰り返し更新できる。 この連続表現での更新を繰り返し、最後のステップで初めて離散トークンをサンプリングする。連続的な生成過程を持つものの、従来の連続拡散モデルのようにODEやSDEを解く必要はない。 学習時には、主に次の2つの工夫を導入している。 ・Multi-pass Training:通常2回のForwardを行い、1回目の予測から得た連続表現を2回目に入力する。勾配は2回目のForwardにのみ流す。 ・Temperature Regularization:Softmaxの温度をランダムに変化させ、異なる確信度を持つ連続表現でも安定して更新できるようにする。 これらの工夫はどちらも重要であり、特にTemperature Regularizationを除くと、GSM8Kの正答率は48.8%から24.8%に低下する。また、2回のForwardの両方に勾配を流すと、42.2%から20.2%に低下した。 さらに、従来の自己条件付けではTransformerの最終隠れ状態を利用していたが、LAWで同じ方法を使うとGSM8Kの正答率は1.5%未満にとどまった。一方、Expected Embeddingを使うと35.8%に改善した。単に過去の隠れ状態を再利用するのではなく、トークン予測分布を経由して連続状態を作り直すことが重要だといえる。 実験では、Sudoku Hardの正答率が従来の91.0%から96.5%、GSM8Kでは34.5%から42.2%に改善した。さらに、既存の離散拡散モデルDuoを追加学習するだけで、GSM8Kの正答率が17.2%から48.1%に改善した。 大規模LLMへの適用は未検証だが、自己条件付けによる反復的な予測改善が拡散言語モデルの生成能力を大きく向上させることを示した、興味深い研究である。 コメント === Looped Transformerや潜在空間での推論と関係するような技術である。これまで成功していなかった部分について、今回の工夫(全て勾配を流すのではなく最後だけ流す、最後の隠れ状態ではなく、トークンの加重平均とする、学習時の温度をランダムにする)であるがこれらが決定的に重要な可能性はある。 この辺の技術がこなれてきた場合、多くの推論が明示的な推論トークン列ではなく途中ステップの連続的な状態の変換(トークンなので分析も可能)で行われるようになるだろうか。
1
8
1
34
5,112
Noiseは、Backpropを使わずにTransformerを事前学習するアルゴリズムを競う公開ベンチマークである。 37.7MパラメータのTransformerを対象に、誤差逆伝播を使わず、前向き計算のみで、同じトークン数の学習でより低い検証損失を達成するか、同じ検証損失により短い時間で到達するかを競う。 すでに通常のBackprop + Adamをわずかに上回る性能を達成するレベルには達しているが、改良されたBackpropを用いる手法には及んでいない。 基本となる考え方は、各層のActivationに摂動を与え、Lossの変化からパラメータの更新方向を推定するNode Perturbationである。 特に、Transformerの残差接続などの構造を利用し、各層が最終的な損失にどれだけ寄与したかを推定するCredit Assignmentを効率化している。 ・ランダムな摂動の代わりに、互いに直交するHadamard摂動を用い、効率的にLossの変化を推定する。 ・学習初期には、各層の出力を共通のLM Headで評価し、後続の層を通さずに局所的な学習信号を得る(Boost)。 ・学習後半では、残差接続の恒等経路や線形演算を利用し、最終的なLossへの影響を近似的に評価する(Cut)。 こうした性能改善では、Transformerの構造を積極的に利用することが重要になっている。 コメント === 完全な誤差逆伝播を行わなくても、ネットワークの構造を利用した近似的なCredit Assignmentによって、効率的に学習できる可能性が示されている。 一方今回の手法がそのまま大規模に適用できるかは不明であり、今後、長期学習時や、モデルの幅や深さに対するスケーリングがどうなるのかが興味深い。 Hadamard摂動の利用は、構造化された乱数を使う準モンテカルロ法やLow-discrepancy列などの考え方も連想させるが、直交性があり効率的に計算できるなど様々な発展が考えられる。 特に学習後半では、多くの更新方向でLossが改善しなくなる中、わずかな改善方向をいかに効率的に見つけるかが重要になる。そのためには、全空間を探索するのではなく、勾配が集中する低次元部分空間を推定し、その中で摂動を与えるといった考え方が必要になるかもしれない。 一般に、誤差逆伝播を使わない学習手法は、微分不可能な計算モジュールを含む場合、例えば離散化、ツールコール、他のエージェントの出力を利用する場合などに有効となる。また、モデルが非常に大きく、層方向の並列化を進めたい場合にも利点が考えられる。 今回の取り組みで得られたCredit Assignmentの効率化手法は、通常の誤差逆伝播による学習にも応用できる部分があるだろう。
1
18
114
10,038
Softmaxとクロスエントロピーの標準的な実装に2種類の数値誤差が存在し、従来のGrokkingの研究で観測されていたSlingshot現象(突然の損失急増)が、この数値誤差によって説明できることが示された。さらに、蒸留で教師と生徒の出力が近い場合にも問題が発生しうることが明らかになった。 “A Fix for Softmax and Cross-Entropy” * Softmax Collapse 通常のクロスエントロピー損失では、正解クラスのLogitに対する勾配は gₜ = qₜ - 1 である。ここでqₜは、Softmaxが正解クラスに割り当てる確率である。 正解クラスに割り当てる確率が大きくなるにつれて勾配は小さくなるが、数学的にはゼロにはならない。しかし、浮動小数点演算ではqₜがちょうど1に丸められ、勾配がゼロになってしまう。 一方、不正解クラスに対する勾配は、より小さな値まで保持される。 本来、すべてのクラスに対するLogitの勾配の総和はゼロ、つまり ∑ₖ gₖ = 0 を満たす必要がある。これは、すべてのクラスのLogitに同じ値を加えてもSoftmaxが変わらないという性質に対応する。しかし、この数値誤差によって勾配の総和がゼロでなくなる。 正解クラスのLogitと、それ以外のLogitのLogSumExpとの差をマージンmと定義すると、FP32ではmが約16.6を超えた時点でこの問題が発生する。 例えば、Logitが [9, -9, -9] であり、最初のクラスが正解の場合、通常実装の勾配は、マージンが18と16.6を超えてしまうので0となり、 [0, 1.523 × 10⁻⁸, 1.523 × 10⁻⁸] となるが、正しい勾配は [-3.046 × 10⁻⁸, 1.523 × 10⁻⁸, 1.523 × 10⁻⁸] である。 この違いは絶対値としては小さいものの、Adamのように勾配の大きさを使って更新量を正規化するOptimizerでは、重大な影響を与えうる。 特にGrokkingでは、ネットワークが訓練データを記憶した後、かなり時間が経ってから汎化性能が急上昇する現象が知られている。その過程では、Lossがゼロ近傍から急激に増大し、再びゼロに戻るSlingshot現象も観測されていた。このSlingshot現象が次の過程で発生することを示している。 1. モデルが訓練データを記憶すると、正解Logitのマージンが増大する。 2. マージンがFP32の約16.6を超えると、正解クラスの勾配がゼロに丸められる。 3. 勾配が極めて小さくなると、Adamの二次モーメント推定値が減衰する。 4. Weight Decayなどによってマージンが再び16.6未満になると、正解クラスの勾配が復活する。 5. 小さくなった二次モーメントによって勾配が正規化されるため、非常に大きなパラメータ更新が発生し、Lossが急増する。 さらに、AdamWのεとWeight Decayが、学習後のLogitのマージンを決めることも分析している。 Weight Decayが働く定常状態では、マージンは近似的に m* ≈ log(1/ε) となる。例えばε = 10⁻⁸の場合、m* ≈ 18.4となり、FP32で数値誤差が発生する16.6を超えてしまう。これもSlingshotが発生する原因となる。 この問題の解決方法は非常に単純であり、クロスエントロピーの計算に必要な L = log(1 + e⁻ᵐ) を、数値的に安定したSoftplus関数 L = softplus(-m) で計算する。 さらに、正解Logitとその他のLogitとの差を使って損失を表現することで、自動微分によって正解クラスの勾配が、その他のクラスの勾配の負の総和として計算されるようにする。 これによって、非常に小さな勾配も保持され、勾配の総和がゼロになる性質も維持できる。 * Softmax Cancellation もう一つの問題は、LLMの蒸留やKL正則化などで発生しうる。 TeacherとStudentの、確率分布をp, q としたとき、Soft-target Cross-EntropyのStudent側のLogitに対する勾配は g = q - p となる。 このときTeacherとStudentが完全に一致すれば、勾配はゼロになるはずである。 しかし、従来の実装ではTeacherとStudentのSoftmaxを別々に計算してから引き算するため、近い値同士の減算による桁落ちが発生する。さらに、前向き計算と後ろ向き計算で演算経路が異なることによって、両者が完全に一致していてもゼロでない勾配が発生する場合がある。 これは、Teacherの重みをコピーしてStudentを初期化する場合などに問題となる。 実際に論文では、Teacherの重みからStudentを初期化して蒸留を続ける実験において、従来実装ではBF16の数値誤差がAdamによって増幅され、StudentがTeacherから乖離して精度が低下することを確認している。 この問題も計算方法の工夫によって解決できる。具体的には、TeacherとStudentの確率を別々に計算して引き算する代わりに、Logitの差から出発し、expm1やlog1pなどの数値的に安定した関数を用いて勾配を計算する。 これによって、TeacherとStudentが一致する場合には勾配が正確にゼロとなる。 コメント ==== 今回指摘された数値誤差の原因となる計算は、PyTorchやJAXを含む現在の標準実装にも残っている。今後、これらの標準実装でも対応が進むと思われるが、大規模LLMの学習で使われるFused Kernelなどについても、同様の問題が存在するかどうかの確認が必要である。 このような一見わずかな数値誤差が、学習性能や学習中の不可解な挙動に大きな影響を与える、数値計算の方法によって勾配が変わり、Optimizerを通じて学習全体の挙動が変化するのは今後注意が必要である。 将来的には、コーディングエージェントが、勾配の総和がゼロになるといった数学的に保たれるべき不変性を理解し、実装の中でそれらが破られる可能性のある箇所を自動的に検出・指摘できるようになることを期待する。
4
115
13
622
73,014
RePRは、生成モデルに適した潜在表現(トークン)を、PMIの因子分解という原理に基づいて学習する手法である。特に、生成モデル向けの表現学習にScaling Lawが成立することを示した点で重要である。 “Generative Representation Learning for Scalable Latent Tokenization” RePRによって得られた潜在表現を使い、同じRAE生成フレームワークで比較したところ、FID-50KでI-JEPAやMoCo v3よりも高品質な画像生成を達成している。 さらに、Encoderの学習に投入する計算量を増やすことで、下流の生成性能が継続的に改善することを示し、また、Encoderと画像生成モデルの両方の学習コストを考慮しても、総学習計算量を減らしながらFIDを改善できることを実証している。 --- 従来の画像生成では、潜在表現を得るためにVAEやDINOv2などのEncoderが使われてきた。しかし、対比学習などの識別的な表現学習では、少数の特徴だけで学習目標を達成できる場合、他の重要な情報が表現から失われる可能性がある(Feature Suppression)。 RePRは、特徴間の類似度を対比学習によって間接的に学習する代わりに、一方の観測によって他方の観測をどれだけ予測しやすくなるかを表すPMIを教師信号として、特徴表現を学習する。 PMI(点ごとの相互情報量)は、2つの観測x, yについて、 PMI(x, y) = log [p(y|x) / p(y)] と定義される。これは、xを知ることによってyの予測にどれだけの情報が追加されるかを測っている。 一方、従来の対比学習では、同じ画像から得られた異なるViewの表現を近づけ、異なる画像の表現を引き離す。この対比学習の目標は、理論的には表現の内積によるPMIの近似と関係していることが知られている。 〈R(x), R(y)〉 ≈ PMI(x, y) RePRはこの関係を利用し、生成モデルによってPMIを推定し、その推定値PMI'と表現の内積が一致するように学習する。基本的な損失関数は次のように定義される。 L = (〈R(x), R(y)〉 - PMI'(x, y))² つまり、生成モデルが推定した情報量を、表現空間の幾何構造に直接埋め込むという発想である。これはWord2VecがPMI行列を暗黙的に因子分解しているという結果にもつながっている。 ここで、PMIの推定には生成モデルを利用する。具体的には、Diffusion Modelで条件なしの場合と条件付きの場合のDenoising誤差を比較し、p(y)とp(y|x)の対数確率の差に対応する量を推定する。 しかし、画像全体のPMIを直接推定することは計算量的に困難である。そこで、画像のパッチを一つ追加したときのPMIの変化量を計算し、それらを積み重ねることで画像全体のPMIを表現する。 PMI(x, y_{≤k}) - PMI(x, y_{<k}) そして、PMIの差分を、特徴表現の差分との内積によって学習する。 ||〈R(x), R(y_≤k) - R(y_<k)〉- (PMI(x, y_{≤k}) - PMI(x, y_{<k})||^2 また、通常のPMIは画像の高周波成分も含め、共有される情報を同等に扱う。RePRではDiffusionのノイズレベルごとに重みを変えることで、生成に有用な情報を重視するようにしている。 さらに、PMIを推定するための生成モデルは、最初の段階ではピクセルを予測するが、次の段階では前の段階で学習された特徴表現を予測対象とする。このIterative Feature Predictionを繰り返すことで、より生成に適した潜在表現を獲得できることが示されている。 コメント === 画像などの表現学習についてはさまざまな研究が提案されているが、RePRは表現学習をPMIと結びつけ、理論的にも実践的にも興味深い結果を示している。 直感的には、一方の観測からもう一方の観測を予測しやすいほど、その特徴表現の内積が大きくなり、逆に予測しにくいほど小さくなるように学習している、と言い換えられる。 また、生成モデル向けの表現学習にもスケーリング則が成り立つことが示されたことも興味深い。さらに、後続の生成モデルの学習コストも含めた総計算量を考慮しても、計算資源の一部を表現学習に振り向けた方が、生成モデルだけを学習するよりも高い性能を達成できることが示されている。 なお、特徴表現の学習過程では、PMIを推定するための生成モデルも既に学習されている。そのため、表現の獲得と生成を別々のモデルで行うのではなく、両者を統一的に実現するモデルを構築できるのではと考えられる。
1
9
56
6,931
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の低い区間を優先して学習するようにしている。例えば、宙返りの着地が難しい場合、動作の最初から繰り返すのではなく、着地直前の状態から重点的に学習できる。 また、今回のタスクは、あらかじめ正解となる動作が存在し、位置や姿勢の誤差から各ステップで滑らかな報酬を計算できるという特徴がある。こうした密な報酬も、微分可能な世界モデルを通じた方策最適化が成功した重要な要因だと考えられる。 そのため、正解動作が存在しないタスクや、物体との複雑な接触を伴う操作タスクにおいても同じような学習効率を達成できるかは、今後の重要な研究課題となる。
3
29
3
206
20,497
FP4による大規模言語モデルの事前学習について、既存手法を体系的に比較した論文が公開されている。 “What Matters for NVFP4 Training - A Comparative Study of Low-Precision Pretraining Recipes” FP4を使った学習は高速な行列演算や必要なメモリ帯域幅を抑えられる一方、BF16と比較すると精度低下や学習の不安定性が残る。 FP4のE2M1形式では、0、0.5、1、1.5、2、3、4、6と、それらの負の値からなる15種類の数値しか表現できない。ただし、NVFP4では16個の値からなるブロックごとにスケール係数を持ち、それを掛けることで表現範囲を調整する。 今回の論文では、8BのDenseモデルと30B-A3BのMoEモデルを最大1Tトークン学習し、既存のFP4学習手法を比較している。 特に提案するMosaic NVFP4では、以下の工夫を組み合わせている。 ・Forwardでは重みを1次元で量子化し、Backwardでは必要な方向に再量子化する。 ・Backwardの2種類の行列積にRandom Hadamard Transformを適用し、外れ値の影響を軽減する。 ・勾配にはMS-EDEN(後述)を使い、Stochastic Roundingより小さな量子化誤差で不偏性を実現する。 ・その他のテンソルにはFour Over Sixによる適応的なスケール選択を適用し、量子化誤差を削減する。 ・最後の15%の層はBF16のまま維持する。 これらを組み合わせることで、30B-A3Bモデルでは1Tトークン学習後のBF16に対する相対学習損失差を0.81%未満まで縮小し、従来のFP4手法から0.36パーセンテージポイント改善した。 特に厄介なのは、すべての層をFP4にすると、8Bモデルでも300Bトークンを超えたところで学習が発散したことである。また、手法間の優劣は150Bトークンを超えるまで安定せず、短期間の小規模実験だけではFP4学習手法を適切に評価できないことも示された。 ・MS-EDEN 勾配の量子化において、バイアスをゼロに保ちながら量子化誤差の分散を小さくする手法。 個々の値を確率的に丸めるStochastic Roundingとは異なり、まずベクトル全体にRandom Hadamard Transformを適用して値を混合し、FP4に量子化する。その後、補正係数を掛けることで不偏性を回復する。 この補正係数はNVFP4のブロックスケールに組み込まれる。スケールはFP8で格納されるため、補正後のスケールを量子化する際にはStochastic Roundingを使用する。 これにより、各要素を確率的に丸める場合よりも小さな量子化誤差で、不偏な勾配推定を実現できる。 ・Four Over Six FP4のE2M1形式では、表現可能な値の4と6の間に5が存在しないため、この領域で量子化誤差が大きくなりやすい。 通常のスケーリングでは、ブロック内の絶対値の最大値をFP4の6に対応させる。しかし、最大値を4に設定した方が、ブロック全体の量子化誤差が小さくなる場合がある。 Four Over Sixは、ブロックごとに最大値を4に対応させる場合と6に対応させる場合の両方を試し、量子化後の二乗誤差が小さい方を選択する。 コメント ==== FP4は行列積の演算性能を向上させ、必要なメモリ帯域幅も削減できるため有望だが、大規模な事前学習での実績はまだ限られている。 今回の研究でも、かなり学習が進んだ後に発散する現象が確認された。最後の15%の層をBF16に維持することで発散を抑えられ、この条件ではBackwardの量子化手法による性能差もかなり小さくなる。 これは、Backwardの量子化バイアスを抑えるだけでは解決できない問題が残っていることを示している。著者らは量子化による表現精度の問題を示唆しているが、発散の詳しい機構はまだ分かっていない。 学習が進むにつれて小さな数値誤差が蓄積し、学習ダイナミクスが変化する可能性も考えられる。Forward側の表現誤差や、量子化とパラメータ更新の相互作用も重要だと思われる。 これはFP4に限らず、誤差逆伝播法を近似したり、一部を別の計算方式に置き換えたりする手法にも共通する課題である。小規模な実験では問題なく動作していても、モデルを大きくしたり学習を長期間続けたりすると、発散や性能劣化が生じる場合が多い。 深層学習はかなり粗い数値計算でも動作する印象があるが、特に長期間の学習における勾配計算や数値誤差の扱いは、想像以上に繊細である。
1
14
2
84
9,897
生成AI向けAIプロセッサーへの光電融合技術の実装を目指し、Tenstorrentと共同で研究開発を進めていきます。 本共同研究開発では、学習・推論の双方を高性能かつ高効率に実行できるAIプロセッサーシステムの実現を目指します。 具体的には、光電融合技術と最先端のパッケージング技術を活用するとともに、これまで培ってきた超高メモリ帯域幅を実現する3次元メモリ積層技術を、最先端の半導体プロセスと組み合わせることでさらに進化させます。加えて、機密計算に対応したCPUチップレットとの連携を進め、システム全体の性能、電力効率、安全性の向上を目指します。
【発表】PFNとTenstorrentの共同研究が、NEDOプロジェクトのテーマ「(f8)光電融合技術を実装したxPU設計開発」に採択されました。これを受け、両社で共同研究開発を開始、学習・推論ともに高性能・高効率に実行可能な生成AI向けAIプロセッサーのシステム化を目指します。preferred.jp/ja/news/pr20261…
1
19
93
12,441
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キャッシュ管理や負荷分散システムなどは公開されていないようである)
3
40
4
188
18,695
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で速度場を条件付き期待値の回帰によって学習するのと同様に、尤度推定に必要な発散も条件付き期待値の回帰に帰着させたという点で、非常に興味深い。
1
12
1
126
19,398
Scallop, Few-Step Boltzmann Generators via Scalable Likelihood Flow Maps tonyauyeung.github.io/scallo…
3
14
3,068