LEAPS: A discrete neural sampler via locally equivariant networks

Thank you for your thorough and valuable feedback on our manuscript. Below we itemize and address your comments and concerns.

Further experiments and benchmarks:

• We have included additional experiments, as possible within the limited time frame. Specifically, we run additional tests of the Ising model at the critical parameters you suggested, and benchmarked this setup against the samplers that reviewer YU3Q suggested (DISCS) and annealed importance sampling (AIS). The results are provided in the anonymous google drive https://tinyurl.com/leapsicml where we see that the LEAPS performs favorably in this setting. To save space to address your other concerns, please see our response to YU3Q for a thorough discussion of these results.

No experiments with locally equiv MLP

• We have tested the MLP architecture and it has performed significantly worse. The convolutional architecture is geometrically translation equivariant, which is known in deep learning to outperform MLPs if the data (e.g. the statistical physics models we consider) is symmetric. We therefore do not consider this a noticable benchmark.

Experimental details. Thank you for letting us know about missing details for replication. Here we supply an exhaustive list:

• Learning rate=5e-3
• Hyperparameter selection: We peform manual optimization of the neural network architectures and hyperparameters.
• Batch size 256 walkers, stored in a replay buffer. New walkers are added to the buffer every 50 training steps
• Training iterations: 30k
• Ising model: periodic lattice.
• External magnetic field: $\mu=0$.
• Comparison to theoretically available solutions: The theoretically available solutions are infinite grid sizes. We run Glauber dynamics that simulate the actual physics of the system for a sufficiently long time to recover the statistics that we compare against.

We will add the above details to the paper. Please let us know if you are missing other details.

Essential References Not Discussed:

• Thank you for pointing us to these citations. We will happily add them to the text when we are allowed to edit it. We note that the paper by Tian et al is cited already in multiple locations of the manuscript.
• One distinguishing feature of our approach to autoregressive models is that we can draw samples with fewer than d steps where d is the dimension. For example, we achieved an ESS of 69% in $<100$ steps for the Ising model. An autoregressive could only achieve the same results in $d=L^2$ steps. Therefore, our effective ESS will be higher.

Non-locally invariant architecture clarification

• Thank you for bringing this up. As stated in section 7, a non-locally equivariant architecture would require $\mathcal{O}(N*d)$ evaluations of the network for computation for $N$ the number of neighbors and $d$ the dimension. This number is prohibitively large for any $d$ that is interesting for applications, e.g. $d=225$ in the Ising experiment. We do not consider a need to give experiments to show that this is unfeasible. Of course, one could train a neural network architecture to be locally equivariant via an additional loss but then lose the computation of exact importance weights, which is one of the main goals of our paper.

Log-variance.

• We thank the reviewer for bringing up the log-variance loss. Evaluating the log-variance divergence is computationally much more expensive here, which is why we do not benchmark against it. Specifically, the log-variance divergence is $K$ times more expensive in memory, where $K$ is the number of simulation steps. The reason for that is as follows: The log-variance loss is not valid when sampled pointwise (it requires the whole trajectory). In contrast, the PINN loss can be evaluated pointwise. However, the PINN objective generally requires a “discrete divergence” computation (the $\mathcal{K}_t$ operator). We get rid of the cost of this discrete divergence with the locally equivariant networks. Note that this is not necessarily true in the continuous case, where the log-var can avoid the computation of a divergence via Ito-integrals.

Proposal distribution.

• We use the model itself as a proposal distribution.

Finally, we like to emphasize that the we consider the the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore stress our methdological and theoretical contributions:

• A new derivation of Radon-Nikodym derivatives of path measures
• Proactive importance sampling as an IS scheme for CTMCs
• A PINN-objective for discrete samplers by bounding the variance of the IS weights
• Local equivariance as a symmetric constraint to enable scalable proactive IS
• Locally equivariant neural network architectures such as convolutional neural networks

Thank you again for your valuable review. We hope that this addresses your questions. We would accordingly appreciate any increase in rating you see fit.

