언어 모델링을 위한 토큰 시간 연속 확산

초록

본 논문에서는 토큰 시간 연속 확산(TTCD)이라는 새로운 확산 언어 모델을 소개한다. 이 모델은 (a) 연속 공간에서 작동하며, 가우시안 노이즈를 추가 샘플링 없이 결정론적으로 최종 토큰 캔버스에 매핑하고, 중요한 점으로 (b) 토큰별 시간이라는 새로운 개념을 도입하여 일부 토큰이 다른 토큰보다 더 빠른 속도로 노이즈에서 토큰으로 진행되도록 한다. 연속 공간 모델링은 TTCD가 여러 토큰의 병렬 샘플링을 피하는 데 도움을 주며, 이는 이산 공간에서만 반복하는 모델의 고속화 시 부정확성의 주요 원인이다. 토큰별 시간 개념은 TTCD가 조건부 생성을 더 잘 모델링하고, 더 확실한 토큰이 더 빠른 속도로 진행되도록 하며, 정제 과정에서 토큰 간의 차별화된 상호 영향을 가능하게 한다. TTCD는 고속화 시 이산 모델보다 우수한 성능을 보인다. 우리는 OpenWebText에서 1억 6천만 개의 파라미터를 가진 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.