Paper: Momentum Particle Maximum Likelihood (ICLR-2024)

Reviewer comment (Novelty):
The idea of adding momentum to gradient updates is very well known and it doesn't require deep theoretical justification to come up with Equations 2a - 2d, which follow quite obviously from PGD and the definition of Nesterov accelerated momentum.

Author response:
While we agree that the idea of adding momentum is well-known, we disagree that it follows obviously from PGD and Nesterov's Accelerated Gradient (NAG) algorithm. What may be immediate from PGD and NAG is a version of MPGD that only accelerates one component, which does not perform as well as MPGD (see replies to yQ9T). Furthermore, there are multiple methods to incorporate momentum into PGD, for example, through the principle of least action [1]. Taking this perspective, the resulting algorithm is difficult to simulate and so was discarded [2, see Appendix D]. In our work, we advocate for the vector field perspective (as seen in [3]) as a suitable perspective for incorporating momentum into gradient flow algorithms in the product space $\mathbb{R}^{d_\theta}\times \cal{P(\mathbb{R}^{d_x})}$.

Reviewer comment (Novelty):
Showing convergence is valuable, but most gradient flow algorithms have a natural recipe to provide convergence under strong convexity and/or smoothness assumptions. So, this is not really a novelty here, although it does contribute to the overall quality of the paper.

Author response:
We agree that arguing convergence for gradient flows under strong convexity assumptions is straightforward and not a substantial contribution. However, our MPGD flow in Equation 2 is not a gradient flow. This is due to the fact that it cannot be written as $\dot{z} = -\nabla F(z)$ for some $F$ and suitable notation of gradient $\nabla$. As a result, proving convergence does not follow from standard gradient flow theory and can be seen in Appendix F to be far from trivial.

We stress that when working in a momentum-enriched system, establishing convergence does not follow the same recipe as for 'pure' gradient flows. One of the key differences between gradient flows and momentum-enriched flows is that we can only show that the objective is non-increasing (which does not establish convergence). A different strategy is needed for momentum momentum-enriched system. One usually has to construct a Lyapunov function that couples the momentum variable with the variable of interest (for example, under convexity, see [4, Theorem 3]; and under the log-Sobolev assumption, see [3, Proposition 1]). We note again that the proof under log-Sobolev assumptions, see [3, Appendix C] is far from trivial. Similarly, the bulk of our appendix is required to prove this fact, see Appendix F. Specifically, this is done by a generalization of the argument in [3] to product space. An example of a unique extension is that we have to control an interaction term that occurs from the interaction in the product space (see Appendix F.4.6).

Reviewer comment (Novelty):
If we were to simply add a momentum step to every paper with a gradient update algorithm and show some results where it performs better, there would be a deluge of papers. For example, in the Coin EM paper there is a SVGD-like gradient step for updating the density represented by particles, and one could certainly add momentum to that — but does that make it a novel contribution?

Author response:
We agree that the general idea of accelerating PGD using momentum is natural and, on its own, need not count as exceedingly novel. The novelty (and the challenge) lies in how exactly this should be done. Should one add momentum to the parameters? To the particles? To both? Should one do so at the level of $(\theta,q)$ or at that of $(\theta,x)$? What discretization should one use? Should one use full or partial updates? And so on.

As such one of our implicit contributions is that we have provided a roadmap for others to incorporate momentum in this space. We have shown how one might proceed to derive and study momentum methods in this space which particularly requires care at all stages. Specifically, the stages where the process can go awry are:

* Derivation of the flow: As mentioned previously, there are alternative recipes that one can employ. These will lead to different continuous-time dynamical systems - some are more practical than others. We show that taking the vector field perspective when deriving a recipe for a momentum-enriched system is a reasonable place to start.
* Proving convergence: As mentioned previously, momentum-enriched algorithms are not gradient flows, and so establishing convergence is not easy, due to the presence of interaction terms between components in different spaces (which we note is not typically a concern for typical momentum-enriched methods), there is a question whether the Lyapunov approach is sufficient to prove convergence of the flow, and if so what will the Lyapunov function look like? We answer in the affirmative and provide hints for future work on what to choose as a suitable Lyapunov function.
* Discretization of the flow: As mentioned by eLbr, discretization of momentum-enriched systems requires care, and, as you have mentioned, the partial update is crucial for stable discretization. This idea may also be applied to other momentum-enriched systems in this product space.

