提出HollowFlow,让生成模型采样更快,大幅降低计算开销。
HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing
- 用新型非回溯图网络构建块对角雅可比矩阵,实现恒定次数反向传播。
- 在大系统上达到100倍加速,理论与实验均验证其线性扩展性。
- 适用于各类图神经网络和注意力模型,适合高维科学建模场景。
基于流和扩散的模型在科学应用中表现出强大能力,尤其适用于采样未归一化概率分布,如玻尔兹曼生成器(BGs)。然而,其部署面临关键挑战:样本似然计算随系统规模 $n$ 呈不可行增长,难以用于大规模问题。为此,我们提出 $ extit{HollowFlow}$,一种基于新型非回溯图神经网络(NoBGNN)的流模型。通过强制块对角雅可比结构,HollowFlow 的似然评估仅需常数次反向传播,实现高达 $/mathcal{O}(n^2)$ 的加速,显著推动 BGs 向更大系统扩展。重要的是,该框架可泛化至任意等变图神经网络或基于注意力的架构。我们在两个不同系统上训练 BGs 验证 HollowFlow,随着系统增大,采样与似然评估时间急剧下降,符合理论预期。在更大系统上实现了 $10^2\times$ 的加速,充分展示其在以往受计算瓶颈限制的高维科学问题中的潜力。
原文摘要 · Abstract (English)
Flow and diffusion-based models have emerged as powerful tools for scientific applications, particularly for sampling non-normalized probability distributions, as exemplified by Boltzmann Generators (BGs). A critical challenge in deploying these models is their reliance on sample likelihood computations, which scale prohibitively with system size $n$, often rendering them infeasible for large-scale problems. To address this, we introduce $\textit{HollowFlow}$, a flow-based generative model leveraging a novel non-backtracking graph neural network (NoBGNN). By enforcing a block-diagonal Jacobian structure, HollowFlow likelihoods are evaluated with a constant number of backward passes in $n$, yielding speed-ups of up to $\mathcal{O}(n^2)$: a significant step towards scaling BGs to larger systems. Crucially, our framework generalizes: $\textbf{any equivariant GNN or attention-based architecture}$ can be adapted into a NoBGNN. We validate HollowFlow by training BGs on two different systems of increasing size. For both systems, the sampling and likelihood evaluation time decreases dramatically, following our theoretical scaling laws. For the larger system we obtain a $10^2\times$ speed-up, clearly illustrating the potential of HollowFlow-based approaches for high-dimensional scientific problems previously hindered by computational bottlenecks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。