Categorical Flow Maps

我们提出了“分类流映射”（Categorical Flow Maps），这是一种基于自蒸馏的流匹配方法，旨在加速分类数据的少步生成。该方法建立在近期关于流匹配的变分建模框架之上，并顺应了扩散模型与流模型领域整体向高效推理演进的趋势。我们定义了一种朝向概率单纯形（simplex）的流映射，使其能够将概率质量持续输运至一个预测的目标分布，从而导出一种天然约束模型输出的参数化形式。由于我们的轨迹是连续而非离散的，分类流映射既可采用现有的蒸馏技术进行训练，也可借助一种基于终点一致性的新型优化目标进行训练。这种连续化的建模方式还自然地支持测试阶段的推理：我们可直接复用现有针对分类数据的引导（guidance）与重加权（reweighting）技术，在采样过程中灵活调控生成结果，以更好地满足下游任务目标。实验表明，我们在图像、分子图和文本三类任务上均取得了当前最优的少步生成性能；即便仅使用单步生成，仍展现出强劲的生成效果。

如何在分类数据（如图像、分子图、文本）上实现加速的少步（甚至单步）生成，同时保持高质量和可控性——这是扩散模型和流模型在离散/分类空间中推理速度慢、采样步数多的核心挑战。

提出Categorical Flow Maps：一种基于连续流匹配的变分框架，将离散类别分布嵌入单纯形（simplex），定义从均匀先验到目标类别的连续概率流；通过自蒸馏训练，联合优化轨迹建模与端点一致性目标，使模型隐式学习约束性预测，并天然兼容连续空间中的指导（guidance）与重加权技术。

在图像（CIFAR-10/100）、分子图（ZINC）、文本（WikiText-2）上均达到少步（1–4步）生成SOTA；首次在分类生成中实现单步高质量采样；无需离散化近似或Gumbel-Softmax等松弛技巧；开源代码；端点一致性损失和单纯形流参数化是新设计；测试时可直接复用Classifier-free Guidance等连续域技术。

Flow Matching for Generative Modeling (Lipman et al., 2023); Rectified Flow (Liu et al., 2023); Discrete Diffusion Models (Hoogeboom et al., 2021; Austin et al., 2022); Categorical Score-Based Generative Modeling (Campbell et al., 2023); Self-Distillation in Diffusion (Salimans & Ho, 2022)