ALoDLM、トークンごとに計算量を変える拡散言語モデル

拡散言語モデル（DLM）は、文章を左から右へ1トークンずつ生成する自己回帰モデルとは異なり、複数の位置を並列に予測しながら文章を段階的に確定します。生成を高速化できる可能性がある一方、同程度の規模の自己回帰モデルと比べると、品質面で差が残ることが実用化の課題でした。

論文「ALoDLM: Adaptively Looped Diffusion Language Models」は、その理由の一つを計算量とトークン難度の不一致に求めます。部分的に観測された配列では、容易に予測できる未知トークンもあれば、より多くの計算を必要とするトークンもあります。しかし従来のDLMは、各デノイズ段階で未知の位置にほぼ一律の計算深度を割り当てます。その結果、簡単な位置に余分な計算を使う一方、難しい位置を十分に精緻化できない可能性があります。

トークン適応型の潜在再帰

ALoDLMの中心的なアイデアは、計算を位置ごとに適応させる潜在再帰です。各デノイズ段階でモデルは潜在表現を繰り返し更新し、十分に確定した位置は離散トークンとして確定します。確定済みのトークンは、後続の予測で利用できる明示的な文脈としてモデルへ戻されます。一方、まだ不確かな位置は潜在状態を保持したまま、追加の再帰パスで更新されます。

つまり、文章全体が同じ回数だけ計算される必要はありません。複数位置を並列に処理する利点を保ちつつ、計算を必要としているトークンに集中させる設計です。

予測と計算配分を同時に学習

この方式では、モデルはトークンを予測するだけでなく、各位置をいつ確定させるかも学習しなければなりません。そこで著者らは、トークン単位の計算スケジュールを潜在変数として扱い、予測と計算配分を同時に学習する条件付き負の証拠下限（NELBO）を導出しました。計算深度を固定の推論設定ではなく、学習される振る舞いの一部として扱う点が特徴です。

ALoDLMは1.7Bと8Bのパラメータ規模で訓練されました。論文によれば、11のベンチマークにおいて、両規模の平均スコアが評価対象のDLMおよび対応する自己回帰ベースラインを上回りました。また並列デコード能力を維持しており、最適化された推論エンジンを使った場合、評価対象の自己回帰モデルと拡散モデルの中で、品質と効率のバランスにも強みを示したとされています。

意義と今後の確認点

ALoDLMが示すのは、拡散モデルにおける並列性と差別化された計算が両立し得るという方向性です。すべての位置を同じ経路で処理するのではなく、確信度の低い位置に追加計算を割り当てることで、自己回帰モデルとの差を縮めようとしています。

ただし、提供された素材にはベンチマークごとのスコア、具体的な遅延値、計算スケジュールに関する詳細なアブレーションは含まれていません。タスクや系列長、ハードウェアが変わった場合の効果については、完全な論文と追加実験の確認が必要です。それでも、トークン単位で計算を配分するという考え方は、拡散言語モデルの効率と品質を同時に改善するための明確な研究方向を示しています。

出典：Hugging Face Daily Papers