技術的サマリー：再学習ではなくアライメント：表現アライメントを介した自己回帰言語モデルの拡散言語モデルへの適応

問題提起

拡散言語モデル（DLM）は、非逐次生成、双方向編集、反復的洗練において、標準的な自己回帰（AR）モデルと比較して明確な利点を提供します。しかし、DLM のスケーリングは計算コストが高価です。理論的には、DLM は AR モデルの単一の左から右への経路と比較して、$L!$ の生成経路を学習する必要があり、しばしばはるかに多くの計算資源を必要とします。

最近のアプローチは、事前学習済みの AR チェックポイントを DLM に変換することでこのコストを軽減しようとしています。これらの方法は AR 重みから DLM を正常に初期化しますが、主に目的関数、アテンションパターン、またはサンプリング規約の変更を伴う継続的なノイズ除去トレーニングに依存しています。決定的な点は、既存のレシピが変換された DLM に対して、元の AR 事前学習中に学習された内部表現幾何学を明示的に保持することを制約していないことです。これにより、核心的な問いが生じます：AR と DLM を、言語表現の再学習を必要とする分離されたパラダイムとして扱う必要があるのでしょうか、それとも AR 事前学習によって学習された意味構造を変換中に明示的に保持できるのでしょうか？

手法：REPR-ALIGN

著者は、DLM トレーニングを表現学習問題ではなくアライメント問題として扱う、ミニマリストな適応手法であるREPR-ALIGNを提案します。核心的な仮説は、言語生成の「難しい部分」は意味的および構文的構造（表現）を学習することであり、これは AR 事前学習ですでに習得済みであるというものです。残りのタスクは、これらの特徴を反復的かつ任意の順序のデコーディング機構に適応させることに過ぎません。

この手法は、同一のアーキテクチャを持つ 2 つのモデルを含みます：

1. ティーチャー：因果的アテンションを持つ凍結された事前学習済みの自己回帰（AR）トランスフォーマー。
2. 生徒：同じ AR 重みから初期化された、双方向アテンションを持つマスク付き拡散言語モデル（DLM）。

トレーニング手順

トレーニング中、生徒モデルは汚染されたシーケンス（トークンがマスクされたもの）を消費し、標準的なマスク付きノイズ除去目的関数を使用してマスクされたトークンを予測しようとします。同時に、凍結されたティーチャーモデルは、因果的アテンションの下でクリーンなシーケンスを処理します。

AR モデルの内部幾何学を保持するために、著者はレイヤーごとの表現アライメント損失を導入します。トランスフォーマースタック内の各レイヤー $\ell$ において、生徒の隠れ状態（$h^{(\ell)}_D$）は、コサイン類似度損失を使用してティーチャーの隠れ状態（$h^{(\ell)}_{AR}$）とアライメントされます。

全体のトレーニング目的関数は以下の通りです：
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
ここで：

• $\mathcal{L}_{diff}$ は標準的なマスク付きノイズ除去の交差エントロピー損失です。
• $\mathcal{L}_{align}$ は、すべてのレイヤーとマスクされた位置にわたる生徒とティーチャーの隠れ状態間の平均コサイン距離です。
• $\lambda$ は重み付けハイパーパラメータ（デフォルト 10）です。

主要な設計選択には以下が含まれます：

• アダプター不要：この手法は、アテンションマスクを因果的から双方向に切り替える以外に、追加のパラメータやアーキテクチャの変更を必要としません。
• 凍結されたティーチャー：AR ティーチャーは凍結されたままであり、安定した表現座標系を提供します。
• シフト規約：この手法は、次のトークンインデックスと位置ごとのノイズ除去を調和させるために、AR から DLM への適応で使用される標準的なシフト規約に準拠しています。

主要な貢献

