Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The authors propose a Path Integral Sampler based on Schrodinger bridge. They turn the sampling problem into a stochastic control problem and use a neural network to parameterize the control policy. In order to find the optimal control policy, several approximation techniques are introduced. The authors further propose an important sampling scheme to handle the case of a suboptimal policy. They provide convergence guarantees under Wasserstein distance and demonstrate the proposed method on synthetic datasets, alanine dipeptide and binary MNIST.
main_review: The idea of using stochastic control for sampling is very interesting and the empirical performance seems improved. I have the below questions.

1.	Compared with previous MCMC methods, one benefit claimed in the paper is that PIS can generate samples with fewer steps. It is not clear to me why this is the case. T seems a hyperparameter in the algorithm and there are no theoretical or empirical supports to show that T is less than the number of steps in MCMC methods. Besides, what does “until converged” mean specifically in Algorithm 1? How do the authors check convergence? From Figures 1 and 2, the main advantage of PIS seems the ability of sampling from multimodal distributions. Why is it the case?

2.	In Theorem 2, the Wasserstein distance between the sampler and the target density increases as T increases which seems weird to me. Should the sampler be closer to the target distribution as the number of SDE time-discretization steps increases?

3.	Since the authors use a neural network to parameterize control policy, how does the expressiveness of neural networks affect the performance? When do the users need to include the importance sampling step? There seems no discussion about how to choose the neural networks.

4.	The empirical comparison between PIS and previous methods generally lacks interpretation. Surely the evaluation metric of PIS seems better, but it will be much helpful to provide an explanation on why it is the case. For example, is it because the target distributions are multimodal or PIS converges faster? This kind of interpretation of empirical results sheds more light on when and how to use the proposed method for users. 

5.	This paper seems closely related to [Deep Generative Learning via Schrodinger Bridge, ICML 2021] which also uses Schrodinger Bridge algorithm for sampling. A discussion with this paper would be helpful.


summary_of_the_review: In summary I think the proposed method is interesting and technically sound. But the benefits over previous methods are not clear both in theory and practice.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper presents a control approach to sampling. It is proposed to use a controlled diffusion initialized at time t=0 to sample at a given time t=T from a given unnormalized target distribution. This is achieved by minimizing a forward KL between two suitable diffusions. The resulting expression is simple as it is given by the integrated squared control plus a final cost term. The authors propose to parameterize the control by a neural network and minimize the KL using stochastic gradient techniques. In experiments, this novel method appears to outperform some recent alternatives.
main_review: The theoretical result giving the KL expression is known as acknowledged by the authors. However, to the best of my knowledge, this has never been exploited to solve sampling problems. Instead it has been proposed to instead directly approximate the Schrodinger-F\"ollmer drift. It is an interesting idea to attempt to exploit the KL expression to obtain a practical sampling algorithm. 
The authors propose to parameterize the control term by some neural networks which involves the gradient of the log-target; see e.g. equation (15). The gradient of the loss is trained using recent neural SDEs techniques. The method is then demonstrated on various examples where it appears to outperform solid baselines.

I find the idea proposed in this paper interesting but I also think that a significant part of the paper was spent discussing unnecessary material and that important details were omitted. The limitations of the methodology should also be better spelled out. 

First, the authors discussed a general class of controlled diffusions, review HJB for such models and Fleming logarithmic transform. It is shown that the the logarithm of the value function satisfies a simple linear PDE which can be solved using Feynman-Kac. However, these results are never used in the paper. The paper only deals with scenarios where we initialize the diffusion at the origin and consider $dX_t=u_t(X_t)dt+dW_t$ as one needs to know the marginal density of the uncontrolled process at time T. So I think it's really unnecessary to spend so much time on the general case, you might want to include this material in the supplementary and instead concentrates on more important algorithmic details. Similarly please simplify the presentation of Algorithms 1 and 2.

