言語モデリングのためのトークン時間連続拡散

要旨

本論文では、トークン時間連続拡散（TTCD）を導入する。これは新しい拡散言語モデルであり、(a) 連続空間で動作し、ガウスノイズをさらなるサンプリングなしで最終的なトークンキャンバスに決定論的にマッピングする。そして重要なことに、(b) トークンごとの時間という新しい概念を組み込んでおり、一部のトークンは他のトークンよりも速い速度でノイズからトークンへと進行する。連続空間モデリングにより、TTCDは複数トークンの並列サンプリングを回避できる。これは、離散空間のみで反復するモデルにおいて高速化時に不正確さの主な原因となる。トークンごとの時間の概念は、TTCDが条件付き生成をより適切にモデル化する助けとなり、より確実なトークンがより速い速度で進行することを可能にし、精緻化の際にトークン間の影響に差別化をもたらす。TTCDは高速化時に離散モデルを上回る性能を示す。我々はOpenWebText上で1億6000万パラメータのTTCDモデルを訓練し、その後自己蒸留を行った。その結果、高速化時において、非条件付き生成品質で同等であり、条件付き生成では、同じデータで訓練・自己蒸留された同程度の規模の既存のいくつかのモデルを上回る性能を達成した。数独解決においても同様の利得を得ている。

English

In this paper we introduce token time continuous diffusion (TTCD), a new diffusion language model which (a) operates in continuous space, deterministically mapping Gaussian noise to a final token canvas with no further sampling, and crucially (b) incorporates a new notion of per-token times, with some tokens proceeding from noise to token at a faster rate than others. Continuous space modeling helps TTCD avoid the parallel sampling of multiple tokens, which is a key source of inaccuracy at high speedups for models that iterate purely in discrete space. The notion of per-token times helps TTCD to better model conditional generation, allows for more sure tokens to proceed at a faster rate, and allows for differentiated inter-token influences during refinement. TTCD outperforms discrete models at high speedups. We train a 160M parameter TTCD model on OpenWebText, and then self-distill it; we find that at high speedups we are comparable in unconditional generation quality, and outperform in conditional generation, several existing models of similar size trained, on the same data, and self-distilled. We achieve similar gains in Sudoku solving as well.