トークン予測の見直し：木構造拡散言語モデル

概要: 離散拡散言語モデルは、自己回帰型言語モデルに対する競争力のある代替として登場してきましたが、限られたパラメータおよびメモリ予算のもとでそれらを効率的に学習することは依然として難しい課題です。現代のアーキテクチャは主に、語彙全体を対象としたトークン予測レイヤーに基づいており、これがモデルのパラメータの大きな割合を占めます（例：小規模なDiTスタイルの設計では20%以上）。さらに、このレイヤーがピーク時のGPUメモリ使用量を支配することが多いのが実情です。これは、制約のある学習リソースの下で、パラメータとメモリの双方を非効率に使うことにつながります。
この問題に対処するために、私たちは明示的な語彙全体予測が本当に必要かどうかを改めて見直し、その代わりにトークン間に内在する構造を活用して、木構造化された拡散言語モデルを構築します。具体的には、あらかじめ構築された語彙ツリーにおけるトークンの祖先ノードに対応する中間潜在状態を用いて拡散過程をモデル化します。この木構造による因数分解は、分類の次元を指数関数的に低減し、予測ヘッドのサイズを実質的に無視できるほど小さくし、パラメータを注意ブロックをより深くするために再配分できるようにします。実験的に、同じパラメータ予算のもとで、提案手法はピーク時のGPUメモリ使用量を半減しつつ、最先端の離散拡散言語モデルと同等のパープレキシティ性能を達成します。