Follow Hamiltonian Leader: An Efficient Energy-Guided Sampling Method

Rebuttal

Thank you for reviewing our paper and providing insightful feedback. We appreciate your positive remarks and constructive criticism, which will help us improve our work. Below, we address your comments and questions:

1. Quantitative Experiments on CIFAR-10:
   We understand your concern about the absence of quantitative experiments on complex image datasets. To address this, we have conducted experiments to evaluate the performance of FHL on the CIFAR-10 dataset, and the results are included in the overall response. These findings will be incorporated into the revised manuscript to highlight the method's applicability to complex distributions.

2. Out-of-Distribution (OOD) Experiments with EBMs or Score-Based Models: We apologize for any confusion regarding out-of-distribution experiments in our context. Our method is fundamentally a sampling technique designed to accurately generate the target distribution, whereas out-of-distribution experiments are typically performed in classification or detection tasks.

   However, as an alternative, we had conducted experiments in scenarios where sampling is challenging, such as sampling from poorly-conditioned functions, as shown in Figure 4. These experiments aim to validate the stability and robustness of FHL during sampling. We are glad to extend the results in the revised paper to provide a more comprehensive analysis of the method's capabilities.

Discussions

1. Integration with Advanced MCMC Methods:
   We recognize that our algorithm could benefit from integration with other advanced MCMC methods. In fact, our method is orthogonal to these techniques and can be combined with methods such as LMC or parallel tempering. However, due to space and time constraints, our current version focuses solely on comparison and integration with HMC. We will extend our work to include integration with other advanced MCMC methods in the future.

2. Quantitative Comparison with Other Methods:
   As suggested, we have included quantitative experimental results for CIFAR10 in the overall response.

Question

• Entropy-SGD

  The key concept of Entropy-SGD is its incorporation of Bayesian optimization elements. Unlike other optimization methods such as Adam or Momentum SGD, Entropy-SGD optimize over a local area (referred to as "local entropy") instead of searching for an isolated optimal point.

  $$F(x, \gamma) = \log \int_{x' \in \mathbb{R}^n} \exp\left( -f(x') - \frac{\gamma}{2} ||x - x'||_2^2 \right) dx'.$$

  The extra term $E(x, x') = \frac{\gamma}{2} ||x-x'||^2_2$ in Entropy-SGD controls the range over which it seeks valleys of specific widths. Another way to write the objective used in the Entropy-SGD (see Equation 3, [2]) is:

  $$f_\gamma(x) := -\log \left( G_\gamma * e^{-f(x)} \right); G_\gamma = (2\pi\gamma)^{-N/2} \exp\left( -\frac{||x||^2}{2\gamma} \right)$$

  The objective above shows that Entropy-SGD optimizes a modified, smoother energy landscape, which can be viewed as a convolution of the original energy landscape with a Gaussian kernel. This approach aids in identifying wider valleys.

• PARLE: The Parallel Variant of Entropy-SGD
  An extension of Entropy-SGD, known as PARLE [2], essentially parallelizes the Entropy-SGD algorithm. This method optimizes the following objective function, which reduces to Entropy-SGD when $n=1$:

  $$\arg \min_{x, x^1, \ldots, x^n} \sum_{a=1}^{n} f_\gamma(x^a) + \frac{1}{2\rho} || x^a - x ||^2.$$

  By deriving this expression, it becomes evident that PARLE encourages particles to sample around $\bar{x}$, where $\bar{x} = \frac{1}{n} \sum_{a=1}^{n} x^a$ is the average of the particles.

• Connection and Comparison for FHL

  Both Entropy-SGD and PARLE are optimization methods that ultimately produce a final point instead of a Markov chain. In contrast, FHL is specifically designed as a sampling method. Furthermore, there are other similarities and differences between FHL and Entropy-SGD/PARLE:

  • (Similarity) Exploring the Landscape: All of three approaches allow to explore the function landscape around the anchor point (Entropy-SGD: the average of the SGD trajectory of a single particle; PARLE: the average of multiple particles; FHL: the leader of the particles).

  • (Difference) Zeroth-Order Optimization: FHL leverage zeroth-order information, and uses the leader as the anchor point instead of the average of particles. Averaging is commonly used in optimization techniques to reduce variance. However, in sampling tasks, noise is typically not introduced, and pulling towards the average may degrade convergence performance. Therefore, FHL could be expected to perform better.

[1] P. Chaudhari, A. Choromanska, S. Soatto, Y. LeCun, C. Baldassi, C. Borgs, J. T. Chayes, L. Sagun, and R. Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In ICLR, 2017.

[2] Chaudhari, Pratik, et al. "Parle: parallelizing stochastic gradient descent." arXiv preprint arXiv:1707.00424 (2017).

Thank you again for your valuable feedback. We are confident that the additional results and planned revisions will address your concerns and enhance the overall quality and impact of our work.

Thank you for your detailed review and for highlighting both the strengths and areas for improvement in our paper.

Explanation

For most sampling algorithms, the objective is to develop a proposal function $Q(x'|x): \mathbb{R}^d \rightarrow \mathbb{R}^d$ that generates a new sample $x'$ from an existing sample $x$, with the goal of having $x'$ fall into a region of high probability density. Provided the sampling algorithm adheres to the principles of detailed balance and ergodicity, it typically converges to the correct distribution.

Therefore, designing an effective sampling method generally involves improving the proposal at each step. FHL introduces a bias for each particle, guiding them towards more accurate results, specifically aiming for a lower value of the objective function $U(x)$, which aligns with common optimization tasks. There are two potential scenarios where our approach can be beneficial:

1. Missing Gradient: When the gradient is unavailable, the leader can be used as a reference point. This concept can be illustrated with a convex function. Consider a convex function $f$, where $x$ represents our particle and $y$ represents its leader. Then, $f((1-\rho) \cdot x + \rho \cdot y) < (1-\rho)f(x) + \rho f(y) < f(x)$, suggesting that $\rho (y-x)$ is a potential loss-descent step as long as $0 < \rho < 1$. Therefore, even when the gradient vanishes, our method still ensures a decrease in loss.

2. Non-Optimal Gradient Direction: As illustrated in Figure 1, the leader's pulling bias can improve the convergence direction when the gradient descent direction and the Newton descent direction are not aligned. It can be demonstrated that as long as the leader lies within the cone $C = \text{cone}(d_N(x), \theta_x)$, the descent direction could be improved. Here, $d_N(x)$ is the Newton descent direction (red arrow in Figure 1), $d_G(x)$ is the gradient descent direction (blue arrow in Figure 1), and $\theta_x = \theta(d_G(x), d_N(x))$.

We hope this explanation clarifies our contributions and provides a better understanding of the potential of the FHL method.

Rebuttal

We have carefully considered your feedback and address your points below:

1. Complexity and Computational Resources:
   We acknowledge that the FHL method introduces modifications that may increase complexity. However, only one additional hyperparameter, the pulling strength $\lambda$, is specifically introduced, while the number of particles per group can be determined based on available computational resources. FHL is designed for large-scale parallel sampling, supported by two key factors:

   1. The computation can be executed in parallel.
   2. The communication cost scales logarithmically.

   As the number of particles increases, the total running time is primarily influenced by the communication cost, making FHL highly scalable with the number of particles. Additionally, techniques like lazy communication can be employed to further reduce communication overhead.

2. Selection of the Leader Particle:
   The leader selection is indeed crucial to the method's performance. However, unlike stochastic optimization, most sampling algorithms require that the zeroth-order information be correct. If this assumption fails, these methods would become ineffective, as the key Metropolis-Hastings step [1] relies on the zeroth-order information to satisfy the detailed balance condition.

   Furthermore, even with suboptimal leader selection, although the convergence rate might be affected, the method should still be able to converge to the correct distribution, as long as the algorithm satisfies the detailed balance condition.

   Even if we are forced to sample from a biased function $\psi(x) = U(x) + \frac{\lambda}{2}||x - z||^2$ instead of the true function $U(x)$, the error can be bounded in some sense. We have provided a theoretical analysis of this in our appendix, Section C.

   [1] Chib, Siddhartha; Greenberg, Edward (1995). "Understanding the Metropolis–Hastings Algorithm". The American Statistician, 49(4), 327–335.

3. Lack of Theoretical Analysis:
   We have included a theoretical analysis in the overall response, demonstrating that our algorithm has the correct distribution (to sample from) as its invariant distribution.

4. Scalability Concerns:
   We have included a large-scale experiment in the overall response, demonstrating that FHL can generate biomolecule conformations despite the high number of degrees of freedom and numerous energy landscape barriers. This sampling experiment is not only extensive but also highly significant for real-world simulations of biomolecules.

Thank you once again for your valuable feedback and insights. We believe the planned revisions and additional analyses will address your concerns and enhance the paper's quality. We appreciate your assessment and are committed to addressing the concerns raised to improve the paper's overall quality and impact.

Dear Reviewer N2ad,

Thank you for your detailed review and valuable feedback on our paper. We appreciate the opportunity to address your concerns and provide further clarification on our rebuttal.

Addressing Your Concerns:

• Complexity and Computational Resources: FHL introduces minimal complexity with only one additional hyperparameter. The method is designed for parallel execution, making it scalable as communication costs scale logarithmically. Techniques like lazy communication further reduce overhead.

• Selection of the Leader Particle: While leader selection is crucial, our method remains effective as long as it adheres to the detailed balance condition. Even with suboptimal selection, convergence to the correct distribution is achievable.

• Lack of Theoretical Analysis: We have included a theoretical analysis in our appendix and overall summary (see pdf), demonstrating the correctness of our algorithm's invariant distribution.

• Scalability Concerns: Our large-scale experiments showcase FHL's capability to generate biomolecule conformations effectively, highlighting its significance for real-world applications.

We are committed to improving our paper based on your insights and believe that the planned revisions will address your concerns. Thank you again for your valuable feedback.

Best regards,

Thank you for taking the time to review our paper and for providing constructive feedback. We appreciate your comments and have addressed your points and questions below:

1. Simplicity and Clarity:
   We are pleased that you found our approach simple and clear. The goal was to design a method that effectively incorporates energy information into gradient-based sampling techniques while remaining intuitive. We believe this clarity helps in demonstrating the benefits of our approach through the examples provided.

2. Overhead and Computational Cost:
   Thank you for pointing out the need to include a discussion on the overhead of our method. We have made a discussion on both communitation and computation cost in the overall response.

3. Tension Coefficient Analysis:
   The tension coefficient is indeed critical to our method, and we apologize for not including an ablation study on this parameter. We have now provided a table of the values we explored and a brief discussion in the overall response.

4. Determination of the Tension Coefficient:
   The tension coefficient is determined empirically based on preliminary experiments. We agree that providing an ablation study for this coefficient across the three scenarios would strengthen our findings, and we are currently conducting this analysis. The results will be included in the revised manuscript.

5. Clarity in Figure 7:
   Thank you for your suggestion regarding Figure 7. We acknowledge that the current presentation could be more clear and we have included quantitative metrics in the overall response.

Limitations

We acknowledge that the limitations section was brief. We will expand this section to discuss the scenarios where our method may face challenges, as well as ethical considerations and potential applications.

Thank you once again for your valuable feedback and constructive suggestions. We are confident that the revisions and additional analyses will address your concerns and improve the paper.

Thank you for your feedback and for acknowledging our additional quantitative evaluations. We appreciate your interest in the ablation studies on the tension coefficient.

As mentioned in the third paragraph of the "3. Sensitivity to Hyperparameters" section of the overall summary above, we have included a table displaying the hyperparameters we explored. Additionally, we have conducted an ablation study to address your concerns.

Ablation Study

Due to time constraints, we will briefly discuss the selection of $\lambda$, and for simplicity, we will consider updating the momentum directly by $- \eta \nabla f(x) + \lambda (x-z)$. Notice that as $\lambda \rightarrow 0$ FHL performs similarly to vanilla HMC method. As mentioned in the overall summary, determining an optimal value for $\lambda$ is challenging because it depends on the interplay of various hyperparameters and the characteristics of the sampled function. Thus, we will consider the case where $- \eta \nabla f(x) + \lambda (x-z)$ is a better choice than $- \eta \nabla f(x)$.

Observation

We have two key observations from both theoretical and experimental results:

• Setting $\lambda$ to a small value may enhance performance.
• Increasing the number of particles in the group, $n$, is likely to improve performance.

Theoretical Results It is well-known that the Newton method typically converges faster than the gradient method (the Newton method exhibits quadratic convergence near the optimal solution while the gradient method has a linear convergence rate for strongly convex function). The Newton method suggests a direction given by $d_N(x) = -(\nabla^2 f(x) + \epsilon I)^{-1}\nabla f(x)$, where $\epsilon$ is a small value added to prevent the inversion of a singular matrix. Hence, we consider it desirable to have search directions that are close to $d_N$.

Let $\theta(u, v)$ denote the angle between vectors $u$ and $v$, and let $\cos(u,v) = \frac{u^T v}{||u|| \cdot ||v||}$. Define the leading direction with respect to a leader $z$ as $d_{\lambda}(x;z) = -\eta \nabla f(x) + \lambda (x-z)$. Consider the positive definite quadratic function $f(x) = \frac{1}{2} x^T A x$, which results in $d_G(x) = -\eta \nabla f(x) = -\eta A x$ and $d_N(x) = -x$. For the upcoming theorems, we define a cone with axis $d$ and angle $\theta_c$ as $\text{cone}(d, \theta_c) = { x : x^T d \geq 0 \text{ and } \theta(x,d) \leq \theta_c }$. Additionally, we define $\theta_c(x) = \theta(d_N(x), d_G(x))$.

Proposition 1: If $\lambda_1 > \lambda_2$ and $\theta(d_{\lambda_1}(x;z), d_N(x)) < \theta(d_G(x), d_N(x))$, then $\theta(d_{\lambda_2}(x;z), d_N(x)) < \theta(d_G(x), d_N(x))$. In particular, if $z \in C$ and $0 < \lambda < 1$, then $d_{\lambda}(x; z) \in C$.

Proposition 2 Define the outward vector for a point $x$ as $N_x = d_N - \frac{<d_N,\ d_G>}{||d_G||}\cdot \frac{d_G}{||d_G||}$. If $z$ satisfies $-N_x^T (z-x) < 0$, then for sufficiently small positive $\lambda$, $d_\lambda (x; z) \in \text{cone}(d_N(x), \theta_c(x))$.

Claim 3 There exists a region that consistently improves performance. If each particle within a group of size $n$ converges to the distribution $\pi(x) = \frac{1}{Z} e^{-\frac{1}{2}x^T A x}$ (where $Z$ is a normalizing factor), then as the number of particles $n$ increases, the probability of finding an improving leader for an arbitrary particle $\tilde{x}$ increases. Specifically, this probability is given by

$1 - \left(\frac{1}{Z} \int_{ z \in (z \mid f(z) < c(\tilde{x}) ) } e^{-\frac{1}{2}z^T A z} \ dz\right)^n,$

where

$c(\tilde{x}) = \frac{|| N_{ \tilde{x} } ||_2^4} {2 \cdot v(\tilde{x})} ,$

and

$v(\tilde{x}) = N^T_{\tilde{x}} A^{-1} N_{\tilde{x}}.$

Experimental Results

We initialize $A$ as a positive-definite diagonal matrix with diagonal elements randomly selected from a uniform distribution in the range $(0.2, 20)$. The dimension of $A$ is $2000$ and we then run FHT for $2000$ iterations, adjusting $\lambda$ by dividing it by the number of particles (using $\lambda_n = \lambda / n$ instead of $\lambda$). We report the value of the leader $z$, $f(z)$, for different values of $\lambda$ and $n$.

The experimental results confirm the consistency with our theoretical observations.

Please let us know if you have further questions or need additional details.

Thank you for your thoughtful and detailed comments on our paper. We appreciate your feedback and are grateful for the opportunity to clarify some aspects of our work. Below, we address your comments and questions:

1. Novelty of Zeroth-Order Energy Information
   We are glad you found the incorporation of zeroth-order energy information into sampling algorithms a novel and promising idea. We believe this approach offers a fresh perspective on improving sampling stability and efficiency for sampling, particularly in challenging scenarios.

2. Sensitivity to Key Hyperparameters
   We agree that exploring the sensitivity to key hyperparameters, such as the number of parallel sampling chains, is essential for a comprehensive understanding of the FHL algorithm. We have now provided a table of the values we explored and a brief discussion in the overall response.

3. Discussion of Computational Cost/Overhead for Parallel Sampling
   We greatly appreciate your recognition of the elegance of the FHL algorithm across parallel sampling chains. Our aim was to develop a principled approach that effectively leverages both types of information to enhance sampling performance. We would like to point out that our method only broadcasts the leader's parameters, making it communication-efficient.

   Furthermore, our method is compatible with advanced techniques like parallel tempering, as they are orthogonal approaches. We are also exploring methods to minimize overhead, such as dimensionality reduction and lazy communication.

4. Theoretical Analysis and Quantitive Results
   Thank you for requesting a theoretical analysis of the convergence properties and sampling quality of the FHL algorithm. We have included convergence guarantees comparable to those offered by most sampling methods. Additionally, we provide further quantitative results for our CIFAR10 experiment (Section 5.2 in the paper).

5. Large-scale Experiment
   Scalability to high-dimensional spaces is an important aspect of our algorithm. We have included a large-scale experiment in the overall response, showing that FHL can generate biomolecule conformations despite the high number of degrees of freedom and many energy landscape barriers. This sampling experiment is not only extensive but also highly significant for real-world biomolecular simulations.

Limitations

• Discussion on Limitations and Societal Impacts
  We acknowledge that our paper did not explicitly discuss limitations or potential negative societal impacts. In the revised version, we will include a dedicated section addressing these concerns.

We sincerely thank you for your constructive feedback, which has provided valuable insights into areas for improvement. We are committed to addressing these points in our revised manuscript and believe these enhancements will strengthen our work.