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

Q1: There is no mention as to how the inverse map is computed in practice.
A1: The computation of the inverse of the mirror map, $(\nabla G)^{-1}$, is discussed in line 225 and further detailed in Appendix D. Specifically, $(\nabla G)^{-1}$ is computed as the convex conjugate of the mirror map $\nabla G$, as defined within the framework of our convex potential flow (CP-Flow). The convexity of CP-Flow ensures both the existence and uniqueness of this inverse. Our approach leverages conjugate convex optimization to efficiently compute $(\nabla G)^{-1}$, enabling CPMLA to smoothly transition between primal and dual spaces without introducing significant computational overhead.

Q2: The theory claims to show convergence of the CPMLA algorithm during training, but does not consider the learnable parameters.
A2: In our theoretical analysis, Assumptions 5.1 through 5.4 ensure that the convergence properties hold even as the parameters of CP-Flow and the EBM evolve during training. Importantly, the training parameters do not alter the underlying network structure or its properties. This is guaranteed by the structure of $G$, as stated in Equation 5. Notably, the convex potential retains strong convexity with respect to the input $x$, independent of the learnable parameters $\vartheta$, ensuring the persistence of the unique inverse map, the retention of optimality via Brenier's theorem, and the continued ability to optimize the loss function without affecting convergence guarantees.

Q3: Critical typos in Algorithm 1: $\nabla f$ is evaluated in dual space. And it should be $\nabla G(\nabla G^*y_t)$ in the sqrt.

A3: These are not typos, indeed these are correct formulation of Mirror Langevin Algorithm. We provide a detailed explanation below.

• Gradient Descent in the Dual Space: The method of evaluating the gradient in the dual space, as presented in Algorithm 1, follows the standard procedure for mirror gradient descent and mirror Langevin Dynamics. Specifically, previous works in this field, including those by Hsieh et al. (2018) and Jiang (2021), have consistently evaluated the gradient in the dual space. This approach aligns with the fundamental design of the mirror Langevin dynamics and is fully consistant with the established literature on this topic.

• Square Root in Equation 8: We would like to clarify that the term $\nabla^2 G(y_t)$, which appears under the square root, is indeed correctly placed according to the derivations outlined in the reference by Jiang (2021), particularly in Section 4.3, which discusses $MLA_{\text{AFD}}$. In this context, $\nabla^2 G$ represents the Hessian of the convex potential and is properly evaluated at each iteration, reflecting the standard formulation for alternative forward discretization.

Q4: Does CPMLA with additional computation outperform standard Langevin methods during test time with a fixed EBM?
A4: CPMLA outperforms CoopFlow in terms of both inference efficiency and sample quality, and CoopFlow itself already provides significant improvements over standard Langevin methods. Therefore, CPMLA represents a meaningful advancement in comparison with the standard Langevin techniques. Standard Langevin sampling methods are computationally demanding and non-mixing. The additional computation introduced by CPMLA is designed to address these limitations and create a better sampling algorithm. Specifically, CPMLA demonstrates a clear performance advantage over standard Langevin Monte Carlo methods under similar test conditions, particularly when using a fixed EBM on CIFAR-10, as shown in Appendix I, Figure 6.

Q5: What is the motivation behind Equation 8? Algorithm 1 seems contradictory with Equation 7 and 8.

A5:

• Motivation behind Equation 8: Equation 8 and Algorithm 1 use a discretization method where the gradient descent and SDE steps are separated. This design is deliberate, as it allows the SDE step’s drift term to become zero, simplifying the implementation and improving accuracy when discretizing. By removing the drift term in the SDE, we achieve a straightforward numerical scheme for the noise term that is computationally efficient. Also the Hessian of the convex potential plays a key role in adapting the sampling steps to the underlying geometry of the data manifold. The motivation for this inclusion is well-documented in the literature, particularly in the work by Jiang (2021), which discusses the alternative forward discretization scheme and highlights its advantages in terms of convergence and bias reduction.

• Clarification on Apparent Contradiction: The points raised in lines 259 and 269 of the paper clearly explain how Algorithm 1 and Equation 8, though differing slightly in form, represent the same underlying discretization approach for MLA. It IS related to the noise distribution $\xi \sim \mathcal{N}(0, \nabla^2 G)$.

Q6: The practical approximations are hidden away in the text, especially the diagonal Hessian approximation.

A6: The diagonal Hessian approximation is a practical approximation employed in our method mentioned on line 263. This simplification allows us to avoid the full Hessian calculation, significantly reducing computational costs without compromising sampling accuracy. This choice is particularly advantageous in high-dimensional settings where full Hessian computations are impractical. We will revise the manuscript to include a paragraph of practical modification near Algorithm 1.

Q7: There is an arbitrarily chosen cut-off $T=21$ in Figure 2. There is no comparison with similar parameter count.

A7:

• Clarification on Iteration Cut-Off: The iteration cut-off in Figure 2 was selected to showcase CPMLA's rapid convergence and inference efficiency compared to CoopFlow. We acknowledge that it may appear that CoopFlow could overtake CPMLA if allowed more iterations. However, our goal was to emphasize the significant efficiency gains CPMLA provides early on. Specifically, as shown in Table 2, CPMLAprt achieves a FID score of 21.43 at $T = 20$, which is comparable to CoopFlow's FID of 21.16 at $T = 30$. This represents a 33% reduction in the number of iterations required to reach similar performance, demonstrating CPMLA's practical advantages in scenarios where faster sampling is critical.

• Memory Efficiency and Parameter Count: The key strength of CPMLA lies in achieving better sampling efficiency with fewer parameters, rather than generating higher-quality images with the same parameter count. The significant difference between 0.26M and 28.78M parameters highlights CPMLA’s efficiency, which is particularly important for applications with limited computational resources. While our primary comparisons focused on CoopFlow, we are open to extending future analyses to include methods with comparable parameter counts to provide a more comprehensive evaluation.

Q8: Please define relaxed log-concavity.

A8: Relaxed log-concavity is a key assumption of target distributions as we have mentioned in Line 75 and 109. It is a less restrictive version of log-concavity, permitting some flexibility while still maintaining key properties like unimodality or a tendency towards a peaked structure in the function. To begin, we first define strongly log-concave. Just as a strongly convex function ensures well-conditioned geometry for fast, unbiased optimization, a strongly log-concave distribution ensures a similar well-conditioned structure, facilitating efficient sampling of true samples. This property is explored in several works, such as [1][2][3][4][5].

A function $f(x)$ is strongly log-concave if the logarithm of the function is strongly concave, i.e., for any two points $x_1$ and $x_2$ and for any $\lambda \in [0, 1]$,

In our work, relaxed log-concavity is a broader condition compared to "strongly log-concave" distributions. Specifically, a distribution satisfying a Logarithmic Sobolev Inequality (LSI) can be considered an example of a "relaxed log-concave" distribution. The LSI condition ensures certain concentration properties similar to those of log-concave distributions, even though the global curvature constraints are less stringent.

[1] Alain Durmus and Eric Moulines. Non-asymptotic convergence analysis for the unadjusted langevin algorithm. Annals of Applied Probability,Annals of Applied Probability, Dec 2016a.

[2] Alain Durmus and Eric Moulines. High-dimensional bayesian inference via the unadjusted langevin algorithm. Bernoulli,Bernoulli, Dec 2016b.

[3] Xiang Cheng, Niladri S. Chatterji, PeterL. Bartlett, and MichaelI. Jordan. Underdamped langevin mcmc: A non-asymptotic analysis. Cornell University - arXiv,Cornell University - arXiv, Jul 2017.

[4] Arnak S. Dalalyan and Avetik Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications, 129(12):5278–5311, Dec 2019. doi: 10.1016/j.spa.2019.02.016. URL http://dx.doi.org/10.1016/j.spa.2019.02.016.

[5] Wenlong Mou, Yuanlin Ma, MartinJ. Wainwright, PeterL. Bartlett, and MichaelI. Jordan. High- order langevin diffusion yields an accelerated mcmc algorithm. arXiv: Machine Learning,arXiv: Machine Learning, Aug 2019.

Thank you for your comment.

