Непрерывная диффузия по времени токенов для языкового моделирования

Аннотация

В данной статье мы представляем модель непрерывной диффузии с временными метками токенов (token time continuous diffusion, TTCD) — новую диффузионную языковую модель, которая (а) работает в непрерывном пространстве, детерминированно отображая гауссовский шум в конечное полотно токенов без дополнительной выборки, и, что важно, (б) включает новую концепцию индивидуальных времен для каждого токена, при которой некоторые токены переходят от шума к токену с более высокой скоростью, чем другие. Моделирование в непрерывном пространстве помогает TTCD избежать параллельной выборки нескольких токенов, что является ключевым источником неточности при высоком ускорении для моделей, работающих исключительно в дискретном пространстве. Концепция индивидуальных времен токенов позволяет TTCD лучше моделировать условную генерацию, обеспечивает более высокую скорость для более уверенных токенов и позволяет дифференцировать взаимные влияния токенов в процессе уточнения. TTCD превосходит дискретные модели при высоком ускорении. Мы обучаем модель TTCD с 160 млн параметров на наборе данных OpenWebText, а затем применяем к ней самодистилляцию; обнаруживается, что при высоком ускорении мы сопоставимы по качеству безусловной генерации и превосходим по качеству условной генерации несколько существующих моделей аналогичного размера, обученных на тех же данных и подвергнутых самодистилляции. Аналогичные улучшения достигаются и в решении судоку.

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.