Thank you for your valuable feedback on our work. Below we address your questions and provide new information on additional experiments.

Updates to experiments: harder sampling problem and comparison benchmarks We obtained new experimental results and benchmarks as possible within the limited time frame. The results can be found at this anonymous link, in which we show that our method performs well at the critical point of the Ising model and compare it extensively against other methods in this hard regime (see more below): https://tinyurl.com/leapsicml

Additional benchmarks (1) - LEAPS vs. MCMC. We use the DISCS benchmark [1] for discrete samplers to compare LEAPS to various MCMC samplers. We also adapt the parameters of the Ising model to the critical temperature for $\beta=0.4407$. Note that this makes the problem harder (both for traditional methods and for LEAPS). The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\sim 70$% vs an ESS of $<0.1$% of MCMC-based samplers. This shows how LEAPS effectively converts a simple distribution into a highly complex distribution that can only be sampled with many MCMC runs. In figure 3, we also show that the samples recover known physical observables of the Ising model. Finally, we note that there are significant difficulties in creating a fair assessment between non-equilibrium (LEAPS) and equilibrium methods (MCMC):

• The ESS is measured differently in DISCS because it focuses on equilibrium MCMC methods, rather than non-equilibrium sampling methods such as LEAPS.
• The ESS quantities for non-equilibrium methods such as LEAPS do not take into account the number of steps used for annealing (i.e. for simulation of the continuous-time Markov chain). To account for this, we have also plotted the same quantities normalized by the number of function evaluations in Figure 1 (ESS/NFEs). However, LEAPS does not require evaluations of the energy at inference time but only of the neural network, while MCMC methods require energy evaluations (see metrics used in DISCS [1]). The best way we account for that now would be to normalize by the NFEs as measured in neural net evaluations. However, note that this does not measure efficiency or wall clock time, which would highly dependent on implementation and is usually not measured in the literature on neural sampling methods. The main goal of LEAPS is to solve the sampling problem, not efficiency.

We hope that our new experiments showcase the efficacy of the LEAPS algorithm, while also explaining the difficulty in benchmarking LEAPS vs traditional equilibrium methods.

Additional benchmarks (2) - LEAPS vs. AIS. We also benchmark LEAPS against annealed importance sampling (AIS) [2]. AIS is a standard non-equilibrium sampling technique. Note that here, ESS is measured in the same way and the two methods are exactly comparable. In Figure 2, we plot the number of AIS steps vs the ESS of AIS. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps. We note that LEAPS is a true extension of AIS and we can run AIS with the learned transport without changing the weights. In addition it can be turned into a sequential monte carlo sampler by performing resampling along the trajectory. However, here, for the purposes of benchmarking we turn off the AIS (here and above) to showcase that the learnt transport via the neural network enables the efficiency boost.

Comparison to other works cited: Many of the works that we cite in our work are about the mathematical realization of discrete diffusion models and sampling algorithms for continuous variables. To the best of our knowledge, we do not know of another paper which learns a neural sampler for discrete distributions via Continuous Time Markov Chain. Therefore, we cannot benchmark against them. If you know of any, please let us know.

Finally, we would like to emphasize that the we consider the the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore emphasize again our methdological and theoretical contributions:

• A new derivation of Radon-Nikodym derivations of path measures
• Proactive importance sampling as an IS scheme for CTMCs
• A PINN-objective for discrete samplers by bounding the variance of the IS weights
• Local equivariance as a symmetric constraint to enable scalable proactive IS without losing representational power
• Locally equivariant neural network architectures such as LE-convolutional neural networks

Thanks again for your valuable feedback. If these results have fulfilled your questions and concerns, we would greatly appreciate any increase in score.