A key component of the paper is the parameterization $u_t(x)=NN_1(t)-NN_2(t)\nabla \log \mu(x)$ in eq. (15) of the control policy using neural networks. This should be better motivated. The use of the gradient of the log-target seems essential for the method to work well but I find it hard to believe it can solve all issues. Clearly if one has a multimodal target, gradient information will not be sufficient to explore all the modes; i.e. Langevin does not mix well in multimodal scenarios. Obviously one has also another component $NN_1(t,x)$ in the control which could deal with scenarios. However, if we are interested in sampling a high dimensional multimodal target, how can the initial control ever visit the region of high probability mass under $\mu$ and learn anything. The forward KL is the criterion been minimized so I would expect you will miss modes in some examples. In this respect, it would be good to discuss how the method is initialized (i.e. what is $u_t(x)$ at the first iteration, how is $T$ selected? etc). 

The failure modes of the proposed method are never mentioned, was it necessary to train multiple times the model with various initializations to obtain good results etc, did you find examples where the method fails miserably? (this would be fine, I don't expect any method to solve all sampling problems). I think these are really important points which are unfortunately not discussed. 

Similarly, I was somewhat frustrated by the lack of details on the gradient computation. Instead of just citing a paper, I would encourage the authors to discuss these details (instead of having a fairly long introduction on stochastic control). It would also be good to explain if there are any gradient issue (gradient vanishing/exploding) as the number of discretization steps increases (you seem to be using 100 in simulations, are there numerical problems when you increase the number of steps) etc.

Other questions:

- As in the AFT paper, the training time is not taking into account when evaluating the performance of the algorithm. 
How large is the training time? Is it better/worse than AFT? If the training time is included, is the path integral sampler competitive to an SMC algorithm using a large number of intermediate temperatures?

- How do you select T and $\delta t$? If you initialize using $u_t(x)=0$, then at time $T$ you will have a state $X_T \sim N(0,T*I)$ but if the mode of the target is really in the tails of this distribution, would not this be a problem?

- Shouldn't the weight of interest be  $dQ^{\star}(\tau)/dQ^u(\tau)$ and not its inverse given in (17). You are sampling from $Q^u(d\tau)$ and are targetting $Q^{d\star}(\tau)$. The ESS is $\mathbb{E}_{Q^u}[{dQ^{\star}(\tau)/dQ^u(\tau)}^2]$ and not what is written in Theorem 3 given in eq. (17).

- It is claimed that a single step of the sampler in Algorithm 2 is similar to a single step of leap frog step of HMC. As the control is computed by a neural network, this really depends on the size of the network. I could not find this point discussed in the paper. 

- There is quite scant information on the tuning of SMC and AFT (what flow was used? how was HMC tuned? was the Adam step size reasonable?)

- 10 temperatures as a default for SMC is certainly not standard practice. In practice, the temperature schedule is also typically selected adaptively.

- In the 2D example (Figure 1), the scale of the target is set to be larger than the initial distribution so no surprise SMC will struggle. I find the example quite misleading. The comment "The small variance is selected deliberately to distinguish the performance of the different method" in F4 should be included earlier. 

- I couldn't find any details on the neural nets used in the experiment. Did you use the same architecture for all examples?

Minor comments

- SMC with annealing trick is due to Del Moral et al. (2006) not Chopin & Papaspiliopoulos (2020).

- page 13: Dai Pra (1991); ?)

- Table 1: How do you estimate the log normalizing constant for LGCP?

- Table 2: caption indicates KL-divergences but you cannot compute it as you don't have the normalizing constant.

Post-rebuttal comment: The authors didn't submit a revised version of their manuscript on time so I revised my score accordingly. I was eventually asked by the AC to consider the revised version that was sent after the deadline. As some of my comments have been taken into account, I have re-adjusted my score. In particular, I do appreciate that some of the limitations of PIS are now detailed. However, a few important points need to be addressed seriously by the authors in their final version of the paper. 

- "Our approach avoid long mixing time theoretically and is more efficient": this claim is not justified and should be removed from the paper. You have no guarantee to be able to compute anything close to the optimal control and, even if you were, the Lipschitz constant of the control might be very high which would require a very fine discretization of the SDEs for numerical stability and hence a long run time. This type of bold claims is not necessary.

