HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing
2510.21542v1
stat.ML, cs.LG
2025-10-28
Авторы:
Johann Flemming Gloy, Simon Olsson
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.
Ссылки и действия
Дополнительные ресурсы: