Gradient-based Discrete Sampling with Automatic Cyclical Scheduling

Summary

The paper presents a novel gradient-based algorithm to sample from complex multimodal discrete distributions based on differentiable energy functions. Overall, the method is based on locally balanced proposals, previously introduced, and instantiates it with parametrized functions and a cyclical schedule for the "learning rate" that promotes the alternation of modes discovery and modes refinement. The authors also provide an algorithm to automatically tune the introduced parameters based on an input acceptance rate and initial and final balancing parameters. The paper features a theoretical analysis that includes concrete convergence rates under (somewhat restrictive) assumptions and finishes with some experiments in learning and sampling from RBMs and EBMs and finally in text infilling with masked language models.

Strengths

- to the best of my (limited) knowledge in this area, the algorithm presented in the paper seem novel and potentially quite impactful (provided the author release an "easy-to-use" implementation) - the illustrative example sets the stage nicely for the need of further development in the field of gradient-based sampling from discrete distribution and provides a very compelling visualization of the efficacy of the method. - the non-asymptotic bounds may offer concrete guarantees when assumptions are met.

Weaknesses

- I think too much of the paper is in the appendix, and the information presented in the paper is not sufficient to fully follow the logical flow (see also points below). In my opinion, some details regarding the development of the method and the theory could be moved to the appendix to make more space for both preliminaries (like how are these methods used in practice) and developing more intuition. - the algorithm is fairly complex and the paper fails in provide enough intuition for some parts of its functioning. Like the authors claim that "it is fairly easy to choose initial and final balancing factors", but do not elaborate why (in the main paper) - the only "real-world" experiment is only sketched in the main paper, not providing enough information to appreciate the task. What is the precise difficulty here? Since roberta is a masked language model I believe it is possible to derive "pseudo-distributions" like in [1] - the theoretical analysis seems to require strong assumptions that may be violated in compelling real-world use cases (like LLMs) [1] Hennigen, Lucas Torroba, and Yoon Kim. "Deriving language models from masked language models." arXiv preprint arXiv:2305.15501 (2023).

Questions

- Can you please discuss in which cases the locally concave hypothesis holds for realistic models such as LLMs? - Even if this is probably more of a sanity check, I'd like to see how the method behave on unimodal, or mildly multimodal distributions - Can you please describe some other applications of the method to real-world problems (like text infilling)

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Partially discussed, see questions

Summary

The paper proposes a solution to the challenge of sampling from high-dimensional discrete spaces, where conventional discrete samplers often get trapped in local modes. To address this, the authors introduce a discrete Langevin sampler with automatic cyclical scheduling. This method comprises three components: a cyclical step size schedule, a cyclical balancing schedule, and an automatic hyperparameter tuning scheme. The authors provide theoretical guarantees for non-asymptotic convergence and inference, and extensive experiments demonstrate the method's superiority in sampling complex multimodal discrete distributions.

Strengths

The paper is well-motivated, and the proposed automatic cyclical scheduling method is presented clearly, making it accessible to readers. The theoretical results, which offer non-asymptotic convergence, support the method's robustness. Additionally, the empirical study is solid, with extensive experiments demonstrating the method's superiority in sampling from high-dimensional spaces.

Weaknesses

- My primary concern lies in the complexity of the proposed methods. The automatic schedule tuning scheme appears to be quite time-consuming, particularly the grid search required for the balancing parameters \beta_i, demanding significant computational resources. - Another concern pertains to the theoretical assumptions underlying the analysis. The non-asymptotic convergence of the proposed samplers relies on the strong convexity of the negative energy function, an assumption that may not hold in practical deep EBM scenarios. Despite this potentially restrictive condition, the analysis provides valuable insights, and empirically, the proposed method demonstrates effective performance, as supported by extensive studies.

Questions

- In equations 8 and 9, how do you estimate the acceptance rate A? Is it estimated by averaging across the training batch? If so, that implies the complexity would increase to n*s times compared to the original DMALA samplers at each step, where n is the number of grids for parameters \beta_i and s is the number of sampling steps per cycle. This would be highly time-consuming. - Could you provide a comparison of the running times between the proposed ACS sampler and the DMALA samplers? - Could you elaborate on the rationale behind setting the target acceptance rate $\rho^*$ to 0.5 in your experiments? What are the implications of setting it to 1 or 0.234 instead?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Despite the strong assumptions in the theoretical analysis and the perceived complexity of the algorithm, this paper presents a substantial contribution to the field by offering a well-explained, theoretically sound, and empirically validated approach to enhancing discrete sampling in high-dimensional spaces.

Summary

The paper introduces a novel method for sampling from multimodal discrete distributions, which presents an innovative approach to address the challenge of local modes trapping in gradient-based discrete sampling, together with non-asymptotic convergence guarantee and empirical validation of the proposed method.

Strengths

1. The proposed method seems novel to address the challenge of sampling from multimodal discrete distributions. 2. The hyperparameter tuning algorithm seems useful for practical use.

Weaknesses

1. In the experiments, there seems to be no error bars in Figure 1 and Table 1. 2. The quality of the samples seems worse than DMALA in Table 2 and Figure 12. Is it possible that the proposed method might sacrifice sample quality to achieve higher diversity?

Questions

1. The paper mentions the proof is not consistent with the specific tuning algorithm used in the experiments. Could you elaborate on the reasons?

Rating

5

Confidence

1

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

Summary

This paper proposes a new discrete sampling method called ACS that addresses a common problem for existing gradient-based approach where they are susceptible to becoming trapped in local modes. ACS combines local-balancing proposals with a cyclic step size to balance local exploitation and global exploration; it is in essence an extension of cyclic stochastic-gradient MCMC to discrete distributions. To ensure proposals are still balanced with a varying step size, ACS uses a cyclic balancing schedule along with an automatic tuning scheme to easily adapt the schedules. Non-asymptotic convergence guarantees are provided. Results demonstrate ACS to outperform prior approaches for sampling from energy based models, training RBMs, and text-infilling.

Strengths

- Using a cyclic step size schedule is a well motivated and effective approach for incorporating global considerations into the original local self-balancing MCMC approach presented in https://arxiv.org/pdf/2109.03867. - ACS is accompanied by an automated tuning scheme to make it easy to configure the two cyclic schedules. - Strong empirical results on EBM tasks. - ACS has non-asymptotic convergence guarantees.

Weaknesses

- Experimental results for RBM have discrepancies wrt results reported in previous papers. In particular the ranges differ for average energy from other papers and the curves for log MMD for ACS show unexpected curvature. Please respond to the questions in the section below to clarify. - Inadequate discussion of text-infilling results. It is not clear why higher perplexity and diversity is good for ACS since the goal is to be able to efficiently sample from a target discrete distribution. - Error bars missing for results in Figure 3 and Table 1. Typos & formatting: - Inline citation format should show author name - Figure 3 caption: cpnvergence -> convergence

Questions

- Why is average energy in Figure 3 negative? Also, the results for ACS on dynamic_mnist and omniglot are unexpected with better performance on fewer iterations before converging. - Why is the scale for average energy in Figure 3 different from that reported in Figure 4 of the [AB paper](https://openreview.net/pdf?id=lEkl0jdSb7B)?

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are adequately discussed.

Thanks

Dear authors. Thank you very much for your rebuttal. I appreciate the additional experiments clarifications on the infilling task. I keep my opinion that this work is a valid addition to the conference.