Convex Potential Mirror Langevin Algorithm for Efficient Sampling of Energy-Based Models

Authors: Yang Zitao, Amin Ullah, Shuai Li, Fuxin Li, Jun Li
Primary area: generative models
Keywords: mirror Langevin dynamics, efficient sampling, energy based models

Abstract:
This paper introduces the Convex Potential Mirror Langevin Algorithm (CPMLA), a novel method designed to optimize sampling efficiency within Energy-Based Models (EBMs). CPMLA employs mirror Langevin dynamics in conjunction with convex potential flow as a dynamic mirror map for sampling in EBMs. By leveraging this dynamic mirror map, CPMLA enables targeted geometric exploration on the data manifold, enhancing the convergence process towards the target distribution. Theoretical analysis proves that CPMLA achieves exponential convergence with vanishing bias under relaxed log-concave conditions, supporting its efficiency and effectiveness in adapting to complex data distributions. Experimental results on established benchmarks like CIFAR-10, SVHN, and CelebA showcase CPMLA's enhanced sampling quality and inference efficiency compared to existing techniques.

Decision: reject
Reviewer scores: 8, 8, 3, 3

Reviews and discussion:
<thread>
creation date:
1729956964289

thread_type:
official_review

<thread_object>
creator: reviewer_67jh

<content>
summary:
The authors consider the specific sampling problem where the potential is given by some learned energy. This work proposes learning a convex mirror potential alongside the energy-based model for generation, motivated by the optimal mappings given by Brenier's theorem. The authors make two key approximations to computing the convex conjugate of an input-convex neural network, and in computing the Hessian of the optimal mapping. The sampling algorithm boils down to constructing a dual mapping, performing mirror Langevin sampling in the dual space from dual noise, and mirroring back to the primal space.

strengths:
* The idea of jointly training an energy model along with another convex model to guide training seems new. There seems to be a hidden adversarial learning perspective, where the convex guide model is learned to have large deviation between clean and generated samples, which allows for an intuitive interpretation as "larger gradient steps towards ground truth samples".

weaknesses:
* (Critical) The algorithm and corresponding theory in its current form are dubious. There is no mention as to how the inverse map $\nabla G^*$ is computed in practice. The theory claims to show convergence of the CPMLA algorithm *during training*, but does not consider the learnable parameters (there is no mention of $\theta$ or $\vartheta$ anywhere), and appears to be a clone of Prop. 3 in Jiang (2021). I suspect it might be possible using techniques from De Bortoli et al. (2021), but this is not considered in this work. 

* (Critical) Algorithmically, there seems to be a critical typo in Algorithm 1 on line 284, on the "gradient descent part", where $\nabla f$ is evaluated on an element of the dual space. This appears also in Equation 8 in the starred equality: it should be $\nabla G(\nabla G^* y_t)$ inside the sqrt. This raises some suspicion about how the authors performed the experiments.
* There is a lack of explanation of how the convex conjugate of the network is actually computed. This is explicitly given in the references but is not found in the main text or appendix.
* For the additional amount of computation that is supposedly needed to perform the mirroring, does the CPMLA method outperform standard Langevin methods during test-time, i.e. with a fixed energy model? This should be addressed to accurately consider the practicality of this method.
* Could the authors clarify the motivation behind Eq. 8 and the surrounding exposition as ``practical application'', if the version of MLA used in Algorithm 1 comes from discretizing Eq. 7? Same with the seemingly contradictory paragraph above Algorithm 1. I believe it has to do with the noise distribution $\xi \sim \mathcal{N}(0, \nabla^2 G)$ but I am not sure.
* While theoretically motivated, the practical approximations are hidden away in the text. For example, diagonal Hessian approximation (l.263) is only mentioned once, but is critical to maintaining computational feasibility. The paper could use a paragraph explaining all the approximations made, perhaps somewhere next to Algorithm 1.  
* There is no comparison towards classical MCMC methods when performing sampling from EBMs. However I am unfamiliar with the energy-based-model sampling literature, so perhaps this is not critical.
* Figure 2 appears to have an iteration cut-off at an arbitrarily chosen $T=21$, where it seems that CoopFlow may overtake CPMLA. In the following table, CPMLA is beaten by CoopFlow at $T=30$. While I appreciate that CoopFlow has significantly more parameters which may cause memory concerns, the iteration count is not enough to discern whether the proposed CPMLA is a practical method. There is no comparison towards methods with similar parameter count either.