Given our work, we believe that questions concerning i) how to incorporate momentum into algorithms in this product space, ii) how to study the convergence of such algorithms, as well as iii) how to discretize such algorithms are now made clearer. As such, we believe that these contributions are worthy of dissemination.

Reviewer comment (Experiments):
Question — How does the partial update discretization compare to a naive momentum step for the Variational Auto Encoder problem?

Author response:
We performed this experiment, keeping all the hyperparameters the same as before, and as suggested by the toy example, we ran into numerical instabilities. This occurred in the first couple of iterations; there is not much to add beyond this.

Reviewer comment (Theory):
I have concerns about whether the discussion on discretization in this paper is sufficient. There are two main types of discretization involved: one for the expected values of the particles, and the other along the time direction. The paper only briefly mentions using Kuntz's method in Section 5. However, it’s important to note that while the proposed method is based on Underdamped dynamics, Kuntz’s approach is for Overdamped scenarios. In usual Langevin dynamics, not specifically Mckean-Vlasov processes, it’s well-known that not paying enough attention to time discretization in Underdamped cases can lead to unstable dynamics, which is a crucial issue for sampling. But this paper does not seem to discuss these potential problems.

Author response:
We agree that naive discretisation of Underdamped Langevin dynamics can fail to yield accelerated rates of convergence. In the underdamped Langevin case, one example of an integrator that yields accelerated rates is the Exponential Integrator (see [2,3]). Our proposed algorithm is a result of utilizing the Exponential integrator with a partial update. As noted by Reviewer 1aSJ, the stable partial update is crucial for a stable discretization. We have further expanded the discussion on this matter in the manuscript.

Reviewer comment (Experiments):
The paper does not provide enough information on how to choose parameters properly. In the field of Langevin dynamics, it’s recognized that using Underdamped dynamics can lead to faster results if the parameters, like the momentum parameter, are chosen correctly. However, there is no clear guidance in the paper on how to achieve this acceleration in the context of the Mckean-Vlasov process.

Author response:
In Appendix J.3, we discuss a heuristic that we used throughout our experiments. It is based on the relationship between our discretization and NAG. While we found that this heuristic was helpful in our problems, we recognize that this is not a complete solution. This is an issue we have inherited from Underdamped Langevin dynamics where it is generally understood to be somewhat more challenging than its overdamped counterpart, as witnessed by the dearth of papers proposing practical tuning strategies for this class of algorithms. [1] describes some strategies for practical tuning of the friction parameter based on heuristics adapted to Gaussian targets, as well as tuning the step size by monitoring energy errors (in their context, acceptance rates in their Metropolis-Hastings step). In our work, we have adopted some ad-hoc strategies based on empirical monitoring of the algorithm; future work could seek to carefully extend the proposed strategies of [1] to the present context. We will expand on this in the manuscript.

Reviewer comment (Experiments):
How should we choose the momentum parameter? Additionally, due to variations in this choice, how much of a difference in performance can we expect with real data, as opposed to a toy model?

Author response:
In Appendix J.3, we discuss a heuristic that we used throughout our experiments. It is based on the relationship between our discretization and NAG. While we found that this heuristic was helpful in our problems, we recognize that this is not a complete solution. This is an issue we have inherited from Underdamped Langevin dynamics where it is generally understood to be somewhat more challenging than its overdamped counterpart, as witnessed by the dearth of papers proposing practical tuning strategies for this class of algorithms. [1] describes some strategies for practical tuning of the friction parameter based on heuristics adapted to Gaussian targets, as well as tuning the step size by monitoring energy errors (in their context, acceptance rates in their Metropolis-Hastings step). In our work, we have adopted some ad-hoc strategies based on empirical monitoring of the algorithm; future work could seek to carefully extend the proposed strategies of [1] to the present context. We will expand on this in the manuscript.

Reviewer comment (Theory):
The momentum method is usually used for acceleration. However, the theoretical advantage, such as an improved convergence rate over PGD, has not been discussed. Indeed, we cannot see the benefit of the convergence analysis (Proposition 4.1) compared to the PGD method.

