Sequential Controlled Langevin Diffusions

6. How sensitive is the procedure wrt the chosen time discretization $[t_n, t_{n+1}]$?

We also would like to note that for our main experiments we did not tune the time discretization, and the baseline choice of a cosine noise schedule with a uniform discretization is a robust choice that allowed SCLD to achieve the best performances on most of the benchmarked tasks.

We would like to additionally emphasize that the choice of time discretization is somewhat redundant as a hyperparameter. A key observation is that CMCD and SCLD depends on time increments $\Delta t$ (note we use $\Delta t=\frac{t_n-t_{n-1}}{L}$ in our experiments) through $\sigma(t)\sqrt{\Delta t}$ after reparametrizing the control (see equations 32-33 in the paper). As such, any tuning of the time discretization (for a fixed step count) can be replaced with equivalent adjustments of the noise schedule $\sigma(t)$. Indeed, we can confirm this experimentally, with training trajectories remaining unchanged when rescaling time and noise levels accordingly.

As such, we do not lose performance by fixing the time discretization, and it is sufficient to just tune the noise schedule. While we expect further performance improvements by further finetuning of the diffusion schedule, we as mentioned leveraged a fixed cosine schedule which has also proven successful in many previous situations such as [d, e].

7. When sampling, is it possible for SCLD to maintain a principled time discretization guidance similar to traditional SMC use of constant ESS? Does the flexibility in proposal distribution come at the cost of losing these methods? If so, how does this trade-off impact the overall performance and reliability of the sampler?

• We have addressed the first question previously in that we can choose the annealing schedule in a principled way by simply learning it end-to-end via gradient optimization

• While we do not think that switching from ESS-based ways of choosing annealing schedules to learning the annealing schedule constitutes a disadvantage, we note that as choosing the annealing schedule (such as by considering ESS) is independent of learning the control, it is also possible to choose the annealing schedule via ESS: First, we set $T=1$ and $\beta_t=t$. Then we can dynamically choose the time discretization $t_n$ based on ESS and learn a drift function that depends on the current timestep $t$ and position $x$. Note that in this case, the total number of discretization steps is variable. We can then apply log-variance-based training in the same fashion as Algorithm 2 from the paper. This constitutes an area for future investigation.

---

References

[a] MCMC Variational Inference via Uncorrected Hamiltonian Annealing: https://arxiv.org/pdf/2107.04150

[b] Langevin Diffusion Variational Inference: https://arxiv.org/abs/2208.07743

[c] Estimation of Thermodynamic Observables in Lattice Field Theories with Deep Generative Models: https://arxiv.org/pdf/2007.07115

[d] Improved Denoising Diffusion Probabilistic Models https://arxiv.org/pdf/2102.09672

[e] Denoising Diffusion Samplers https://arxiv.org/pdf/2302.13834

4. Are there more complex and less toy examples where SCLD would be applicable?

We first emphasize that our tasks are considered challenging benchmarks, covering up to $1600$-dimensional densities from Bayesian inference, $50$-dimensional multimodal densities with $40$ modes, as well as problems from robotic navigation where most existing methods fail. We would say that these are arguably among the most complex tasks used to benchmark other samplers, such as those presented in our related work section.

Nevertheless, we want to illustrate the application of SCLD to the simulation of a statistical lattice field theory. We study the lattice $\phi^4$ theory in $D = 2$ spacetime dimensions (however, note that the dimension of our density will be $d=128$, see below). The random variables in this circumstance are field configurations $\phi \in \mathbb{R}^{W\times L}$, and their density is given by

$p_{\mathrm{target}}(\phi)=\frac{e^{-U(\phi)}}{Z},$

where $U(\phi)=-2\kappa\sum_x\sum_{\mu}\phi_x\phi_{x+\mu}\ +\ (1-2\lambda)\sum_x \phi_x^2 \ +\ \lambda \sum_x \phi_x^4$

