拡散言語モデルにおけるトークン編集のための自己生成エラートレーニング

拡散言語モデルは、近年注目を集めているテキスト生成手法であり、自己回帰モデルとは異なり、並列デコードが可能で、デコード中に編集モジュールを用いて生成トークンを修正できるという特徴を持っています。LLaDA2.1はその代表的なモデルの一つで、特にブロック拡散デコード中にコミットされたトークンを修正するためのトークン間（Token-to-token, T2T）編集モジュールを備えています。しかし、本研究が指摘するように、このT2Tエディタのトレーニングと推論の間には深刻なミスマッチが存在します。トレーニング段階では、エディタはランダムに語彙を破損させたトークン列を入力として受け取り、それらを修正するように学習します。ところが、実際の推論時には、エディタはモデル自身が生成した流暢で高確信度のドラフトエラーに直面します。このようなトレーニングと推論の分布の違いが、エディタの性能を制限し、特に正確な数値の出力や短い事実回答を必要とするタスクにおいて、最終桁の転写誤りや過剰な自己修正といった問題を引き起こします。

この課題に対処するため、Lin YaoらがarXivに提出した論文では、自己生成エラートレーニング（Self-Generated Error Training）という新しい手法が提案されました。この手法の核心は、エディタを推論時と同様のエラー分布で訓練することです。具体的な手順は以下の通りです。まず、モデルは勾配を計算せずに通常の順伝搬でドラフト出力を生成します（ノーグラディエントドラフトパス）。次に、その出力の一部のトークンを、モデル自身の予測トークンで置き換えることで、推論時に見られるような自己生成エラーを作り出します。そして、2回目の順伝搬で、モデルはこれらの自己生成エラーから正しい出力に回復するように教師付き学習されます。このプロセスは勾配計算を必要とせず、LLaDA2.1-miniモデルに対する短いLoRA（Low-Rank Adaptation）継続事前学習によって実装されます。LoRAはモデルのパラメータの一部のみを更新する効率的なファインチューニング手法であり、計算コストを抑えつつ新しいタスクに適応できます。

実験では、複数のベンチマークタスクにおいて、公式のQ-Mode T2T手順を採用し、推論パラメータは一切変更せずに評価が行われました。その結果、自己生成T2T手法はベースラインと比較して総合的な精度を向上させ、T2T編集強度（モデルが生成済みトークンを修正する頻度）を低減できることが示されました。特に注目すべきは、二つの典型的な失敗モードを効果的に緩和した点です。第一に、モデルが正しい推論を行ったにもかかわらず、最終的な数字の出力に誤りが生じる「最終桁転写誤り」が減少しました。第二に、短い事実回答の前にモデルが過剰に自己修正を行い、正しい答えを誤ったものに書き換えてしまう問題が軽減されました。これらの改善は、トレーニングデータの分布を推論時の実際の分布に近づけることの重要性を明確に示しています。

本研究は、拡散言語モデルにおける編集モジュールの信頼性を大幅に向上させる可能性を秘めています。将来的には、より大規模なモデルや多様なタスクへの適用が期待されます。研究者らは、再現性を高めるためにコードとモデルを公開する予定であると述べています。これにより、コミュニティ全体でのさらなる研究発展が促進されるでしょう。