Author response:
Since the convergence analysis of PGD does not extend to MPGD, the role of Proposition 4.1 is largely to establish convergence, rather than to provide evidence that MPGD is more performant than PGD. As noted by 1aSJ, gradient flow algorithms have a straightforward recipe for convergence, but since MPGD is not a gradient flow algorithm, the convergence proof is considerably more involved. If interested, please see our reply to 1aSJ where we are more explicit about the differences between establishing convergence for momentum-enriched methods and gradient flows.

Reviewer comment (Theory):
In light of the theoretical work on sampling and particle-based optimization methods, the provided analysis seems somewhat weak. For instance, the existence and smoothness of the solution of SDE (2a)-(2d), and any guarantees of the discretization (in time and space), are not provided.

Author response:
In the revised manuscript, we provided additional theoretical results proving the existence and uniqueness of the solutions with the appropriate extension, as well as asymptotic guarantees of the space discretization. The proof of existence and uniqueness of the solutions (as well as the space discretization) require an extension of McKean-Vlasov SDE arguments to the product space that is not immediate from the current theory. The statement of the results can be found in Proposition 3.1 and Proposition 5.1 with the proofs in Appendix G and H.

Reviewer comment (Theory):
Is it possible to discuss the time and particle complexities to achieve a given optimization precision? Such an analysis is quite standard in the optimization/sampling literature.

Author response:
Sharp analyses of PGD and MPGD are not yet available, so we cannot theoretically guarantee this at present. While such an analysis may be standard in the optimization/sampling literature, it is not standard for algorithms that operate in the product space (see [1,2,3,4,5] to name a few).

[1] Kuntz, J., Lim, J. N., & Johansen, A. M. (2023, April). Particle algorithms for maximum likelihood training of latent variable models. In International Conference on Artificial Intelligence and Statistics (pp. 5134-5180). PMLR.  
[2] Han, Tian, et al. "Alternating back-propagation for generator network." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 31. No. 1. 2017.  
[3] Sharrock, Louis, Daniel Dodd, and Christopher Nemeth. "CoinEM: Tuning-Free Particle-Based Variational Inference for Latent Variable Models." arXiv preprint arXiv:2305.14916 (2023).  
[4] De Bortoli, Valentin, et al. "Efficient stochastic optimisation by unadjusted Langevin Monte Carlo: Application to maximum marginal likelihood and empirical Bayesian estimation." Statistics and Computing 31 (2021): 1-18.  
[5] Nijkamp, Erik, et al. "Learning multi-layer latent variable model via variational optimization of short run mcmc for approximate inference." Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VI 16. Springer International Publishing, 2020.  
[6] Ma, Yi-An, et al. "Sampling can be faster than optimization." Proceedings of the National Academy of Sciences 116.42 (2019): 20881-20885.  
[7] Holley, Richard, and Daniel W. Stroock. "Logarithmic Sobolev inequalities and stochastic Ising models." (1986).

Reviewer comment (Experiments):
Motivation for this particular extension. While it is natural to consider "momentum enriched" optimization algorithms in this context, I am left wondering whether there is a **strong** motivation, precisely: are there problems where PGD actually **struggles**, *and* the new MPGD (even if well-tuned) really resolves the situation? This is not very clear from the current version of the paper. Yes, we "can" do this extension and it's natural, and it seems to perform better (focussing for now on the toy example, where metrics are more interpretable and relate directly to the MLE problem) than PGD, but the toy example only shows a (relatively minor) *improvement* over PGD which is already working well. Could you think of an example where the benefits of MPGD vs PGD really show, i.e., PGD is converging very slowly and MPGD "fixes" the problem? Again in other words, the authors claim "we derive a discretization of the flow that can achieve better performance than PGD ...", and this is supported by the experiments. It "can be better" indeed, but how relevant is that? It would be stronger to show cases where PGD doesn't do well and MPGD really does (as opposed to cases where they both do well and MPGD does a bit better).

Author response:
Thank you for your comments and questions.

> While it is natural to consider "momentum enriched" optimization algorithms in this context, I am left wondering whether there is a strong motivation, precisely: are there problems where PGD actually struggles, and the new MPGD (even if well-tuned) really resolves the situation?

