连续扩散模型与流匹配模型有望成为语言建模（LM）中自回归方法之外的一种强大替代方案，因为它们解锁了一系列目前仅为连续模态所保留的优势，包括加速采样与倾斜。近期，多项工作已经证明，可以通过一个简单的高斯分布与 one-hot 编码数据分布之间的流匹配过程来连续生成离散数据。
他们进一步展示了通过 Categorical Flow Maps（CFMs）实现加速采样的可行性，从而在少步采样机制下获得了具有竞争力的样本质量。然而，该方法此前仅在相对较小的规模（< 1B）上进行了评估，其可扩展性问题完全悬而未决。
在本文中，我们在 2.1T tokens 上训练了一个 1.7B 参数的基础流模型，并将其自蒸馏为一个 CFM，该 CFM 能够在少至 4 个推理步骤内生成多样化的高质量文本，同时保持接近数据级别的 token 熵。此外，我们引入了半离散设置下 CFM 的似然界，并表明它们可用于在标准 LM 基准上对模型进行评分，取得与离散扩散方法同一区间的结果。
最后，我们揭示了在大规模训练这些模型时出现的一些挑战，并就损失加权与时间调度提供了规范性见解。
- † 牛津大学
- ** 在 Apple 期间完成的工作
Continuous diffusion and flow matching models could represent a powerful alternative to autoregressive approaches for language modelling (LM), as they unlock a host of advantages currently reserved for continuous modalities, including accelerated sampling and tilting. Recently, several works have demonstrated the possibility of generating discrete data continuously by a simple flow matching process between a Gaussian and the one-hot encoded data distribution. They have further shown the feasibility of accelerated sampling via Categorical Flow Maps (CFMs), resulting in competitive sample quality in the few-step regime.
However, this method had only been evaluated at relatively modest scales (< 1B), leaving the question of its scalability completely open. In this article, we train a 1.7B-parameter base flow model on 2.1T tokens and self-distill it into a CFM that generates diverse, high-quality text in as few as 4 inference steps while maintaining near-data-level token entropy. Furthermore, we introduce a likelihood bound for CFMs in the semi-discrete setting, and show that they can be used to score the model on standard LM benchmarks, achieving results in the same range as discrete diffusion methods.
Finally, we uncover some of the challenges that arise from training these models at scale, and we provide prescriptive insights on loss weighting and time scheduling.
- † University of Oxford
- ** Work done while at Apple