[1] De Bortoli, V., Durmus, A., Pereyra, M., & Vidal, A. F. (2021). Efficient stochastic optimisation by unadjusted Langevin Monte Carlo: Application to maximum marginal likelihood and empirical Bayesian estimation. Statistics and Computing, 31, 1-18.

questions:
* (Critical) Clarify that the convergence result in Section 5 does not apply to the training, but rather to sampling for a fixed $G$.
* The assumptions required for Section 5 are not complete: $G$ needs to be $\mathcal{C}^3$.
* Please define relaxed log-concavity.
* Please add a motivation as to where $\log \det$ comes from when training CP-Flow, see Huang et al. (2020).
* (Eq 2) The stated ``Langevin Monte Carlo'' is more commonly called the unadjusted Langevin algorithm (ULA). I suggest using ULA and citing the required references. 
* How is $\nabla G^*$ computed? The authors initially claim that you can compute the conjugate using convex optimization, but this does not give you the gradient.
* (l.242) ``conugate''
* Please weaken the claim at the start of Section 5 about convergence of DNNs. This is a special case of general theory.
* (Appendix F) "our work introduces" -> "our work considers", mirror log-Sobolev has already been introduced in Jiang 2021.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677465989

thread_type:
official_review

<thread_object>
creator: reviewer_sitl

<content>
summary:
This paper expects to enhance sampling quality and inference efficiency by introducing the convex potential flow (CP-Flow) to energy-based models (EBM). Different from standard EBM, this paper not only approximates the gradient of log density of target distribution but also considers the SDE path toward the target. The benefits of enforcing the SDE path to optimal transport in this paper lie in decreasing the mixing time of the convergence expected for the continuous version.

strengths:
1. The motivation for introducing the SDE path requirement to EBM to enhance the sampling efficiency is interesting and reasonable from an intuition perspective.
2. This paper is well written; the intuition, implementation, and experiments are clear.

weaknesses:
1. The novelty of this paper is not enough; controlling the path of generation with additional models beyond EBM is not novel and is highly similar to the baseline Coopflow. This paper only replaces the normalizing flow (in Coopflow) with mirror Langevin Dynamics.
2. The analysis of this paper does not match their claim:
    1. Theorem 5.5 is only established when the approximation error of EBM and CP-Flow are both 0. That means this analysis does not consider how the approximation error affects the final inference efficiency. Since this paper considers a generative model rather than a sampling algorithm, I think the approximation error is necessary to be considered, just as shown in [Chen2022].
    2. The analysis of Theorem 5.5 does not include any novel technique, which just utilizes the framework proposed in [Vempala2019]  to analyze a sampling algorithm. Moreover, the Fokker-Planck equation is actually given by [Jiang 2021]
    3. Line. 37—Line. 38, authors argue that MCMC tends to get trapped in local modes and has a slow convergence. Actually, the analysis of this paper shows that CPMLA achieves KL convergence with the same gradient complexity as that in unadjusted Langevin Algorithm [Vempala2019] w.r.t. the dimension $d$ and error tolerance $\epsilon$ dependence. Besides, it is worse than underdamped Langevin Dynamics [Zhang2023]. 
3. The experiment does not show a significant advantage compared with Coopflow. I suggest that the authors provide an FID curve when T is larger than 20. Since the decrease in the tendency of Coopflow is sharper than that of CPMLA, in Table 2, I suggest that the authors also show the performance of CPMLA when T=30.

I am willing to raise my score if the authors solve my concerns.



[Chen2022] Chen S, Chewi S, Li J, et al. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions[J]. arXiv preprint arXiv:2209.11215, 2022.

[Jiang2021] Qijia Jiang. Mirror langevin monte carlo: the case under isoperimetry. Neural Information Processing Systems, 2021.

[Vempala2019] Vempala S, Wibisono A. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices[J]. Advances in neural information processing systems, 2019, 32.

[Zhang2023] Zhang S, Chewi S, Li M, et al. Improved discretization analysis for underdamped Langevin Monte Carlo[C]//The Thirty Sixth Annual Conference on Learning Theory. PMLR, 2023: 36-71.

questions:
Please check the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701545468

thread_type:
official_review

<thread_object>
creator: reviewer_wyv9