We find that in many cases, MPGD does indeed perform substantially better than the vanilla PGD. Originally, we aimed to demonstrate this through the VAE example (for instance, Table 1), which other reviewers agree demonstrates its performance gains

> [1aSJ] The second dataset is a more realistic problem of building a Variational Auto Encoder where the momentum method described here seems to work very effectively.

The toy hierarchical example was intended to illustrate various aspects of MPGD (e.g., its regimes) and validate its design (in particular, the specific discretization we use and the gradient correction).

Still, you are entirely correct: by simply increasing the model's variance, the performance gains for MPGD relative to PGD again become large for this example; see Appendix K. As the variance of the model increases, PGD struggles. Additionally, we have demonstrated empirically that PGD is more sensitive to initialization of the cloud and thus degrades with poor initialization. Thank you very much for this suggestion. As further evidence, we have added a new example (Appendix K.2). It is density estimation in $1$-d where the true density is a bi-model Gaussian. In 1-d, we can estimate the Wasserstein-$1$ ($W_1$) distance easily. In this problem, MPGD shows a significant improvement over PGD - the final $W_1$ of PGD is around $0.3$ and MPGD is around $0.1$ - a three-fold decrease.

Reviewer comment (Experiments):
This point is related to the above; in the toy experiment, there is not a significant (in the sense of practically, not statistically) difference between MPGD and PGD (Figs 1(a) and (b) ) except for the underdamped and overdamped regimes (which presumably are less desirable). Can you try with different parameters (e.g. variances of the hierarchical model) to see if there is a setting where MPGD really shows its strength?

Author response:
Thank you for your comments and questions.

> Can you try with different parameters (e.g. variances of the hierarchical model)

Thank you for your suggestion, this is a good idea. The results can be found in Appendix J.1. As the variance increases, it can be seen that PGD struggles to converge while MPGD still converges, thus performing significantly better. This can be seen in Figure 6, where we provide the requested experiment comparing PGD and MPGD on the Toy HM with different variances.

Reviewer comment (Writing):
In the last paragraph before 2.3 starts, what does the symbol $\bigotimes$ denote?

Author response:
> What does the symbol $\otimes$ denote

$p \otimes q$ refers to the product measure $p(\mathrm{d}x) \cdot q(\mathrm{d} y)$. Thank you for pointing this out, we have now defined it when it first appears.

Reviewer comment (Experiments):
You enrich both components of the space with momentum variables but why? Is one of the two not enough?

Author response:
Thank you for your comments and questions.

> You enrich both components of the space with momentum variables but why ? Is one of the two not enough?

We found that MPGD performed better than algorithms that only enrich one component (or none at all). The ordering can be seen to be MPGD > Momentum in $\theta$ $\approx$ Momentum in $X$ > PGD. The $\approx$ indicates that there are settings where one is better than the other and vice versa. To illustrate this observation, we provide two additional experiments (see Figure 7 and Figure 8 in Appendix K).

In Figure 7, we use the Toy HM example with $\sigma=12$ and compare MPGD (and PGD) with algorithms that enrich one of the spaces with momentum. In Figure 7, we show the effect of initialization of the particle cloud on each method. Overall, it can be seen that the presence of momentum results in a better performance. Furthermore, we observe that the algorithms which only enrich one component can exhibit problems which are not present for MPGD. For instance, in Figure 7, poor initialization of the cloud drastically lengthens the transient phase of the $\theta$-only enriched algorithm, whereas MPGD can recover more rapidly. This is a typical setting in high-dimensional settings (see Figure 3). Conversely, $X$-only-enriched algorithms are observed to suffer from slow convergence in the $\theta$ component (intuitively, one pays a larger price for poor conditioning of the ideal MLE objective). Overall, it can be observed that MPGD noticeably outperforms PGD and all other intermediate methods on all settings.

In Figure 8, as a more realistic example, we provide an additional experiment on a density estimation task using VAEs on a $1$-dimensional problem. Figure 8b shows the empirical Wasserstein-$1$ distance between the model and the data throughout the training process. It can be seen that by doubly enriching the space (i.e. using MPGD), the transient period is noticeably shortened, attaining a lower $W_1$ than both PGD and the other algorithms which enrich only one of the spaces.