Diffusion Continue Temporelle de Tokens pour la Modélisation du Langage

Résumé

Dans cet article, nous introduisons la diffusion continue par temps de jeton (TTCD), un nouveau modèle de langage par diffusion qui (a) opère dans un espace continu, en mappant de manière déterministe un bruit gaussien à un canevas de jetons final sans aucun échantillonnage supplémentaire, et surtout (b) intègre une nouvelle notion de temps par jeton, certains jetons progressant du bruit au jeton à un rythme plus rapide que d'autres. La modélisation en espace continu aide TTCD à éviter l'échantillonnage parallèle de multiples jetons, une source majeure d'imprécision à des accélérations élevées pour les modèles qui itèrent purement dans l'espace discret. La notion de temps par jeton aide TTCD à mieux modéliser la génération conditionnelle, permet aux jetons plus certains de progresser à un rythme plus rapide, et autorise des influences inter-jetons différenciées lors du raffinement. TTCD surpasse les modèles discrets à des accélérations élevées. Nous entraînons un modèle TTCD de 160 millions de paramètres sur OpenWebText, puis l'auto-distillons ; nous constatons qu'à des accélérations élevées, notre qualité de génération inconditionnelle est comparable, et notre génération conditionnelle est supérieure à celle de plusieurs modèles existants de taille similaire, entraînés sur les mêmes données et auto-distillés. Nous obtenons également des gains similaires dans la résolution de Sudoku.

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.