<content>
summary:
This paper introduces a novel approach for efficiently training and sampling from energy-based models, addressing key limitations such as the challenges in obtaining high-quality samples, which are essential for the training of energy-based models to converge. Additionally, the authors provide valuable theoretical justifications for the convergence of mirror Langevin algorithms within the context of deep neural networks. They demonstrate the effectiveness of their method across several established benchmarks, including CIFAR-10, SVHN, and CelebA, and carefully compare their results with existing approaches.

strengths:
* The proposed method is well-presented, featuring convincing empirical evaluations and solid theoretical justifications.
* The paper is clearly written and effectively organized.
* The proposed algorithm enhances sampling efficiency from energy-based models without significantly increasing the number of parameters in the CP-Flow component.
* The authors include a comprehensive appendix that provides detailed information on the experimental setups and the theoretical proofs.

weaknesses:
* The parameterization of the CP-Flow model in the proposed algorithm appears too constrained, which may complicate the optimization process. Additionally, the authors approximate the Hessian matrix during training using only its diagonal components, and I believe this estimation will become less accurate as the dimensionality increases.

questions:
I am curious about the efficiency of the proposed method in comparison to the approaches presented in [1] and [2], which are relatively easy to implement for both the toy problems and the image generation tasks discussed in the paper. Notably, these methods seem more relevant to the proposed algorithm than some of the literature cited.

[1] Carbone, D., Hua, M., Coste, S., & Vanden-Eijnden, E. (2023). Efficient training of energy-based models using Jarzynski equality. Advances in Neural Information Processing Systems, 36, 52583-52614.

[2] Du, Y., Li, S., Tenenbaum, J., & Mordatch, I. (2020). Improved contrastive divergence training of energy-based models. arXiv preprint arXiv:2012.01316.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730767161455

thread_type:
official_review

<thread_object>
creator: reviewer_kvdi

<content>
summary:
This paper studies mirror Langevin descent algorithm, where the mirror map is the gradient of a convex function. Novel theoretical results are proven for the convergence of this algorithm under weaker assumptions than previously available in the paper "The Mirror Langevin Algorithm Converges with Vanishing Bias", and in a stronger metric (KL divergence). The mirror map is modelled using an input-convex neural network (Amos et al, 2017), which is formulated in a way to guarantee invertibility of the mirror map. This is called CP-Flow. By jointly optimizing the parameters of this network, and performing sampling, an efficient algorithm is obtained for generative modelling, with dramatically fewer parameters compared to other models (this is due to faster convergence because of the preconditioning of the Langevin dynamics with the mirror map).

strengths:
The paper proposes a very interesting new algorithm for generative modelling based on a combination of input convex neural networks with mirror Langevin descent. There are some new theoretical results showing vanishing bias and bounding convergence rates under weak assumptions. The empirical performance is convincing, especially given the fact that the total number of parameters is much lower than some competing models.

weaknesses:
The FID scores on the example are not yet competitive with the latest score-based generative models such as NCSN++. Nevertheless, there is a lot of potential in this method, and simply modifications might still significantly improve performance.

In terms of presentation, a better explanation of exactly how is the mirror map function G is found in practice would be useful. We have found section 3.2 too short and difficult to fully grasp.

questions:
Could you please give more details on exactly how the mirror map function G is found? Do you alternate between exploring the target distribution, and updating G in an adaptive way, or do you fix G in advance before applying mirror Langevin? You seem to mention cooperative learning on page 6, do you mean that G is updated several times after sampling?

In particular, Section 3.2 should be clarified. Section 4.2 is also not sufficiently clearly explained, especially this cooperative learning aspect.
It would be important to be precise about the implemented algorithm, if there is not sufficient space in the main paper, you could add this discussion in the appendix.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733310841762

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response to all reviewers and ACs

comment:
We sincerely thank the reviewers for their positive feedback on the novelty and effectiveness of the proposed CPMLA and their constructive suggestions on the theoretical and experimental improvement. We summarize the reviewers' comments and our rebuttals as follows.

### Summary of Reviewer Feedback:

#### Strengths:
- **Novelty**: Reviewers acknowledged the innovative approach of combining mirror Langevin dynamics with convex potential flow, as well as the novel theoretical results that prove exponential convergence under relaxed conditions.
- **Empirical Results**: Our empirical evaluation, including comparisons with established benchmarks, was positively noted, especially regarding CPMLA’s ability to achieve high sampling efficiency and quality with fewer parameters compared to other methods.
- **Presentation**: The clarity and organization of the paper were appreciated, with particular praise for the clear motivation and sound theoretical foundations.

