LEAPS: A discrete neural sampler via locally equivariant networks

Authors: Peter Holderrieth, Michael Samuel Albergo, Tommi Jaakkola

ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental As a demonstration of the validity of LEAPS in high dimensions, we benchmark it on models of statistical physics, namely the Ising model and the Potts model. We apply LEAPS on each model on a 15 × 15 lattice corresponding to a d = 15 × 15 = 225 dimensional space. To construct ρt, we use temperature annealing making the inverse temperature βt a linear function of time. Ablation Ising model. To establish the validity of our method also experimentally, we compare our results against a ground truth of long-run Glauber dynamics, an efficient algorithm for simulation in this parameter regime. In Figure 8, one can see that LEAPS recovers the expected physical observables. Further, in Figure 8 we compare three different realizations of our method, one using LEA, and the other two using deep LEC that vary in depth. The deep LEC network performs best and we use it for all subsequent experiments.
Researcher Affiliation Academia 1Massachusetts Institute of Technology 2Society of Fellow, Harvard University 3Institute for Artificial Intelligence and Fundamental Interactions. Correspondence to: Peter Holderrieth <EMAIL>, Michael S. Albergo <EMAIL>.
Pseudocode Yes Algorithm 1 LEAPS Sampling 1: Require: N time steps, M walkers, model Fθt, replay buffer B, MCMC kernel Mt, density ρt, coeff. ϵt ≥ 0, resample thres. 0 ≤ δ ≤ 1 2: Init: Xm0 ∼ ρ0, Am0 = 0 (m = 1, . . . , M) 3: Set h = 1/N 4: for n = 0 to N − 1 do 5: for m = 1 to M do 6: Xm t ← Mt(·, Xm t ) with prob. hϵt else Xm t 7: Xm t+h ∼ (1Xt=y + h Qθt(y, Xt))y∈S 8: Am t+h = Am t + h Kθtρt(Xm t ) 9: end for 10: t ← t + h 11: if ESS(At) ≤ δ then 12: Xt = resample(Xt, At) (m = 1, . . . , M) 13: At = 0 (m = 1, . . . , M) 14: end if 15: end for 16: Optional: Store {(Xmt , Am t , t)}t,m in B. Algorithm 2 LEAPS training with optional replay buffer Require: B batch size, N time steps, model Gθt, free energy net Fϕt , learning rate η, replay buffer B. 1: while not converged do 2: if use buffer then 3: (Xmtm, Amtm, tm)m=1,...,B ← Sample Batch(B) 4: else 5: (Xmtm, Amtm, tm)m=1,...,B ← Algorithm 1 6: end if 7: L(θ, ϕ) = 1 B P m=1 |Kθtmρtm(Xmtm) − tmFϕtm|2 8: θ ← θ − η∇θL(θ, ϕ) 9: ϕ ← ϕ − η∇ϕL(θ, ϕ) 10: end while
Open Source Code Yes We provide code in https://github.com/malbergo/leaps/.
Open Datasets Yes As a demonstration of the validity of LEAPS in high dimensions, we benchmark it on models of statistical physics, namely the Ising model and the Potts model. Next, we use the DISCS benchmark (Goshvadi et al., 2023) to benchmark LEAPS against MCMC samplers (see Figure 5). We also apply the LEAPS method on the Potts model following Goshvadi et al. (2023).
Dataset Splits No The paper uses well-known models (Ising and Potts) and refers to the DISCS benchmark. It mentions comparing results against a 'ground truth' attained from Glauber dynamics but does not specify traditional train/test/validation splits for these models or any other dataset.
Hardware Specification No The paper does not explicitly mention any specific hardware (e.g., GPU models, CPU types, or cloud computing instances with specifications) used for running the experiments.
Software Dependencies No The paper does not explicitly provide specific software dependencies with version numbers (e.g., Python version, PyTorch version, specific library versions).
Experiment Setup Yes LEAPS on each model on a 15 × 15 lattice corresponding to a d = 15 × 15 = 225 dimensional space. To construct ρt, we use temperature annealing making the inverse temperature βt a linear function of time. LEAPS is run for 100 steps. For the locally equivariant attention (LEA) mechanism, we use 40 attention heads, each with query, key, and value matrices of dimension 40x40. As such, there are about 350,000 parameters in the model. In addition, the locally equivariant convolutional net (LEC) of depth three uses kernel sizes of [5, 7, 15], while the depth five version uses [3, 5, 7, 9, 15], amounting to around 100,000 parameters. Figure 7. Ablation experiments on simple Ising model with L = 15, J = 0.4, β = 0.7. Figure 8. Repeat experiments from Figure 8 for Ising model with critical temperature. Left: Comparison of the 2-point correlation function against the Glauber dynamics ground truth. Right: Histograms of the magnetization M(x) of configurations as compared to the ground truth set attained from a Glauber dynamics run of 25,000 steps. LEAPS recovers the expected observables. β = 0.4407, J = 1