Improving Gradient-guided Nested Sampling for Posterior Inference

We thank the reviewer for their very helpful feedback.

We have tried to highlight more clearly the original contributions from this paper: Firstly, we combine existing ideas that have never been combined before, resulting in a performance in sampling tasks that beats all existing state-of-the-art methods. Secondly, we use Nested Sampling with GFlowNets to generate a large number of samples from the target distribution. As the reviewer highlighted, we believe that this opens the door to many interesting research directions.

To clarify the value of the individual contributions, we have added ablation studies (Appendix G) to elucidate the value of each component. We have also highlighted more clearly how we use the gradient in guiding the choice of new live points in Section 3. We have also added an intuitive argument about why this update only requires $O(1)$ reflections, as opposed to baseline slice sampling.

Regarding your questions:

In Figure 1, instead of plotting the error in the estimate of log Z versus the number of dimensions, how does the error grow versus the number of likelihood evaluations?

We have added a plot of log Z vs. the number of likelihood evaluations in a separate appendix (Appendix J). GGNS needs fewer likelihood evaluations to reach the same error.

In Figure 1, the error bar (i.e., standard deviation across 10 runs) are much bigger for GGNS compared with the baselines, does it mean it is trading off variance for bias?

Interesting point. This is true inasmuch as the other methods do seem to have less variance while being more biased. At the same time, this is not an instance of the typical bias/variance tradeoff in probabilistic modeling: the other methods are biased because they are not correctly sampling the space (e.g., by missing modes, leading to an easier sampling problem). Therefore, it is not a matter of more model complexity or less regularization leading to less variance, but a matter of the simpler sampling methods being inaccurate.

Table 1, all methods except GGNS are under-biased for the Gaussian mixture example, is it just by chance?

As far as we understand, the reason other methods underestimate the normalization constant for the Gaussian mixture is that they are likely to have missed, or in some way underestimated, a mode. This leads to these methods thinking that there is less total posterior mass and, therefore, that logZ is smaller. We have added a comment about this in the table caption.

In the 'Many Wells' experiment (Figure 4), are you comparing against FAB or FAB with buffer (which should give better results than FAB)?

This is FAB with buffer, we have now clarified this in the text, thank you for pointing this out.

Minor: Formatting under Figure 1 needs to be modified. page 7, "due to the high its high dimensionality" -> "due to its high dimensionality"

Thank you; these have been fixed.

We thank the reviewer for their encouraging feedback.

To address the first point, we have added an appendix (Appendix H) with a complete GGNS algorithm. We have also added an ablation study (Appendix G) to analyse how our different improvements affect the algorithm’s performance.

We have also corrected the formatting issue with the caption of Figure 1.

We thank the reviewer for their very helpful feedback. We have tried our best to address these points with changes to the text.

lack of clarity regarding the proposed method

We appreciate the difficulty in understanding how all the components of the algorithm come together. We have now added an appendix with a concrete GGNS algorithm (Appendix H).

I didn’t understand what was the advantage of the method for Section 5, it seems a nice application that you can apply to generative flow networks, but what have we gained in training of the drift with your method, is there a baseline we can compare your method against?

In summary, Nested Sampling provides only a limited number of samples, but by training a GFlowNet, we can get an arbitrarily large number of samples by amortization.

GFlowNets are off-policy deep reinforcement learning algorithms. They are trained to sample a target density, but work best when the training is guided by a dataset of meaningful ground truth samples (to aid mode discovery), which need not be large or unbiased. The aim of Section 5 is thus twofold: we show that

• By using samples from GGNS in the training policy of a GFlowNet, we promote rapid mode discovery and convergence;
• After the GFlowNet has been trained (with GGNS guidance), we are able to use the trained model to obtain nearly-unbiased samples in a fixed number of sampling steps.

Page 5 at the top there is some overlapping text below Figure 1.

We have corrected the formatting error.

I don’t really understand how you took advantage of ‘GPU interoperability’?

We have tried to clarify this point. The point here is that while Nested Samping algorithms had been written in differentiable programming languages before, none of those implementations took advantage of the gradient of the likelihood – their main advantage over implementations in non-differentiable programming languages arises from GPU interoperability. Therefore, what we meant to say (and we have tried to clarify), is that we do not take advantage only of GPU interoperability, but also of efficient parallel gradient computation.

I don’t quite understand why there error bars reduce with the number of dimensions in Figure 3.

This is a very good question, which we had not fully explored in the original submission. The reason for this is that $\mathcal{D}_{\rm KL} (\mathcal{P} | \Pi)$ decreases as we increase the dimensionality, as we show in appendix F.2, meaning that the information gain when going from prior to posterior goes down. This could occur because of the cancellations between sine and cosine times, as the number of terms in the sum increases.