The Langevin step in Algorithm 1 should be $y - \eta \nabla f_{\theta}(\nabla G^*y) + \sqrt{2\eta}\xi$. To address the gap introduced by this adjustment, we have revised Theorem 5.5 accordingly.

We prove that the total variation distance satisfies $d_{TV}(\rho_t, p_{\text{data}}) < \delta$, where $\delta$ is a small constant. A sketch of the proof is provided below.

We decompose the total variation distance between $\rho_t$ and $p_{\text{data}}$ into three terms, providing controls for each. This revised formulation guarantees that $\rho_t$ converges to the data distribution, thereby addressing the reviewer's concerns.

Specifically, let $q_{\vartheta^*}$ denote the target distribution of CP-Flow and $p_{\theta^*}$ the target of the EBM. For the total variation distance $d_{TV}(\rho_t, p_{\text{data}})$, we decompose it into three terms:
$d_{TV}(\rho_t, p_{\text{data}}) \leq d_{TV}(\rho_t, p_{\theta^*}) + d_{TV}(p_{\theta^*}, q_{\vartheta^*}) + d_{TV}(q_{\vartheta^*}, p_{\text{data}}).$

• The first term is controlled under the original Theorem 5.5 (substituting $\delta$ with $\delta_1$). Using Pinsker's inequality we establish: $d_{TV}(\rho_t, p_{\theta^*}) \leq \sqrt{\frac{1}{2} D_{KL}(\rho_t \| p_{\theta^*})} < \sqrt{\frac{\delta_1}{2}}.$
• The second term can be controlled because Langevin MCMC only runs for $T$ steps rather than converging to its stationary distribution. By analyzing the Fokker-Planck equation of Langevin sampling: $\frac{\partial p_t(x)}{\partial t} = -\nabla_x \cdot \left(p_t(x) \frac{\eta^2}{2} \nabla_x f_\theta(x)\right) + \frac{\eta^2}{2} \nabla_x^2 p_t(x),$
  we estimate the incremental change as $d_{TV}(p_t, p_{t-1}) \leq \sqrt{\frac{1}{2} D_{KL}(p_t \| p_{t-1})} \sim O(\eta)$. Summing over $T$ steps gives:
  $d_{TV}(p_{\theta^*}, q_{\vartheta^*}) \sim O(\eta \sqrt{T}).$
  Thus, the second term can be bounded by $\delta_2$ by balancing $\eta$ and $T$, particularly in the later stages of training.
• The third term leverages the universality property of CP-Flow (Theorem 3 in [11]). Given that the initial noise distribution is absolutely continuous with respect to the Lebesgue measure, there exists a sequence $q_{\vartheta_n}$ such that $d_{TV}(q_{\vartheta_n}, p_{\text{data}}) < \delta_3 \text{ as } n > N.$The optimality of CP-Flow (Theorem 4 in [11]) further guarantees almost sure convergence in distribution of $q_{\vartheta_n}$ to the optimal Brenier map $q_{\vartheta^*}$, ensuring that $d_{TV}(q_{\vartheta^*}, p_{\text{data}}) < \delta_3$.

Combining these results, we conclude that $d_{TV}(\rho_t, p_{\text{data}}) < \delta = \sqrt{\frac{\delta_1}{2}} + \delta_2 + \delta_3$. While this conclusion is weaker than our original claim, it remains sufficient to establish the convergence of Algorithm 1.

We will revise the theorem with detailed proof.

We greatly appreciate the reviewer’s thoughtful comments and constructive feedback. We are especially grateful for the positive remarks regarding the clear and well-organized presentation of our work, the convincing empirical evaluations, and the solid theoretical justifications. Below, we provide detailed responses to the raised questions.

Q1: The parameterization of the CP-Flow is too constrained, which may complicate the optimization process.

A1: The constraints on the CP-Flow parameters are necessary to ensure that the entire neural network is strongly convex with respect to the input $x$, thereby satisfying Brenier’s theorem and ensuring optimality toward the target distribution while meeting the conditions for the mirror map [1]. Regarding the optimization process, we follow the approach outlined in [1], which involves conjugate gradient descent and quadratic optimization methods that are computationally efficient and not overly complex. Specifically, we adopt the "input-augmented" version of the Input-Convex Neural Network (ICNN), which enhances the architecture by directly connecting half of the hidden units to the input. This modification introduces a form of skip connection, improving the propagation of gradients during training and alleviating optimization difficulties.

