HollowFlow: Efficient Sample Likelihood Evaluation using Hollow Message Passing

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.

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.

We 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

The performance penalty is larger than one would hope, calling into question whether the method would work for more difficult systems.

We 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.

Additionally, 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.

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.

By 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.

The kNN graph construction is non-differentiable. Can the authors comment on how this changes the validity of the divergence computation or the CNF?

It 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.

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.

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?

In 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. If 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.

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?

Naively, 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.

As 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. 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.

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.

The 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$.

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.

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. Future 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.

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?

For 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.

Limitations:

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.

While 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.

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.

We are currently working on running HollowFlow on Alanine Dipeptide and plan to include the results in the camera ready version.

We thank the reviewer for the valuable and very detailed feedback. Answers to the reviewer's questions and comments can be found below.

1. 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.

2. 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.

   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$.

   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.

3. 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.

4. 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.