- Figure 1 is an unncessary ``clickbait" figure which, if kept in the manuscript, will just convince any AIS/SMC expert that you haven't run serious comparisons. The caption mentions that "The underlying distribution is chosen deliberately to distinguish performance of different methods". This is not quite correct, the problem is not the "underlying" (i.e. target) distribution but the fact that, instead of picking a diffuse initial distribution compared to the target as recommended in the SMC iterature, you pick a standard normal distribution whereas the modes of the target are at (-5,0), (-5,-5) etc. If you were using an initial distribution N(0,\sigma^2 I) with sigma=5, I am sure it'd work very well. Also the selection of 10 intermediate distributions is just arbitrary and, contrary to what the authors state, is not a recommendation made in [Arbel, 2021]. I recommend the authors to read and cite Zhou, Yan and Johansen, Adam M and Aston, John AD, Toward automatic model comparison: an adaptive sequential Monte Carlo approach, JCGS 2016. Ideally you'd run simulations with this method.

- I still believe that having a general drift and volatility in eq (2) is just unncessary when you use f=0 and g=1. Please simplify the presentation. Algorithm 1 and Algorithm 2 could be similarly simplified. It's not like you can use very general f and g in the first place as you need to know the marginal distribution of X_T for the uncontrolled SDE.

- Please cite standard references in importance sampling instead of Au & Beck, 2003. Also please do cite Neal, Annealed Importance Sampling, 2001.

- Figure 3 is another unnecessary "our method works best" figure which tells us nothing given no detail is provided (i could not find any) in the main text.

- You mentioned it in your rebuttal that "We admit a better temperature schedule may improve the performance of SMC": this is certainly true and should be stated explicitly in the paper and you should include here the reference to  Zhou, Yan and Johansen, Adam M and Aston, John AD  which addresses this problem. This is followed by a claim that you follow "the setup for SMC and AFT recommended by AFT and its official codebase". I checked the code (which is a codebase for SMC and not the official codebase) and the paper [Arbel, 2021]. You should use obviously the same settings as in this paper when you present similar examples but, if 10 temperatures is written in a code, this is not a recommendation for all examples. I think including more thorough comparisons (i.e. more temperatures, optimized schedules) would make the paper better.  You compute the "ground thruth" for the LGCP using SMC (something which should be added in the paper too and not only mentioned in the rebuttal), so it must work well sometimes.

-Section F.1 didn't address my concerns and has to be rewritten/completed. 
"We note PIS training algorithm is different from AFT, where the AFT algorithm requires optimization during generating samples and
retraining the network for the next batch of samples. In contrast, the PIS control policy is trained once and used everywhere."
Once you run AFT to learn the flows, then you can rerun it with the trained flows (this is what the authors appear to be doing from my reading - "We concentrate our empirical value evaluation on the learnt flow, which is equivalent to using the test set particles." ), so it's not different from PIS. 
I want to see a table with two columns, one for training where the training times for all methods are given (what is diplayed currently is the execution time for PIS and the training time for AFT) and another column where the execution times for all methods are presented  (AFT will correspond here to SMC+learned flows).




summary_of_the_review: Overall, this is an interesting approach and the authors have demonstrated empirically on a variety of models that it appears to perform well. 

Two points should be addressed:

- I believe that it would be good to give much more details about neural network architectures, initialization, gradient computation etc. 

- The limitations of the methodology should be better spelled out. The authors should for example illustrate the failure modes of the algorithm.  

More generally, it remains unclear whether any of the recent proposed Monte Carlo methods relying on training flows/neural networks by minimizing KL/maximizing ELBO are competitive with AIS/SMC when taking into account the training time.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper shows that sampling can be treated as a stochastic
optimal control program and introduces a novel sampling algorithm
that works by simulating a stochastic differential equation under
optimal control with appropriately chosen dynamics and cost function.

main_review: This paper is excellently written and great care is taken
to highlight exactly what properties of each formalism are
used to get the nice closed-form equations. The work is
highly novel and one that would be greatly valued by the
community. The proofs are clear and easy to follow as well.

