HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing

TL;DR

HollowFlow tackles the computational bottleneck of evaluating sample likelihoods for flow-based Boltzmann Generators in high-dimensional systems. It introduces Hollow Message Passing with a non-backtracking GNN to enforce a block-diagonal/hollow Jacobian, enabling a constant number of backward passes and large speed-ups ($O(n^2)$) when using a $k$NN graph. The approach is compatible with any equivariant GNN or attention architecture and is demonstrated on LJ13, LJ55, and Alanine Dipeptide with substantial gains in sampling and likelihood evaluation. This work enables scalable, symmetry-aware BGs for complex scientific problems, improving achievable system sizes and unbiased observable estimation.

Abstract

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.

Figures (5)

• Figure 1: Summary of HollowFlow: (a) A vector field $b$ is parametrized with an (equivariant) non-backtracking graph neural network (NoBGNN) and a series of (equivariant) neural networks $\tau_i$, such that its Jacobian $\mathbf{J}_b$ (b) can be decomposed into a block-hollow and block-diagonal part with block size $d$ (in this example, $d=3, n=5$). (c) Efficient evaluation of the trace of the Jacobian with only $d$ instead of $nd$ backward passes (BP) through the network. The vector components are one where the entry is green and zero elsewhere.
• Figure 2: Schematic example of line graph message passing. Blue: original graph ($G$), orange: line graph ($L(G)$), gray: information flow between $G$ and $L(G)$. \ref{['fig:LG_a']} Construction of the line graph nodes. Every node in $L(G)$ corresponds to one edge in G. The dashed arrows indicate information flow before the message passing from $G$ to $L(G)$. \ref{['fig:LG_b']} Construction of the edges of $L(G)$. \ref{['fig:LG_c']} Information flow after the message passing (dashed arrows) from $L(G)$ to $G$.
• Figure 3: Runtimes of HollowFlow and baseline for generating $10^5$ LJ13 equilibrium samples. All runtimes are multiplied by the GPU usage. \ref{['fig:speedup']} Total runtime $\text{RT}(L(G))$ and runtime of the forward passes $\text{RT}^f(L(G))$ and backward passes $\text{RT}^{\nabla}(L(G))$ of HollowFlow. In light blue: speed-up (ratio of the runtimes of the fully connected baseline $\text{RT}(G)$) and HollowFlow $$. \ref{['fig:runtime']} Total runtime and runtime of the forward and backward passes for the non-hollow fully connected baseline. The forward and backward runtimes of the baseline were extrapolated from other runs, see \ref{['app:runtimes']} While the total runtime of the baseline is massively dominated by the backward pass, the opposite is true for HollowFlow.
• Figure 4: Schematic example of the graph $G$ (in blue) and its line graph $L(G)$ (in orange): The line graph nodes encircled with a continuos red line correspond to $\mathcal{N}^{lg}(j,i)$ while the node encircled with a dashed red line corresponds to $\mathcal{N}^{lg}(i,j)$. Two edges of $L(G)$ are labelled explicitly to illustrate our labelling convention.
• Figure 5: Average number of connections left in the line graph of a $k$NN graph $G$ as a function of the number of message passing steps. $G$ has 55 nodes whose coordinates are sampled from a standard Gaussian. The average is taken over 10 independently sampled graphs. The larger $k$, the faster the line graph disconnects.

Theorems & Definitions (6)

• Theorem 1: Block-hollowness of HoMP
• Theorem 2
• Theorem \ref{th:hmp}: Block-hollowness of HoMP
• proof
• Theorem \ref{th:complexities}
• proof