We thank the reviewer for their very helpful and through feedback. We have tried our best to address the reviewer’s concerns.

the exact details of the implementations are not self-contained within the paper

First, we appreciate the comment about the lack of details of the algorithm. Therefore, we have added an Appendix H with the details of our implementation. We hope this will add clarity and transparency to the paper.

"Adaptive Time Step Control": "This time step is adjusted dynamically ... increase or decrease ..." How is it adjusted exactly?

The details of how adaptive time step happens are now in lines 21-24 of Algorithm 3.

"Trajectory Preservation": What is exactly a "trajectory" here? It is the states of the Markov chain after multiple Markov chain transitions? Or the intermediate states of a Hamiltonian trajectory as in recycled MCMC methods [1]? "Trajectory preservation": "select a new live point at random from the stored trajectories" how random? Uniformly at random? Or weighted resampling as in typical Hamiltonian Monte Carlo implementations [2,3]?

By trajectory, we mean the steps of the Markov chain, we have clarified it in a footnote in Section 3. The points are sampled fully randomly, as we need to sample uniformly from the prior. We have also clarified this.

"Pruning Mechanism": "This mechanism significantly improves the computational efficiency" How/why does it improve the computational efficiency exactly?

The pruning mechanism improves the computational efficiency because, without it, we would continue to evaluate the likelihood for points that have drifted far away from the region of interest, leading to a waste of computational resources. We have tried to clarify this.

"Differentiable Programming": How is differentiable programming used here? Why does it help?

Differentiable programming is key here, in that our method requires calculating the gradient of the likelihood. Therefore, it will only work for likelihoods for which we can compute gradients

second paragraph: "the fact that gradients guide the path means one no longer requires ": I'm not sure if this is obvious. Is there a proof for this statement?

We have added a intuitive argument for why we get a linear scaling with dimensionality for our method. The idea is that because the gradient is $SO(n)$-equivariant, HSS can explore the principal directions in a number of steps that is independent of the dimension of the ambient space. Please also see the answers to Reviewers skHM and Ec9X for intuition regarding this.

"Mode Collapse Mitigation": "... preventing them from converging prematurely to a single model": Is there theoretical/empirical evidence for this?

We have added empirical evidence for this in our ablation study (Appendix G.3).

"Robust Termination Criterion": "has decreased by a predetermined fraction from its maximum value": What is the theoretical principle behind this termination criterion?

We have added a more detailed motivation for our termination criterion in appendix I, as well as empirical evidence of it working better than the commonly used in in appendix G.4.

"gradient-guided nested sampling ... makes use of the power of ... parallelization for significant speed improvements": I couldn't find any empirical evidence on much this method takes advantage of parallelization. Did the authors measure the strong scaling of this method? Until how many cores does this scale? What is the efficiency?

It is hard to do an empirical study of how much we benefit from parallelization, as it would require rewriting the algorithm in a non-parallel way. Our argument is that, because our nested sampling steps are all run in an embarrassingly parallel way, our algorithm will benefit from an increased number of cores, up to $n_{live} / 2$, which is the number of points we can update simultaneously.

A more detailed study with different computing configurations would be interesting but beyond our capabilities for the response period. Beyond that, we cast our contribution, in part, as an adaptation of nested sampling algorithms to hardware intended for modern machine learning workflows, featuring massive parallelization on GPUs. This is particularly important in data processing settings that combine nested sampling with deep learning, such as when the prior or likelihood models are given by deep neural networks. We have added this to the paper.

I sincerely thank the authors for addressing my comments. In fact, I am impressed by the effort they have put in. Unfortunately, I have some major concerns that remain, that prevent me from increasing my score. In particular, one of the main claims of the paper that the method is scalable in terms of dimensionality, is still not clearly demonstrated.

• Pooled Baselines? For objectively judging the performance in a controlled manner, I believe the experiments producing Figure 1 and Figure 2 are the most important. However, the baselines are "pooled" and not properly evaluated on the same problems. While the authors replied that, after showing that their method performs well on one experiment, they chose to "move on" to a different set of baselines, I do not quite agree that this is scientifically rigorous. Is it obvious that we will see the same experimental result despite a different setup?
• Lack of High-Dimensional Experiment Furthermore, the experiment for Figure 2 is way too weak. For instance, take a look at Table 1 by Zhang and Chen (2022). There, they included a Log-Gaussian Cox process experiment, which is both geometrically challenging and also actually high-dimensional. On the other hand, the current paper only considers the 10-D funnel and a 2-D mixture. The remaining experiments have a different set of baselines for some reason, which, again, I do not find that it provides clear evidence of superiority.
• Paper Organization Lastly, while I appreciate the authors' effort to improve the paper and incorporate the comments of the reviewers, the paper is still not self-contained. This unfortunately, does not address the concern that other reviewers have also raised. I expect some major reorganization to make the paper meet the bar for the quality expected for an ICLR paper.