The experiments are rigorous and persuasive. One thing I
don't understand is why was the work not compared against
neural density estimators? There is a rich literature of
flow models that is referenced in the paper but aside from
SNF not really compared against. Some of it even uses
neural networks that backpropagate through a SDE solver!

summary_of_the_review: This work ntroduces a novel sampling algorithm with lots of interesting
properties. This paper makes a solid contribution with the strengths and
limitations of the approach very clearly articulated.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper proposes an algorithm called Path Integral Sampler (PIS) for sampling from unnormalized distributions by parameterizing the control policy in the Schrödinger bridge problem with neural networks and supplying it with the gradient information from the target density. Performance guarantees are provided, and experiments on several illustrative toy problems as well as on sampling molecular structures and latent space of VAEs are conducted. The method is shown to produce high-quality samples from the posterior, demonstrating state-of-the-art performance.
main_review: Strengths: clear exposition, theoretical guarantees provided, informative experiments and ablations

Weaknesses:
0) It would be informative to see a comparison to Stein-Variational Gradient Descent (SVGD) or other ParVI methods [1]
1) The paper says "To calculate gradients, we rely on backpropagation through trajectory." Are there any issues associated with that? E.g., exploding/vanishing gradient. How long can trajectories be? How much memory/time is required?
2) In Algorithm 1, specify what is input and output of the algorithm
3) In Algorithm 2, it is a bit confusing that 'w' is used as the argument of the exponential, which is different from Eq. (17), where it is used as the exponential itself
4) In Theorem 2, Condition 1 is only defined in the Appendix. Either state the condition before the theorem or expand it in the statement of the theorem
5) The paper says "training consumption is not included in the comparison". Could the authors state at least the order of magnitude of how much time the proposed method takes compared to the baselines?
6) In Table 1, highlighting PIS_RW-GT may be seen as a distraction, since it is using the ground truth. Therefore, perhaps one could either separate that line or set those numbers in italic and use bold on PIS_RW-Grad for MG(d=2) columns
7) "more clearer" -> "clearer"; "as it in a toy example" -> "as in a toy example"
8) What are the limitations? Adding a paragraph on drawbacks and limitations would be helpful.

[1] Liu, C., Zhuo, J., Cheng, P., Zhang, R., & Zhu, J. (2019, May). Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning (pp. 4082-4092). PMLR.
summary_of_the_review: The paper is overall good. The method is based on taking the established framework of inference as control via Schrödinger bridge problem and utilizing NNs as a parameterization for the policy. Exposition is clear, convergence is characterized, the properties of the algorithm are highlighted on relatively low-dimensional problems.

====

After rebuttal

I thank the authors for addressing my comments. I trust that the authors will include the promised changes and additions regarding the training time and the comparison to ParVI into the final version of the paper. I maintain my score "8: accept, good paper".
review_title: null

Meta-review: This paper introduces a control-based approach to sampling. All of the reviewers found the idea interesting. There were serious concerns by some of the reviewers regarding how the paper positioned itself relative to the literature, how it designed baselines for experiments, and how it compared itself to existing methods. There was vigorous rebuttal phase. The authors submitted a slightly late revision based on a procedural misunderstanding, and I decided to incorporate their late revision.

Based on the revision, the majority of reviewers felt that the paper was at least above the bar for acceptance and some of the more positive reviewers stood strongly by the paper. I believe that this paper is of value to the community, so I will recommend that it is accepted, but I want to be very clear about something: the authors **must** incorporate the late revision as the basis of their camera ready and I **strongly recommend** that they address the concerns of reviewer d7Mk, including but not limited to:

- "Our approach avoid long mixing time theoretically and is more efficient": this claim is too strong.
- Figure 1 is not particularly informative and the authors should reconsider it.
- The presentation of Eq (2), Algorithm 1, and Algorithm 2 should be simplified. 
- Section F.1 should incorporate the comments of Reviewer d7Mk.
- Citing standard references mentioned by Reviewer d7Mk.

The reason I highlight these recommendations is that I believe they will greatly improve the quality, longevity, and impact of this paper. Slightly overselling ideas feels like a good strategy, but it is a bad long-term strategy. I believe addressing these points is in the interest of the authors.
Decision: Accept (Poster)