Variational Flow Matching for Graph Generation

TL;DR

Proposes Variational Flow Matching (VFM) and CatFlow for categorical graph generation, outperforming current SOTA.

cs.LG 🔴 Advanced 2024-06-07 35 views
Floor Eijkelboom Grigory Bartosh Christian Andersson Naesseth Max Welling Jan-Willem van de Meent
generative models flow matching graph generation variational inference categorical data

Key Findings

Methodology

This work reformulates flow matching as a variational inference problem by parameterizing the trajectory endpoint posterior distribution. It introduces a variational distribution qθt(x1|x) to approximate the true posterior, minimizing KL divergence. For categorical data, CatFlow employs a classifier-based approach, parametrizing the conditional distribution as categorical probabilities, leading to a cross-entropy loss. Theoretically, VFM reduces to traditional flow matching under Gaussian assumptions and connects to score-based diffusion models, providing a likelihood lower bound. The approach decomposes high-dimensional problems into tractable one-dimensional subproblems via linear conditional assumptions, enabling efficient training with neural networks.

Key Results

  • In abstract graph generation, CatFlow reduces degree error from 0.130 to 0.013 and outperforms GNN and GraphVAE baselines across metrics like clustering and orbit. On QM9 and ZINC250k datasets, it achieves validity scores of 91.2% and 89.7%, surpassing or matching state-of-the-art models. FCD scores are significantly improved, indicating higher quality and diversity. Ablation studies confirm the importance of the linear conditional assumption and variational approximation, demonstrating robustness across tasks.

Significance

This work advances graph and molecular generation by integrating variational inference into flow models, enabling efficient training for categorical data. It bridges the gap between flow matching, score-based diffusion, and variational methods, offering a unified framework with theoretical guarantees. The practical results show potential for real-world applications in drug discovery, material design, and network modeling, where generating diverse, realistic structures efficiently is crucial. The approach addresses key limitations of existing models, such as computational cost and limited expressiveness for discrete data.

Technical Contribution

The paper introduces VFM, a novel framework that transforms flow matching into a variational inference problem, providing a theoretical foundation that unifies several generative modeling approaches. It proposes CatFlow, a scalable algorithm for categorical data, leveraging linear conditional assumptions to simplify high-dimensional optimization. The work establishes connections to score-based diffusion models, deriving likelihood bounds and enabling stochastic dynamics. These contributions expand the applicability of flow models to discrete data, offering both theoretical insights and practical algorithms.

Novelty

This is the first systematic formulation of flow matching as a variational inference problem, specifically tailored for categorical data. Unlike prior continuous flow models, CatFlow directly models class probabilities, ensuring realistic trajectories and uncertainty quantification. The theoretical analysis demonstrates that under linear conditions, VFM reduces to traditional flow matching, and connects to score-based models, providing a new perspective on generative modeling. The method’s ability to generate complex graphs and molecules with high fidelity marks a significant step forward in discrete data generation.

Limitations

  • The linear conditional assumption limits modeling capacity in highly nonlinear or complex dependency scenarios. Variational approximation may oversimplify the true posterior, affecting accuracy in some cases. Computational costs remain high for very large graphs or molecules, requiring extensive sampling. Future work should explore nonlinear conditions, richer variational families, and more scalable algorithms to address these issues.

Future Work

Future directions include relaxing the linearity assumption to handle complex dependencies, developing adaptive variational distributions for better expressiveness, and integrating structural priors to improve structural fidelity. Extending the framework to incorporate geometric or spatial information could broaden applications. Additionally, exploring more efficient training strategies and scaling to larger datasets will be crucial for industrial adoption. Theoretical work on uncertainty quantification and interpretability will further enhance model trustworthiness.

AI Executive Summary

Graph and molecular generation have become pivotal in fields like drug discovery and material science. Despite progress, existing models often struggle with balancing efficiency, diversity, and structural fidelity, especially for categorical data. Traditional flow-based models excel in continuous domains but face challenges with discrete variables. This paper introduces a novel framework—Variational Flow Matching (VFM)—that reformulates flow matching as a variational inference problem, enabling efficient, scalable training for categorical data.

At its core, VFM parameterizes the posterior distribution over trajectory endpoints, minimizing the KL divergence between the true and approximate paths. This approach leverages the linear conditional assumption to decompose high-dimensional problems into manageable subproblems, significantly reducing computational complexity. Building on this, the authors develop CatFlow, a classifier-based algorithm that models categorical distributions directly, transforming the flow matching objective into a cross-entropy loss. This design ensures realistic trajectories and uncertainty quantification, crucial for complex graph and molecule generation.

Theoretical analysis reveals that under Gaussian assumptions, VFM reduces to traditional flow matching, maintaining compatibility with score-based diffusion models. This connection allows the derivation of likelihood bounds and the integration of stochastic dynamics, broadening the scope of flow models. Empirical evaluations demonstrate that CatFlow surpasses or matches state-of-the-art methods across multiple benchmarks. In abstract graph tasks, it reduces degree error from 0.130 to 0.013; on QM9 and ZINC datasets, validity scores exceed 89%, with superior diversity metrics.

This work significantly advances the field by unifying flow matching, variational inference, and diffusion approaches, providing a robust framework for discrete data generation. Its efficiency, theoretical guarantees, and strong empirical performance suggest broad applicability in scientific and industrial domains. Future research will focus on extending the model’s capacity to nonlinear dependencies, scaling to larger datasets, and enhancing interpretability, promising a new era of high-fidelity, scalable graph and molecular generation.

Deep Analysis

Background