and where summation over $x$ denotes summation over lattice sites and summation over $\mu$ corresponds to summing over neighbours. For our experiments we pick $W=16$ and $L=8$, leading to a target in $d=16\cdot 8 = 128$ dimensions. We also pick $\lambda=0.022$ and $\kappa=0.3$ (corresponding to a so-called critical point) as in [c].

A physical question of interest is to estimate the free energy $F=-\frac{1}{d}\log Z$ of the system. We used long-run HMC as in [c] to obtain groundtruth free energy values. We apply both the SMC-ESS method considered previously and SCLD. We tune adaptive tempered SMC (SMC-ESS) in the same manner as before, maximizing ELBOs. For SCLD we use the same design choices as we did for the main experiments except for the network, where we use an equivariant version of PISGradNet that satisfies $u(\pmb{x})=-u(-\pmb{x})$ to account for the symmetry in the potential. For both SCLD and SMC-ESS we estimate the free energy with 2000 particles.

We run SCLD with 4 training seeds for 2000 training iterations and evaluate both SCLD and SMC-ESS using 50 evaluation seeds. We compute the average error $\Delta F$ of our free energy estimates across seeds and obtain the following results:

SCLD is able to obtain on average better estimates of free energy than SMC-ESS. We will add a comparison to all baselines in our final version. We anticipate that the application of SCLD to further problems in the natural sciences is a promising area of future work.

5. When moving to the continuous time/path space measure setting of SCLD it seems to me that we loose these principled ways to choose the time discretization

We assume here that "time discretization" for traditional discrete-time SMC algorithms refers to the chosen annealing schedule of the densities (i.e the $\lambda_t$ in [1], which we denote $\beta_t$ in our paper). We will refer to this subsequently by the annealing schedule. In our algorithm, we would like to highlight that our choice of $\beta_t$ (which determines factors like ESS) is entirely decoupled from how we discretize time.

In SCLD, we are able to, and choose to, learn $\beta_t$ (Section 2.4 - Annealing path) via direct gradient optimization of the divergence objective. Gradient optimization is a principled and desirable (see [a]) way of selecting algorithm parameters, so while we don't choose the $\beta_t$ with ESS in mind, our method of choosing the annealing schedule is also principled. In fact, our method of choosing $\beta_t$ exhibits several benefits, the most notable being that we can draw samples using a fixed, finite number of sampling steps, which is not the case with ESS-based discretization schemes. We have added a section in the appendices (A6.5, in orange) which visually compares the learned annealing schedule with a fixed, linear annealing schedule, and shows that the learned annealing schedule is able to produce a more gradual transition from the prior to target density.

We will also talk more specifically about choosing time discretization in the next point.

(response continued in next post)

Dear Reviewer ZuHY,

Thank you for your considered review. We are pleased that you appreciate our pedagogical style of writing, and are very grateful for your added insights, especially on the constant ESS idea. We respond to your questions and raised issues below.

1. The reader struggles to properly get the contributions

We provide several novel ideas in our paper. We introduce a new SMC framework for combining SMC with diffusion samplers by adopting a continuous path space perspective to SMC and formulating a novel family of training objectives, culminating in the highly effective Sequential Controlled Langevin Diffusion (SCLD) method. Please refer to our General Response for a more detailed discussion of our contributions.

2. A comparison with more traditional SMC samplers to justify the need for the control term and the training procedure would be a nice thing to have

We will jointly answer this in the next point.

3. Are there clear examples where the more flexible proposal is worth the added need for training over running a slightly more sophisticated SMC sampler (eg [1], [3])?

We compare SCLD against the Adaptive tempered SMC implemented in [2], which utilizes the constant ESS method for choosing discretization steps as seen in [1], [3] (we identify it as the closest relative among the Blackjax library of the methods presented). As for SCLD, we utilize a single HMC step for the SMC kernel with 10 leapfrog integration steps, and apply the same tuning procedure for step size as we did for our own SMC method. We additionally sweep over the ESS threshold $\alpha \in \{0.3, 0.5, 0.75, 0.9, 0.95, 0.99\}$ and use systematic resampling. Otherwise, we apply the same benchmarking procedure as for our other experiments.

