Momentum Particle Maximum Likelihood

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 then we would have a deluge of papers!

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.

How does the partial update discretization compare to a naive momentum step for the Variational Auto Encoder problem.

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.

How does this algorithm compare to the performance of other current particle-based approaches such as SVGD-like approaches.

In our experiment, we compare against other particle methods such as Alternating Backpropagation and Short Run MCMC. We chose not to compare SVGD since it incurs an $\mathcal{O}(M^2)$ cost instead of $\mathcal{O}(M)$ like ours and our other baselines.

[1] Wibisono, Andre, Ashia C. Wilson, and Michael I. Jordan. "A variational perspective on accelerated methods in optimization." proceedings of the National Academy of Sciences 113.47 (2016): E7351-E7358.
[2] Taghvaei, Amirhossein, and Prashant Mehta. "Accelerated flow for probability distributions." International Conference on Machine Learning. PMLR, 2019.
[3] Yi-An Ma. Niladri S. Chatterji. Xiang Cheng. Nicolas Flammarion. Peter L. Bartlett. Michael I. Jordan. "Is there an analog of Nesterov acceleration for gradient-based MCMC?." Bernoulli 27 (3) 1942 - 1992, August 2021. https://doi.org/10.3150/20-BEJ1297
[4] Su, Weijie, Stephen Boyd, and Emmanuel Candes. "A differential equation for modeling Nesterov’s accelerated gradient method: theory and insights." Advances in neural information processing systems 27 (2014).

Thank you for your suggestions and comments. Here is our response:

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.

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

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. It does go to overall quality of the paper though.

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

Thank you for your comments.

I have concerns about whether the discussion on discretization in this paper is sufficient ... But this paper does not seem to discuss these potential problems.

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.

Additionally, the paper does not provide enough information on how to choose parameters properly. ... How should we choose the momentum parameter?

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.

[1] Riou-Durand, Lionel, et al. "Adaptive Tuning for Metropolis Adjusted Langevin Trajectories." International Conference on Artificial Intelligence and Statistics. PMLR, 2023.
[2] Cheng, Xiang, et al. "Underdamped Langevin MCMC: A non-asymptotic analysis." Conference on learning theory. PMLR, 2018.
[3] Sanz-Serna, Jesus Maria, and Konstantinos C. Zygalakis. "Wasserstein distance estimates for the distributions of numerical approximations to ergodic stochastic differential equations." The Journal of Machine Learning Research 22.1 (2021): 11006-11042.

It is better that the authors provide an example that satisfies all assumptions to convince the readers that the theory is not vacuous.

There is already an example in the paper that satisfies the assumptions: the toy hierarchical model in Section 6.1. Its log-likelihood is a quadratic polynomial of $(\theta,x)$, so the log-likelihood's gradient is linear and trivially Lipschitz. The log-likelihood is also strongly concave (cf. Appendix J.4.1), and the extended log Sobolev-Inequality holds for all models whose log-likelihoods are strongly concave (i.e., whose Hessian is upper bounded by $-\lambda$ times the identity matrix for some constant $\lambda>0$). The argument for this is as follows: if $(\theta_t,q_t)$ denote the solution to $\cal{E}$'s gradient flow (see Section 2.3 in our paper), then Eqs. (28,30) in [1] respectively read

I(\theta_t, q_t) := -\frac{\mathrm{d}}{\mathrm{d}t} \mathcal{E}(\theta_t, q_t) =\|\nabla \mathcal{E}(\theta_t, q_t)\|^2_{q_t},\qquad \frac{\mathrm{d}}{\mathrm{d}t}I(\theta_t, q_t) \le -2\lambda I(\theta_t, q_t) .$$ (In [1] the minus sign is missing from (30)'s RHS, but this is a typo as can be easily verified from the equation's proof.) The inequality implies that $I(\theta_t, q_t)\leq e^{-2\lambda t}I(\theta_0, q_0)$. Plugging this into (28) and integrating it over time we find that

\mathcal{E}(\theta_0, q_0) - \mathcal{E}(\theta_t, q_t) \le |\nabla \mathcal{E}(\theta_0, q_0)|^2_{ q_0}\left (\frac{1- \exp(-2\lambda t)}{2\lambda} \right ).\qquad(*)$$

In [1, Theorem 3], it was argued that $(\theta_t,q_t)$ converges to $(\theta^*, p_{\theta^*}(\cdot|y))$ in a certain sense as $t$ tends to infinity. It is possible to strengthen this statement to $\lim_{t\to\infty} \mathcal{E}(\theta_t,q_t)=\mathcal{E}(\theta^*, p_{\theta^*}(\cdot |y))$. By $\cal{E}$'s definition, $\mathcal{E}(\theta^*, p_{\theta^*}(\cdot |y)) = - \log p_{\theta^*}(y)$ and, because the flow's initial condition $(\theta_0,q_0)$ was arbitrary, the LSI follows by taking the limit $t\to\infty$ in $(*)$.

However, formally establishing the limit $\lim_{t\to\infty} \mathcal{E}(\theta_t,q_t)=\mathcal{E}(\theta^*, p_{\theta^*}(\cdot |y))$ requires a somewhat involved and lengthy study of the PGD gradient flow which lies beyond the scope of this paper (we are aware of an upcoming publication that covers this). Nevertheless, we hope the above sketch will persuade you that there are indeed many models that satisfy our assumptions and that our theoretical analysis is not ultimately vacuous.

Aside: We conjecture that the extended log-Sobolev inequality is more general than strong log-concavity (as is the case for the classical LSI). This is motivated by the hints in [6, Proposition 1], which established log-Sobolev inequality for potentials that are strongly concave outside a ball. However, to establish this fact would require a non-trivial generalization of the widely-used Holley-Stroock comparison Theorem [7].

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.

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

We thank the reviewer for their comments.

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.

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.

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.

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.

The important assumptions should be exposed in the main text, and their reasonability should be discussed. However, Assumptions 1 and 2 (Lipschitz smoothness and LSI) are hidden in the Appendix.

We agree with this and have made the necessary changes (see Section 4).

The authors missed the series of mean-field optimization works

Thank you for your references. We have added a section in the Appendix dedicated to related work which included your reference (see Appendix B).

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.

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.

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.