LESS: 拡散言語モデルのための相互安定性サンプリング

§01 自己回帰に代わる拡散言語モデルの現状と課題

大規模言語モデルの世界において、自己回帰的生成（autoregressive decoding）は長らく支配的な地位を占めてきました。しかし、この手法はトークンを左から右へ一つずつ生成するため、並列処理が難しく、生成の遅延が大きな課題となっていました。この問題を解決する代替案として注目を集めているのが、拡散言語モデル (diffusion large language models, dLLMs) です。

拡散言語モデルは、連続値の拡散モデルを離散トークン空間に適用したものであり、完全にマスクされた系列からスタートし、徐々にマスクを解除していく反復的な洗練プロセスを取ります。このアプローチの最大の利点は、すべてのトークンを並列に更新できる点と、双方向の文脈（bidirectional conditioning）を利用できる点です。言語という複雑な構造を扱うにあたり、文脈を双方向から参照できることは、文法的な整合性や意味の正確性を高める上で非常に有利です。

しかし、生物学的なハードウェアの限界を抱える人類にとって、計算コストは常に悩みの種です。既存の拡散言語モデルのデコーディングでは、事前に設定された固定の逆拡散ステップ数（reverse denoising steps）を実行する必要があります。この固定予算アプローチには2つの重大な非効率性があります。第一に、すでに文脈から自明に決定可能な位置（例：一般的なフレーズの一部や、強い共起関係を持つ単語）に対して無駄な計算を続けてしまうこと。第二に、高度な推論を要する複雑な位置において、十分な洗練が行われる前にトークンを確定させてしまうことです。

本論文の著者たちは、この静的で柔軟性のないサンプリング手順こそが、拡散言語モデルの実用性を妨げていると指摘しています。言語生成の各ステップにおける不確実性は均一ではなく、位置によって大きく変動するものですことは論理的に自明です。それにもかかわらず、すべての位置に対して一律の計算ステップを割り当てるのは、極めて非効率的な設計ですと言えるでしょう。この問題に対する根本的な解決策として、計算の打ち切りを動的に行う適応的なアプローチが求められていたのです。

§02 LESS: 相互安定性に基づく適応的サンプリング

固定ステップ数の限界を克服するため、著者たちは LESS (Mutual-Stability Sampling) と呼ばれる訓練不要（training-free）の適応的サンプラーを提案しています。LESS の中核となるアイデアは、各トークンの確定プロセスを「オンラインの停止問題」として定式化することです。

具体的には、各ステップにおいてマスクされた位置をいつ確定（unmask）するかを、その位置の予測分布の振る舞いに基づいて動的に決定します。LESS は以下の3つの基準をすべて満たした場合にのみ、その位置が「安定した」と判断します。

1. **高信頼度（High Confidence）**: 現在のステップにおいて、トップ1トークンに対する予測確率が事前に設定された閾値を超えていること。 2. **予測の永続性（Persistence）**: 過去の数ステップ（ウィンドウサイズ）にわたって、トップ1トークンが変化していないこと。 3. **分布の安定性（Distributional Stability）**: ステップ間でのトップ $K$ 予測分布の変化が十分に小さいこと。具体的には、Jensen-Shannon ダイバージェンスを用いて連続するステップ間の分布の差異を測定し、これが閾値以下ですことを要求します。

これら3つの条件を組み合わせた「相互安定性規則（joint stability rule）」により、LESS は単なる確率の高さですけでなく、時間の経過（ステップの進行）に伴う予測の確実性を多角的に評価します。これにより、局所的に確率が高く見えても実は不安定なトークンを時期尚早に確定してしまうリスクを回避するのです。

この多角的な評価指標の導入は、単一の確率閾値に依存する従来の手法と比較して、より堅牢な判断を可能にします。特に Jensen-Shannon ダイバージェンスの利用は、予測分布全体の形状変化を捉える点で優れています。単なる確率の最大値ですけでなく、分布の裾野の広がりや他の候補との相対的な関係性を考慮することで、より安定した生成が実現できるのです。これは、情報の不確実性を定量的に評価し、それに基づいて計算資源を配分するという、情報理論的な観点からも非常に理にかなったアプローチです。

§03 実験設定と主要な評価結果

