HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing
Authors: Johann Flemming Gloy
Simon Olsson
Venue: NeurIPS 2025 poster ()
Keywords: Generative Models
Flow-based models
Graph Neural Networks
Boltzmann Generators
Primary area: machine_learning_for_sciences
TL;DR: We present HollowFlow, a flow-based generative model based on a GNN that makes sample likelihood computations drastically more efficient.

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.

HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing
Johann Flemming Gloy! Simon Olsson'
Department of Computer Science and Engineering,
Chalmers University of Technology and University of Gothenburg SE-41296 Gothenburg, Sweden. {gloy, simonols) @chalmers.e.
Abstract
Flow and diffusion-based models have emerged powerful tools for scientific applications, particularly for sampling non-normalized probability distributions, 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 HollowFlow, a flow-based generative model leveraging a novel non-backtracking graph neural network (NoBGNN). By enforcing plock-diagone Jacobian structure, HollowFlow likelihoods are evaluated with constant number of backward passes n, yielding speed-ups of up to Onn)): significant step towards scaling BGs to larger systems. Crucially, our framework generalizes: any equivariant GNN or attention-bass baated architecture can be adapted into 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 102 x speed-up. clearly illustrating the potential HollowFFow-basee approaches for high-dimensional scientific problems previously hindered by computational bottlenecks.
Introduction
Efficiently sampling bigh-dimensional, non-normalized densities - cornerstone Bayesian infer- ence and molecular dynamics remains computationally intractable for many scientific applications. Boltzmann generators (BGs) [1] address this by training a surrogate model Pi (x) to approximate
Z-1 e--1exp(-B0r exp(-Bu(x)), where u(x) is the potential energy of the system configuration x E Ω C RN B the inverse temperature, and Z the intractable partition function. While BGs enable unbiased estimation of observables through reweighting x; Pr with weights wi & (uxx)// p1 (x;), their practicality is limited by a fundamental trade-off: likelihood computations of expressive surrogates rely N automatic differentiation backward passes with system dimensionality N, rendering them infeasible for large-scale problems.
the Boltzmann distribution u1
To address this problem, we outline framework combining non-backtracking equivariant graph neural networks with continuous normalizing flows (CNFs), breaking this trade-off. This strategy enforces a Jacobian that can casily be split into block-diagonal and block-hollow part, reducing the number of backward passes required for likelihood computation to scale as ((1) in N. Our contributions include:
Hollow Message Passing (HoMP): A general scheme for message passing with a block- diagonal Jacobian structure.
"Corresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

c(x)
T(h, x)
(a)
X1
h1 (x-1)
b1
J b(x)
b(xx=
X2
h2 (x-2)
b2
J(xx)
r(x)
NoBGNN
(block- hollow)
(block-
diugonal)
Pnn (x-m) m
b
(b)
Tr (T(J)x))
(c)
BP
BP3
Figure I: Summary of HollowFlow: (a) A vector field is parametrized with (equivariant) non-backtracking graph neural network (NoBGNN) and a series of (equivariant) neural networks T;, such that its Jacobian J, (b) can be decomposed into block-hollow and block-diagone part with block size d (in this example, d = 3, 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 one where the entry is green and zero elsewhere.
Theoretical Proofs and Guarantees: proof of the block-diagonal structure HoMP and scaling laws of the computational complexity of forward and backward passes.
Scalable Boltzmann Generators: HollowFlow-based BG implementation achieving 10' x faster sampling and likelihood evaluation, e.g., 10 x for LJ55, system of 55 particles in 3-dimensions interacting via pair-wise Lennard-Jones potentials.
Experiments on multi-particle systems validate HollowFlow's scalability, bridging a critical gap in high-dimensional generative modelling for the sciences.
2 Background
2.1 Boltzmann Generators and Observables
Computing observables by averaging the Boltzmann distribution is central challenge in chemistry, physics, and biology. The Boltzmann distribution is given by,
μ(x) Z exp( (0x(()Bux))
(1)
where 8 = (k is the inverse thermal energy and u Ω + IR is the potential energy of a system configuration
C R.V.
Observables are experimentally measurable quantities, such as free energies 2 Generating ensembles which align with certain observables give us fundamental insights into the physics of molecules, including mechanisms, rates, and affinities of drugs binding their targets [5] 6]. protein folding
phase transitions
12]
The primary strategies for generating samples from the Boltzmann distribution involve Markov Chain Monte Carlo (MCMC) or molecular dynamics (MD) simulations. These approaches generate unbiased samples from the Boltzmann distribution asymptotically, however, mix slowly on the complex free energy landscape of physical systems, where multiple, metastable states separated by low-probability channels. In turn, extremely long simulations are needed to generate independent sampling statisties using these methods. In the molecular simulation community, this problem has stimulated the development of a host of enhanced sampling methods, for recent survey, see [1]
Boltzmann Generators [ instead try to learn an efficient surrogate from biased simulation data, and recover unbiased sampling statisties through importance weighting. The key advantage BGs that sampling using modern generative methods, such as normalizing flows, is much faster than the conventional simulation based approaches, which would enable amortization of the compute cost per Boltzmann distributed sample.

---

2.2 Normalizing Flows
To enable reweighting, evaluating the sample likelihood under a model, Pt. critical. A normalizing flow (NF) [14] 15) allows efficient sampling, and sample likelihood evaluation, by parameterizing a diffeomorphism, Φv R + R-V, with a constrained Jacobian structure. The diffeomorphism (or flow) Φu maps an easy-to-sample base distribution, /o to an approximation Pr of the empirical data density, pD. c.g., by minimizing the forward bullback-Leible divergence, DKL (PDllp) = Ex-pp [logpp(x)) - log P1 (x)]. The change in log-density when transforming samples with is given by
0 -
A log p (x) == log (x) log Pa p = log det
(2)
0x
Computing the log-determinant of a general N N Jacobian is computationally expensive (O(N))),
d0 (x) with structure, i.e., triangular 8x
however, endowing the Jacobian
18). can make the
likelihood evaluation fast.
Continuous Normalizing Flows A flow can also be parametrized as the solution to an initial value problem, specified by a velocity field bo(x, t) RRN [0 - RN of ordinary differential equation (ODE)
dx(t) = by (x((), dt, x(0) ~ po+ x(1) ~ P1.
The general solution to this problem is flow. (x(),,, R [0. 1| + +R which satisfies the
initial condition, c.g., do (x))))))) = x(0). Such a model is known as a contimuuus normalizing flow (CNF). The ODE eq. B together with the prior Po gives rise to time-dependent distribution p(x(t), t) given by the continuity equation
a = (prbg(x,
(4)
with O, and V. denoting the partial derivative w.r.t. t and the divergence operator, respectively. Solving eq. 4 get the change in log- probability for the CNF,
A log p CNF == log p1 (x(1)) log Pu (x(0) = -- V (o(())), dt.
(5)
In principle, CNFs allow for much more expressive architectures to used to parametrize by-
The currently most scalable approach to train CNF's is conditional flow matching (CFM) or related objectives 22 In CFM, the intractable flow matching loss, where by regressed directly against the unknown true vector field uy (x). of a flow é, is replaced by the conditional flow matching loss which has the same gradient w.r.t. the parameters [0
CCFM
(xpp (x, ue (x(2))11
(6)
The true vector field un (x) and its corresponding probability density (x( are given as marginals over the conditional probability and an arbitrary conditioning distribution p(z):
px(x) = f Pe(x)=))(( (((x)==(<((())))))
Pr (x(x)z) (x/z)
dz and ux(x)
p(z) dz.
(x(x)
where we parameterize pt (x|z) and un (x(2)) with linear interpolant with connections to optimal transport theory [22]:
X1), p(z)= = (Xx. x1) (x/2 N(xx + (1 t)xo, o),
Wy (x/z)) X1 X
where the coupling π (x1, x4) minimizes the 2- Wasserstein optimal transport map between the prior
Pu(Xo) and the target μ(x;) using mini-batch approximation [20 22 23
Unfortunately, akin to general Jacobian log-determinant calculation of normalizing flows, computation of V bp(x() (x(())) !) (eq. (5)) requires O(N) backward passes, for each numerical integration step along
I, making the evaluation of sample likelihoods impractical high-dimensions, and consequently, their use for BGs limited.

---

Equivariant flows Symmetries can be described in terms of a group G that acts on a finite- dimensional vector space V = R through a group representation p : + (d R). function f: IRd R is called G-equivariant if f((()))) = f(g) and h R - V called G-invariant if h((gg)))) h(x) for all g E 9 and x E Rd where V is another vector space. If the vector field be of a CNF H--qquivariant, where 7( < 9 and the prior density Po is G-invariant it has been shown that the push forward density, of the flow P1 is H(-invariant (23 In practical applications, symmetries under the Euclidean group E3) are of particular interest as these are the symmetries that naturally obeyed by many physical systems such molecules. Thus, several GNN architectures have been proposed that are equivariant or invariant under the action of E(3) or some subgroup of E(3) 26 27
2.3 HollowNets
Reverse-mode automatic differentiation, or back propppoppggttion allows for the computation of vector- Jacobian products with a cost which is asymptotically equal to that of the forward pass (28). In practice, we use this mode to compute the divergence of velocity fields parametrized with neural networks. However, isolating the diagonal of the Jacobian to compute the divergence, for general, free-form vector fields, requires ((N) backward passes [29 30).
Using a special neural network construction HollowNet, Chen and Duvenaud showed, that one can access the full diagonal of the Jacobian, J E IR-XN, with single backward pass. A HollowNet depends on two components:
Conditioner: h4 = (x--),, where o; :RN-1 + IR"h Transformer: b, (x,, h.), where T IR R + R
(9)
(10)
where h; is latent embedding which is independent of the i'th dimension of the input, x. Applying the chain rule to by = (r(r,, c((-))). the Jacobian J db decomposes as follows:
dx
db
Ōr Oh
(11)
dx
dx Jdiag
Th
J hallow
which splits the Jacobian into diagonal Jdiag and 'hollow' Jhollow contributions [2 The divergence
can then be evaluated in ((1) backward passes, through rewiring the compute graph and computing 1 ddiag
3 Hollow Message Passing
is straightforward to construct HollowNets for regular forward neural networks, and convo- lutional neural networks, through masking operations (29). However, extending these principles to cases where the data or its domain are subject to other symmetries not as straightforward.
To address this problem, we develop Hollow Message Passing (HoMP) to work with graph data, subject to permutation equivariance. We further specialize this approach work with data which subject to common Euclidean symmetries, useful in molecular applications. Here, molecules and particle systems are represented as points in Rd (typically particles, making N dn. HoMP depends the following developments:
3), one for each of the n atoms or
Permutation Equivariant Conditioner: We make the conditioner functions J. per- mutation equivariant, using a non- backtrackiknn GNN (NoBGNN) in combination with a changing underlying graph structure such that it fulfils the same constraints the original conditioner. Namely, 0hg 0 vi E nt).
Dx.
We lift the HollowNet from working with scalar features, i.e., X; E R to work with Euclidean features, e.g., X, E Rd Moreover, NoBGNN allows for the use of message passing architectures that are invariant or equivariant under the group action p(g), g E G of a group G. Furthermore, for each there can be additional scalar input data Z, E IR. In practice (see section while Z; encodes the atom type.
Euclidean Symmetries of Node Features:
X; E R3 are the coordinates of an atom

---

(a)
(b)
(c)
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). 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). (b) Construction of the edges of L(G). (c) Information flow after the message passing (dashed arrows) from L(G) to G.
Permutation Equivariant Conditioner We design GNN for the conditioner such that information from any node never returns to itself. A GNN where this holds for one message passing step is called non-backtracking GNN (31] line graph GNN (32 and can be constructed as follows:
Let = (N, E)) be directed graph with nodes N := [I,...,,,) n.. and edges E' S {( E N, W. + j}. To each node i, we assign a node feature n4. Given G, we construct the non-backtracking line graph L(G) = (Nla, Els), where
N'a = E, E = ((..i.)) E (i,)) E E and + k).
(12)
The node features of (G) are n are defined to be
niy
Ig
Ti.
(13)
The hidden states h L(G) are initialized to be hig n. Using this setup, standard message passing scheme 33 is applied L(G):
mag = (h//. mig),
mly =
mag-
h = (hi mLy);
(14)
(EN
bj Σ R(hi n2))
(15)
VEN())
d, sb and R are all learnable functions, (i. (ki i, E')). (k](k E E}} and
Tle is the number of message passing steps the line graph. This construction is illustrated in fig. 2
The blue nodes and edges are part of G while the orange nodes and edges are part of L(G)) The
). gray dashed arrow in fig. (( illustrate the projection of the node features onto L(G') (eq. [3 The
orange arrows in fig. the projection of the node features from L(G) back G summarized eq. (15). Equations (13
the message passing on L(G) (eq. 14 Figure 2(c) visualizes
and (14) take the role of the conditioner e(x) while 15 takes the role of the transformer 7 (h, x) (see appendix B.
So far, this GNN is, except in special cases that we detail in appendix B only non- backbtrkta just message passing step is performed. In general, to ensure the non-bacctracking nature of the GNN, some connections passing steps. This be done by keeping track of which nodes of L(G) received information of which nodes of G after each message passing step. To this end, we define the three dimensional time-dependent backtracking array B(t) Ruxnxn
L(G) need to be removed after an appropriate number of message
{
ah dxt
+
B(t)ijA=
(16)
else.
ah O whenever (i,j) & E. Notice that whenever B(t) =
with the additional convention that Tx2 information that started in node j will return to node j during readout. Algorithm shows in detail how B() is calculated (steps 8 and 11) and used to remove edges in L(G) (step 10). A more detailed
explanation and proof that this construction does indeed result in a non bankbbkkkaakkingggg GNN for an
arbitrary number of message passing steps is given in appendix B
5

---

In typical molecular applications, nodes are represented in Euclidean space (e.g.. x; E Rd where d and may have additional features Z, (such as atom type). Our goal is to embed (x,, 7:) E Rd x IR into a set of equivariant (v;) and invariant (s;) node features, (U:, 8; = embed( x; Z,) E R x I which serve as inputs for message passing on (see eq. (33 ) For this task, we can use any Geequivariant embeddings inherent to GNN
Euclidean Symmetries of Node Features:
architectures such as PaiNN or E3NN (for d = 3 and G < E(3)) [26 27).
The message passing and readout functions (eqs. [4] and (5) are made G-equivariant by using the message and update functions of any G-equivariant GNN. As we only modify the underlying graph of the GNN, this trivially ensures equivariance.
By construction, the d-dimensional node inputs x; allow the Jacobian J, to be easily split into a block-diagona and block-hollow part with block size d x d. This observation and its consequences are summarized in algorithm|| [ and theorem I (see appendix B for a proof and additional explanations). Algorithm D can also
the pairwise euclidean differences of the node features x, (pd = true in algorithm TT This is common practice in molecular applications to achieve translation invariance. brief illustration of algorithm D can be found in fig. T
Algorithm 1 Hollow message passing
Input: {xi, 2=)==1 where x; E Id and Z, IR, pd E (true, J salse) Returns: {b:)==1 b, E Rd
1: if' pd then
2:
di = X; Xj.
Gjj = (dbe 2,)
3: else
4: n1 = emmbed(x; Z,)
5: end if
6: Calculate G (N, E), e.g., as a kNN graph. Self-loops, i.e., (i, E are prohibited.
7: Calculate the line graph L(G)- = (N's Ela) and its node features "ig"
ety if pd
N'a E, E (1) E E, (J. E and + k). n
n4 else
8: Calculate initial node features f.g and initial backtracking array B(0) V Σ ((n(2)) if pd
(1ifk E (i. else) if pd
h
(E.)
B(0)
lg nig
else
(1 = i, 0 else) else
9:
for + 0 to (T 1) do
10:
Remove appropriate edges of LG); El9 + E'9 \ {(i,j,)) E E's B(t) =1}
11:
Update backtracking array: Bt + 1) maxieN ((t)) (B() B(t).jk
12:
Do message passing:
mky ohh Maa), m = ExENiO(Ss) mlags Mt (h m!y)
13:
end for
14:
Perform readout and project back to G:
Tmmeed(x; detach () Xi, pd
by EiENGy R(h!)) miz), where ni
ng
else
Theorem (Block-hollowness of HoMP), Algorithm I defines a function b Rddn Rdn whose Jacobian Jo(x) Rdnxxd can be split into blolk-hollow and a block-diagonal part, i.e.,
Jb(x) Jo(x) + Jr(x)
(17)
where Je(s) is block-hollow, while J Jx) is block-diagonal with block size d respectively. This structure of the Jacobian enables the exact computation of differential operators that only include diagonal terms of the Jacobian with instead of nd backward passes through the network.
In appendix A we show how HoMP can be adapted to attention mechanisms well.

---

Computational Complexity
4.1 Graph Connectivity
The number of backward passes through a vector field b parametrized by algorithm that is required compute its divergence is reduced by factor of n compared to standard GNN. However, the forward pass through b is slower than through a standard GNN, as is operating on the line graph. This overhead will scale with the number of edges, which we denote # for G and # 1 for L(G). In table [ we show how the number of edges varies with graph types. For a fully connected graph, # E scales n² resulting in impractical overhead in high-dimensional systems. To address this, we consider & nearest-neit neighess (kNN) graphs instead which scale as # E = Onnk²). Choosing S (v)) will leave the cost of a forward-pass through comparable to a standard GNN with fully connected graph. To test whether this restriction in connectivity results in a sufficiently expressive model, we empirically test different values for k in section
Using a ANN graph poses a locality assumption that might limit the model's ability to learn the equi- librium distribution of systems with long-range interaction (e.g., systems with coulomb interactions such as molecules), One possible strategy to circumvent this problem is described in appendix C
Table 1: Number of edges of L(G) and G
fully connected
h neighbors
#E #E's
(n (nn))
O(nk)
O(nk²)
4.2 Divergence Calculation Speed-up
We summarize our runtime estimations for one integration step during inference RT including likelihood calculations, for HollowFlow and a standard GNN based flow model in the following theorem:
Theorem 2. Consider GNN-based flow model with fully connected graph G and T message passing steps and HollowFlow model with a k neighbors graph Gk and T message passing steps. Let both graphs have n nodes and d-dimensional node features.
The computational complexity of sampling from the fully connected GNN-based flow model including sample likelihoods
RT (G -
(18)
while the complexity of the HollowFlow model for the same task is
(Lstep(L(Gg.)) On + dk)).
(19)
Moreover, the speed-up up of HollowFlow compared to the GNN-based MoNe-baedffor flow model is
RT (G O (Tsstep(T[Gaa
Tn²d
(20)
T'gk2 + dk
derivation of these estimates can be found appendix B.2
Experiments
We test HollowFlow on two Lennard-Jonee systems with 13 and 55 particles, respectively, (LJ13,
LJ55). We also present results on Alanine dipeptide (ALA2) in Appendix
To this end, we train CNF (for training details and hyperparameters see appendix D) using the conditional flow matching loss 200 to sample from the equilibrium distribution of these systems. We parametrize the vector field of our CNF using the 0(3) -(3)-eqquva PaiNN architecture (26). We
additionally use minibatch optimal transport [22 and equivariant optimal transport (23 For each system we train multiple different models:
7

---

Baseline:
equivvaaiaa GNN on fully connected graph with 3-7 message passing steps.
HollowFlow: eqqiivvava HollowFlow on kNN graph with different values for kt with two message passing steps.
In the case of LJ13, we systematically test all values of k from 2 to 12 in steps of 2 to gain insight into how k affects the performance and runtime the model. Based on our findings and the theoretical scaling laws, pick & {7,27) for LJ55.
Motivated by the symmetry of pairwise interactions between atoms, all KNN graphs are symmetrized. i.c. whenever the edge (i.j) is included in the kNN graph the edge (j. is included as well, even if it is not naturally part of the &NN graph.
All training details and hyperparameters are reported in appendix D
Criteria to Assess Effectiveness We measure the sampling efficiency through the Kish effective samples size (ESS) (35]. As the ESS is highly sensitive to outliers, we remove a right and left percentile of the log importance weights, log w,, to get more robust, yet biased, metric, ESSrem (see appendix D.7p [36]. As HoMP is, in general, a restriction compared to normal message passing, expect the ESS to be lower. However, long the speed-up is sufficiently large, our method can still be an overall improvement.
To gauge the improvement we introduce compute normalized, relative ESS, i.e., the number of effective samples per GPU second and GPU usage during inference (see eq. 21)). This gives us fairer grounds for comparison of HollowFlow against our baseline, standard O 0 -e]-quiivaii CNE Further, by dividing the relative ESS for all hollow architectures by the baseline relative ESS, we get the effective speed-up up (EffSU) and correspondingly the EffSUrem calculated from the relative ESSrem in the same way. To ensure fair comparison, we ran all inference tasks on the same type of GPU for every system, respectively. Details can be found in appendix D
The results for LJ13 and LJ55 are listed in tables and 3 For both systems we observe that the ESS and the ESS are highest for the baseline and significantly smaller for HollowFlow. However, the relative ESS,ee is significantly higher for all Hollow models. This is reflected the effective speed-up EftSU gauging how many times more samples we can generate per compute compared to our non-hollow baseline. While the maximal EftSUrem that we observe for LJ13 is about 3.3 for k = 6, for LJ55 we observe maximal EffSU rem of about 93.7 for k = 7. This shows that HollowFlow can reduce the compute required for sampling by about two orders of magnitude on our larger test system.
In fig. we provide a detailed analysis of the runtime of different parts of models trained on LJ13 during sampling. Figure ((b) demonstrates that the vast majority of the compute in the baseline model is spent on the divergence calculations while only tiny fraction spent on the forward pass. fig. 3(a] shows that exactly the opposite is true for HollowFlow, resulting in a decrease of runtime by a factor of up to ~ 15, depending on
ESS x #samples
relative ESS(.)
relative ESS
EASUU(-) = relative ESS ESS(baseline)
(21)
RT (GPU usage)*
6 Related Works
Efficient Sample Likelihoods in Probability Flow Generative Models Evaluating the change in probability associated with probability flow generative models is subject to research but broadly fall into two different categories: 1) approximate estimator by invoking stochastic calculus (37) or
38 39 which are efficient, but prone high variance 2) NFs with structured possibly combined with annealing strategies (42-455 which offer efficiency, but limited expressivity or an inability to exploit certain data symmetries such as permutation equivariance.
trace estimators 30
Jacobians [16 18
Boltzmann Emulators and Other Surrogates Several works forego quantitative alignment with physical potential energy function u(x) in favor of scaling to larger systems through coarse-gran -rrainiggg or enabling faster sampling. These qualitative and semi-quantitative Boltzmann surrogates are often

---

15
AT(L(G))
1000
(LGGH
10
runting (8)
spend-up
runtime
500
RT(G) RZYL(C)
5
10a
6
9
12
RT(G)
RT (G)
RTY(G)
(b)
Figure 3: Runtimes of HollowFlow and baseline for generating 109 IJJ13 equilibrium samples. All runtimes are multiplied by the GPU usage.
(a) Total runtime (T(I(G)) and runtime of the forward passes RT' (L(G) and backward passes RT (L(G)) of HollowFlow. In light blue: speed- up (ratio of the runtimes of the fully connected baseline RT(G)) and HollowFlow (b Total runtime and runtime of the forward and backward passes the non-hollow fully connected baseline. The forward and backward runtimes of the baseline were extrapolated from other runs, appendix While the total runtime of the baseline is massively dominated by the backward pass, the opposite is true for HollowFlow.
[.6) D.6
Table 2: Comparison of HollowFlow trained using different &NN graphs and the non-hollow fully connected baseline, all trained on LJ13. The two rightmost columns show the effective speed-up compared to the baseline (1) with and without weight clipping. Bold indicates best performance.
ESS (%)
ESS (1) (%)
EffSU rem (t)
EffSU
k 2
0.054
0.039
2.922.85
1.059 070
1.099 009
1.16
0.041 0.110.102
10.3410.40
2.649 2.6492.666
0.64]
3.300 3.30
20.20 29
3.260 3.243
3.278
31.8753.80
= 8
2.745 250
3.453.200 236
20.64 20.60
2.310 3233
18.26
2. 297
k 10
1.057]
16.5016.58
1.484 1.441.775 175
492
5.51 1.61
12
4.0691.929
3.20 200
19. 19.7219.80
1.627 1.636
20.1 1. 35.55
5.66
baseline
2.1322.231
40.7340.5.
1
Table 3: Comparison of HollowFlow trained using different kNN graphs and the non-hollow fully connected baseline, all trained on LJ55. The two rightmost columns show the effective speed-up compared to the baseline eq. (2I) with and without weight clipping. Bold indicates best performance.
ESS (%)
ESS (t) (%)
EffSU rem (1)
EffSU
1: 7
0.006
0.530.50
93.7377 99. 99.071
13.26136.60
88.484
A: 0.0070.012
0.64
9.4668 .88
8.4614.90
A: =
55
0.0200.023
0.74 0.
4.365 5883
9.11
13.18 13.18
4.57
baseline
0.0480.051
2.963.03
3.003
2.89
referred to as Boltzmann Emulators [46]. Some notable examples include, SMA-MD [47]. BioEmu 48). AlphaFlow [49]. JAMUN (50). and OpenComplex2 [5T]. Other approaches, instead aim to
mimic MD simulations with large time-steps including ITO (2)54). TimeWarp 41]. ScoreDynamics
55
and related approaches 606-60).
Non-backtracking Graphs Non- Non-backtracking graphs are applied in a variety of classical algo- rithms, e.g., [61] 62). NoBGNNs that are non-backtracking for only one message passing step have been applied in a range of other context including for community detection (32), as way to reduce over-squashing (31] and alleviate redundancy in message passing (63). However, nobody has yet exploited these methods to construct NoBGNN that is backbrtktaakknngg for an arbitrary number of message passing steps to give the Jacobian structure that allows for cheap divergence calculations.

---

7 Limitations
We demonstrate good theoretical scaling properties of HollowFlow large systems and support this experimentally on systems of up to 165 dimensions. However, the kNN graph used in the current HollowFlow implementation involves an assumption of locality, which may break down in molecular and particle systems with long-range interactions, such as electrostatics. Future work involves exploring strategies to incorporate such long-range interactions, which may involve other strategies construct the graph used in HollowFlow, wwwald-summation-basss
methods [64] 65 or linear/fast attention [66] [6]. Also, the fact that some edges in the line graph underlying HollowFlow need to be removed after each message passing step might be problematic at scale. Despite that, there still a wide range of physical systems where long-range interactions are effectively shielded, such as metals and covalent solids, which might enjoy the properties of Hollow models even without these technical advances.
Overall, the performance of our baseline is not in line with state-of-the in terms of ESS for the studied systems. However, since our HollowFlow models build directly on the baseline we anticipate any improvement in the baseline to translate into the HollowFlow models as well.
8 Conclusions
We introduced HollowFlow, a flow-based generative model leveraging novel NoBGNN based on Hollow Message Passing (HoMP), enabling sample likelihood evaluations with constant number of backward passes independent of system size. We demonstrated theoretically that we can achieve sampling speed-ups of up to (n²)). if the underlying graph G of HollowFlow is a &NN graph. By training a Boltzmann generator on two different systems, we found significant performance gains in terms of compute time for sampling and likelihood evaluation, the largest system enjoying 10 x speed-up in line with our theoretical predictions. Since HoMP can be adopted into any message- passing-base based architecture (68), our work can pave the way to dramatically boost the efficiency of expressive CNF- Boltzmann Generators for high-dimensional systems.
Acknowledgements
This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. Results were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Alvis (projects: NAISS 2024/22-688 and NAISS 2025/22-841), partially funded by the Swedish Research Council through grant agreement no. 2022-06725. The authors thank Ross Irwin and Selma Moqvist and other team members of the AIMLeNS lab at Chalmers University of Technology for discussions, feedback and comments.
References
[1] Frank Noć et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning". In: Science 365.6457 (2019), eaaw 1147.
[2] Jan Hendrik Prinz, Bettina Keller, and Frank Noé. "Probing molecular kinetics with Markov models: Metastable states, transition pathways and spectroscopic observables". In: Physical Chemistry Chemical Physics 13.38 (2011), pp- 16912-16922 ISSN: 14639076. Dor: 1039/
c1cp21258c
[3] Simon Olsson et al. "Combining experimental and simulation data of molecular processes via augmented Markov models". In: Proceedings of the National Academy of Sciences of the United States of America 114.31 (2017), pp. 8265-8270. ISSN: 10916490.
[4] Christopher Kolloff and Simon Olsson. "Rescuing off-equilibrium simulation data through dynamic experimental data with dynAMMo". In: Machine Learning: Science and Technology 4.4 (2023). ISSN: 26322153, DOL: 10. 1088/26322222 2153/ ad10ce
[5] Ignasi Buch, Toni Giorgino, and Gianni De Fabritiis. "Complete reconstruction of an enzyme- inhibitor binding process by molecular dynamics simulations".. In: Proceedings of the National Academy of Sciences 108.25 (June 201 1), pp. 101844 14119999
TSSN: (091-6490. DOT: 10. 1073/
pnas 1103547108
10

---

[6] Nuria Plattner and Frank Noé. "Protein conformational plasticity and complex ligand binding kinetics explored by atomistic simulations and Markov models". In: Nature Communications
6.1 (July 2015). ISSN: 2041-1723. DOT: 10. 1038/ ncomms8653)
[7] Kresten Lindorff-
et al. "How fast-folding proteins fold". In: Science 334.6055 (2011),
pp. 517-520.
[8] Frank Noć et al. "Constructing the equilibrium ensemble of folding pathways from short off-equilibrium simulations". In: Proceedings of the National Academy of Sciences 106.45 (Nov. 2009), pp. 19011-19015 ISSN: 1091-6490. DO1: 10. 1073 3003//pas..ar 0905466106
[9] M. Parrinello and A. Rahman. "Polymorphic transitions in single crystals: A new molecular dynamics method". In: Journal of Applied Physics 52.12 (Dec. 1981). pp. 7182-7190. ISSN:
1089-7550. DOL: 10. 1063/1011111 328693
[10] William C. Swope and Hans C. Andersen. -arrt-lrticlee molecular-dynamic study of homo-
geneous nucleation of crystals in a supercooled atomic liquid". In: Physical Review B 41.10
(Apr. 1990), pp. 7042-7054. ISSN: 1095-3795. DOE: 10. 1110///yyrrebbb....m 41 7042
[11] Pieter Rein ten Wolde, Maria J. Ruiz-Montero, and Daan Frenkel. "Numerical Evidence for bec Ordering at the Surface of a Critical fee Nucleus". In: Physical Review Letters 75.14 4 (Oct. 1995), pp. 2714-2717. ISSN: 1079-7114, 14. DOI: 10. 10t11100/phhh whtppphllsttttttt 75......22 77.244
[12] Hendrik Jung et al. "Machine-guider" path sampling to discover mechanisms of molecular self-organizationn. In: Nature Computational Science 3.4 (Apr. 2023), pp. 334-345. ISSN:
2662-8457. DOT: 10. (00003//438802220 023- -00-5888---2-002----
[13] Jérôme Hénin et al. "Enhanced Sampling Methods for Molecular Dynamics Simulations [Article v1.0]". In: Living Journal of Computational Molecular Science 4.1 (2022). DOI:
10 00.33011/1ivecom
4. 1. 1583)
[14] Esteban G. Tabak and Eric likelihoodd'. In: Communications in Mathematical Sciences 8.1 (2010), pp. 217-233.
"Density estimation by dual ascent of the log-
[15] Danilo Jimenez Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows.
2016. arXiv: 1505. 05770
[16] Laurent Dinh, Jascha Sohk-Dicks and Samy Bengio. Density estimation using Real NVP.
2016. eprint: arxiv: 1605. 08803
[17] George Papamakarios, Theo Pavlakou, and lain Murray. Masked Autoregressive Flow for Density Estimation. 2017. eprint: arXiv: 1705. .70557
[18] Durk P Kingma and Prafulla Dhariwal. "Glow: Generative flow with invertible 1xl convolu- tions". In: Advances in neural information processing systems 31 (2018).
[19] Qiang Liu. "Rectified flow: A marginal preserving approach to optimal transport". In: arXiv
preprint arXiv:2209. 14577 (2022).
[20] Yaron Lipman al. Flow Matching for Generative Modeling. 2023. arXiv: 2210 02747
[21] Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. "Stochastic Interpolants: A Unifying Framework for Flows and Diffusions". In: arXiv (2023).
[22] Alexander Tong et al. "Improving and generalizing flow-based generative models with mini- batch optimal transport". In: Transactions Machine Learning Research (2024). Expert Certification. ISSN: 2835-8856.
[23] Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. 2023. aivi 2306 [st
15030
[24] Jonas Köhler, Leon Klein, and Frank Noe. "Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities". In: P Proceedings of the 37th International Conference on Machine Learning. Ed. by Hal Daumé III and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, 13-18 Jul 2020, pp. 5361-53770
[25] Danilo Jimenez Rezende et al. Equivariant Hamiltonian Flows. 2019, arxiv: 1909 13739 [stat.ML)
[26] Kristof T. Schütt,, Oliver T. Unke, and Michael Gastegger. "Equivariant message passing for the prediction of tensorial properties and molecular spectra". In: CoRR abs/2102.03 02.00 150 (2021).
arXiv: 220.011220000 03150|
[27] Mario Geiger and Tess Smidt. e3nn: Euclidean Neural Networks. 2022. arXiv: 09453
LG)

---

[28] Andreas Griewank and Andrea Walther. Evaluating Derivatives: Principles and Techniques of Algorithuic Differentiation, Second Edition. Society for Industrial and Applied Mathematics. Jan. 2008. ISBN: 9780898717761
DOT: 10. 1137/1 1 9780898717761
[29] Ricky T. Q. Chen and David Duvenaud. "Neural Networks with Cheap Differential Operators".
In: CoRR abs/1912.00 2.93599 (2019), arXiv: 1912.03579
[30] Will Grathwohl et al. "Scalable Reversible Generative Models with Free-form Continuous
Dynamics". In: International Conference Learning Representations. 2019.
[31]
Seonghyun Park et al. Non-backtracking Graph Neural Networks. 2024. arXiv: 2310.07430
[cs. LG]
[32] Zhengdao Chen, Xiang Li, and Joan Bruna. Supervised Community Detection with Line Graph
Neural Networks. 2020. arXiv: 1705. 1705.08415 [stat.Mi.
[33]
Justin Gilmer et al. Neural Message Passing for Quantum Chemistry, 2017. arXiv: 1704
01212 [C]
[34]
Aram- -ram-Alexandre Pooladian et al. "Multisample Flow Matching: Straightening Flows with Minibatch Couplings." In: ICML. Ed. by Andreas Krause al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 28100-28127.
[35]
Leslie Kish. Survey Sampling. John Wiley & Sons, 1965.
[36] Selma Moqvist et al. "TThemmdyyammic Interpolation: A Generative Approach to Molecular Thermodynamics and Kinetics". In: Journal of Chemical Theory and Computation 21.5 (Feb. 2025), pp. 2535- 2545. ISSN: 1549-96228 DOT: 10. 1021/acs aac...c.. /e2///cs.cccc................... 4c011577
[37] Marta Skreta et al. The Superposition of Diffusion Models Using the Itô Density Estimator. 2024. eprint: 2412.
[38] M.F. Hutchinson. "A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines". In: Communications in Statistics Simulation and Computation 19.2 (Jan.
1990), pp. 433-450. ISSN: 1532-4141. DO1: 10. 1080 0000000310000 1011 1109999008111 2866
[39] Ryan P. Adams et al. Estimating the Spectral Density of Large Implicit Matrices. 2018. eprint:
arXiv: 1301.
[40] Laurence Illing Midgley et al. "SE(3) Equivariant Augmented Coupling Flows". In: Thirty- seventh Conference on Neural Information Processing Systems. 2023.
[41] Leon Klein et al. "Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Coare-Correeeed Dynamics". In: Advances in Neural Information Processing Systems. Ed. by A. Oh et al. Vol. 36. Curran Associates, Inc., 2023, pp. 52863-52883,
[42] Laurence Illing Midgley et al. "Flow Annealed Importance Sampling Bootstrap". In: The Eleventh International Conference on Learning Representations. 2023.
[43] Michael S. Albergo and Eric Vanden-Fijnden. NETS: A Non-Equilibrium Transport Sampler. 2024. eprint: :exivv24100 241.:.22110001111111
[44] Marta Skreta et al. Feynman-Kar Kac Correctors in Diffusson: Annealing, Guidance, and Product
of Experts. 2025. eprint: arxiv: 2503. 02819
[45] Charlie B. Tan et al. Scalable Equilibrium Sampling with Sequential Boltzmann Generators.
2025. eprint: arxaa: 2502 18462
[46] Bowen Jing et al. Torsional Diffusion for Molecular Conformer Generation. 2022. eprint:
arXiv:2206.0117:
[47] Juan Viguera Diez et al. "Generation of conformational ensembles of small molecules via surrogate model-assisted molecular dynamics". In: Machine Learning: Science and Technology
5.2 (Apr. 2024), p. 025010. ISSN: 2632-2153. DOI: 10. 1ww1100882222222 2///2/2222222222222222 ad3b64
[48] Sarah Lewis et al. "Scalable emulation of protein equilibrium ensembles with generative deep learning". In: (Dec. 2024). DOI: |0. 1101/ 12. 05.6888
[49] Bowen Jing, Bonnie Berger, and Tommi Jaakkola. AlphaFold Meets Flow Matching for Generating Protein Ensembles. 2024. eprint: arxi: 2402 240.244455
[50] Ameya Daigavane et al. JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensembles. 2024. eprint: arxiv: 2410. 14621.
[51] Qiwei Ye. "Towards Unraveling Biomolecular Conformational Landscapes with a Generative Foundation Model".. In: (May 2025). DOI: 10. 1101/2025 05. 01.01......0..1111....0..0000000 65 651643
[52] Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit Transfer Operator Learning: Multiple Time-Resolutie Surrogates for Molecular Dynamics. 2023. arXiv: 2305 18046
[physics. chem ph]
12

---

[53] Juan Viguera Diez et al. "Boltzmann priors for Implicit Transfer Operators".. In: The Thirteenth International Conference on Learning Representattonss 2025.
[54] Juan Viguera Diez, Mathias Schreiner, and Simon Olsson. Transferable Generative Models Bridge Femtosecond to Nanosecond Time- Step Molecular Dynamics. 2025, eprint: arXiv 2510. .75889
[55]
Tim Hsu et al. "Score Dynamics: Scaling Molecular Dynamics with Picoseconds Time Steps via Conditional Diffusion Model". In: Journal of Chemical Theory and Computation 20.6 (Mar. 2024), pp. 2335-2348. ISSN: 1549-9626. DO.:|10.
[56] Fabian L. Thiemann et al. Force-Free Molecular Dynamics Through Autoregresii Equivariant Networks. 2025. eprint: arXiv 2503, 23794
[57] Magnus Petersen, Gemma Roig, and Roberto Covino. DynanicsDiffusion: Generating and
Rare Event Sampling of Molecular Dynamic Trajectories Using Diffusion Models. 2024.
[58] Bowen Jing et al. Generative Modeling of Molecular Dynamics Trajectories. 2024. eprint:
arXiv: 2409. 17808)
[59] Pantelis R. Vlachas al. "Accelerated Simulations of Molecular Systems through Learning of Effective Dynamics". : In: Journal of Chemical Theory and Computation 18.1 (Dec. 2021), pp.
ISSN: 1549-9626. DOI: 10. 1021/ http//ccscctc.......o......o.c.c jctcjccocc..0.9
[60] Mhd Hussein Murtada, Z. Faidon Brotzakis, and Michele Vendruscolo. MD-LLL-11: A Large Language Model for Molecular Dynamics. 2025, eprint: arXiv: 2508 03709
[61] Mark Kempton. Non-bucktracking random walks and weighted Ihara's 's theorem. 2016. arXiv:
1603. 05553 [math.CO
[62]
M. E. J. Newman. Spectral community detection in sparse networks. 2013. arXiv: 6494
[63] Rongqin Chen al. "Redundancy-Free Message Passing for Graph Neural Networks". In: Advances in Neural Information Processing Systems. Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., 2022, pp. 4316-4327.
[64] Arthur Kosmala et al. "Ewald- based
Message Passing for Molecular Graphs". In: International Conference on Machine Learning (ICML). 2023.
[65] Bingqing Cheng. "Latent Ewald summation for machine learning of long-range interactions".
In: npj Computational Materials 11.1 (Mar. 2025). ISSN: 2057-3960. DOl: 10. 1038/ 1038//415544 #41224
025-01577 -1577777
[66] J. Thorben Frank et al. Euclidean Fast Attention: Machine Learning Global Atomic Represen- tations at Linear Cost. 2024. eprint: arxiv: 2412.08541
[67] Alessandro Caruso et al. Extending the RANGE of Graph Neural Networks: Relaying Atlention
Nodes for Global Encoding. 2025. eprint: arXiv: 13797
[68] Petar Velivckovié. "Message passing all the way up"" In: ICLR 2022 Workshop on Geometrical and Topological Representation Learning. 2022.
[69] Michael M. Bronstein et al. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges. 202 arXiv: 2104. 13478 [cs. LG]]
[70] Petar Velivcković et al. Graph Attention Networks. 2018. arXiv: 1710. 10903 [stat. ML]
[71] Jason Ansel al. "PyTorrh 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation". In: 29th ACM International Conference on Architec- tural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '24).
ACM, Apr. 2024. DOI: 10 1145/3620666 0tt55333626655334433666 3640366
[72] Michael Poli al. TorchDyn: Implicit Models and Neural Numerical Methods in PyTorch.
[73] Kristof T. Schütt et al. "SchNetPack: A Deep Learning Toolbox For Atomistic Systems". In
Journal of Chemical Theory and Computation 15.1 (2019). pp. 448 455. DOI: 10. 1021/acs
jctc. 8b00908, eprint: https: //ttps:///oi... .rgg 10. 1021 acs jctc 8b00908
[74] Kristof T. Schütt et al. "SchNetPack 2.0: A neural network toolbox for atomistic machine learning". In: The Journal of Chemical Physics 158. 14 (Apr. 2023), p. 144801. ISSN: 0021-
9606. DOI: 10. 1063/5. 0138367 eprint: https::/pubs aip wwtp:/ogg/aaaapeee http:/aaa///ppp/tillleee article- pdf/doi/10. 1063/5..0050 0188011113877 0003866776622667/14811111818 013.l6665..13367.pp.
[75] Stefan Elfwing, Eiji Uchibe and Kenji Doya. Weightee Linear Units for Neural Network Function Approximation in Reinforcement Learning. 2017,
eprint: arXiv: 1702
03118
13

---

Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 2017. arXiv: [76]
1412. 6980 .G]
14

---

A Adaptation of HollowFlow to Attention
attention can be considered special case of message passing neural networks. Thus, the adaptation is straightforward, we outline two possible strategies in this section.
As illustrated 69
First, replace the general message function & (eq. (14)) by product of the importance coefficients from self-attention mechanism
R x R R and a learnable function R - R
m a(hi (e)))
mly E
maag
hij (((h)).m))
(22)
(EJ)
N (i.j) is defined in eq. 28
Second, for graph-attention-t like mechanism [70], we define the line graph softmax function for graph G9 with the corresponding notion of nearest neighbors Nis.9(i,
exp (Visk)
softmaxiy(Ust Ele exp (yert)
(23)
The updates of HollowFlow then become
Oki
f.q
softmax (())))))
(24)
1.l W kikik Ojjj
(25)
(EA
where W ij E RaaX% is learnable weight matrix, : R - IR is a nonlinearity (applied element wise) and a : R" R - RR is a non-linear self-attention function, typically feed-forward neural network. One also generalize multi-head attention by independently executing eqs. (24)
and (55 possibly even on differently connected graphs 1 potentially circumventing expres-
siveness problems, and then concatenating the result:
f.g kij
softmax
by (a (a(W"h W)h(a)
(26)
||
Σ
aksy
Wifhh
(27)
q=1
(EA)l
Here, W and w E R Rm,.,n, are learnable weight matrices and H is the number of attention heads.
These formulations demonstrate that both self-attention and graph-attention layers be integrated into HollowFlow without altering its block-hollow Jacobian structure, enabling efficient divergence computations within attention- based architectures.
B Derivations and Proofs
B.1 Hollowness of HollowFlow
We first reintroduce some notation and do then restate algorithm a and theorem and proof it.
Notational remarks
By (1. and (t) we mean the following notions of nearest neighbors: ( j) =
E' Ai (k) E}.
(28)
(29)
Equation 28 is illustrated in fig. 4 The functions Φ. 3. tb, 3 and R used in algorithm | are all dependent back tracking array
learnable functions. B(t) E R is the ftree-dimensional (see eq. (6) and algorithm D. the integer / indicates the current message passing step. The binary variable pd [Irue, false] indicated whether the algorithms runs on the (translation invariant) pairwise differences or not. The detach() method means that derivatives w.r.t. the detached variable ignored (i.e., equal zero).
15

---

jk
(1,
kj
il
ja
jm
(,,,,)
m.j
778
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 AA (j. i) while the node encircled with a dashed red line corresponds to (i, Two edges of L(G) are labelled explicitly to illustrate our labelling convention.
Algorithm 1 Hollow message passing
Input: {xi, 2i) where x, IRd and Z, IR pd E (true, false)
Returns: (b:) E Rd
1: if pd then
2:
dig x; xi -xj
Eij = embed Z.)
3: else
4: 71, embed(x, 2,)
5: end if
6: Calculate G (N, E). e.g., as a kNN graph. Self-loops, i.c., (i, E are prohibited.
7: Calculate the line graph L(G)- = (N's, El9) and its node features lg.
Eij if pd n else
Nls E. E (i. (i, E, kj E and + kJ. n
8: Calculate initial node features h. and initial backtracking array B(0)
J( Σ )(n())) if pd kEA)
(1 (i. else) if pd
B(0)
n else
(1 =,, 0 else) else
9:
for +- 0 to - do
Remove appropriate edges of (G): El < E ((i,j,)) E'9 B())))) 1}
11:
Update backtracking array: (t + 1)ijk maxieN B(t B())ss})
12:
Do message passing:
mkk = o(h hlg), mly EKENTO(US) mayg
end for
14:
Perform readout and project back to G:
(xbbed(x,detacll () xj, Zi) if id
by EieNG) R(IT") mig) where nij
nj
else
Theorem (Block-hollowness, of HoMP). Algorithm 1 defines a function b Rn Rd whose Jacobian Jaxe bxx) E Rdnxdn can be split into a block-hollow and a block-diagonal part, i.e.,
Jucx) = Jeck) + Jr(x)
(30)
where J Jeix) block-hollow, while Jryx) is block block-diagonal with block size d respectively. This structure of the Jacobian enables the exact computation of differential operators that only include diagonal terms of the Jacobian with instead of nd backward passes through the network.
Remark: Figure 4 might be helpful illustration to keep track of the indices when reading the proof.
16

---

Proof. By the chain rule step 14 gives
dbs Σ
aR
0h T
0R Ōnij
(31)
dx iEW(s)
OX
Omay
8x
If show that 2 [I,., n the Jacobian decomposes as stated. We split the proof into two cases:
hig
0 vt 2 0 and
0f +
for all (i,j) e E and for all & E
OxL
Case I: pd = false:
That
Ong 7hc
if a & is trivial for all ni as Ray = embed(x,, 2j).
We show dx,
M = 0 for all t 2 0 and for all (i.j) E by induction on t:
Base case:
1. h exbed Z;) is by definition independent of x; for all (i,j) E.
2. Furthermore, B(0)a is by definition one exactly when hi depends and zero other- wise (see step 8).
Induction step: Assume
1. big is independent of x, for all (i,j) E E and
2. B(t)i is zero exactly when hy is independent of X2-
We want show that these conditions also hold for t + 1.
To see that h4t1 is independent of Xj for all (i,j) E E, by step 12 of algorithm and the first induction hypothesis, only need to show that hig does not depend Xj for all k E (i.)) Assume h, does depend on Xj. By the second induction hypothesis, this implies (tt) = 1 and by step 10 of algorithm thus (Is, i. & E This means & (i,j), a contradiction.
if and only if hij" independent of x&) by proving both Next, we prove (B(t + 1)ijk implications:
B(t + 1)sk =0 = x++ is independent of xx:
hg Assume Biju (t 1) 0. By step 11 of algorithm this implies B(t)aa = OVIE (is j) and B(t)ijk = 0. By the second induction hypothesis Tollows that hi and hiy are independent of x for all 1 (i.(() which by step 12 of algorithm implies that 1+1 is independent of x
h is independent of X B( 1)izk 0:
Assume h4t1 is independent of X By step 12 of algorithm 1 this implies that hi and Mi are independent of x₁ for all E (i) By the second induction hypothesis and step 11 of algorithm ] we get Bt + 1)ikk = O as desired.
Case 2: pd = true:
On. 0 j * le is trivial for all {1,..., n. as nij embed (x;.ddxccdeeech () X,, Z,)) and That 8xg the detach() method means that we ignore derivatives w.r.t. the detached variable.
We again show 2x, compared to case is the base case, h is initialized differently.
_h hi for all 2 0 and for all (i,j) E E by induction on The only difference
17

---

Base case:
1. b is function of nig where & (i,(i. As k j and i j. Hl is independent of x₃-
2. We need to show B()) is zero exactly when hi is independent Xx- Assume B(0) =
0. Then & & (i. by step 8. Thus, by definition of kij (step 8), 8), A2 is independent of
X5. Conversely, n. be independent of Xg. Then k & (i, meaning B(0) by step 8.
Induction step: The induction step is the same as in case 1.
Jacobian diagonal with d backward passes: We can now exploit the decomposition eq. 31 of
the Jacobian to get its diagonal terms with only d vector-Jacobian products (backward passes).
This can be achieved by applying the detach() () operation to h in step 14 of algorithm such
that the block-diagonal part of the decomposition eq. (31 remains unchanged while the hollow part vanishes. we construct d column vectors {rz E Rdn as
{
if = i d(k 1), k E (1......)
(ot)
(32)
else.
The vector-Jacobian products JJ do then give us access to all diagonal terms of Ja, e.g.. 1 for illustration in the
tr(J;) E of Jon (see fig.
5, d = 3).
Further remarks The conditioner c(x) and the transformer Th,, x) shown fig. can explicitly be written component-wii wise follows:
(x = (he heNv)) == =h
(33)
(),,,,,))= Σ
R(K T
(34)
(ENU)
where hTy"" and niy are defined in algorithm
Whenever directed graph G acyclic, its directed line graph T.(G) will also be acyclic which implies that even without edge removal in L(G) (step 10 in algorithm message passing on the line graph ensures that information starting from node i of G can of message passing steps. If the shortest cycle in G' has length (i.e., G has girth g), one can do at most g 2 message passing steps without returning information if edges of LGG) are removed. The shift -2 comes from the fact that initialization of the node features (step 7 in algorithm I and the readout (step in algorithm effectively bridge one edge of G each.
return itself after arbitrary number
B.2 Computational Complexities
Theorem { Consider a GNN-based flow model with a fully connected graph G and message passing steps and a HollowFlow model with a neighbors graph Gg and T message passing steps. Let both graphs have nodes and d-dimensional node features.
The computational complexity of sampling from the fully connected GNN-based flow model including sample likelihoods is
( =
18
while the complexity of the HollowFFow model for the same task is
RTsteP (L(G)) = (n 222 + dk)).
19
Moreover, the speed-up of HollowFlow compared the GNN-based flow model is
RT (G = =0 0 Tn?d
20
RT (T(te((((Gr)
T + dk
18

---

Proof. The sampling procedure involves a forward pass and one divergence calculation for each integration step. We start by estimating the runtimes of the forward pass RT'. The key assumption is that the computational complexity of the forward pass is proportional to the number of edges of the graph (see table D and the number of message passing steps.:
RT (c) = O(T# = O(Tn²)
(35)
= ( nk
(36)
The computational complexity of the divergence calculation RT will be proportional to the number of backward passes necessary to compute the divergence times the complexity for each backward pass. For the fully connected GNN we need dn backward passes at the cost of Eyo:
RT (G)) OT Fyydn) OTn''d).
(37)
For HollowFlow we only need backward passes by exploiting the structure of the Jacobian (See the- orem[]. []. Each of these backward passes has a complexity O Ek), we only need to backpropagate through the transformer, not the conditioner:
RT (L(Gx)) = O## Egd) = (dnk))
(38)
The runtime of one integration step. RT that is proportional to the total runtime, will scale as RT + RTV Thus, we get
(TTte = (GJO) + = + O(Tná))) + O(dnk) = =n + dk)),
(39)
(I((k)) = ( (I(G)) RT
(40)
where the right term the sum in eq. (39 dominates as (n > (n and constant
independent of n and T. The speed-up during sampling with likelihood evaluation when using HollowFlow compared to a standard, fully connected GNN thus
RT (G O
Tnad
Tnd T
(41)
RTStep (I(G))
Tlgk2 + dk
where the last equality assumes k = (v)) (see section 4.1 In conclusion, we expect a speed-up
of On²²) for constant k and a speed-up of ((n) for k = O(vn)
B.3 Memory scaling of B(t)
From the definition of (tt) (eq. (16)) one might expect a cubic memory scaling in However, it is sufficient to only save information about from which nodes G all the nodes of L(G) have received information from. If G has nodes and L(G) has Mig nodes, the size of the array will effectively be nng- For a &NN graph, nig = O(nk) so the memory requirements of Bt)) in our implementation scale as n n2
C Additional Results
C.1 Non-equivariant HollowFlow
Adaptation of HollowFlow to non-equivariant GNN architectures is straightforward as HollowFlow only affects the construction of the underlying graph while preserving the message and update functions of the GNN including their equivariance properties. Relaxing permutation symmetry can be done by using unique embeddings for all nodes the graph (as done for Alanine Dipeptide, see appendix
C.2 Multi-head HollowFlow
Using a ANN graph poses a locality assumption that might limit the models ability when learning the distribution of systems with long-range interactions (e.g., systems with coulomb interactions such as molecules). To circumvent the locality assumption while still keeping the computational cost manageable, we additional consider multi-head strategy with H heads. The central idea is to run HoMP on H different graphs parallel and sum up the result. These graphs can be constructed in different ways, explore the following two strategies:
19

---

Table LJ13
ESS (%)
ESSrem (1) (%)
EffSU (1)
EffSU
0.054 00.00 068
2.92 95
1.059 .770
1.16
.049
=
0.041 0.17
0.001
10.3410:29
10.28
2.649 6363
667
0.64
6
3.300 301
20.20 29
3.260 9.278
31.877.77
2.745 250
20.6420072
2.310
18.26
1.0570.803
16.50
1,484|4492
5.519.19
1.475
12
4.0693.260 926
19.7219.63
1.627 636
20.11 35.55
H = 2
2.199
1.783
17.55 17.64
17.47
1.613
12.08 20.68
H 3
0.741
13.78 13.7813:75
1.339
4.17
H
742 0.334
11.32
1.064
1.85
Н = 2, =
2.86 2044
23.3523.25
23.44
0,864
16640.868
8.53 1.98
H = 3, 3
2.484
19.0618:95
19.14
1.206 1.200
1221
9.23
15,60
H 4,1 =1
0.322
11.52]1.58
11.5211.46
0.845
1.42
И 4, =
2.2397.541
15.3215.40
0.849
0.811
7.40
baseline k = 6
0.002
003
43.5243.63
1.010
0.00
baseline
2.132
40.73
1
Non-overlapping scale separation: Consider fully connected graph G embedded in euclidean space with F. edges and sort the edges by length. Divide this list into H equally sized chunks of length #E/F with no overlap. If there is a remainder, distribute the remaining edges the heads equally possible.
2. Overlapping scale separation with overlap number J: Consider fully connected graph G embedded in euclidean space with # E edges and sort the edges by length. Divide this list into H equally sized chunks of size #E/(H I) such that the centers of these chunks are spaced evenly. If there is remainder, distribute the remaining edges to the heads as equally possible.
Both strategies ensure that there is no locality assumption. The first strategy assumes scale separability while the second one allows for (limited) scale overlap. In terms of scaling, observe that # On²)) for a fully connected graph. Thus, each head has (nn2/ (H 7)) edges giving us on average Onn/ (( - I)) edges per node. Finally, the line graph of each head has thus O(n"/ - 2) edges, giving us O(Hn² (H ))²) line graph edges total. Thus, if we, e.g., choose the number of heads H and the overlap number I proportional to the scaling is (n²)).
We both of the aforementioned multi-heat strategies on LJ13, LJ55 and Alanine Dipeptide. The results together with the single-head results from the main paper can be found in tables 4 to
C.3 Alanine Dipeptide
We additionally trained HollowFlow and corresponding baseline on Alanine Dipeptide using number of different graph connecting strategies (multi headed and single headed). For all models, we break permutation equivariance. Interestingly. the overlapping multi-head strategy with H and I = seems to perform best in terms of ESS,em out of all HollowFlow models, suggesting that multi-head approaches might be a promising direction to alleviate limitations of HollowFlow. However, we do not observe effective speed-up of HollowFlow compared to the hollhollow baseline. We believe that further hyperparameter tuning and engineering efforts might close this performance gap.
20

---

Table LJ55
ESS (%)
rem (1) (%)
EffSUrem (t)
EffSU
k 7
0.006 10060.00 010
0.53g.51
03.7378 48.484
82.26
003
k = 27
0.007 012
0.64
9.466 986
8.463.31
003
8.965
k = 55
0.020 025
0.74
0.77
4.3654 144
683
9.11 13.18
H 2
0.008 003
0.53
5.259
6.442.07
H
0.006 009
0.71
7.198 7.576
4.861.86
003
H 6
0.009 013
0.480.45
5.901 6224
9.39 33933.8
H = 8
0.015 020
0.53
7.144 7-599
15.41
23.41
6.711
6.96
H 10
0.006
0.53
7 7.604
7.02 .252
baseline 27
0.0070.012
0.64
0.3240-342
0.28
baseline
0.048
0.0480.051
2.96
1
1
Table ALA2
ESS (%)
ESS (t) (%)
EftSU (t)
EffSU
2
0.008 0.0089.011 011
005
0.688.70
0.65
0.566 589
0.544
3.63
H = 3
0.01
0.23 0.22
0.198 0-211
4.97
H 4
0.006
0.11
0.1149.125
2.71
5
0.003
0.04
0.042
1.433.24
= 3, =1
0.017
0.73
0.3990
3.64
k = 11
0.010
0.52
0.588 0.550
4.52
0.29
baseline
0.135
224
16. 16.14.0.29
15.97
1
1
D Experimental Details
D.1 Code and Libraries
All models and training is implemented using PyTorch with additional use the following libraries: bgflow [1, 24). torchdyn (72). TorchCFM 20 and SchNetPack 73 74] The conditional
flow matching and equivariant optimal transport implementation based on the implementation used
in (23) J. The implementation of the PaiNN (26) architecture is based [73 74). The code and the
models are available here: http::///gttttth...oooooooooooo wtwp:/ssssooooooooooooooooooo
D.2 Benchmark Systems
The potential energy of the Lennard-Jones systems LJ13 and LJ55 is given by
12
6
(x)
Σ 1.j
Tim dig
T'm
(42)
2T
di
where the parameters are as in [33 Im = 1, c and = 1. dej is the pairwise euclidean distance between particle i and j. Detailed information about the molecular system (Alanine Dipeptide) can be found in (23].
D.3 Training Data
used the same training data as in (23 available at https:///of...ooo///////eeeee fnll
The training data has been generated using MCMC,
details can found in the appendix [23] For systems we used 10 randomly selected samples
for training as in 23 Validation was done using 10º of the remaining samples.
21

---

D.4 Hyperparameters
Choice of the Hyperparameter k Heuristically, we expect maximal expressiveness for mid-range values of k: By construction, some of the connections in the line graph need to removed after certain number of message passing (MP) steps (See appendix B for more details). The larger k, the more connections need to be removed after only one MP step up the point where (almost) no connections are left after only one MP step. We thus have two counteracting effects when & increases: The first MP step gains more connections and thus expressiveness while the second MP step and all later MP steps loose connections and expressiveness. This effect is illustrated in fig. 5 Eventually, latter seems to take over limiting the expressiveness of the model for large values of k. This non-linear qualitative scaling of performance with le is experimentally supported by the results in table
10a
k
k
k
k=
10a
k
10 k
==
#connsctions
11
10a
10
0
T
Figure 5: Average number of connections left in the line graph of a ANN graph G as a function the number of message passing steps. G has 55 nodes whose coordinates sampled from standard Gaussian. The average is taken over 10 independently sampled graphs. The larger k; the faster the line graph disconnects.
Discontinuities Caused by ANN Graph The vector field bz(z) (eq. (3) is constructed using KNN graph. This might cause discontinuities on the decision boundaries of the ANN construction. As the divergence of the vector field is only defined if it is differentiable, these discontinuities can cause theoretical issues, including but not limited to the continuity equation not being well-defined everywhere eq. (4). We have observed any numerical instabilities other issues in practice. To further validate empirically that the discontinuities do not cause problems, we tracked the divergence of trained, fully connected well trained, neirest- HollowFlow IJ13 model during sampling using a forward Fuler scheme with 20 integration steps. We monitor the maximum and average absolute difference of the divergence of consecutive steps, computed over 1000 initial conditions, to assess whether there are significant discontinuities in the divergence of the nearest- neighbors model compared to the fully connected one. Results can be found in table 7 While both values are slightly higher for the &NN model, there are no significant differences further supporting our claim that the discontinuities do cause issues in practice.
Table 7: Maximal and average absolute difference of the divergence of consecutive steps, computed over 1000 initial conditions, for two different LJ13 HollowFlow models.
model
maximal difference
average difference
fully connected
72.5
5.3
t = 6
74.2
5,8
22

---

Hyperparameters Used in Experiments: Depending on the system, we used a different number of message passing steps and a different value for the hidden dimension n2, following (23). For IJJ3, we chose nt = 32 for all runs. All HollowFlow experiments with IJ13 and kNN graphs used two message passing steps while the baselines and the multi-head experiments used three. For IJ55, we chose n4 = 64 for all runs. All HollowFlow experiments with LJ55 and &NN graphs used two message passing steps while the baselines and the multi-heal experiments used seven. For Alanine Dipeptide, we chose n = 64 and five message passing steps for all runs. All neural networks in the PaiNN architecture use the SiLU activation function [75].
The training details for LJ13, LJ55 and Alanine Dipeptide are reported in tables g to 10 For LJ13, we selected the last model for all &NN experiments and the fully connected baseline experiment, while we selected the model with the lowest validation loss for all remaining experiments for inference. For IJ55 and Alanine Dipeptide we always selected the model with the lowest validation loss for inference. All training was done using the Adam optimizer [76].
Generally, we observed that HollowFlow can be trained with larger learning rate compared to the baseline without running into instability issues. Nevertheless, the effect on model performance seemed limited.
Table 8: Training details LJ13
batch size
training time (h)
learning rate
epochs
b
1024
10
1000
5.8
h 4
1024
5 X 10
1000
6.3
=
1024
X 10
1000
6.7
1024
10
1000
7.3
è = 10
1024
5 X 10
1000
7.5
A = 12
1024
X 10
1000
7.9
H 2
1024
10
1000
9.9
H = 3
1024
5 X 10
1000
10.3
= 4
1024
X 10
1000
10.7
Н 2, I 1
1024
10
1000
13.1
H 3,
1024
5 X 10
1000
11.0
= 4,1
1024
5 X 10-4
1000
11.3
H =441
1024
5 10 10
1000
12.7
bascline k = 6
256
5 X 10
1000
2.6
baseline
1024
5 X 10-4
1000
2.5
Table 9: Training details LJ55
batch size
training
learning rate
epochs
time (h)
256
(5 10 to 10 x 10-5
(150, 850)
10.5
27
70
(5 x 10 to 10-5,5 10 50-5)
(150, 69)
21.9
55
5 10
110
32.0
Н 2
128
(5 10-3 10 10
(150, 850)
16.7
И = 4
128
10=0 3
10 10"4)
(150,
16.8
Н =
80
10= 10-4)
(150,260)
11.8
H
128
x 16x
(150,850)
14.0
H = 10
128
(5 x 16x1103105xx000..,
(150, 850)
14.0
baseline k
1024
10001055 10-5)
(150, 850)
6.6
baseline
256
10 1o 10 10x 10-5)
(150, 844)
24.0
All inference details for IJ13, LJ55 and Alanine Dipeptide are reported in tables 11 to 13 The
integration of the ODE (eq. (3) was performed using a fourth order Runge Kutta solver with fixed step size. We used 20 integration steps.
23

---

Table 10: Training details Alanine Dipeptide
batch
training
learning rate
epochs
size
time (h)
H 2
100
(5 x 10 to to
10 x 1 10-4)
(150, 850)
20.3
H = 3
100
(5 x 10 to 10xx1000xx00 10-4)
(150, 850)
17.2
= 4
128
(5 x 10 t0 10 10"4)
to5
(150, 850)
16.7
H 5
128
(5 x 10-3 10-1.xx 10-1
(150, 850)
16.8
H = 3, =
80
(5 X t 10-1,5 10-1)
(150, 260)
11.8
k = 11
128
(5 X 10 to 10 x 10=4)
(150, 850)
14
baseline
1024
X 10-4
1000
6.6
Table Sampling details LJ13
batch size
GPU
GPU mem.
RT (s)
RT (s)
RTV (s)
# samples
usage (%)
usage (%)
k
407
367
29
40000
71
66
2 x 10
k 4
674
577
83
1024
61
92
2 x 10a
k 6
887
790
71
3500
73
67
2 x 10a
k
1141
1023
86
3500
46
2 10a
A 10
1411
1275
97
3500
82
86
2 10a
k = 12
1532
1389
103
3500
83
84
2 x 10a
H
1723
1527
159
1024
73
2 10a
H
1886
1682
165
1024
57
42
2 10S
Н 4
2050
1833
174
1024
54
26
2 2 10S
H 2,
3581
3261
242
1024
64
2 10
H 3,
2488
2221
207
1024
66
96
2 10
H 4
2395
2135
202
1024
60
53
2 10
H = =4,1=
3197
2856
257
1024
52
10a
baseline k =
4599
36
4450
12000
98
64
2 103
baseline
2508
29
2425
1000
85
18
1 1 10a
Table 12: Sampling details LJ55
batch size
GPU
GPU mem.
RT (s)
RT (s)
RTV (s)
# samples
usage (%)
usage (%)
k 7
1370
1147
171
256
73
56
10a
k 27
12609
11820
599
70
96
61
4 X 10
k 55
31439
30028
1114
30
96
83
x 10*
22602
20239
2173
79
67
104
H
20048
17897
1902
15
87
87
4 X 10
H
16586
14618
1707
18
87
84
x 10ah
H
14960
13097
1596
20
87
4 10a
H 10
13725
11982
1445
25
90
80
4 10a
baseline k = 27
353693
348
351221
256
99
65
4 X 10a
baseline
529793
596
525997
250
99
82
4X 10a
D.5 Errors
All errors were obtained 68% symmetric percentile around the median from bootstrap sampling using 1000 resampling steps. The reported value is the mean of these bootstrap samples.

---

Table 13: Sampling details Alanine Dipeptide
batch size
GPU
GPU mem.
RT (s)
RT (s)
RT (s)
# samples
usage (%)
usage (%)
H
2791
2382
370
100
70
79
5 x 10*
H
2532
2123
369
100
75
81
5 X 10
H =
2167
1837
284
128
72
92
5 x 10a
H
2182
1848
285
128
85
5 10*
H 3, =
3734
3236
431
80
79
91
5 x 104
= 11
1838
1633
178
128
78
68
5 x 104
baseline
26554
71
26094
1024
60
5 10*
D.6
Runtimes
The runtimes were computed by directly measuring the wall-clock time of the computation of interest. Due to missing data, the runtimes of the forward and backward pass in fig. [(b] needed to be extrapolated from different run with the same model by assuming the same ratios to the known total runtime.
D.7
Weight Clipping
As described in section of the log importance weights, log to make the ESS and EffSU estimations more robust. The corresponding quantities are named ESSrem and EffSUrem- While these metrics biased, this procedure was necessary to obtain a reliable and robust estimate of the effective speed-up of hollow flow as reflected the errors of the right most column of tables 6o
remove left and right percentile (i.e., one percent each side)
D.8 Computing Infrastructure
All experiments were conducted on GPUs. The training and inference for LJ13 and Alanine Dipeptide was conducted on NVIDIA Tesla V100 SXM2 with 32GB RAM. The training for LJ55 was conducted on a NVIDIA Tesla A100 HGX with 40GB RAM, inference was performed on a NVIDIA Tesla A40 with 48GB RAM.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main claims are that the number of backward passes to evaluate the diagonal of the Jacobian HollowFlow is constant in and that this results in dramatic speed-up when evaluation sample likelihoods. This is shown theoretically in section B and demonstrated in practice in section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We included an extra section where we discuss limitations and potential issues of our newly introduced method (section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
26

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We have two theoretical results (theorems appendix
| and 2 that are both proved in
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We provide all details about the experiments in appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
27

---

reproducible, they 5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: The training data is publicly available as stated in appendix D The code, including instructions how to use it, will be made public.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc.
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https for more details.
/nips. wttpcc/pcccccpcccccooooooop httpi//gggudeee/ddSSuuubbiiiii wwweiisssonPoPiiiiooooo
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All details can be found in appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Errors are provided in all tables except for runtimes that would require several runs which is impractical due to the huge computational overhead. Details how those errors were calculated can be found in appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
28

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Information about all compute resources that have been used are provided in appendix D
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't itdn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpps://eeuriss..ccppbbb////ees (et.eiieeeeEiiccccccdeeeeieee
Answer: [NA]
Justification: The paper describes methodological contributions of broad interest and with no immediate ethical concerns.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper describes methodological contributions broad interest and with immediate societal impact.
Guidelines:

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper describes methodological contributions broad interest with immediate concerns about misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Details about the datasets and code used, including creators and license, are reported in appendix
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
30

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thood cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our implementation of HollowFlow and the models we trained are detailed in No new datasets have been created in the course of this work.
appendix D
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer:
[NA]
Justification:
No crowdsourcing nor research with human subjects is involved in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No crowdsourcing nor research with human subjects is involved in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
31

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: All method development was done without the help of LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
32


Reviews:
[{"review_id": "Fic1tIfFsn", "reviewer": "Reviewer_PymQ", "summary": "This paper proposes HollowFlow, a framework that combines an non-backtracking GNNs and some NNs to speedup sampling and likelihood evaluation in previous BG methods. Specifically, HollowFlow's non-backtracking GNNs enforce block-diagonal Jocabian structure, which effectively reduces the number of backward pass. Experiments conducted on LJ13 and LJ55 further validate the effectiveness of HollowFlow.", "questions": "See above", "limitations": "See above", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. HollowFlow innovatively introduces HoMP to effectively reduces the number of backward passes required for the likelihood evaluations. KNN graphs are also incorporated to reduce the graph complexity.\n\n2. The paper develops theoretical scaling laws that aligns well with experimental results.\n\n\nWeaknesses:\n1. Although the paper states that:\n\n> the performance of our baseline is not in line with state-of-the-art in terms of ESS, however, since our HollowFlows are built directly upon this baseline, we anticipate that any improvement in the baseline will be reflected in the HollowFlow models as well.\n\nNevertheless, the lack of direct comparison with state-of-the-art methods reduces the persuasiveness of the results regarding HollowFlow's performance. I suggest that the authors supplement their experiments with additional benchmarks against leading approaches.\n\n2. While HollowFlow achieves remarkable speed-ups compared to the baseline, the ESS and ESS_rem is noticably lower. Furthermore, increasing the value of $k$ does not seem to improve HollowFlow's ESS and ESS_rem beyond a certain upper bound, and this upper bound remains significantly lower than that of the baseline. What is the rationale of this phenomenon? and How to tune HollowFlow in practice?\n\n3. The introduction of KNN assumes the locality of the graph, as stated in the limitation section. How does HollowFlow perform with a fully connected graph? Additionally, I recommend including an ablation study to disentangle the contributions of HoMP, KNN, etc.\n\n4. Experiments are conducted solely on LJ systems. Thus, the generalizability of the proposed model seems not sufficiently validated.", "was_revised": true, "final_justification": "most concerns have been solved. i raise the score.", "rebuttal": "We thank the reviewer for the valuable and very detailed feedback. Answers to the reviewer's questions and comments can be found below.\n\n1. Thanks for your suggestion, a direct comparison against a state-of-the-art method might indeed be interesting. However, we want to emphasize that our goal is not to beat the state of the art in terms of ESS but to demonstrate a significant improvement in terms of ESS per compute during sampling for a given architecture. In our experiments, we decided to use the PaiNN architecture as a base to implement HoMP even though state-of-the-art results have been achieved with other architectures like the architecture, that we will refer to as EGNN in the following, used in the Equivariant Flow Matching paper (Klein et al., 2023). HoMP can also be implemented with EGNN but in contrast to PaiNN it requires some arbitrary symmetry breaking on the line graph as all terms in the message passing will cancel out otherwise. While a base line using EGNN achieves an ESS that is close to the state-of-the-art, it seems to lower performance of HollowFlow, presumably due to the aforementioned symmetry breaking. Thus, we stuck with PaiNN for our baseline and HollowFlow, even though the PaiNN baseline is not in line with state-of-the-art methods.\n    \n\n2. HollowFlow uses a restricted architecture that we expect to have a limited expressiveness resulting in a lower ESS. However, as pointed out above, our goal is to improve the ESS per compute during sampling which we still achieve due to enormous computational savings. The fact that the ESS does not increase further with $k$ beyond some upper bound can heuristically be explained as follows: By construction, some of the connections in the line graph need to be removed after a certain number of message passing (MP) steps (See appendix B for more details). The larger $k$, the more connections need to be removed after only one MP step up to the point where (almost) no connections are left after only one MP step. We thus have two counteracting effects when $k$ increases: The first MP step gains more connections and thus expressiveness while the second MP step looses connections and expressiveness. Eventually, the latter seems to take over effectively limiting the expressiveness of the model.  We will include a brief summary of this explanation in the camera ready version to improve clarity.\n\n   As stated in table 3, we found empirically that HoMP with relatively small $k$'s ($k=6, k=8$ for LJ13) seems to work about as well as fully connected HoMP ($k=12$). To (heuristically) tune $k$ in practice, one can efficiently numerically work out how many edges of the line graph are left after $T$ message passing steps without training a model. There should be a significant amount of edges left after one MP step and the graph should not be extremely sparse in the beginning. We will include a plot showing the percentage of edges remaining as a function of the number of message passing steps for different $k$’s in the appendix of the camera ready version together with an instruction on how to tune $k$.\n\n   In practice, one could also use a multi-head strategy where several relatively sparse HoMP with different connectivities (e.g. separated length scales) are run in parallel and combined in the end (see appendix A for a brief discussion). This makes sure that each network does not get disconnected as fast while not having to do any locality assumption. This introduces an additional scaling with the number of heads while the scaling with the number of particles remains the same for each head. We will include a more detailed discussion of multi-head HollowFlow in the camera ready version.\n\n3. While we indeed assume locality of the graph, it is possible to circumvent this assumption by the multi-head approach described above. In Table 3 we show how HollowFlow performs with a fully connected graph on LJ13 (see row labelled $k=12$). This can already be seen as an ablation to dissentangel the contribution of the locality assumption. We will add a similar study for LJ55. Running HollowFlow without HoMP is not possible as HoMP is inevitable to achieve the hollow structure of the Jacobian unless we only do one message passing step. We will include an additional ablation study in the camera ready version where we run the baseline with a KNN graph.\n\n4. It would indeed be desirable to test the method on other systems. We are currently working on running HollowFlow on Alanine Dipeptide and will include the results in the camera ready version."}, {"review_id": "nIlKoDKvM5", "reviewer": "Reviewer_tKoR", "summary": "This paper introduces \"HollowFlow,\" a novel continuous normalizing flow (CNF) architecture designed to drastically reduce the computational cost of likelihood evaluation for generative models, particularly in the context of Boltzmann Generators (BGs). The authors' key contribution is a new message passing scheme called Hollow Message Passing (HoMP) addressing the prohibitive scaling of computing the divergence of the vector field in CNFs, which traditionally requires $O(N)$ backward passes.\n\nHoMP is implemented using a non-backtracking graph neural network (NoBGNN) that operates on a line graph representation of the system's connectivity. This construction splits the Jacobian of the neural network vector field into a block-diagonal part and a \"block-hollow\" part, which allows for the exact computation of the Jacobian's trace (the divergence) with a constant number (dimension of coordinates) of backward passes.\n\nThe authors provide a theoretical analysis of the computational complexity, showing that HollowFlow can achieve a speed-up of up to $O(n^2)$ over standard equivariant GNN-based flows. They validate their approach by training E(3)-equivariant HollowFlow models on two Lennard-Jones (LJ) particle systems, LJ13 and LJ55.", "questions": "- The appendix clarifies the update rule for the backtracking array B(t), which is of size $n^3$. Could you comment on the practical computational and memory cost of storing and updating this array in your implementation? Could it become a practical bottleneck for systems with thousands of nodes, and do you see a path to optimizing it?\n- Your results show a non-monotonic dependence of model performance on the hyperparameter 'k'.  Do you have any intuition or a proposed heuristic for selecting an optimal 'k' for a new system without having to train multiple models? Is there a risk that for a complex landscape, the optimal 'k' would need to be so large that it negates the computational advantage of using a sparse graph?\n- Could you elaborate on the training of the baseline model? Given its very low ESS, do you believe it was fully converged? How does its ESS compare to other published flow-based Boltzmann Generators for these LJ systems? A stronger baseline would make the \"effective speed-up\" claims even more compelling.\n- This work is motivated by achieving exact divergence calculation efficiently. How do you see your method comparing to those that use stochastic trace estimators (e.g., Hutchinson's estimator)? While approximate, they are also computationally cheap. In your opinion, what are the key applications or systems where the exactness of the divergence provided by HollowFlow is critical and approximation is insufficient?", "limitations": "The authors provide an honest discussion of the limitations in the main paper, which I find accurate. Most of my main concerns align with theirs, and are:\n- The current kNN-based implementation assumes locality and is not well-suited for systems with critical long-range interactions (e.g., electrostatics).  This significantly limits its immediate applicability to a wide range of important molecular systems.\n- The HoMP algorithm relies on updating a backtracking array, B(t), to dynamically remove edges from the line graph at each message passing step.  This array is of size $O(n^3)$. This could become a new computational or memory bottleneck for very large systems (e.g., n > 1000), potentially undermining the overall speed-up. The appendix provides the algorithm but does not analyze the practical cost of this step.\n- The core idea of HoMP is to restrict information flow to gain computational efficiency.  This is an explicit trade-off. It remains an open question whether this restriction will prevent the model from learning the complex energy landscapes of more challenging systems, leading to a sample quality (ESS) so poor that even a massive speed-up cannot compensate for it.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths:\n- The proposed HoMP algorithm is a clever adaptation of the \"hollow network\" concept to the crucial domain of equivariant GNNs. This offers a direct and effective solution to the divergence calculation bottleneck in CNFs.\n- The paper is well-supported by theory and offers step-by-step proofs in the appendix.\n- The experimental results are impressive, demonstrating a speed-up of nearly 90x for the LJ55 system. \n- The paper provides code and all training details for reproducibility.\\\n- The paper provides an extension to attention architectures offering integration with various SOTA architectures.\n\n### Weaknesses:\n- A notable weakness is the performance of the baseline model, particularly in terms of the effective sample size (ESS). For LJ13, the baseline achieves an ESS of only 0.001%, and for LJ55, it's 0.167%.  These are very low values, suggesting the baseline model is not particularly strong. While the authors acknowledge this, it tempers the impact of the \"effective speed-up\" metric, as the comparison would be more compelling against a state-of-the-art baseline with a higher ESS.\n- The experiments are confined to two relatively small, well-behaved systems (LJ13 and LJ55) that are dominated by short-range interactions. The paper's claims of general applicability to \"high-dimensional scientific problems\" would be much stronger if tested on more diverse and challenging systems.\n- The use of a kNN graph is a necessary compromise to avoid the poor scaling of the line graph construction on a fully connected graph.  However, the choice of 'k' is a critical, non-trivial hyperparameter. Figure 3(c) shows that performance is quite sensitive to 'k', and the appendix does not offer a heuristic for selecting it, adding a layer of tuning complexity.", "was_revised": true, "final_justification": "The authors address most of my concerns. However, I'm still concerned about the scalability of the method.", "rebuttal": "We thank the reviewer for taking the time to give valuable and very detailed feedback. Answers to the reviewer's questions can be found below. Finally, we additionally comment on some limitations that the reviewer pointed out.\n\nQuestions\n\n1.\n>The appendix clarifies the update rule for the backtracking array B(t), which is of size $n^3$. Could you comment on the practical computational and memory cost of storing and updating this array in your implementation? Could it become a practical bottleneck for systems with thousands of nodes, and do you see a path to optimizing it?\n\nIn practice, it is sufficient to only save information about from which nodes of $G$ all the nodes of $LG$ have received information from. If $G$ has $n$ nodes and $LG$ has $n_{lg}$ nodes, the size of the array will effectively be $nn_{lg}$. For a KNN graph, $n_{lg} = \\mathcal{O}(nk)$ so the memory requirements of $B(t)$ in our implementation scale as $n^2k$. In practice, only indexing operations, element wise binary logical functions and PyTorch scatter-add operations are applied to $B(t)$. Those do, to the best of our knowledge, not scale more than linearly with the size of $B(t)$ and can be (and probably are internally in PyTorch) heavily parallelized. We never experienced any computational bottlenecks during these calculations.\nIf it does become a problem, one possible strategy to circumvent it is to use a graph $G$, that has a fixed girth (length of the shortest cycle) of at least $T^{lg} + 2$, where $T^{lg}$ is the number of message passing (MP) steps on the line graph. This ensures that the Jacobian stays hollow even without removing connections rendering $B(t)$ unnecessary.\n\n2.\n>Your results show a non-monotonic dependence of model performance on the hyperparameter 'k'. Do you have any intuition or a proposed heuristic for selecting an optimal 'k' for a new system without having to train multiple models? Is there a risk that for a complex landscape, the optimal 'k' would need to be so large that it negates the computational advantage of using a sparse graph?\n\nNaively, one might expect that the model performance (ESS) would monotonically increase with $k$ as there are more connections in the graph for larger $k$. However, this is not the case in practice (see e.g. table 3). The fact that the ESS does not increase further with $k$ beyond some upper bound can heuristically be explained as follows: By construction, some of the connections in the line graph need to be removed after a certain number of MP steps (See appendix B for more details). The larger $k$, the more connections need to be removed after only one MP step up to the point where (almost) no connections are left after only one MP step. We thus have two counteracting effects when $k$ increases: The first MP step gains more connections and thus expressiveness while the second MP step looses connections and expressiveness. Eventually, the latter seems to take over effectively limiting the expressiveness of the model.  We will include a brief summary of this explanation in the camera ready version to improve clarity.\n\nAs larger $k$'s are not necessarily better, we do not think that for complex landscapes the optimal $k$ will be so large that it negates the computational advantage.\n   To (heuristically) choose $k$, one can efficiently numerically work out how many edges of the line graph are left after $T$ message passing steps without training a model. There should be a significant amount of edges left after one MP step, and the graph should not be extremely sparse in the beginning. We will include a plot showing the percentage of edges remaining as a function of the number of message passing steps for different $k$'s together with an instruction on how to tune $k$ in the appendix of the camera ready version.\n\n3.\n>Could you elaborate on the training of the baseline model? Given its very low ESS, do you believe it was fully converged? How does its ESS compare to other published flow-based Boltzmann Generators for these LJ systems? A stronger baseline would make the \"effective speed-up\" claims even more compelling.\n\nThe baseline model was trained in a similar way as in the Equivariant Flow Matching paper (Klein et al., 2023). However, there are some significant differences. Beside a different learning rate scheduler (see appendix C4), most importantly, we use a different architecture (PaiNN) instead of their architecture that we will refer to as EGNN. The ESS for e.g. LJ13 that was reported in (Equivariant Flow Mathching, Klein et al., 2023) is almost 58\\% while our ESS is close to zero if we do not remove a small number of outliers and comparable (49\\%) if we do remove these outliers. As the two architectures are different (different update rules, different embeddings etc.), there might be non-trivial effects that could explain the difference in ESS, similar to the non-linear scaling with $k$.  \n\n   While we do believe that our baseline models are fully converged, it might be possible to further increase performance with different training strategies and choice of hyperparameters.\n\n   It is important to keep in mind that the ESS is extremely sensitive to outliers of the log importance weight histogram. To illustrate this effect, we included $ESS_{rem}$ where these outliers are removed in a controlled way.\nFuture work will be necessary to fully understand how different architectures used in HollowFlow and in the non-hollow baseline do affect the ESS and especially how different architectures influence the amount of outliers that significantly lower the ESS.\n\n4.\n> This work is motivated by achieving exact divergence calculation efficiently. How do you see your method comparing to those that use stochastic trace estimators (e.g., Hutchinson's estimator)? While approximate, they are also computationally cheap. In your opinion, what are the key applications or systems where the exactness of the divergence provided by HollowFlow is critical and approximation is insufficient?\n\nFor Boltzmann generator applications, variance in the likelihood estimator leads to biased weights under self-normalized importance sampling. Thus, exact likelihoods and exact divergences are critical and even an unbiased estimator like the Hutchinson estimator is insufficient. Furthermore, considering how sensitive the ESS is to even a small number of outliers, it is important to get the importance weights exactly right.\n\n\nLimitations:\n\n>The current kNN-based implementation assumes locality and is not well-suited for systems with critical long-range interactions (e.g., electrostatics). This significantly limits its immediate applicability to a wide range of important molecular systems.\n\nWhile we agree with the reviewer's observation that we currently assume locality, we want to emphasize that this can be circumvented by a multi-head strategy, where several relatively sparse HoMP with different connectivities (e.g. separated length scales) are run in parallel and combined in the end (see appendix A for a brief discussion). This makes sure that each network does not get disconnected as fast while not having to do any locality assumption. This introduces an additional scaling with the number of heads while the scaling with the number of particles remains the same for each head. We will include a more detailed discussion of multi-head HollowFlow in the camera ready version.\n \n>The core idea of HoMP is to restrict information flow to gain computational efficiency. This is an explicit trade-off. It remains an open question whether this restriction will prevent the model from learning the complex energy landscapes of more challenging systems, leading to a sample quality (ESS) so poor that even a massive speed-up cannot compensate for it.\n\nWe are currently working on running HollowFlow on Alanine Dipeptide and plan to include the results in the camera ready version."}, {"review_id": "zhr3GOM7zm", "reviewer": "Reviewer_KaKL", "summary": "This paper addresses the problem of expensive divergence computations for likelihood evaluations in continuous normalizing flows. The solution adapts a previously proposed network architecture which can report its Jacobian trace in a single backward pass. The present authors describe the construction of a similar architecture with equivariant GNNs. The resulting network is less powerful than standard GNNs, but the fast likelihood computation enables fast ESS/s on the well-studied LJ-13 LJ-55 systems.", "questions": "See above.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n* The paper treats an important and very unappreciated problem in the Boltzmann sampling literature, which is the difficulty of scaling exact likelihood CNFs to high dimensions despite their superior performance to coupling flows. There are few obvious ways of tackling this challenge, resulting in fewer papers published in this area. The authors' present work is among the more creative recent contributions to Boltzmann sampling.\n* Although the basic idea of HollowNet is pre-existing, the authors present a highly nontrivial adaptation to GNNs operating equivariantly via careful construction of the message passing graph. The runtime characteristics of this network are carefully considered and formalized.\n* The empirical results indeed show that trading off accuracy for fast likelihoods is sometimes worthwhile, supporting the promise of the research direction.\n\n**Weaknesses**\n* It is not acceptable to manipulate the ESS metrics by excluding outliers, a protocol which has no justification in importance sampling. **The authors must correct this, or I will lower the score.**\n* The performance penalty is larger than one would hope, calling into question whether the method would work for more difficult systems.\n* The authors should discuss considerations, if any, in adapting the framework to non-equivariant GNNs or transformers, as these are increasingly the architectures of choice for molecular diffusion or flows.\n* The kNN graph construction is non-differentiable. Can the authors comment on how this changes the validity of the divergence computation or the CNF?", "was_revised": true, "final_justification": "I maintain my positive assessment of the technical aspects of the work from my original review.", "rebuttal": "We thank the reviewer for the valuable, detailed feedback and thoughtful questions. Answers to the reviewer's questions and additional comments can be found below.\n\n1.\n>It is not acceptable to manipulate the ESS metrics by excluding outliers, a protocol which has no justification in importance sampling. The authors must correct this, or I will lower the score.\n\nWe agree that 'manipulating' benchmark statistics is not acceptable. We however emphasize that we include these to illustrate the impact of the outliers on the ESS, and how controlling these is a path forward for improving the presented approach\n\n2.\n> The performance penalty is larger than one would hope, calling into question whether the method would work for more difficult systems. \n\nWe agree that the penalty is large, however, we expect this to be linked to the locality assumption which we use throughout the manuscript. To test this, we will include an additional ablation study in the camera ready version where we use a KNN graph for the baseline to disentangle the contributions of locality and HoMP on the performance penalty.\n\nAdditionally, for the camera ready version, we will implement multi-headed HoMP similar to the multi-head attention described in appendix A , where several relatively sparse HoMP with different connectivities (e.g. separated length scales) are run in parallel and combined in the end. This makes sure that each network does not get disconnected as fast while not having to do any locality assumption. Furthermore, we are currently trying HollowFlow on Alanine Dipeptide and will include the results in the camera ready version.\n\n3.\n>The authors should discuss considerations, if any, in adapting the framework to non-equivariant GNNs or transformers, as these are increasingly the architectures of choice for molecular diffusion or flows.\n\nBy non-equivariant, we assume the reviewer means, not equivariant with respect to the euclidean groups, e.g. E(3), as explored in our examples. This can be done by simply adopting a non-equivariant MP/transformer architecture as is, but ensuring the MP/attention is done in a manner that is faithful to the HoMP recipe, as outlined in the paper. Relaxing permutation symmetry can be done by avoiding using the same embeddings for the same particles/nodes in the graph.\n\n4.\n>The kNN graph construction is non-differentiable. Can the authors comment on how this changes the validity of the divergence computation or the CNF?\n\nIt is unclear what the reviewer means here. The KNN graph is used to limit the connectivity of the MP graph, and the divergence calculation does not depend on the graph construction once it is set."}]

Author rebuttal:

Meta-review: The authors propose a flow-based generative model based on a GNN that makes sample likelihood computations more efficient, using a novel non-backtracking graph neural network (NoBGNN) enforcing a block-diagonal Jacobian structure. They showed speedups on two Lennard-Jones particle systems. Reviewers found the problem to be important and the experimental results significant. Most questions/concerns were resolved during discussion, after which all reviewers leaned towards acceptance. One major point of discussion was the reporting of ESS metrics including outliers (KaKL), which I remind the authors to incorporate into the paper.
Decision: Accept (poster)