We have added a section (A6.8, in orange) in the appendices of our paper concerning this comparison. While the new SMC implementation (which we refer to as SMC-ESS) significantly outperforms the baseline, it is still worse than SCLD on 9 out of 11 benchmark tasks, while also sacrificing a key benefit of SCLD which is that of fixed sampling times once trained. For instance, SMC-ESS was only able to attain better performance on LGCP in part due to taking 789 SMC steps on average, compared to SCLD's 128. We display the results below for reference as well.

Table: Comparison of SCLD against SMC-ESS in terms of ELBOs

Comparison of SCLD against SMC-ESS in terms of ELBOs.

---

Table: Comparison of SCLD against SMC-ESS in terms of Sinkhorn distances

Comparison of SCLD against SMC-ESS in terms of Sinkhorn distances.

Finally, we note that such improved SMC methods can also be used for SCLD to further improve its performance. We will discuss this in a subsequent section.

(response continued in next post)

Dear Reviewer ZuHY,

Thank you for your response and your additional feedback. We are pleased that you recognize that SCLD and its possible derivative approaches "have great potential."

SMC experiments: We are very sorry to hear that our additional experiments and revised paper could not convince you. We agree with your suggestion on the presentation and added the SMC-ESS results (and our new SMC-FC results outlined below) to our main tables.

We would also like to emphasize that the HMC kernel for SMC-ESS is not fixed across all iterations of SMC but is instead determined by the chosen annealing schedule. We define step sizes $\varepsilon_t = f(\beta_t)$, where $\beta_t$ is the annealing schedule and $f$ is a piecewise linear function over $[0,1]$ which we select using a large grid search. As such, in terms of experimental fairness and the original review, we consider SMC-ESS to be our best interpretation of the baseline that was first requested.

Nevertheless, we are happy to present an additional baseline (SMC-FC) that is taken as is from [1] (as you suggested). SMC-FC utilizes the full-covariance tuning approach for Independent Rosenbluth Metropolis-Hastings (IRMH) proposals (on top of using adaptive tempered SMC). We employ $100$ MCMC steps per step and use the same evaluation protocol as for the other methods. In particular, we use $2000$ particles for all methods (which we clarified in the appendix). We tabulate the results below.

Table: Comparison of SCLD against various SMC methods in terms of ELBOs

---

Table: Comparison of SCLD against various SMC methods in terms of Sinkhorn distances

---

We definitely agree that sophisticated SMC samplers have the potential to offer strong performance. However, while we see improvements of SMC-FC over SMC-ESS for several tasks (Robot1, Robot4, Credit, Sonar), SCLD still exhibits an overall advantage on 9 out of 11 tasks.

Concluding remarks: We would like to reiterate that our main contribution is a general continuous-time framework for end-to-end learning of diffusion-based proposals. We understand the great advantages of SMC methods and embrace them in our work, offering a simple way to merge them with diffusion-based samplers. While we avoid more sophisticated design choices to have a fair comparison to previous methods, such as CRAFT and PDDS, we also note that there is much potential in applying more complex SMC ideas to our framework. We significantly outperform previous deep learning-based sampling methods (in terms of cost and accuracy), and as such believe our work has a place in the literature of this growing field.

---

We have already revised our paper to reflect the new experiments above. We thank you again for your suggestions and hope that this response clarifies your remaining concerns.

Dear Reviewer TwR8,

Thank you very much for your extensive review. We appreciate that you value our novel sampling framework that integrates SMC and diffusion-based sampling and achieves state-of-the-art performance on high-dimensional examples. We are in particular happy that you highlight our application of the log-variance divergence, which indeed is an essential ingredient that has not been used for resampling-based algorithms before. Let us address your questions and comments in the sequel.

Contribution