1. 表現保持の特定：本論文は、現在の AR→DLM 変換レシピにおいて欠けている要素が、AR モデルの内部表現幾何学の明示的な保持であることを特定しました。
2. ミニマリストな変換レシピ：REPR-ALIGN は、レイヤーごとのコサインアライメントを介して DLM 生徒を凍結された AR ティーチャーに固定する、シンプルでアダプター不要な手法を導入します。このアプローチは、モデルスケール全体にわたって変換の品質とサンプル効率を向上させます。
3. データと計算の効率性：著者は、AR→DLM 変換が本質的に大量のデータを必要とするものではないことを実証しました。表現アライメントを用いれば、データ全体のサブセット（0.8B トークン）でのトレーニングが、同じステップ予算の下での全データセット（50B トークン）でのトレーニングを上回る性能を発揮します。さらに、大規模なブロック（埋め込みと MLP）を凍結することは、品質を低下させることなくトレーニングのスループットを最大 2 倍向上させます。
4. スケールアップの検証：著者は、このレシピを使用して 4B パラメータの「oDLM」をトレーニングしました。Dream-7B（AR 初期化に基づいて構築された強力な公開 DLM）と比較して、oDLM はより少ないパラメータとはるかに軽量なデータおよび計算予算にもかかわらず、HumanEval および HumanEval+ の pass@10 スコアでより良い結果を達成しました。

実験結果

この手法は、0.6B、1.7B、および 4B スケールの Qwen3 チェックポイントを使用して、コード生成ベンチマーク（HumanEval、HumanEval+、MBPP、MBPP+）で評価されました。

• 適応速度：REPR-ALIGN は一貫して適応を加速しました。HumanEval pass@10 において、アライメントされたモデルは、変換の初期段階を通じて AR 微調整およびスクラッチトレーニングの両方を上回りました。200k ステップにおいて、0.6B モデルは 6.1 ポイント、1.7B モデルはベースラインに対して 9.4 ポイントの向上を示しました。
• スケーリング：アライメントの恩恵はモデルサイズとともに増大しました。1.7B モデルは 0.6B モデルよりも大きな絶対的な改善を示し、より大きな AR 事前学習幾何学を保持することがより価値があることを示唆しています。
• データ効率性：アライメントを用いた 0.8B トークンのランダム部分集合でのトレーニングは、同じステップ数における全 50B トークンストリームでのトレーニングよりも高い pass@10 スコアをもたらしました。
• 凍結：トークン埋め込みと MLP ブロックを凍結することは、パフォーマンスを維持またはわずかに向上させながらスループットを約 2 倍に増加させ、主要な学習信号が表現の再学習ではなくノイズ除去機構の適応にあるという見解を支持しました。
• アブレーション：
  • 指標：コサインアライメントは L2 距離を上回りました。これは、有用な信号がメトリック（大きさ）ではなく幾何学的（方向）であることを示唆しています。
  • 重み：$\lambda=10$ が固定と適応の間の最適なトレードオフを提供しました。
  • レイヤー：すべてのレイヤーをアライメントすることが pass@1 において最良の結果をもたらしましたが、pass@10 に対しては中間および上位レイヤーが特に重要でした。

意義と主張

本論文は、言語表現が生成順序に関わらず普遍的であると主張しています。DLM トレーニングをアライメント問題として再定義することで、著者は、拡散事前学習の高コストの大部分が、AR 事前学習ですでに提供された表現の不要な再学習に起因すると論じています。

この研究の意義は、以下の点において拡散言語モデルのトレーニングのための「定番の技術」を提供するところにあります：

• 計算の削減：標準的な事前学習コストの断片で競争力のある性能を達成します。
• トレーニングの簡素化：複雑なアダプターや重微調整の必要性を排除し、代わりに単純なアライメント目的関数に依存します。
• アクセスの民主化：低データ環境および制約された計算資源において、高品質な DLM トレーニングを実行可能にします。

著者は、この研究を「再学習」から「アライメント」への転換として位置づけ、離散生成における将来の進歩は、ゼロからトレーニングパラダイムを構築することではなく、より良い機構アライメントから生じる可能性があると示唆しています。