拡散言語モデルのためのアンマスキングポリシーの学習

要旨

拡散型（大規模）言語モデル（dLLM）は、多くのタスクにおいて自己回帰型モデルと同等の下流性能を達成しつつ、推論時の効率性向上が期待される。特に成功しているバリアントとして、特殊なマスクトークンで満たされたバッファをモデルの語彙からサンプリングされたトークンに段階的に置換する「マスク離散拡散」がある。複数のトークンを並列にアンマスクすることで効率化が図れるが、一度に過剰なトークンを処理すると生成品質の劣化を招くリスクがある。このため、dLLMの重要な設計要素は、拡散過程の各ステップでどのトークンを置換するかを決定するサンプリング手順である。実際、最近の研究では、信頼度閾値処理などのヒューリスティック戦略がランダムなアンマスクに比べて品質とトークン処理量の両方で優れることが明らかになっている。しかし、こうしたヒューリスティック手法には欠点もある：手動調整が必要であり、我々の観察ではバッファサイズが大きくなると性能が低下する。本研究では、強化学習を用いてサンプリング手順を訓練する手法を提案する。具体的には、マスク拡散サンプリングをdLLMを環境とするマルコフ決定過程として定式化し、dLLMのトークン信頼度をアンマスク判断に変換する単層Transformerベースの軽量ポリシー構造を設計する。実験により、訓練されたポリシーは半自己回帰生成と組み合わせた場合に最先端のヒューリスティック手法と同等の性能を発揮し、完全拡散設定ではそれを上回ることを示す。また、これらのポリシーが新たな基盤dLLMや長い系列長へ転移可能であることを確認する。一方、ドメイン外データへの適用時には性能劣化が生じ、精度と効率性のトレードオフの微調整が課題となることも明らかになった。

English

Diffusion (Large) Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks, while holding the promise of being more efficient during inference. One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary. Efficiency can be gained by unmasking several tokens in parallel, but doing too many at once risks degrading the generation quality. Thus, one critical design aspect of dLLMs is the sampling procedure that selects, at each step of the diffusion process, which tokens to replace. Indeed, recent work has found that heuristic strategies such as confidence thresholding lead to both higher quality and token throughput compared to random unmasking. However, such heuristics have downsides: they require manual tuning, and we observe that their performance degrades with larger buffer sizes. In this work, we instead propose to train sampling procedures using reinforcement learning. Specifically, we formalize masked diffusion sampling as a Markov decision process in which the dLLM serves as the environment, and propose a lightweight policy architecture based on a single-layer transformer that maps dLLM token confidences to unmasking decisions. Our experiments show that these trained policies match the performance of state-of-the-art heuristics when combined with semi-autoregressive generation, while outperforming them in the full diffusion setting. We also examine the transferability of these policies, finding that they can generalize to new underlying dLLMs and longer sequence lengths. However, we also observe that their performance degrades when applied to out-of-domain data, and that fine-grained tuning of the accuracy-efficiency trade-off can be challenging with our approach.

拡散言語モデルのためのアンマスキングポリシーの学習

Learning Unmasking Policies for Diffusion Language Models

要旨

Support