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.

翻译：基于流和扩散的模型已成为科学应用中的强大工具，特别是在采样非归一化概率分布方面，以玻尔兹曼生成器（BGs）为例。部署这些模型的一个关键挑战在于其对样本似然计算的依赖，其计算复杂度随系统规模 $n$ 呈禁止性增长，通常导致其在大规模问题上不可行。为解决此问题，我们提出了 $\textit{HollowFlow}$，这是一种基于流的生成模型，其利用了一种新颖的非回溯图神经网络（NoBGNN）。通过强制雅可比矩阵具有块对角结构，HollowFlow 的似然评估仅需常数次关于 $n$ 的反向传播，从而实现了高达 $\mathcal{O}(n^2)$ 的加速：这是将 BGs 扩展到更大系统的重要一步。至关重要的是，我们的框架具有普适性：$\textbf{任何等变 GNN 或基于注意力的架构}$ 都可以适配为 NoBGNN。我们通过在两个规模递增的系统上训练 BGs 来验证 HollowFlow。对于这两个系统，采样和似然评估时间均大幅减少，符合我们的理论缩放定律。对于较大的系统，我们获得了 $10^2\times$ 的加速，这清楚地展示了基于 HollowFlow 的方法在解决先前受计算瓶颈制约的高维科学问题方面的潜力。