We provide several novel ideas in our paper. We introduce a new SMC framework for combining SMC with diffusion samplers by adopting a continuous path space perspective to SMC and formulating a novel family of training objectives, culminating in the highly effective Sequential Controlled Langevin Diffusion (SCLD) method. Please refer to our General Response for a more detailed discussion of our contributions.

Scaling analysis of log-variance divergence

Thank you for raising the scaling properties of the log-variance divergence, which are indeed very relevant to our algorithm. Contrary to the KL divergence, which scales exponentially in the dimension (see Proposition 2.4), the relative error of the log-variance divergence can be shown to not have this unfavorable scaling behavior. This has been investigated in a slightly different setting already in Proposition 5.7 in [1]. For convenience, we have stated the proof in our revised version, adapted to our setting, see Appendix A.2 (highlighted in green). In addition, we compare the performance of KL and log-variance divergence-based training in Appendix A.6.10. Please let us know if you have further questions.

Thank you also for your suggestion to add the log-variance loss to Algorithm 1. However, note that this algorithm only describes the sampling and not the training phase and is therefore independent of the loss. We have previously presented the training algorithms in our appendix and referenced them from the main part. However, we agree with your suggestion and moved the algorithm to the main part; see Algorithm 2 in the revised version (highlighted in green). Moreover, note that the training scheme with replay buffers can be found in Algorithm 4 in the appendix. Please let us know if anything is unclear or if you have further suggestions to improve the writing.

Off-policy training

We agree that the term off-policy training originates from the reinforcement learning community. However, several recent works, e.g., [2, 3, 4], have adopted this term in the same context to denote training the sampler with modified trajectory distributions (such as by using replay buffers). As such, we consider off-policy training to be established terminology. Nevertheless, we agree that the term should have been better clarified due to its relative novelty in this context, and thank the reviewer for bringing this to our attention. We have added Remark A.3 in our revised paper regarding this and referenced it in Section 2.3 (highlighted in green).

Table 1 and Differences from alternative methods

Thank you for your suggestions on the writing of this part, which has helped us improve our paper. We, for instance, added lines 125-126 and referenced Appendix A.1 in our Table 1. Due to the large amount of related literature present on sampling problems, SMC methods, and probabilistic modeling, we chose to defer discussion on specific methods and their differences to SCLD to the appendix, mainly due to space limitations in the main part of the paper. We hope that our "related works" section in Appendix A.1 addresses the author's concerns and we also significantly extended it in our new revision. As for PDDS, it is also due to our lack of space that the discussion was deferred to the appendix. However, we now added an extensive comparison between SCLD and PDDS that demonstrates the superiority of SCLD on our set of benchmark tasks (see Appendix A.6.7, highlighted in blue).

We hope that the added content proves satisfactory, and would be happy to work further with you to improve our writing quality.

(response continued in next post)

Dear Reviewer Kjto,

Thank you for your helpful review and for appreciating our clear and detailed writing as well as the effectiveness of our proposed method. We comment on your questions and concerns in the following:

Extended literature review

Thank you for the suggestion. Since our combined related work sections have already been quite long and generative modeling focuses on a different task, we did not mention these approaches in our submission. However, we agree that we should discuss the most related approaches and have added a detailed discussion in Appendix A.1 (highlighted in brown), also covering (generalized) Schrödinger bridges. Based on your suggestion, we also added another section on posterior sampling with diffusion-based models and connections to stochastic optimal control.

Posterior sampling with diffusion models (and SMC)

Since posterior sampling is a special case of the sampling problem, it can be directly tackled with SCLD without reliance on any samples. However, for complex and very high-dimensional prior distributions, e.g., image/audio/video distributions, it becomes crucial to leverage data. In particular, most popular diffusion-based sampling approaches leverage a pre-trained diffusion prior, i.e., a diffusion model trained on samples from the prior distribution. During inference, i.e., the generative process, additional guidance terms are then approximated on the fly to obtain approximate samples from the posterior. This is also the setting of [4], which additionally leverages SMC to obtain asymptotic (in the number of particles) unbiased samples. We present details in Appendix A.1 (highlighted in green).