Lastly, I have some additional minor comments. The paper states, "Differentiable Programming Whilst nested sampling algorithms written in differential programming languages exist," for which the authors clarified the meaning. I believe this wording is quite misleading. The sampler itself does not internally use differentiable programming. (There are, in fact, samplers that internally use differentiable programming in a non-trivial manner, for example, by differentiating through the Matropolis-Hastings correction step. So it's worth clarifying that the method is not doing something like that here.) All that it operates with, are the provided gradients right? I think it would be worth clarifying the language that the proposed methodology takes advantage of gradients not necessarily differentiable programming. Furthermore, for the experiment in Figure 2, the authors mention that what is called "HMC" here used the harmonic mean estimator following Zhang and Chen. Unfortunately, the harmonic mean estimator is known as "the worst Monte Carlo estimator ever" due to its potentially infinite variance. Therefore, it cannot be considered to be a valid baseline for estimating marginal likelihoods. I recommend the authors look for better baselines, for instance, thermodynamic integration.

Dear Authors,

Paper Organization

We are confused about this point, as we feel like we have addressed the points raised by the other reviewers in this regard. The paper now contains an implementation with all the algorithm details, an ablation study, and theoretical arguments for improved scaling. Therefore, we are not sure what the paper is lacking in terms of reorganization.

I am specifically referring to my original comment about "self-containment." By self-containment, I am saying that the details currently deferred to the appendix should have been blended in the main text. I believe readers should be able to have a rough idea of the proposed algorithm after having read only the main text, which does not seem to be the case. Recall that other reviewers have also expressed similar concerns, and I feel the main text still does not properly address this. Furthermore, given the current organization of the paper, it appears to me that this will require a major restructuring of the paper, which would require an additional round of thorough review, unfortunately.

Note that all of this is for the benefit of the paper, I strongly believe the technical contribution of the paper deserves better presentation.

Nevertheless, given many of my original concerns have been addressed, I'll raise my score.

We thank the reviewer for their feedback. We have tried to address all four points that were brought up:

1. We have tried to further clarify our main contributions. Please see response to all reviewers for a summary.

2. We thank the reviewer for bringing up this point. We have now added an appendix (Appendix H) with a concrete GGNS algorithm, including all the details.

3. We appreciate this point. To attempt to address it, we have tried to provide a justification for what is, in our opinion, the most important part of this work, which is the fact that each step only needs O(1) bounces; as this is what allows the algorithm to have a better scaling with the number of dimensions. This has been added to Section 3. In summary, while slice sampling only has information about one dimension per step (needing $O(n)$ steps to explore an $n$-dimensional space), when we use Hamiltonian slice sampling, the gradient provides information about all $n$ dimensions on each step. The $SO(n)$-equivariance of Hamiltonian slice sampling implies that exploration of a density concentrated around a low-dimensional manifold is less sensitive to the dimension of the ambient space. Please also see the answer to Reviewer skHM for intuition regarding this.

4. This is a very good point. We have added an Appendix G with a detailed ablation study. Please see the response to all reviewers for discussion.

We thank the reviewer for their very helpful feedback and comments. We particularly appreciate the suggestion of an ablation study. We have now added an Appendix G with such a study. Please see the response to all reviewers for discussion.

With respect to the $O(1)$ bounces, the reviewer makes a good point that our original submission did not justify this claim. We have now added a theoretical justification for this point in Section 3. The main idea is that, while other methods like slice sampling only have information about one dimension per step (needing $O(n)$ steps to explore an $n$-dimensional space), when we use Hamiltonian slice sampling, the gradient provides information about all $n$ dimensions on each step, meaning we need only $O(1)$ steps.

For intuition, consider a density concentrated around a $d$-dimensional linear subspace of an $n$-dimensional space ($d\ll n$). Usual (coordinatewise) slice sampling would require $O(n)$ bounces to mix within the subspace despite the space having only $d$ degrees of freedom. On the other hand, Hamiltonian slice sampling, which is equivariant to orthogonal transformations of the space, would take steps within the low-dimensional subspace and does not depend on the dimension of the ambient space.

We have also fixed multiple typos and misspellings.

We have just uploaded a revised pdf with changes marked in $\color{red}\text{\sf red}$.