[1] Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. Learning, 2020.

Q2: The approximation of Hessian will become less accurate as the dimensionality increases.

A2: The diagonal approximation of the Hessian matrix during training was chosen to balance computational tractability with accuracy. Empirically, we found that this approximation still provides sufficient information for updating the mirror map $\nabla G$, particularly when combined with stochastic optimization techniques. This is supported by the FID scores and inference speed presented in Table 2 and Figure 2.

Q3: What is the efficiency of CPMLA in comparison with the approach presented in [2] and [3]? These methods are more relevant to CPMLA.

A3: Empirically, we have shown that our method achieves superior FID scores in real-world image generation tasks compared to both [2] and [3], indicating that CPMLA can outperform the generative quality of models trained using the techniques in those works. Intuitively, our approach operates as a second-order sampling algorithm, which provides both higher sampling efficiency and greater representational power compared to methods relying on pre-determined data augmentations or hierarchical architectures, as used in [2] and [3]. We appreciate the reviewer’s observation that [2] and [3] are closely related to our work. We will add comparisons in the revised manuscript.

[2] 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.

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

We sincerely thank the reviewer for their thoughtful and constructive feedback. We will address reviewer's questions as below.

Q1: 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.
A1: We respectfully disagree with the assessment that our work “only replaces the normalizing flow (in CoopFlow) with mirror Langevin Dynamics.” CPMLA is the first to introduce the Mirror Langevin Algorithm (MLA) as a sampling technique within deep neural networks, whereas previous MLAs lacked large-scale experimental validations and only used predefined fixed mirror maps. Compared to CoopFlow, CPMLA is a second-order sampling algorithm that utilizes the Hessian of CP-Flow to better capture the underlying geometry of the data manifold, thereby achieving more efficient sampling. Experimental results demonstrate that CPMLA outperforms CoopFlow in terms of inference efficiency and FID scores, while using only 0.9% of CoopFlow’s parameter count. This highlights the effectiveness of our dynamic mirror map and the practical advantages of CPMLA in terms of computational efficiency and model complexity.

In summary, CPMLA introduces a different approach to sampling in EBMs through a dynamic, geometry-aware mirror map. We believe these innovations clearly differentiate CPMLA from CoopFlow.

Q2: The approximation error is necessary to be considered.
A2: In line 401 and 456, empirically, our excellent FID scores demonstrate that CPMLA provides a reliable approximation of the target distribution, without introducing bias.

Q3: The analysis of Theorem 5.5 does not include any novel technique compared with [Vempala 2019] and [Jiang 2021].
A3: Vempala's convergence framework is the analysis of the Unadjusted Langevin Algorithm (ULA), whereas we focus on the analysis of the Mirror Langevin Algorithm (MLA). In contrast to Vempala's framework, which uses the standard norm $||z||$, we introduce second-order information by incorporating a geometry-aware mirror flow. Specifically, we define the extended norm as $||z||_{\nabla^2 G}^2 = z^{\top} {\nabla^2 G} z$ in the dual space. This specific adaptation is crucial for analyzing the dynamics in our setup, as CPMLA utilizes a dynamically adjusted mirror map within the CP-Flow structure.

In contrast to Jiang's work, our approach leverages the bounded eigenvalues of the CP-Flow Hessian to enhance theoretical guarantees. This enables us to exert more precise control over the behavior of the dynamics in dual space, ensuring more robust convergence, particularly in high-dimensional settings. Notably, both Vempala's and Jiang's work are foundational and widely recognized in the field for analyzing the convergence of sampling algorithms, and in fact, Jiang’s analysis of continuous Mirror Langevin Dynamics (MLD) is itself inspired by Vempala’s methods.

Below, we summarize the additional experiments we have included: a comparison of wall clock time, a comparison of FID scores in further iterations and a comparison of FID scores and the number of parameters across different EBMs and flow-based frameworks.