However, in our setting, we do not have a pretrained diffusion prior (and for many of our tasks not even a prior distribution in the Bayesian sense). Without a pretrained diffusion prior, the method of [4] cannot be applied or (when using an uninformed prior) is effectively just a version of SMC, against which we already compare in our experiments. Note that we also added a comparison in Appendix A.6.8 (highlighted in orange) showing that SCLD can also outperform more advanced SMC methods. The most related method leveraging diffusion-based samplers and SMC is PDDS; see Appendix A.1 for a description and comparison. However, SCLD offers several advantages, leading to better performance (see Appendix A.6.7, highlighted in blue).

Nevertheless, we believe that (non-trivial) adaptations of SCLD are possible to apply its main idea to posterior sampling problems and reward fine-tuning of diffusion models. We mention an outlook at the end of Appendix A.1 and thank the reviewer for bringing up this exciting avenue for future research.

Dear Reviewer Kjto,

We hope this message finds you well. As the review period is coming to an end, we wanted to know if you have further questions or concerns which we can address?

Dear Reviewer 6QXs,

Thank you very much for your careful review. We are very pleased that you enjoyed reading our paper and found it a useful contribution. We would like to address your concerns below.

Overall contributions

We provide several novel ideas in our paper. We introduce a new SMC framework for combining SMC with diffusion samplers by adopting a continuous path space perspective to SMC and formulating a novel family of training objectives, culminating in the highly effective Sequential Controlled Langevin Diffusion (SCLD) method. Please refer to our General Response for a more detailed discussion of our contributions.

The combination of continuous and discrete time seems somewhat arbitrary

Overall, we believe that our framework is a principled way to combine continuous-time diffusion samplers and traditionally discrete-time SMC approaches to take advantage of the strengths of both methods. We agree with the reviewer that if a perfect control is learned then resampling wouldn't be needed, and that when the learned control isn't perfect, SCLD can correct for "imperfections of the control function via resampling". This is a key benefit of our method, leading to significantly improved convergence and performance.

The number of $t_n$'s, as the reviewer points out, is an important design choice in SCLD, and one which we study in Section 3.4 (the part on the choice of number of SMC steps). This must be determined by the available computational budget (a larger number of discretization steps take longer to simulate) as well as the specific target at hand. We found that on most tasks using large numbers of resampling steps (for a fixed total discretization step count) worked well and we did not optimize the position of the $t_n$ to have a fair comparison to our baselines. However, it is possible that performance could be further improved by using more elaborate strategies: as mentioned in the discussion with Reviewer ZuHY, the number of resampling steps and their locations can also be chosen systematically and adaptively as in [3], and this constitutes an area of future work.

We will further discuss the choice of the number of $t_n$'s in your subsequent question regarding the downsides of excessive resampling.

Equation 12

Thank you for identifying our typo, we have amended Equation 12 (in red).

Figure 5 seems to suggest that the more SMC steps, the better the performance, even though the paper noted potential downsides of excessive resampling (line 521).

Yes, we have found that excessive resampling can impact performance on highly multimodal tasks. In Figure 4 (previously Figure 5), you can see that on the Robot4 task, the method performs worse when 128 resampling steps are performed at training and evaluation, compared to using 4 resampling steps at training and none at evaluation. This is the result of mode collapse. We also found that in the GMM40 (50d) task, using 128 resampling steps led SCLD to collapse onto one mode. As such, we used fewer (4) resampling steps on multimodal tasks as opposed to 128. However, our results were all obtained using multinomial resampling, which is well known to be suboptimal [1], in the interests of giving a fair comparison to methods like CRAFT. It is likely that mode collapse and loss of sample diversity can be mitigated by using better resampling schemes such as systematic or parallel resampling [2].

References

[1] https://arxiv.org/pdf/2402.06320

[2] https://arxiv.org/abs/1903.12583

[3] https://arxiv.org/abs/1808.07730