Difusión Continua de Tiempo de Token para el Modelado del Lenguaje

Resumen

En este artículo introducimos la difusión continua de tiempo por token (TTCD, por sus siglas en inglés), un nuevo modelo de lenguaje de difusión que (a) opera en espacio continuo, mapeando determinísticamente ruido gaussiano a un canvas de tokens final sin necesidad de muestreo adicional, y crucialmente (b) incorpora una nueva noción de tiempos por token, donde algunos tokens progresan desde el ruido hasta el token a una velocidad más rápida que otros. El modelado en espacio continuo ayuda a TTCD a evitar el muestreo paralelo de múltiples tokens, una fuente clave de inexactitud a altas aceleraciones en modelos que iteran exclusivamente en espacio discreto. La noción de tiempos por token ayuda a TTCD a modelar mejor la generación condicional, permite que tokens más seguros avancen a una velocidad más rápida, y permite influencias diferenciadas entre tokens durante el refinamiento. TTCD supera a los modelos discretos a altas aceleraciones. Entrenamos un modelo TTCD de 160 millones de parámetros en OpenWebText y luego lo auto-destilamos; encontramos que a altas aceleraciones somos comparables en calidad de generación incondicional y superamos en generación condicional a varios modelos existentes de tamaño similar entrenados con los mismos datos y auto-destilados. También logramos ganancias similares en la resolución 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.