Table 1: A comparison of wall clock time for inferring 1,000 images of similar quality on the CIFAR-10 dataset

As shown in Table 1, our model, CPMLAprt, outperforms CoopFlow in terms of inference speed.

Table 1 compares the wall clock time required to generate 1,000 images of similar quality (with FID scores of 21.43 vs. 21.16) on the CIFAR-10 dataset for CPMLAprt (T = 20) and CoopFlow (T = 30). To account for potential initialization and other sources of error, we calculated the average wall clock time by generating 50,000 images when computing the FID score. The results show that, although CPMLA involves additional computations, the lower iteration count and fewer parameters more than compensate for this, leading to a superior inference speed compared to the baseline.

Table 2: A comparison of FID scores under different iteration counts

As shown in Table 2, the FID score of CPMLA is lower than that of CoopFlow when the iteration count is 20, and is comparable when the iteration count reaches 30. We would like to emphasize that the key strength of CPMLA lies in achieving better sampling efficiency with fewer parameters. When comparing the difference in FID scores relative to the differences in parameter counts, the former is minimal, as shown in the following Table 3.

Table 3: A comparison of FID scores and the number of parameters across different EBMs and flow-based frameworks

Table 3 demonstrates that parameter count of CPMLA is significantly fewer than other methods, while still maintaining high image quality. This highlights the efficiency of CPMLA in terms of both parameter usage and performance.

We appreciate the reviewer's comments and would like to clarify the following points:

Q1: The approximation error of EBM and CP-Flow should be clarified. I hope they can provide some evidence, such as the approximation in some synthetic data.

A1: We have included empirical evidence in Figure 1 with synthetic data. CPMLA accurately estimates the eight-Gaussian mixture distribution density, demonstrating both approximation quality.

We have revised Theorem 5.5 and proven that the total variation distance $d_{TV}(\rho_t, p_{\text{data}}) < \delta$ with a small constant $\delta$. We provide the proof sketch as follows.

We decompose the total variation distance between $\rho_t$ and $p_{\text{data}}$ into three terms, providing controls for each. This revised formulation guarantees that $\rho_t$ converges to the data distribution, thereby addressing the reviewer's concerns.

Specifically, let $q_{\vartheta^*}$ denote the target distribution of CP-Flow and $p_{\theta^*}$ the target of the EBM. For the total variation distance $d_{TV}(\rho_t, p_{\text{data}})$, we decompose it into three terms:
$d_{TV}(\rho_t, p_{\text{data}}) \leq d_{TV}(\rho_t, p_{\theta^*}) + d_{TV}(p_{\theta^*}, q_{\vartheta^*}) + d_{TV}(q_{\vartheta^*}, p_{\text{data}}).$

• The first term is controlled under the original Theorem 5.5 (substituting $\delta$ with $\delta_1$). Using Pinsker's inequality we establish: $d_{TV}(\rho_t, p_{\theta^*}) \leq \sqrt{\frac{1}{2} D_{KL}(\rho_t \| p_{\theta^*})} < \sqrt{\frac{\delta_1}{2}}.$
• The second term can be controlled because Langevin MCMC only runs for $T$ steps rather than converging to its stationary distribution. By analyzing the Fokker-Planck equation of Langevin sampling: $\frac{\partial p_t(x)}{\partial t} = -\nabla_x \cdot \left(p_t(x) \frac{\eta^2}{2} \nabla_x f_\theta(x)\right) + \frac{\eta^2}{2} \nabla_x^2 p_t(x),$
  we estimate the incremental change as $d_{TV}(p_t, p_{t-1}) \leq \sqrt{\frac{1}{2} D_{KL}(p_t \| p_{t-1})} \sim O(\eta)$. Summing over $T$ steps gives:
  $d_{TV}(p_{\theta^*}, q_{\vartheta^*}) \sim O(\eta \sqrt{T}).$
  Thus, the second term can be bounded by $\delta_2$ by balancing $\eta$ and $T$, particularly in the later stages of training.
