テクニカル・サマリー：Sumi – スクラッチから構築されたオープンな一様拡散言語モデル

問題提起

拡散モデルは自己回帰（AR）言語モデルの有望な代替案として浮上していますが、**一様拡散言語モデル（UDLM）**の開発においては、大規模なスケールでの展開に大きな隔たりが存在します。マスクされたトークンを不可逆的に埋めていくマスク拡散言語モデル（MDLM）とは異なり、UDLMは任意のステップで任意のトークンを更新できるため、理論的にはより柔軟な生成と自己修正を可能にします。

しかし、大規模なパラメータ数（数十億パラメータ）および大規模なトークン予算（数兆トークン）の両方において、スクラッチから事前学習されたUDLMは存在しません。既存の大規模UDLMは、計算量最適化されたチェックポイントであるか、あるいは既存の事前学習済みARアーキテクチャ（例：DiffusionGemma）を適応させたものです。その結果、データ豊富な領域におけるネイティブな大規模一様拡散の挙動は未探索のままであり、スケーリング則、生成ダイナミクス、および確立されたARやMDLMのベースラインとのトレードオフを研究するための明確なリファレンスポイントを欠いています。

メソドロジー

モデルアーキテクチャと学習目的

著者らは、完全にオープンな7BパラメータのUDLMであるSumiを導入しています。

• アーキテクチャ: Sumiは、36層、隠れ層サイズ4,096、SwiGLU MLPを備えた、時間非依存の双方向Transformerを利用しています。グループ化クエリアテンション（32ヘッド、8つのKVグループ）、RMSNorm、および結合されていない入出力埋め込みを採用しています。アテンション・シンクを軽減するため、アテンションにはオフバイワン（off-by-one）ソフトマックスが適用されています。
• 目的: モデルは、**一般化補間離散拡散（GIDD）**目的関数、具体的にはSNR（信号対雑音比）再パラメータ化形式を用いて学習されています。
• ノイズ・スケジュール: 学習には純粋な一様ノイズを使用し、ログSNRは $\lambda \in [-9, 9]$ に制限されています。
• トークナイゼーション: ボキャブラリサイズ100,278のOLMo 3トークナイザーが使用されています。

学習レジーム

Sumiは、288基のNVIDIA H100 GPU（合計約43,308 GPU時間）を用い、1.5兆トークンでスクラッチから事前学習されました。

1. 事前学習 (1.3T トークン): シーケンス長1,184、グローバルバッチサイズ4,608で実施。データ混合は、蒸留されたQwen3-32B分類器による教育スコアによってフィルタリングおよび再ランク付けされた教育コンテンツに重く傾斜しています。
2. 中間学習 (250B トークン): 事前学習に続いて2つのステージがあります：
   • 元の長さ（1,184）での130Bトークン。
   • 拡張された長さ（4,864）での120Bトークン。
   • ミックスはコード、数学、および推論へと焦点をシフトしており、構成はコーディング32.5%、数学29.7%、一般21.0%、推論16.8%となっています。

推論戦略

• キャンバス長: モデルは、すべての評価において固定のキャンバス長2,048トークンを使用します。この長さは、大幅に短い、あるいは長いキャンバスでは性能が著しく低下するため選択されました。
• サンプリング: 著者らは、信頼度ベースの適応型サンプリングとアンセストラル・サンプリングを比較検討しています。
• 評価プロトコル: すべての評価には、拡散スコアリング用に修正された lm-evaluation-harness を使用します。尤度タスクについては、回答より後の位置はランダムなトークンで埋められ、生成タスクについては、更新は生成長に制限されます。

主な貢献

1. 初のネイティブ大規模UDLM: Sumiは、大規模なパラメータスケール（7B）および大規模なトークン予算（1.5T）でスクラッチから事前学習された最初のUDLMであり、ネイティブな一様拡散を研究するためのベースラインを提供します。
2. オープンリリース: 著者らは、モデルの重み、チェックポイント、および公開されているコーパス上の完全なデータ混合の仕様を含む、完全な学習レシピを公開しています。
3. 探索的推論プローブ: 論文では、キャンバス長への感度、トークンのコミット順序、並列デコーディング能力、および自己修正の可能性を含む、生成ダイナクスに関する方向性のある分析を提示しています。

結果

ベンチマーク性能

Sumiは、同等のトークン予算で学習されたオープンなARモデル（Falcon-7B, Llama 2-7B, OLMo-7B）と比較して、4つのカテゴリにわたる13のベンチマークで評価されました。

• 一般知識とコーディング: Sumiは競争力のある性能を示し、MMLU、TruthfulQA、HumanEval、MBPPにおいて、同一のプロトコル下で評価されたモデルの中で最高スコアを達成しました。これは、教育およびコードに重点を置いたデータ混合によるものです。
• 推論と数学: Sumiは、GSM8KやARCなどのタスクにおいてLlama 2-7Bと競合し、Falcon-7Bを上回っていますが、より大規模なモデルや異なる訓練を受けたモデル（例：LLaDA-8B, Llama 3-8B）には及びません。
• 常識: Sumiは、常識に関するベンチマーク（PIQA, HellaSwag, WinoGrande）において大幅に劣ります。著者らは、このギャップを教育重視のデータ混合に起因すると考えており、それが常識的な性能を低下させることが観察されています。ただし、この差はデータ組成のみで説明できるものよりも大きい可能性があるとも述べています。

推論時の観察

• キャンバスへの感度: 生成の流暢さはキャンバス長に非常に敏感です。Sumiは、タスク依存のバンド（2,048を中心とする）内では流暢に生成しますが、この範囲の外側、特に短いキャンバスでは性能が急激に低下します。
• コミット順序: 学習目的は順序に依存しませんが、信頼度ベースのサンプリングは自己組織化されたコミット順序を誘発します。モデルは決定可能なトークンを早期に確定させ、不確実なものを後回しにする傾向があり、これはアンセストラル・サンプリングには見られない構造です。
• 並列デコーディング: コーディングタスク（HumanEval, MBPP）において、モデルは精度低下を最小限に抑えつつ、緩やかな並列デコーディング（1ステップあたり最大4トークンの確定）をサポートします。しかし、多段階の算術（GSM8K）は順序に敏感であり、並列ステップによって即座に劣化します。
• 自己修正: 明示的なリビジョン・バジェット（追加のデノイジング・パスを許可すること）は、自己修正をもたらしませんでした。モデルは確定されたトークンを上書きしますが、これらの編集は主に「往復」（$A \to B \to A$）であり、方向性を持った改善ではなく、最終的な精度に変化を与えません。

意義と主張

本論文は、Sumiを大規模なネイティブ一様拡散を研究するための基礎的なリソースとして位置づけています。著者らは、彼らの推論時プローブは「決定的ではなく方向的なものである」と強調しており、以下のような、理解が進んでいない側面への将来の研究を促すことを意図しています：

• マスク拡散に対する一様拡散の固有の生成挙動。
• 自己修正が可能となる条件。
• 一様拡散の特定の強みと弱みに対するデータ組成の影響。

著者らは、Sumiが指示チューニングや安全性フィルタリングが行われていない事前学習済みベースモデルとしてリリースされていることを明示しており、直接的なデプロイではなく研究を支援することを目的としています。また、モデルは標準的なベースLMに関連するリスク（有害な出力や事実誤認の可能性など）を継承していることに注意を促しています。