We introduce Categorical Flow Maps, a flow-matching method for accelerated few-step generation of categorical data via self-distillation. Building on recent variational formulations of flow matching and the broader trend towards accelerated inference in diffusion and flow-based models, we define a flow map towards the simplex that transports probability mass toward a predicted endpoint, yielding a parametrisation that naturally constrains model predictions. Since our trajectories are continuous rather than discrete, Categorical Flow Maps can be trained with existing distillation techniques, as well as a new objective based on endpoint consistency. This continuous formulation also automatically unlocks test-time inference: we can directly reuse existing guidance and reweighting techniques in the categorical setting to steer sampling toward downstream objectives. Empirically, we achieve state-of-the-art few-step results on images, molecular graphs, and text, with strong performance even in single-step generation.
翻译:本文提出类别流映射,一种通过自蒸馏实现分类数据加速少步生成的流匹配方法。基于近期流匹配的变分公式化以及扩散模型与流模型加速推理的广泛趋势,我们定义了一种朝向单纯形的流映射,该映射将概率质量输送到预测的终点,从而产生一种自然约束模型预测的参数化。由于我们的轨迹是连续而非离散的,类别流映射可以使用现有的蒸馏技术以及一种基于终点一致性的新目标进行训练。这种连续公式化还自动解锁了测试时推理:我们可以在分类设置中直接复用现有的引导和重加权技术,以引导采样朝向下游目标。实证结果表明,我们在图像、分子图和文本数据上实现了最先进的少步生成效果,即使在单步生成中也表现出强劲性能。