Simplifying Flow Matching Transformations with Low-Rank Mixture Models

技术摘要：利用低秩混合模型简化流匹配变换

问题陈述
归一化流（Normalizing Flows），特别是连续归一化流（CNF），学习从复杂的数据分布到简单的基分布（通常是标准正态分布）之间的可逆映射。虽然流匹配（Flow Matching, FM）通过直接回归 ODE 向量场，实现了 CNF 的可扩展训练，但选择标准正态作为基密度往往会导致与目标数据分布在拓扑结构上的不匹配。这种不匹配迫使学习到的流变换必须异常复杂，以弥合简单潜空间与复杂数据流形之间的差距。因此，这种复杂性导致了训练收敛变慢、推理时的计算成本增加（需要更多的 ODE 求解器步数）以及生成性能下降。

方法论
作者提出使用**概率主成分分析混合模型（MPPCA）**来取代标准正态基密度。MPPCA 是一种低秩高斯混合模型，具有两个关键优势：

1. 解析似然性： 与某些学习到的基分布不同，MPPCA 保留了精确的似然表达式，从而保留了归一化流进行密度估计所需的易处理性。
2. 高效拟合： 通过将协方差矩阵约束为低秩（即在学习到的低维子空间上建模协方差），MPPCA 模型可以使用期望最大化（EM）算法快速且廉价地拟合高维数据。

所提出的工作流包含一个“热启动”策略：

1. 预训练： 使用 EM 算法将 MPPCA 模型拟合到训练数据中，以紧密逼近目标数据分布。
2. 流训练： 随后使用流匹配目标训练连续归一化流，并以预拟合的 MPPCA 作为初始基分布（$p_0$）。
3. 集成： 该方法在标准的方差保持（VP）流匹配和基于最优传输（OT）的流匹配目标上均进行了测试。

核心贡献
本文概述了三个主要贡献：

• 富有表现力的基构造： 作者利用 MPPCA 模型构建了具有表现力的基分布，并证明其拟合过程相对于后续的流训练而言，计算开销极低。
• 与最优传输的协同作用： 他们展示了 MPPCA 基分布可以有效地与基于 OT 的流匹配目标相结合。这种结合能生成更高质量的样本，同时减少推理时所需的积分步数。
• 实证验证： 该方法在具有挑战性的密度估计任务（UCI 表格数据集）和高维生成任务（包括 FashionMNIST、CIFAR-10 和 CelebA 在内的图像数据集）上进行了定量评估，证明了其相对于基准模型的持续增益。

实验结果
作者在多个数据集上将该方法与使用标准正态基（$N(0, I)$）的基准模型进行了对比评估：

• 训练效率： 在表格数据集（Hepmass, Miniboone, Bsds300）上，使用 MPPCA 基的流模型需要显著更少的训练轮数才能收敛。例如，在 Miniboone 数据集上，VP-MPPCA 仅需约 24 个 epoch，而 VP-Normal 则需要约 54 个。尽管存在拟合 MPPCA 的开销，但总训练时间仍有所减少，因为 EM 拟合时间仅占总训练时间的不到 3%。
• 推理复杂度： 自适应步长 ODE 求解器所需的函数计算次数（NFE）在基于 MPPCA 的模型中一致较低。这表明学习到的向量场更简单且复杂度更低，减轻了生成过程中的计算负担。
• 生成质量： 在图像数据集上，OT-MPPCA 模型实现了最低的 Fréchet Inception Distance (FID) 和统计差异分箱数 (NDB) 分数。值得注意的是，通过视觉观察生成的样本（图 3）可以发现，MPPCA 初始化的模型在仅仅第一个训练 epoch 后就能产生有意义且具有结构的样本，而标准正态基线则需要更多轮次才能达到类似的结构。

意义与主张
本文声称，学习一个具有表现力的基分布可以作为一种强大的归纳偏置，从而简化流变换。通过使用接近目标数据的分布来初始化流（最小化基分布与数据之间的 KL 散度），模型可以减轻学习复杂全局变换的负担。相反，流只需要对预设结构的基分布进行精细化调整。

作者强调，由于 MPPCA 模型避免了与全秩高斯混合模型相关的内存和条件数问题，因此该方法在处理高维任务时具有实用性。这项工作表明，使用低秩混合模型进行流训练的“热启动”是一种有效的策略，可以在不牺牲生成质量的前提下，提高训练效率和推理速度。未来的工作建议探索将 MPPCA 拟合扩展到更高维度，并将此方法与其他流匹配目标（如动作匹配和薛定谔桥流匹配）相结合。