[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for discrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.

We thank the reviewer for their valuable feedback on our work. Below, we provide answers to your questions and concerns.

Additional benchmarks and experiments. We obtained new experimental results. The results can be found under this anonymous link, in which we show that our method performs well at the critical phase of the Ising model and compares favorably against existing samplers (see more below): https://tinyurl.com/leapsicml We use two benchmarks:

1. LEAPS vs. MCMC. We benchmark LEAPS against a diversity of MCMC samplers on the critical temperature of the Ising model. The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\sim 70$% vs ESS$<0.1$% of MCMC-based samplers.
2. LEAPS vs. AIS: We also benchmark LEAPS against annealed importance sampling (AIS), see figure 3. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps.

Please see our response to YU3Q for a thorough discussion of these results.

Questions. We address your remaining questions and comments in the following:

• What does "no transport" refer to?: This refers to not using the neural network at all. We show via this ablation that it is truly the neural network (and not the AIS) that enables the sampling from the distribution with only a few samples. We ran additional benchmarks on this that we explained above (AIS vs LEAPS) not using AIS at all for LEAPS.
• LEAPS uncorrected vs. LEAPS: For the simulation of the continuous-time Markov chain, we always use the rate matrix $Q_t$ represented by our network. Due to the local equivariance, we get importance weights for free. "LEAPS uncorrected" means that we do not use the importance weights (at perfect training, this would not be necessary). "LEAPS corrected" means that we use the importance weights to reweight samples towards the correct distribution. Both parts are essentials as our ablations show.
• More datasets and benchmarks: We provided other benchmarks, in particular against a series of MCMC samplers. We further actively run experiments on the Potts model, another model from statistical physics. The only other established benchmark for discrete sampling is combinatorial optimization. However, we note that metrics used for combinatorial optimization do not aim to faithfully sample from a distribution or any observable thereof, where importance weights could be used. Rather it tries to find low energy states. This puts our method fundamentally at a disadvantage.
• Clarifying equivariance: We note that for an input $x$ the neural network outputs $F_t^\theta(y^i,i|x)$ for every $y^i,i$ for token $y^i$ and index $i$ (this is also true for discrete diffusion models, it returns a rate per neighbor). Therefore, this requires only one forward pass. Thank you for bringing it up that this needs further clarification. We will highlight the distinction between input and output of the neural network more in future versions.
• Connection with NETS/ALD: Here is the analogy. Because ALD is like running Langevin dynamics locally on some $\rho_t \sim e^{-U_t}$ for fixed $t$ and then iterating $t \rightarrow t + dt$, so too can we add any local MCMC kernel to the dynamics of the CTMC that preserves detailed balance. This would "simulate approximately the prescribed path" just like Langevin does in the continuous setting. Does that make sense? We will add a remark about this in the text.
• Question about the PINN loss: The goal of LEAPS (and of all importance sampling schemes) is that the weights have as low variance as possible. In our case, this would mean that $A_t$ is just a constant in time and space. In other words, this means $\partial_sA_s=0$. As we show in the paper, this is equivalent to $\mathcal{K}_s\rho_s=\partial_sF_s$. Thank you for bringing this up. We will highlight this in future versions of our work.
• Interpolation schemes: The interpolation scheme is given by a time-dependent Ising model for coupling constant $J_t=tJ_1$ where here $J_1=1$. As you point out, this is equivalent to temperature annealing, i.e. $\rho_t(x)\propto \exp(-tU_1(x))$. As also pointed out, discrete diffusion models use arithmetic means (or probabilistic mixtures) between a uniform (or mask) and the target distribution. In principle, one could use such an interpolation. We focuses on temperature annealing because it is common in the sampling literature and it allows to sample from the Ising model for higher temperatures on the fly (by simply stopping at earlier time points). Therefore, this is a natural and physical way. We have explored various annealing speeds and schedules (i.e. time parameterizations) for the annealing that is happening during training and found the schedule of $t\to \sqrt{t}$ to perform best.

We thank the reviewer for their valuable feedback on our work. We are glad to read that the reviewer considers our work a "sound" and "novel" contribution. Below, we provide answers to your questions and concerns.

We understand that the reviewer sees our experiments as the main area of improvement. Addressing this, we obtained new experimental results - as possible within the limited time frame. The results can be found under this anonymous link, in which we show that our method performs well at the critical phase of the Ising model and compares favorably against existing samplers (see more below): https://tinyurl.com/leapsicml

We would also like to emphasize that the we consider the main contribution of this work the introduction of a new paradigm for learning to sampling from discrete distributions. We therefore stress our methodological and theoretical contributions:

• A new derivation of Radon-Nikodym derivatives of path measures
• Proactive importance sampling as an IS scheme for CTMCs
• A PINN-objective for discrete samplers by bounding the variance of the IS weights
• Local equivariance as a symmetric constraint to enable scalable proactive IS
• Locally equivariant neural network architectures such as convolutional neural networks

Additional benchmarks (1) - LEAPS vs. MCMC. Thank you for sharing the DISCS benchmark [1] for discrete samplers with us. To use this benchmark, we adapt the parameters of the Ising model to match theirs, i.e. we use the parameters for the critical temperature for $\beta=0.4407$. Note that this makes the problem harder (both for traditional methods and for LEAPS). We then ran benchmarks of LEAPS vs other discrete samplers (including the ones in the DISCS benchmark). The results can be found in figure 1 in the google drive folder. As one can see, LEAPS achieves an effective sample size (ESS) of $\sim 70$% vs ESS$<0.1$% of MCMC-based samplers. This shows how LEAPS effectively converts a simple distribution into a highly complex distribution that can only be sampled with many MCMC runs. In figure 3, we also show that the samples recover known physical observables of the Ising model. Finally, we note that there are significant difficulties in creating a fair assessment between non-equilibrium (LEAPS) and equilibrium methods (MCMC):

• The ESS is measured differently in DISCS because it focuses on equilibrium MCMC methods, rather than non-equilibrium sampling methods such as LEAPS.
• The ESS quantities for non-equilibrium methods such as LEAPS do not take into account the number of steps used for annealing (i.e. for simulation of the continuous-time Markov chain). To account for this, we also plot the same quantities normalized by the number of function evaluations in Figure 1 (ESS/NFEs). However, LEAPS does not require evaluations of the energy at inference time but only of the neural network, while MCMC methods require energy evaluations (see metrics used in DISCS [1]). The best way we account for that now would be to normalize by the NFEs as measured in neural net evaluations. However, note that this does not measure efficiency or wall clock time. The main goal of LEAPS is to solve the sampling problem, not efficiency.

We hope that our new experiments showcase the efficacy of the LEAPS algorithm, while also explaining the difficulty in benchmarking.

Additional benchmarks (2) - LEAPS vs. AIS. We also benchmark LEAPS against annealed importance sampling (AIS). AIS is a standard non-equilibrium sampling technique. Note that here, ESS is measured in the same way. In Figure 2, we plot the number of AIS steps vs the ESS of AIS. As one can see, even with $100k$ simulation steps, AIS only achieves an ESS of $<30$% vs an ESS of $69$% of LEAPS with 100 simulation steps. We note that LEAPS is a true extension of AIS (as well as SMC) and we can run AIS with the learned transport without changing the weights. However, for the purposes of benchmarking we turn off the AIS (here and above) to show that the learnt transport via the neural network is what enables the performance boost.

Other benchmarks. Thank you for bringing up other possible benchmarks. We note that combinatorial optimization does not aim to faithfully recover a distribution or any observable thereof, where importance weights could be used. This puts our method fundamentally at a disadvantage. Please let us know if there are other benchmarks that come to mind that we can run to convince you of the potential of the method.

References. Thank you for highlighting the references that we included in the updated draft of our work. We have benchmarked against the suggested methods (see above).

We thank the reviewer again for the insightful response. If these results have addressed your concerns, we would greatly appreciate any increase in score.

[1]. Goshvadi K, Sun H, Liu X, et al. DISCS: a benchmark for discrete sampling[J]. Advances in Neural Information Processing Systems, 2023, 36: 79035-79066.