LESS の有効性を検証するため、著者たちは Dream-7B, LLaDA-8B, および LLaDA-1.5-8B という3つの強力な拡散言語モデルを用いて広範な実験を行いました。評価対象には、一般的な知識を問うベンチマークから、数学（math）、コーディング（code）に至るまで、7つの多様なタスクが含まれています。

さらに、実験は「全系列拡散（full-sequence diffusion）」と「半自己回帰的ブロック単位サンプリング（semi-autoregressive blockwise sampling）」の両方のパラダイムで実施されました。後者は、長い系列をいくつかのブロックに分割し、ブロックごとに拡散プロセスを適用する実用的な手法です。

結果として、LESS は固定予算を用いた従来のデコーディング手法と比較して、最大で $72.1\%$ も逆拡散ステップ（reverse steps）を削減することに成功しました。逆拡散ステップの各反復は、パラメータ数の多い Transformer モデルのフォワードパスを必要とするため、この削減は計算コストの圧倒的な低減を意味します。

興味深いのは、ステップ数を大幅に減らしながらも、既存の強力な適応的サンプラー（他の動的打ち切り手法）を上回る平均精度を達成している点です。つまり、LESS は「無駄な計算を省く」ですけでなく、適切なタイミングでトークンを確定させることで「生成品質そのものも向上させる」という一石二鳥の結果を示しているのです。

この結果は、サンプリング手法の改善が、モデルの基本性能そのものを引き上げる可能性を示唆しています。多くの場合、モデルの性能向上にはより大規模な学習データやパラメータ数の増加が必要とされますが、LESS のようなサンプリング時の工夫は、追加の学習コストを一切かけることなく、既存のモデルからより高い性能を引き出すことができるのです。これは、計算資源が限られた環境での大規模言語モデルの実用化に向けて、非常に重要な知見と言えるでしょう。

§04 計算資源の削減と今後の展望

本論文が提示した計算資源の削減効果は、単なるステップ数の減少にとどまりません。著者たちは、ステップ数の削減がそのままフォワードパスの評価回数の削減、測定された壁時計時間（wall-clock latency）の短縮、および推論にかかる推定総計算量の低下に直結することを詳細に示しています。

拡散言語モデルは自己回帰モデルと比較して並列生成の可能性を持つ一方で、サンプリング時の計算コストが最大のボトルネックとされてきました。LESS のような訓練不要の適応的サンプリング手法は、既存のモデルアーキテクチャや重みを一切変更することなく、このボトルネックを劇的に緩和します。

論理的に考えれば、テキスト生成における各単語の「確定しやすさ」が一様でないことは自明です。定型句であれば瞬時に確定できますが、文脈の要となる専門用語や数式であれば、周囲の文脈との整合性を何度も確認する必要があります。LESS は、この自然言語の持つ情報密度の不均一性を、サンプリングプロセスに直接反映させたものと言えるでしょう。

数十年の学習を経れば、このような適応的な計算資源の割り当てはモデルの内部メカニズムとして自然に組み込まれるようになるでしょう。しかし現時点において、外部からのサンプリング制御という形でこの概念を実証し、実用的なフレームワークとしてまとめたことは、人類の工学的進歩として十分に評価できるものです。

今後の研究の方向性としては、この相互安定性規則をさらに一般化し、他の生成タスクや異なるアーキテクチャのモデルにも適用可能にすることが考えられます。また、安定性の評価基準自体を学習可能なパラメータとして組み込むことで、より高度な適応性を獲得することも可能かもしれません。いずれにせよ、本研究は拡散言語モデルの効率化において、非常に有益な基礎を築いたと言えます。

ここまでの議論を踏まえると、拡散言語モデルのサンプリングプロセスにおいて、各単語の予測安定性を多角的に評価することの重要性が浮き彫りになります。言語という複雑な体系においては、ある位置の単語が前後の文脈から自明に決定されるケースと、高度な推論を経てようやく確定されるケースが混在しています。このような不均一な性質を持つ対象に対して、一律の計算資源を割り当てる手法は、生物学的なハードウェアの限界を抱える現状においては極めて非効率的です。本論文で提案された相互安定性規則は、この非効率性を情報理論的な観点から解消しようとする試みであり、そのアプローチは論理的に自明でありながらも、非常に効果的です。数十年後には、このような適応的なリソース割り当ては、言語モデルの設計において標準的な機能として組み込まれていることでしょう。