• The third term leverages the universality property of CP-Flow (Theorem 3 in [11]). Given that the initial noise distribution is absolutely continuous with respect to the Lebesgue measure, there exists a sequence $q_{\vartheta_n}$ such that $d_{TV}(q_{\vartheta_n}, p_{\text{data}}) < \delta_3 \text{ as } n > N.$The optimality of CP-Flow (Theorem 4 in [11]) further guarantees almost sure convergence in distribution of $q_{\vartheta_n}$ to the optimal Brenier map $q_{\vartheta^*}$, ensuring that $d_{TV}(q_{\vartheta^*}, p_{\text{data}}) < \delta_3$.

Combining these results, we conclude that $d_{TV}(\rho_t, p_{\text{data}}) < \delta = \sqrt{\frac{\delta_1}{2}} + \delta_2 + \delta_3$. While this conclusion is weaker than our original claim, it remains sufficient to establish the convergence of Algorithm 1.

We will revise the theorem with detailed proof.

Q2 & Q3: ULA uses comparable assumptions with Assumption 5.1 in Theorem 1 of [Vempala2019]. I am not convinced that the mixing time is improved compared with ULA.

A2 & A3: We will remove the claim that CPMLA about weaker assumptions and the mixing time. Instead, we will emphasize that our CPMLA demonstrates a non-bias advantage, while ULA has been shown to exhibit bias in practical settings [1,2].

[1] On the Anatomy of MCMC-Based Maximum Likelihood Learning of Energy-Based Models [2] Improved Contrastive Divergence Training of Energy-Based Models

We thank the reviewer for their insightful comments and positive feedback on our work. We are gratified that the reviewer recognizes the novelty of both our algorithm and theory, and the compelling empirical performance in terms of both efficiency and generation quality. Below, we provide detailed responses to the raised questions.

Q1: The FID is not yet competitive with NCSN++.

A1: While NCSN++ is a score-based model, CPMLA is an energy-based model (EBM). As a result, a direct comparison between the two may not be quite fair. As shown in Table 2, CPMLA outperforms in terms of FID scores and demonstrates superior efficiency in comparison with previous EBM sampling methods.

Q2: Is $G$ updated in an adaptive way, or is it fixed in advance before MLA sampling?

A2: As mentioned in Lines 70 and 287, the mirror map $\nabla G$ in our method is not fixed in advance. Instead, it is optimized jointly with the energy-based model (EBM) during training. Additionally, as noted in Line 389, CPMLA operates with two settings: one where both CP-Flow and the EBM are trained from scratch, and another where CP-Flow is pretrained before training both components together. In both settings, $G$ is updated adaptively. This adaptive update process, which we refer to as "cooperative learning," is a key innovation of our method. Since $\nabla G$ acts as a dynamic mirror map, it evolves to better capture the geometry of the target distribution throughout the training process.

Q3: How exactly is the mirror map $\nabla G$ found in practice?

A3: As outlined in Lines 259 and 276 of Algorithm 1, the update of the mirror map $\nabla G$ alternates between two key steps:

• Sampling: We use the mirror map $\nabla G$ and the EBM obtained from the previous iteration to perform sampling via the mirror Langevin algorithm (MLA).
• Updating $\nabla G$: After sampling, the mirror map $\nabla G$ is updated based on feedback from the target distribution (via the EBM), improving its alignment with the target geometry. This dynamic evolution of $\nabla G$ ensures that the sampling and optimization reinforce each other throughout the training.

By tightly coupling $\nabla G$ and the EBM, our method efficiently models complex distributions without relying on a pre-defined or static mirror map.

Q4: Sections 3.2 and 4.2 should be explained in more detail, especially the cooperative learning aspect.

A4: We sincerely thank the reviewer for their thoughtful suggestions. In line 291 of Section 4.2, we describe the cooperative learning scheme in which we update the parameters $\vartheta$ and $\theta$ using both the original examples $\{x_i\}$ and the synthesized samples $\{x_i^{\text{out}}\}$ based on Equations 6 and 3. We will provide more detailed information, including the practical steps for optimization, the frequency of $G$ updates relative to the sampling steps, and the role of the feedback loop between $G$ and the EBM in enhancing sampling performance. Additionally, we will include a diagram and an algorithm to illustrate the cooperative learning process.