#### Suggested Revisions:
We have carefully considered the points raised by the reviewers and present the key revisions that will be incorporated in the revised manuscript:

**Clarity on Mirror Map Calculation**:
  We will enhance Section 3.2 by adding a detailed explanation of how the mirror map $\nabla G$ is computed, including how the optimization is performed. We will elaborate on how the mirror map $\nabla G$ is updated adaptively during the training process, as mentioned in Lines 70 and 287, with references to the cooperative learning mechanism described in Line 389.

**FID Scores and Comparisons**:
  We will include further experiments comparing CPMLA with more relevant models and present results with extended iterations to highlight CPMLA's competitive advantages over wall clock time. 

  **Table 1:** A comparison of wall clock time for inferring 1,000 images of similar quality on the CIFAR-10 dataset  
|                    | CPMLAprt (T = 20) | CoopFlow (T = 30) |
|:------------------:|:-----------------:|:-----------------:|
| Time (s/1k images) | 16.14             | 16.84             |

**Table 2:** A comparison of FID scores under different iteration counts  
|                   | 10    | 20    | 24    | 27    | 30    |
|:-----------------:|:-----:|:-----:|:-----:|:-----:|:-----:|
| CoopFlow (T = 30) | 57.26 | 30.74 | 25.03 | 22.45 | 21.16 |
| CPMLA (T = 20)    | 56.40 | 26.92 | 24.76 | 23.47 | 22.35 |

  We will also include a comparison with methods of similar parameter counts to further demonstrate the efficiency of our approach.

**Table 3:** A comparison of FID scores and the number of parameters across different EBMs and flow-based frameworks  
| Methods       | Number of Parameters | FID   |
|:-------------:|:--------------------:|:-----:|
| NT-EBM        | 23.8M                | 78.12 |
| GLOW          | 44.2M                | 45.99 |
| EBM_FCE       | 44.9M                | 37.30 |
| CoopFlow      | 45.9M                | 21.16 |
| Flow++ only   | 28.8M                | 92.10 |
| VAEBM         | 135.1M               | 12.16 |
| **CPMLA**     | **17.39M**           | **22.35** |

**Theoretical Extentions**:
  We will revise Section 5 to provide a more nuanced explanation regarding the approximation error and its impact on inference efficiency. We will clarify that our advantages on relaxed log-concavity assumption and mixing time, and update Theorem 5.5 to prove that the total variation distance $d_{TV}(\rho_t, p_{\text{data}}) < \delta$.

**Hessian Approximation and Optimization**:
  We will add a discussion detailing the advantages and limitations of the diagonal Hessian approximation used in our method, including its benefits for computational efficiency, particularly in high-dimensional settings where full Hessian calculations are not feasible. This will be addressed near Algorithm 1.

**Typo Clarification**:
  We will ensure that all mathematical formulations in Algorithm 1 are accurate and clearly explained. We will correct the identified typo in Line 284, where the gradient is evaluated in the dual space $y - \eta \nabla f_{\theta}(\nabla G^* y) + \sqrt{2\eta}\xi$, and revise notations in Equation 8 $\nabla^2 G(\nabla G^* y_t)$. These clarifications will help ensure that readers can easily understand the algorithm's implementation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors introduced a new method for sampling from energy based models, where the algorithm trains a dynamic mirror map, before using a mirror based Langevin algorithm to sample from the energy. The reviews here have very high variance, and made the decision somewhat unclear. On the positive side, the algorithm appears to be novel, the authors performed extensive empirical evaluation, and the performance is good albeit not quite state of the art. On the negative side, an extended discussion with a reviewer led to a lot of confusion on both the implementation and the theoretical aspects of this paper. 

Most importantly, this discussion did not resolve all of the original concerns. And without access to the significant revisions to the manuscript on both the theoretical front, and clarity into the implementation of the algorithm, many of these concerns cannot be resolved at this point. Unlike the review process for journals, the reviewers at ICLR cannot request a major revision to further continue the discussion and improve the manuscript. Therefore, I believe the paper in its current form is not ready to be accepted. 

Given that I also personally find many components of this paper confusing to read myself, I would encourage the authors try to reformulate the story and technical explanations of this paper before resubmitting. There is a potentially very interesting algorithm with theoretical results here, and an improved manuscript should be a welcomed addition to the energy based models literature.