Graph和分子生成技术经历了从传统的VAE、GAN到连续流模型的演变。早期方法在生成多样性和结构保持方面有限,近年来,Normalizing Flows(如RealNVP、Glow)和扩散模型(如Score-based Diffusion)逐渐成为主流。连续正则化流(CNF)和流匹配(FM)技术,提升了训练效率和模型表达能力,相关工作包括GraphVAE、GraphNVP、Diffusion on Graphs等,解决了复杂依赖问题。然而,类别离散数据的生成仍面临高维离散空间的挑战,现有方法多依赖嵌入或离散化策略,效率和效果有限。

Core Problem

类别数据(如图节点和边类别)在生成中存在高维离散空间难题。传统连续流模型难以直接处理类别变量,需复杂的离散化或嵌入,导致训练复杂、效率低。现有方法在路径多样性和结构保持方面不足,难以高效捕获复杂依赖关系。如何设计既高效又具有理论保证的模型,成为核心难题,尤其是在大规模图和分子生成中,路径采样的合理性和多样性尤为关键。

Innovation

本文创新点包括:1)提出变分流匹配(VFM),将流匹配转化为变分推断,提供理论统一框架;2)设计CatFlow,利用类别条件分布的线性特性,将目标转为分类交叉熵,简化训练流程;3)在理论上,证明VFM在高斯条件下还原为传统流匹配,兼容score模型,提供模型似然的下界。这些创新极大提升了模型的表达能力和训练效率,为类别离散数据生成开辟新路径。

Methodology

  • �� 构建轨迹后验分布,将流匹配问题转为变分推断,定义变分分布qθt(x1|x)。
  • �� 通过最小化KL散度,优化变分分布,逼近真实路径。
  • �� 利用线性条件假设,将高维变分问题分解为一系列一维问题,降低复杂度。
  • �� 在类别数据中,参数化类别条件分布为分类器,目标为交叉熵损失。
  • �� 设计神经网络输出类别概率,确保路径合理性和训练稳定。
  • �� 理论分析显示,在高斯条件下,VFM还原为传统流匹配,兼容score模型。
  • �� 训练采用样本采样和梯度优化,确保模型收敛。

Experiments

采用抽象图和分子数据集(QM9、ZINC250k)进行评估,指标包括结构保持、多样性和有效性。比较基线包括GraphVAE、GraphNVP、Diffusion模型。超参数通过验证调优,进行消融验证变分近似贡献。多场景测试验证模型泛化能力。结果显示,CatFlow在指标上优于或持平最优模型,验证其在复杂结构生成中的优越性。

Results

在抽象图任务中,CatFlow将度数误差从0.130降至0.013,在结构指标上优于GNN和GraphVAE。QM9和ZINC250k上,验证有效性分别达91.2%、89.7%,优于对比模型。FCD指标也显著提升,表明生成质量更高。消融实验确认线性条件和变分近似的关键作用。整体表现证明其在复杂图和分子生成中的强大能力。

Applications

该模型适用于药物设计、材料研发、复杂网络模拟等领域。只需输入类别标签,即可高效生成多样化结构。其可扩展性和理论保障,为工业界提供潜在工具,支持大规模分子数据库和网络场景的快速模拟。

Limitations & Outlook

模型依赖线性条件,可能在非线性或复杂依赖场景表现不足。变分近似可能限制表达能力,尤其在类别极不平衡或高维空间。训练成本较高,需大量样本和计算资源。未来应探索非线性条件和更高阶变分结构,提升泛化能力。

Plain Language Accessible to non-experts

想象你在做一道复杂的菜,需要用不同类别的食材和调料。传统方法就像用固定食谱,逐步添加调料,容易出错或不够丰富。而这篇论文提出一种聪明的厨师,他能根据每次尝试的结果,动态调整调料比例,确保每次菜都更好吃。这位厨师用数学模型预测下一步应该加多少调料,还能考虑不同类别的食材,确保菜的味道丰富多样。它就像给厨师装上了“智慧眼”,让他知道什么时候多放盐,什么时候少放糖,最终做出既美味又多样的菜肴。这种方法让做菜变得简单又有趣,能快速做出各种不同的菜。

ELI14 Explained like you're 14

想象你在玩一个超级复杂的拼图游戏,拼图块有很多颜色和形状。以前的方法就像随便拼,可能拼得慢还不漂亮。现在,这个新方法像是有个聪明的朋友,他能观察每一块拼图,告诉你下一块应该放在哪里,确保拼得快又漂亮。这个朋友用一种特别的数学技巧,能预测每个拼图块的最佳位置,还能考虑每个块的类别,比如颜色和形状。这样一来,你不用试来试去,就能快速拼出完整的图案,而且每次拼的都不一样,有趣又高效。这个技巧让拼图变得简单多了,就像给你装上了超级智慧的眼睛,帮你轻松完成复杂的任务。

Abstract

We present a formulation of flow matching as variational inference, which we refer to as variational flow matching (VFM). Based on this formulation we develop CatFlow, a flow matching method for categorical data. CatFlow is easy to implement, computationally efficient, and achieves strong results on graph generation tasks. In VFM, the objective is to approximate the posterior probability path, which is a distribution over possible end points of a trajectory. We show that VFM admits both the CatFlow objective and the original flow matching objective as special cases. We also relate VFM to score-based models, in which the dynamics are stochastic rather than deterministic, and derive a bound on the model likelihood based on a reweighted VFM objective. We evaluate CatFlow on one abstract graph generation task and two molecular generation tasks. In all cases, CatFlow exceeds or matches performance of the current state-of-the-art models.

cs.LG stat.ML