From Predictors to Samplers via the Training Trajectory
Authors: Soumya Ram
Akhila Ram
Venue: ICLR 2026 Poster (2026)
Keywords: sampling
energy based models
discrete sampling
synergistic interactions
markov chain monte carlo
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Use a predictor's training checkpoints as an annealing schedule to improve sampling

Abstract: Sampling from trained predictors is fundamental for interpretability and as a compute-light alternative to diffusion models, but local samplers struggle on the rugged, high-frequency functions such models learn. We observe that standard neural‑network training implicitly produces a coarse‑to‑fine sequence of models. Early checkpoints suppress high‑degree/ high‑frequency components (Boolean monomials; spherical harmonics under NTK), while later checkpoints restore detail. We exploit this by running a simple annealed sampler across the training trajectory, using early checkpoints for high‑mobility proposals and later ones for refinement. In the Boolean domain, this can turn the exponential bottleneck arising from rugged landscapes or needle gadgets into a near-linear one. In the continuous domain, under the NTK regime, this corresponds to smoothing under the NTK kernel. Requiring no additional compute, our method shows strong empirical gains across a variety of synthetic and real-world tasks, including  constrained sampling tasks that diffusion models are unable to handle.

Published as a conference paper at ICLR 2026
FROM PREDICTORS TO SAMPLERS VIA THE TRAINING TRAJECTORY
Soumya Ram Akhila Ram²
1 Independent Researcher
2 Department of Computer Science, Massachusetts Institute of Technology
ABSTRACT
Sampling from trained predictors fundamental for interpretability and as a compute-light alternative to diffusion models, but local samplers struggle on the rugged, high-frequency functions such models learn. We observe that standard neural-network training implicitly produces coarse-to-fine sequence of models. Early checkpoints suppress high-degree/ high-frequency components (Boolean monomials; spherical harmonics under NTK), while later checkpoints restore de- tail. We exploit this by running a simple annealed sampler across the training trajectory, using early checkpoints for high-mobility proposals and later ones for refinement. In the Boolean domain, this can turn the exponential bottleneck aris- ing from rugged landscapes or needle gadgets into near-lineal one. In the con- tinuous domain, under the NTK regime, this corresponds to smoothing under the NTK kernel. Requiring no additional compute, our method shows strong empiri- cal gains across a variety of synthetic and real-world tasks, including constrained sampling tasks that diffusion models are unable to handle.
INTRODUCTION
In contrast to the trend toward billion-parameter Transformer LLMs, model deployments for medicine, recommendation systems, and decision support based on structured data continue to be dominated by small CNNs/MLPs. These dominate Al in medical devices (Singh et al. 2025|MMing jet al. 2025), production models for personalized recommendations (Feng et al., 2024), and decision support models for credit scoring, recidivism risk, insurance underwriting, and hospital operations
triage (eba 2023 Grinsztajn et al. 2022 McEifresh al. 2023 Holzmüller et al. 2024).
Despite their relative simplicity, they can be opaque and encode brittle shortcuts. For example, dermatology CNN approved for EU clinical use was shown post-hoc to over-weight surgical skin markings/rulers rather than lesion content; adding simple violet marker to the same benign lesion skyrocketed its melanoma probability (Winkler et al., 2019 Bevan & Atapour-Abarghee Sampling minimal counterfactual edits could have revealed this shortcut. Such cases underscore the importance of sampling from the trained predictor.
2022).
Apart from interpretabilityy candidates for active learning. For instance, many works sample fit DNA sequences from models predicting -NN-ssrancciittion factor (TF) affinity (de Almeida et al., [222 Reddy et al. 2024). This works well as DNA-TF affinity assays test millions of sequences per experiment, often from largely random libraries-enti enabling near-unbiased exploration of sequence space (Gallego Romero Leal )023). Similarly, for protein engineering, many methods sample sequences from a learned fitness model for bayesian optimization (Hu et al., 2022] Ren et al., (2022).
often want to sample from trained predictor pick high-value
However, these sampling tasks can become difficult for certain common landscapes. When the landscape is rugged, high-frequency, high-magnitude fluctuations create many sharp local optima. Another key culprit is synergy outcomes depend on rare combinations of variables- so the indi- vidual effects look innocuous while the rare, joint effect is large. Such concealed interactions evade single-step proposals and short-horizon heuristics, which only "see" myopic gains. As result, dis- covering the right multi-varian variable change becomes essentially a needle-ine in-a-naystacck problem, with
"Corresponding author; soumya 000ma coo

---

Published as a conference paper at ICLR 2026
success of random exploration drops exponentially as the number and order of synergistic interac- tions grow, along with the number of spurious variables.
option is to train award-condiitoned diffusion discrete-walk jump sampler. These are powerful generative approaches, but in the settings we focus on they have three drawbacks: (1) They require training separate generative model in addition to the predictor, which can demand substantial extra compute compared to reusing an already-trainee predictor with test-time trajectory-anneant
To deal with this difficulty,
MCMC. In many domains, strong predictors have already been trained on large, unbiased data, and practitioners operate in compute-constraing
regime where training an
additional generative model is not feasible. (2) Implementing hard constraints such as a Hamming- radius ball or minimal counterfactual edits typically requires additional machinery (e.g., auxiliary guidance networks [Shen et al. (2024) or SMC-style schemes Wu et al. (2023)), rather than a simple modification of the sampling rule. (3) They do not directly support sampling from deployed model for interpretability, whereas our method operates on the existing predictor without any additional training.
We study the plug-and-play test-time sampling problem for trained scalar predictor f": X I. and we aim to draw from the Gibbs density it induces, x*(x) xx exp(/"(x:) optionally under hard constraints. Our solution is trajectory annealing: rather than sample only from IT, we traverse training checkpoints {ft) and run brief MCMC updates targeting ne (x) (x)( before arriving at n²
This exploits the coarse-to-fine learning dynamics of predictors: early checkpoints suppress high- degree components, smoothing the landscape for rapid mixing. This yields two wins. For π (x) with nigh-magnitude, high-frequency variation, early checkpoints bypass the rugged barriers that cause exponential mixing. For synergistic interactions where only higher order, rare combinations of variables are predictive, we show that early checkpoints correspond to low degree projections that reveal modes of π (x), converting sundom-wall behavior with exponential mixing times to near-linear. Our method works as-is, requiring no additional compute or training changes. We demonstrate strong empirical gains on synthetic tasks, sampling from discrete masey-bassed models, and challenging real-world DNA and materials design tasks. the best of our knowledge, this the first work to leverage a neural network's 's training trajectory to improve sampling.
2 RELATED WORK
Smoothing for sampling There is long line of work that smooths the target to accelerate sampling. Prominent examples include reward-conditioned diffusion and discrete walk-jump schemes that walk smoothed manifold and jump back to the discrete space (Yuan et al. 2023 Frey et al,, 2024]. Kirjner et al. (2024) train graph-smoothed protein fitness models ( 250K node sequence graphs) and Zhu et al. (2025) show this smoothing induces spectral bias that disproportionately damps high-degree Boolean monomials. Our approach instead leverages the natural smoothing in network's training trajectory; focus on the no-extra-compute regime and therefore do not benchmark against explicit smoothing methods.
Interpretability Many approaches probe a trained predictor via sampling/optimizzaion. Minimal counterfactual edits are typically posed as constrained optimization on the fixed model,with spar- sity or proximity regularization (Verma et al., (024). A complementary line samples the predictor to quantify rare events- using MCMC to estimate the mass of inputs that violate a property or elicit outlier predictions (Wehb et al., 2019). However, these approaches inherit the exponen- tial barriers from rugged, higuenryy landscapes and random-war exploration for synergistic interactions- issues our method mitigates when substituted for the sampling/optimizatinn.
Test- Time MCMC Sampling Existing test-time, plug-and-play MCMC methods use gradient- guided methods for search and temperature annealing to overcome barriers. However, such tempera- ture annealing methods (e.g- parallel tempering. annealed importance sampling, etc.) cannot bypass the random walk exploration that occurs from rare synergies (Hénin et al. (022). Furthermore, for functions with high barriers, tempering relaxes barriers but offers little directional guidance leading back to the random walk exploration. Thus, mixing time in the above settings remains exponential. Recent methods such Diffusive Gibbs Sampling introduce an auxiliary noisy variable and alter- nate Gaussian noising with gradient-based denoising steps in Gibbs scheme to improve mixing
2

---

Published as a conference paper at ICLR 2026
on multi-moda targets CChnn et al., 2024). Similarly, Iterative Reasoning through Energy Diffusion (IRED) learns annealed energy landscapes but still depends on local energy gradients at test time (Du al. 2024). However, all test-time gradient based methods are limited by the informativeness of local energy gradients.
Discrete sampling Grathwohl et al. (2021) introduced Gibbs-with-Gradien (GWG), using model gradients to choose which coordinate to flip rather than sampling indices uniformly. Since then, flurry of work has pushed discrete MCMC forward along complementary axes: locally bal-
anced/informed proposals that improve Metropolis-Hasting tuning Zanella, 2017 Sun et al.)
2022; non-local or parallel gradient moves via discrete Langevin and related formulations (Zhang et al. 2022; Sun et al. 2023); automatic cyclical scheduling of gradient- -basd updates for better mixing and reduced tuning (Pynadath et al.. 2024); MAL-A-inspired discrete kernels with auxiliary- variable preconditioning (Rhodes & Gutmann 2022); and reheated gradient-based samplers tailored to difficult combinatorial objectives (L.i & Zhang 2025). Our method is compatible with all of these gradient-based discrete kernels and could be combined with their proposal mechanisms or sched- ules; for simplicity, we use GWG throughout our discrete experiments.
Coarse-to- Fine Learning A growing body of theory suggests that gradient- based training and sam- pling in high-dimensional models proceeds in a coarse-to-fine manner. In diffusion models, linear and Gaussian analyses show that high- variance or low-frequency modes of the data covariance learned and expressed in samples much earlier than low-variance, fine-detail modes, leading to an ordered emergence of global structure before local detail (Wang) 2025 Wang & Vastola 2024). Re- lated analyses of SGD on neural networks reveal multi-phase,
dynamics in which
low-complexity or small-support features are acquired first, progressively enabling the learning of
higher-order interactions (Abbe et al. (023).. Similar spectral decompositions of the NTK further
indicate that only a few dominant eigendirections are amplified early in training, biasing learning toward coarse structure before finer modes are fit (Murray et al., 2022).
3 METHODS
Test-time setting. We work in plug-and-pul regime with trained predictor X (MSE- trained on {(r,,,:)). and we sample from its induced density r"(x) x exp(f"(x)), optionally under hard constraints. We do not compare against setups that modify training or fit auxiliary gener- ative/score models (e.g., diffusion); our contribution entirely in the test-time sampling procedure.
Trajectory annealing. Rather than run MCMC only on π², traverse checkpoints along the training trajectory (fe)/_o with fr = f", defining intermediate targets n (x) x exp (f{(z)). Start- ing from & = 0, we apply short Markov kernel for N, steps targeting Te, earry the resulting state forward as the initializer for Tt+1, and continue this coarse-to-fine progression until t = T.. For kernels, we use GWG+MH (Gibbs w/ Gradients + Metropolis Hastings) for discrete and MALA cittropolisAddjstter Langevin Algorithm) for continuous z2
Neural networks learn coarse-ffine: low-frequency structure emerges early, highffrequency later. In discrete models trained with SGD, gradients align more with lower-degree monomials, so those coefficients converge first. In continuous models in the NTK regime, kernel eigenvalues decay with cateericl-harmonne degree; the predictor is f convolved with the kernel, giving strong early smoothing that relaxes over time.
3.1 BOOLEAN VARIABLES
Abbe al. (2023) show SGD learns Boolean functions hierarchically: low-degree monomials learned first as fewer variables leads to greater gradient alignment. We leverage this to turn expo- nential sampling into polynomial time. We study two hard classes: (i) f that are hard because of high-frequency, high-magnitude variation-herr, early checkpoints haven't 't learned the high-degree spikes yet, so the landscape smooth and mixes quickly; and (ii) f* that are hard because there's no variation (the needle gadget). For needles, local sampling is random walk exponential in However, mixing on /" projected to monomials of degree S mixes in (d log d), with mixing worsening the largest degree increases. The low-degree projection of f* also acts as an associative memory that can store many needles.
the needle dimension
3

---

Published as a conference paper at ICLR 2026
3.1.1 BACKGROUND: LOWER DEGREE MONOMIALS ARE LEARNED FIRST
Abbe al. (2023) formalize hierarchy in how SGD fits sparse Boolean targets. Writing the target as a sum of Boolean monomials, they define the leap as the smallest & for which one can order the nonzero monomials so that, when adding the next monomial in that order, the union of involved variables introduces at most k new variables. A pure "staircase" target each term extending the previous by one fresh variable (e-g. I +1+ 2 (122))) -has leap
This notion predicts hierarchical learning under SGD. For laie functions, lower-degree monomials are learned first and higher-degree monomials later: initial gradients correlate more strongly with terms that require fewer new variables, so SGD first aligns small set of coordinates; that alignment then amplifies gradients toward the next monomial, and so on. The trajectory passes through saddle-to-saddle plateaus; phase that requires acquiring L new variables at once takes 6 (ddmax((-111)) steps, the training time is dominated by the largest leap (i.e., the hardest stage).
They prove this in a restricted two-layr fully connected networks with smooth activations, trained on i..dd data using modified SGD (layer-wise updates plus a projection step)-and are complemented by empirical evidence: loss curves for deeper networks on hypercube data exhibit clear plateaus and drops consistent with learning across successive leaps (Abbe et al. 2023).
We provide additional empirical evidence for the hierarchical-leamingg picture across fully connected and convolutional networks, spanning variety of activations, widths, and depths in Appendix C Two regularities emerge: (i) lower-degree Fourier-Walsh components finish aligning with the target function earlier than higher-degree components, and (ii) the degree- wise mass grows only after all its monomials are fully aligned. See Fig. [ for an example. A caveat is transformers, where we observe experimentally they satisfy (i) but (ii). A degree-2 monomial could become aligned and grow in mass before all degree- 1 monomials were aligned.
Core assumption (degree-wise checkpoints). We assume the setting in Abbe et al. (2023) holds for the larger networks we consider. Specifically, along the training trajectory (f. with fr f", there exist increasing checkpoints TO < т1 < .... < <TK- I T such that at Tx the model has effee- tively learned all interactions up to degree k while higher-degree components are still negligible. Equivalently, we may treat fr. as the degree-k projection of the final model:
fck
fex(z) = Σ (S) II Z₁-
SCI
iES
SI
Between these checkpoints, higher-degree terms may be partially learned; we assume only the exis- tence and monotone ordering of (Tk)-
3.1.2 MAGNITUGN HIGH- FREQUENCY VARIATION
Early checkpoints in the training trajectory suppress high-degree terms, smoothing the landscape and making it easy to mix. We exploit this to handle targets with large high-degree components.
As running example, consider
{+1]d with
Σ
T., x exp
7
T3
i=1
where the linear term favors many +1 entries and the parity term II. x. creates a high barrier when |y| is large.
At low temperature, vanilla Gibbs the full objective mixes in exponential time Ô exp(clr once a random start flips to satisfy the parity term, any move that increases the number +1 bits must cross lnlesized barrier, the chain gets stuck near suboptimal states.
Our trajectory sampler avoids this. We first run a short chain at checkpoint T1; this mixes in Odd d) under Gibbs (see App. A and quickly reaches states with many +1 entries. We then continue the chain at the final checkpoint to adjust the parity. Thus, we are able to hit the global maxima in near-linear time sidestepping the exponential barrier.
4

---

Published as a conference paper at ICLR 2026
Fourier-Wald mass by degree during training (10-D indicator) For 3-Layer FCNN of width 2.5
22
deg
3
deg
4eg
deg
deg
deg
2.0
deg
deg
Fourier mass 5 itsy
deg
1.-
deg
deg
deg 10
1.0
0
50 100
150
200
250
300
350
400
Epoch
Figure 1: Fourier- L mass by degree during training d-10 indicator. The tar-
get is f(x) = 1(x } with T,, 2 E {+1}. In -coorditddss this expands as f(x) = Escia (les((:t))))) so the degree-k component consists of all k-way monomials [2
aligned with the pattern =² Curves show Eis|= I(()) over epochs for 3-layer FCNN (width 128). Vertical dashed lines mark the alignment epoch for each degree k: the first epoch at which the sign of every degree-k coefficient agrees with the sign implied
aligned monomial (i.e., all signs point toward >²). We observe consistent ordering: each degree first becomes aligned and then its Fourier mass rises, with lower degrees aligning and growing earlier than higher degrees. Although some final masses are larger at higher degrees (there are (+) degree-k monomials), this effect does not change the early-epoch ordering-loo-degeee components begin to align and increase first.
3.1.3 SYNERGISTIC INTERACTIONS
Needle-like synergistic interactions. our setting, variables interact synergistically; payoff arises only when a specific joint configuration is met (a "needle gadget"). Let 2* E {+1) denote that pattern and define the indicator
(x) 1xx=2") x {=1
Over the Boolean (Walsh) basis, f decomposes into 2d monomials aligned with =*:
(*)x) = 2-d II (1+ 2- Σ (=ta:)
1=1
SCI RES
so the degree-k component is the sum of all k- way monomials =1,) with |S]
Because the density is flat off of a tiny set C' (e.g., single configuration), a local Markov chain be- haves like a simple random walk on the 2d-vertex hypercube until it enters the 1 I-Hamming neigh- borhood of C. This is exponential in d.
Intermediate checkpoint (degree-2 aligned) Now suppose we are at checkpoint 72. Let aligned spins y xzzE. The degree- 2 surrogate can be written as
d
(v) R (5e(1) 2-d
Σ U1
Σ
y:Ul
ISisiEd
which is the Curie Weiss Hamiltonian with a positive external field up to scaling. Existing results show that in the low-temperature regime, hit z* with high probability after Odd steps
5

---

Published as a conference paper at ICLR 2026
with constant number of parallel chains. See App. for details. By exploiting the checkpoint 72, we have converted the exponential random walk into near-linear mixing.
Multiple Needles Even with multiple needles, show theoretically (via connections to binary Hopfield models) and experimentally that low-degree monomials are sufficient to store and retrieve needles. Further details are in App. D
All checkpoints help Sampling against fsk shows a monotone pattern: as k grows during training, the landscape sharpens and the needle hitring time increases. Empirical evidence is in App. Sampling is fastest with k = 2 but sampling from any intermediate checkpoint with & < d helps.
3.2 CONTINUOUS VARIABLES
Let f S + R and let denote the time parameter. Gaussian (diffusion) smoothing on the sphere acts degree-wise degree-k: coefficient of I* multiplied by M&(() = exp(- k(k + 2)} (larger = more smoothing; higher k = stronger damping). NTK training (idcalized FCNN: infinite width, zero init, uniform data) also acts degree-wise on spherical harmonics: the scaling (t decays with degree as determined by the activation, @k--)) for ReLU and 6(k- vit for Tanh (Murray al. 2022). Takeaway. The NTK training trajectory {fe} already provides a continuum of smoothed versions of the the same coarse-to- sineine effect as heat-kernel smoothing. unlike diffusion which learns these smoothed functions explicitly. See App. F Flor more details.
spherical harmonics: time i, the degree-k coefficient of ft equals
4 RESULTS
We evaluate our method under matched compute sampling from discrete functions synthetic boolean functions (ruggedness, synergy), binary MNIST-EBM, DNA design task (including con- strained sampling), and continuous functions Ackley 10-D, and superconductor design task.
4.1
DISCRETE SAMPLING EXPERIMENTS
Experimental details are in App.
4.1.1 SYNTHETIC BOOLEAN EXPERIMENTS
We conceptually show how our method can turn exponential mixing into near-linear on synthetic targets. It succeeds by leveraging (1) fast mixing on low- low-degree surrogates and (2) knowledge the function's 's support gleaned from those surrogates.
We evaluate four functions: (i) targets dominated by high-degree components, (ii) indicator func- tions of increasing size, (iii) indicators with adversarial non-convex linear terms, and (iv) multiple indicators (synergistic interactions). To stress (2), every task includes 500 spurious variables.
To make the contrast stark, we run GWG on the final checkpoint for up to 2000 steps, while our method uses only 40 steps total. We report the percentage of runs that reach the global optimum. Even with 50x more steps (2000 vs 40), GWG's s hit rates remain very low (e.g. S 3%), whereas our method near-perfect on most tasks.
This is because GWG is a random walk that succeeds mainly when it starts near the target, and that probability collapses exponentially with growing indicator size, more distinct synergies to satisfy, and many spurious variables. In contrast, our method's 's (1) fast mixing and (2) support knowledge do degrade these regimes.

---

Published as a conference paper at ICLR 2026
Table I: Sampling with 500 spurious variables on 8-variable polynomial target dominated by high-frequency variation f() 0.1 Σi= 0.2 Σ 22-22 0.4 Eie I Zi+* +
0.8 Eie(0.2) Io Titk 3.0 EiE Ih Litk. Success hits all +1 global max. on variables. Result: Ours outperforms both baselines (0.52 vs. 0.04) with fraction of steps.
Ours: Success (95% % CT) 0.5200 (0.45-0..) 0.599
GWG: Success (95% CI) 0.01 (0.00-0.02)
GWG Temp Anneal: Success (95% CI) 0.04 (0.01-0.06)
High-magnitude, high-frequency variation We construct synthetic function whose coefficients increase with degree. With only 40 steps (vs 2000 for the baseline), we hit the global maxima 52%, compared to 4% for remperature-annealed GWG on the final checkpoint. See Table This is because the early checkpoint allows for fast mixing without the high- figh-frequency variation.
Table 2: Sampling with 500 spurious variables and d indicator variables. Ours 20 GWG steps at epoch 25 plus 20 at final. Baseline = GWG on the final checkpoint for 2000 steps. All runs use GWG with 10. Reported are success fractions with 2SD Cls; Success hits indicator. Baseline also reports median steps to first hit (conditional on success; CI lower bounds clipped at 0).
Success probability (95% CT)
GWG: median steps to first hit (given success) [95% CI)
Ours
GWG
Med. steps
0.98 (0.96-1.00)
0.47 (0.40 0.54)
[1-1]
1.00 (1.00-1.00)
0.21 (0.15 0.27)
[1-1]
1.00 1.0000000)
0.17 (0.11-00.2)
4 [0-922.9
10 0.99 (0.98-1.00)
0.12 (0.07-0.16)
2 [0-56.8)
Indicator function For an indicator function with 500 spurious variables, a random walk takes, on average, 500 2d steps. GWG's conditional median steps (given hit) are 1-4 steps, indicating it mostly succeeds when the initialization close to the target. Thus, its hits degrade as d increases (12% for d 10) whereas method remains perfect (despite having only 50 steps). See Table 2
Indicator function with adversarial non-convexity We add an adversarial degree-1 terms that opposite to the indicator pattern. However, because the indicator dominates the stationary mea- sure, the local field is still dominated by the indicator's s low-degree expansion. Thus, the adversarial linear term has minimal impact and our method has perfect hit rate. See Table
Table 3: Sampling in non-convex binary landscape. The objective is indicator 10 des- ignated variables that yields 10 only at the all-ones pattern (and otherwise), plus a linear term -0.1 10 24 on the same variables that pulls toward all ls; 500 additional variables are spurious (no effect). Success hitting the indicator. Result: Ours near perfect as the non-convexity is not able to dampen the signal from the intermediate checkpoint
Ours: Success (95% CT) GWG: Success (95% CT)
1.00 (1.00-1.00)
0.080 (0.0416-0.1184)
Table Sampling with 500 spurious variables and 3 non-overlapping length-5 indicators. Success hits all 3 indicators at once. Result: Ours is near-perfect, while the baseline is near chance. With only one length 5 indicator (see Table 3). success rate is 0.21, drops with more indicators.
Ours: Success (95% % CI)
GWG: Success (95% CT)
1.00 (1.00-1.00)
0.0250 (0.0029-0.0471)
Multiple indicator functions GWG for a length-5 indicator has a modest hit-rate of 21% (see Table []. however, performance collapses 3% when have three length-5 indicators (on

---

Published as a conference paper at ICLR 2026
overlapping subsets). However, method remains perfect. This is because GWG depends on starting near a good basin (which becomes exponentially unlikely as synergies compound). How- ever, our method leverages (i) fast mixing lower-degree surrogates and (ii) knowledge of the support to consistently (100% vs 3%) find the solution a fraction of steps. See Table
4.1.2 SAMPLING FROM MNIST ENERGY- BASED MODEL
The phenomena emphasized above frigh-frequenc variation and many synergistic interactions are common to real-world data. Thus, we test our method's 's ability to efficiently sample from energy based models (EBMs) trained on binary MNIST.
We train a binary EBM with GWG using the implementation from Grathwohl et al. (2021). At test time, we compare used in Grathwohl al. (202T), which we denote Temp-GWG) vs. annealing along the training trajectory (our method). For both methods, we report FID after either 1K or 10K GWG sampling steps. For our method, we evenly distribute the steps across 500 evenly spaced checkpoints.
-annalllng from the final checkpoint (the baseline inference method
Table 5: FID (L) on binary MNIST using LeNet features. Mean (std) over 10 bootstraps
Method
IK steps
10K steps
Temp-GWG
29.61 (0.239)
21.12 (0.138)
Ours
11.73 (0.284)
5.49 (0.119)
Ground-truth
0.01 (0.013)
Controlling for the number of steps, we observe substantially better FIDs in Table We provide ablations on the number of checkpoints in App. G which show significant improvements over the baseline across a wide range of checkpoint counts. App. contains random samples; ours are substantially sharper.
4.1.3 DNA DESIGN EXPERIMENT
Recently, CNNs have matched or exceeded transformers for both protein and DNA language mod- oling-consisse with the fact that local, repeated motif patterns align naturally with convolutional filters Yang et al. 2024 Bo et al. 2025). Moreover, standard TFF assays reach million scale because short oligo are mass-proos processsd,eeedd, yielding abundant data with oigher-order motif structure, so trained regression models are both accurate and difficult to sample from (Berger et al. 2006).
We train CNN on a dataset that measures binding affinity to the transcription factor (TF) MAX (Badis et al.) 2009). The dataset consists of 42K 60-mer sequences. For sampling, we sampled a random length-60 sequence and performed 60 steps of GWG (on the final checkpoint for the baseline, and across training checkpoints for our method).
Table 6 shows that our median samples are over 10 more performant (fitness is on log-scale), partially driven by the ease of sampling the TF's motif (79% for ours vs. 39% for baseline).
We also show our method's robustness in the constrained sampling setting where we want to sample within Hamming ball of fixed starting point. This mirrors common biological use cases such as identifying minimal gene perturbations, minimal amino-acid mutations, or minimal CRISPR/hase-cditoo edits to boost activity.
Notably, this setting is not straightforward to handle with standard diffusion pipelines. Diffusion models support inpainting tasks conditioned fixed portions of the final output, but enforcing an exact constraint such as remaining within a Hamming ball of fixed sequence typically requires additional machinery (e.g., auxiliary guidance networks Shhnn et al. (2024) or SMC-style schemes
Wu et al. (2023)) rather than single pre-trained denoiser. By contrast, our method handles this naturally by enforcing the Hamming-ball constraint throughout the sampling process, simply by restricting the MCMC chain to the constraint set. This blug-and-play ability to impose new hard constraints test time, without training any additional generative or reward model, practical advantage of our predictor- based approach in this application.
8

---

Published as a conference paper at ICLR 2026
Table Per-run metries with 95% bootstrap percentile CIs (B=500) for constrained sampling. Sam- pling trajectory is restricted to always stay within hamming distance of 7 from the starting point. Each run starts from a random length-60 DNA sequence, runs 60 mutation steps, and keeps the best; repeated 300 times. Diversity/novelty are recomputed per bootstrap. Percentiles are vs. the training set restricted to y > 0.
Method
Fitness (median)
Pet.
Diversity
Novelty
Motif (96)
Ours
7.42 [7.03, 7.63] 2.09 [0.98, 3.06]
99.45 199.38, 99.47]
45 [45, 45]
33 [33, 33
63.0 [58.2, 68.0]
GWG
91.89 [66.66, 96.65]
45 [45, 45]
33 [33, 34]
31.3 [26.0, 37.3]
Temp
-0.32 [-0.79. -0]10]
0.00 [0.00. 0.00]
45 [45, 45]
33 [33, 34]
3.3, 5.0]
PT-GWG
1.92 [1.51, 2.58]
90.18 [83.25, 95.03]
45 [45, 45]
33 [33, 34] 24.3 (19.3, 29.3]
Table 8: Main and secondary metrics. Bracketed values are 95% CIs computed as Student-r intervals across seeds on per-seed means. Secondary metrics for Superconductor are reported as median [IQR]. Refs: Ackley 0.0; Superconductor 185.0.
Novelty Diversity [IQR]
Experiment
Method
Best
Mean [95% CT]
MCMC-Final
8,5628
16.2164 [56.1559, 16.2770]
SMC-Temp
7.8595
16.3141 [16.2735, 16.3547]
Ackley (10D, 4)
Aemm
8.8096
16.3095
[16.2679, 16.3512]
PT
13.6225
[19.4593 20.2803] 19.8698
SMC-Train
3.6942
13.3311
[12.0867, 14.5755]
MCMC- Final
107.4
76.68 (76.42, 76.94]
17.30 [4.98] 11.76 [7.47]
SMC- SMC-Temp
80.7
21.98 [21.39, 22,56]
34.39 [2.42] / 12.76 [2.69]
Superconductor (1)
AISAutoTTemp
107.4
24.42 [23.88, 24.97)
35.66 [5.18/ / 26.05 [5.28]
PT
107.4
25.97 [24.31, 27.62]
35.29 [5.01] / 25.40 [4.91]
SMC-Train
318.4
155.2 [105.6, 204.8]
20.86 [3.64] 16.60 [6.87]
Table illustrates the challenging task of finding a length-7 Hamming perturbation to random length-60 DNA sequence to improve the sequence's fitness. Again, our method finds samples that 10 more performant, and twice likely (63% vs 31 1)) to contain the motif.
Table Per-run metrics with 95% bootstrap percentile CIs (B=500). Each run starts from a random length-60 DNA sequence, runs 60 mutation steps, and keeps the best; repeated 300 times. Diver- sity/novelty are recomputed per bootstrap. Percentiles are vs. the training set restricted to y > 0.
Method
Fitness (median)
Pet.
Diversity
Novelty
Motif (%)
Ours
10.04 [9.74, 10.20)
99.78 [99.77, 99.80)
45 [45, 45]
34 [33, 34]
74.3 [69.7, 79.3]
GWG
2.72 [1.66 60, 5.53]
95.54 [85. 24, 98.831
45 [45, 45]
33 [33, 34]
38.7 [33.3, 44.0]
ATSAutoTemp
GGG GWG
0.56 10.16, 0.85]
42.86 [12.99, 59.772
45 [45, 45]
33 [33, 33]
4,0 [2.0, 6.3]
PT-GWG
-0.79 [-1.17, -0.25]
0.00 [0.00, 0.00]
45 [45, 45]
33 [33, 34]
13.0 [9.3, 16.7]
4.2 CONTINUOUS SAMPLING EXPERIMENTS
Setup. We compare four samplers under matched compute: Sequential Monte Carlo with tempera- ture annealing (SMC-Temp). Annealed Importance Sampling (AISAutoTemp), Parallel Tempering (PT), and our Sequential Monte Carlo with training-ting checkpoints (SMC-Train). All share the same compute budget, with results averaged over seeds. Sampling/budget details are App. K SMC and checkpointing specifics are in App.
Ackley (10D). Rugged continuous optimization on [-10, 10 SMC-Train attains the top mean and best-of-set under matched compute with non-overlapping CIs (See Table 8).
Superconductor. A real- world benchmark in high-D materials design taken from the design benchmark in Trabucco al. (022).. It has rugged, non-convex, heavy-tailed landscape. Inputs

---

Published as a conference paper at ICLR 2026
E IR encode element composition; the target y is the critical temperature T. (K). SMC-Train achieves the highest mean and best-of-set Te, exceeding the reference and all baselines (See Table
DISCUSSION
Discrete vs. Continuous In continuous domains, early training in the NTK (linearized) regime induces trequency-selectio smoothing; training leaves the NTK regime, this smoothing fades. In discrete domains, the effect stronger. Because boolean targets are learned low-to-high degree, there are less variables in the high degree term, causing higher gradient alignment. This allows for faster learning in fact, the number of steps is asymptotically optimal, matching Correlational Statistical Query (CSQ) lower bounds (Abbe et 2023). Other work suggests that SGD learns with an optimal number of steps in more general settings (Barak et al. 2022). Thus, hierarchical learning arises naturally from SGD's inherent efficiency.
Architectures Continnous For FCNNs, the NTK eigenfunctions are spherical harmonics, so spectral bias aligns directly with smoothness (low degree ++ larger eigenvalues); CNNs/ResNets inherit this (Geifman et al. [022,] Belferr et al. (024). Transformers have different eigenfunctions, so this does not apply. (Hron al., (020). Discrete Hierarchical interaction learning under SGD FCNN- specific and transfers to CNNs/ResNets (we report results for all three). Transformers, by contrast, learn interactions via the attention matrix, which is qualitatively different.
I Limitations Our method does not apply to transformers. However, commonly deployed models in medicine, personalized recommendations, and decision support tend to be CNN/MLPs, as detailed in Section interpretability is crucial in these domains. In addition, for protein/DNA predictor models, CNNs outperform transformers low-N fitness tasks (Dallago et al.) 2021), and have re- cently exceeded transformers in pretrained protein/DNA language models (Yang et al. 2024 [B
et al. 2025). This is because local, repeated motif patterns align naturally with convolutional filters.
6 CONCLUSION
Sampling from trained predictor f is important for interpretability and compute-efficiet design. However, rugged and needle-gadget landscapes lead to exponential mixing times that standard, cemperature-annellie based MCMC methods cannot overcome. We demonstrate our trajectory- annealed samplers bypass this barrier across (1) three common architectures (FCNNs. CNNs, and ResNets) ranging from 2-20 layers and across (2) diverse tasks such as synthetic stress-tests, real- world design tasks, and EBM sampling. We theoretically characterize our method's benefits, show- ing exponential -> near-linear sampling improvements under idealized conditions. To our knowl- edge, we are the first to identify and exploit this training-trajectrry lens for neural network sampling. We hope our analysis can spark further research on this topic. including extensions to transformers. Given the method's 's simplicity, we hope can become a useful tool for efficiently probing predictors.
ETHICS STATEMENT
This work introduces a sampling procedure that reuses training checkpoints to improve efficiency when exploring trained predictor's landscape. When paired with interpretability workflows, this can help surface spurious correlations, biases, and failure modes prior to deployment.
At the same time, any method that accelerates sampling optimization over model scores has dual- use potential: could make it easier to construct high-confidence but misleading inputs (adversarial examples jailbreak prompts), search for harmful designs, or probe models in ways that risk model inversion or privacy leakage if training data contain sensitive information.
8 REPRODUCIBILITY STATEMENT
The code and data for all experiments are contained the supplementary zip file. The only exception is the MNIST-EBM sampling experiments. Here, the GWG repo was used as-is, with minimal changes for our inference method, Those minimal changes are explained in I.3.2
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Machine learning for irb models: Follow-up report from the consultation on the discus- sion paper. Technical Report EBA/REP/2023/8, European Banking Authority. August
2023.
URL
library
20on 2atogfahaccciiiseeeerarno222222d
Emmanuel Abbe, Esther Adserå, and Thibault Misiakiewicz. Sgd learning on neural networks:
leap complexity and saddle-to-saddle dynamics. In Conference on Learning Theory (COLT), 2023.
Daniel J. Amit, Hanoch Gutfreund, and H. Sompolinsky. Storing infinite numbers of patterns in spin-glass model of neural networks. Physical Review Letters, 55(14):550-1153, 1985. doi: 10.1 1000:/hyyyeeLeet.55..5555
Gwenael Badis, Michael Berger, Anthony A Philippakis, Shaheynoor Talukder, Andrew Gehrke, Savina A Jaeger, Esther T Chan, Genita Metzler, Anastasia Vedenko, Xiaoyu Chen, et al. Diversity and complexity in dna recognition by transcription factors. Science, 324(5955) 2009.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hid- den progress in deep learning: Sgd learns parities the computational limit. Advances in Neural Information Processing Systems, 35:21750-211 2022.
Yuval Belfer, Amnon Geifman, Meirav Galun, and Ronen Basri. Spectral analysis of the neural tangent kernel for deep residual networks. Journal of Machine Learning Research, 25(184):1-49 2024.
Michael F. Berger, Anthony A. Philippakis, Aaron M. Qureshi, Fangxue S. He, Preston W. Es- tep, and Martha L. Bulyk. Compact, universal dna microarays to comprehensively determine transcription franseription-fcttor factor binding site specificities. Nature Biotechnology, 24(11):1429-1435 2006. doi: 10. 1038//ttt 246.
Peter Bevan and Amir Atapour-Abarghouei Skin deep unlearning: Artefact and instrument debias- ing in the context of melanoma classification. In Proceedings of the 39th International Confer- ence on Machine Learning (ICML), PMLR 162, 2022. URL http::::pppppeeeeeea.a.a .mlr..innn...................... press/vl 62 htwp:///eeeeee/e///eeeaaaa. .df
Yu Bo, Weian Mao, Yanjun Shao, Weiqiang Bai, Peng Ye, Xinzhu Ma, Junbo Zhao, Hao Chen, and Chunhua Shen. Revisiting convolution architecture in the realm of dna foundation mod- cls. In International Conference on Learning Representations (ICLR), 2025. URL https: //openieviee.net//orraaaae BO77LLLWW--
Anton Bovier and Beat Niederhauser. The spin-glass phase transition in the hopfield model with p-spin interactions. Advances in Theoretical and Mathematical Physics, 5(6), 2001.
Benjamin Bowman. A brief introduction to the neural tangent kernel. 2023.
Wenlin Chen, Mingtian Zhang. Brooks Paige, José Miguel Hernández-Lobate and David Barber. Diffussie gibbs sampling. arXiv preprint arXiv:2402.03003 2024.
Christian Dallago, Jody Mou, Kadina E Johnston, Bruce Wittmann, Nicholas Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks in fitness landscape inference for proteins. bioRxiv, pp. 2021-11, 2021.
Bernardo P. de Almeida. Franziska Reiter, Michaela Pagani, and Alexander Stark. Deepstarr pre- dicts enhancer activity from dna sequence and enables the novo design of synthetic enhancers. Nature Genetics, 54:613-624, 2022. doi: 211088//1188-022-00000
Jian Ding, Eyal Lubetzky, and Yuval Peres. Censored glauber dynamics for the mean field ising model. Journal of Statistical Physics, 137(3-4):407-455 2009. doi: 10.1007/ 110955-009-958 1.
11

---

Published as a conference paper at ICLR 2026
Yilun Du, Jiayuan Mao, and Joshua B Tenenbaum. Learning iterative reasoning through energy diffusion. arXiv preprint arXiv:2406.114 11179, 2024.
H. Feng, B. Zhang, F. Ye, M. Si, C.-H. Chu, J. Tian, C. Yin, S. Deng, Y. Hao, P. Balaji, T. Geng, and D. Tao. Accelerating communication in deep learning recommendation model training with dual- level adaptive lossy compression. In Proceedings of SC '24, 2024. URL https: //ttpsi//arxiv.
org/ aba/ 2407 04272
Nathan C. Frey, Daniel Berenberg, Karina Zadorozhny, Joseph Kleinhenz, Julien Lafrance- -afrance-Vanasse, Isidro Hotzel, Yan Wu, Stephen Ra, Richard Bonneau, Kyunghyun Cho, Andreas Loukas, Vladimir Gligorijevie, and Saeed Saremi. Protein discovery with discrete walk-jump sampling. In Proceedings of the Twelfth International Conference on Learning Representations (ICLR), 2024.
URL wtt.https://att /ttp//rxiiii www.wwwwww....................ee aww.////.66... 12360
Irene Gallego Romero and Amanda J. Lea. Leveraging massively parallel reporter assays for evolu- tionary questions. Genome Biology, 24(26), 2023. doi: 111111///135--0251 023-02856-6.
Amnon Geifman, Meirav Galun, David Jacobs, and Basri Ronen. On the spectral bias of convolu- tional neural tangent and gaussian process kernels. Advances in Neural Information Processing Systems, 35:1 1253-11265, 2022.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops i took a gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning, pp. 3831-3841. PMLR, 2021.
Léo Grinsztajn, Edouard Oyallon, and Gaël Varoquaux. Why do tree-based models still outper- form deep learning on typical tabular data? In Advances in Neural Information Processing Sys- tems (NeurIPS) Datasets and Benchmarks, 2022. URL
08815
Jérôme Hénin, Tony Lelièvre, Michael R. Shirts, Omar Valsson, and Lucie Delemotte. Enhanced sampling methods for molecular dynamics simulations. arXiv preprint arXiv:2202.04164 2022.
David Holzmüller, Léo Grinsztajn, and Ingo Steinwart. Better by default: Strong pre-tuned mips and boosted trees tabular data. In Advances in Neural Information Processing Systems (NeurIPS), 2024. URL httphttss://proceeiiiiiii ngs htwp://eeeeeinn.....nee...eeeeee http://////6622222222///a/333332
Jiri Hron, Yasaman Bahri, Jascha Dick-Dicks,,in, and Roman Novak. Infinite attention: Nugp and ntk for deep attention networks. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Muchine Learning Research, pp. 4376-4386 PMLR, /viiss/vvii 19/hr1120aa. html
2020. URL
Ruyun Hu, Lihao Fu, Yongcan Chen, Junyu Chen, Yu Qiao, and Tong Si. Protein engineering via bayesian optimization-guidedd evolutionary algorithm and robotic experiments. Briefings in Bioinformatics, 24(1):bbac570, 2022. doi: 0..19//bib/bbbc570.
Andrew Kirjner, Jason Yim, Raman Samusevich, Shahar Bracha, Tommi Jaakkola, Regina Barzilay, and Ila Fiete. Improving protein optimization with smoothed fitness land- scapes. In International Conference on Learning Representations (ICLR), 2024. URL 22022224922000077777776555066220 htppe//aalla/pp///////aaaaaaa html See also
ngs .ic.gi..ccliiio icc..
arXiv:2307.00994
Muheng Li and Ruqi Zhang. Reheated gradient-basse discrete sampling for combinatorial optimiza- tion. arXiv preprini 20ti:25033000447 2025.
Duncan McElfresh, Sujay Khandagale, Jonathan Valverde, C. Vishak Prasad, Benjamin Feuer, Chinmay Hegde, Ganesh Ramakrishnan, Micah Goldblum, and Colin White. When do neural
nets outperform boosted trees on tabular data? In Advances in Neural Information Processing
Systems (NeurlPS), 2023. /tppattppp///rrccednggs...o http:///so.....ooor///oooonnnnnn htt.://eeee//eeeeeeee44000000000 05200000855111100000E40000006333
hwwp:aaaappp__ files/
and Sw___Bennhmarrs...
12

---

Published as a conference paper at ICLR 2026
Robert J. McEliece, Edward C. Posner, Eugene R. Rodemich, and Santosh S. Venkatesh. The capacity of the hopfield associative memory. IEEE Transactions on Information Theory, 33(4): 461-482, 1987. doi: 10.1
I. D. Mienye, T. G. Swart, G. Obaido, M. Jordan, and P. Ilono. Deep convolutional neural networks in medical image analysis: A review. Information, 16(3):195, 2025, doi: 10.3390/info10030) 66000195. URL https: www /0/B:/c/c////002228///////////
Michael Murray, Hui Jin, Benjamin Bowman, and Guido Montufar. Characterizing the spectrum of the ntk via power series expansion. arXiv preprint vrXiv:2211.07844, 2022.
Patrick Pynadath, Riddhiman Bhattacharya, Arun Hariharan, and Ruqi Zhang. Gradient- dis- crete sampling with automatic cyclical scheduling. Advances in Neural Information Processing Systems, 37:46728 46763, 2024.
Aniketh Janardhan Reddy, Xinyang Geng, Michael H. Herschl, Sathvik Kolli, Avi- ral Kumar, Patrick D. Hsu, Sergey Levine, and Nilah M. Ioannidis. Design- ing cell-type-specffic promoter sequences using conservative model-based optimiza- tion. In Advances in Neural Information Processing Systems (NeurIPS), 2024. httpa///epeeeeeeee///e//////////
URL a961 wot0000oobaaaaaff33333do
https://papers:
Zhizhou Ren, Jiahan Li, Fan Ding, Yuan Zhou, Jianzhu Ma, and Jian Peng. Proximal exploration for model-guided protein sequence design. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 18520-18536. PMLR, 2022.
pttttps::/proocediigg. www.nggsmm......... pre.trres//ee 62/ http//en22a.. html
Benjamin Rhodes and Michael U. Gutmann. Enhanced gradient-based meme in discrete spaces. arXiv preprint ar0.v:2208. 00040,, 2022. URLptttps:/arxxiv.orr
2208 00040
Yifei Shen, Xinyang Jiang. Yifan Yang. Yezhen Wang, Dongqi Han, and Dongsheng Li. Understand- ing and improving training-free free loss-based diffusion guidance. Advances in Neural Information Processing Systems, 37:108974 109002, 2024.
R. Singh, M. Bapna, A. R. Diab, E. S. Ruiz, and W. Lotter. How ai is used in fda-authorized medical devices: taxonomy across 1,016 authorizations. npj Digital Medicine, 8:388, 2025. doi: 10. 17400444417--00-18800000000 URL hollhttps:/ /wwww/wwwww... nat ure cow/ hoam/articles//// s41746 025 01800 p.f
Amos Storkey and Romain Valabrègue. The basins of attraction of a new hopfield learning rule. Neural Networks, 12(6):869-876, 1999. doi: 10.1066/10899-660099)000333
Haoran Sun, Hanjun Dai, and Dale Schuurmans. Optimal scaling for locally balanced proposals in discrete spaces. arXiv preprint 28X11:2209.08/883, 2022. URL.
2209.05 08183
Haoran Sun, Hanjun Dai, Bo Dai, Haomin Zhou, and Dale Schuurmans. Discrete langevin samplers via wasserstein gradient flow. In Proceedings of AISTATS, volume 206 of PMLR, pp. 6290-6313, 2023. UuRLLattss//proceeeee
mlr pwtss/p20e/sss2222222222
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. In International Conference Machine Learn- ing. pp. 21658-21676. PMLR, 2022.
Sarthak Verma, Raphael Boonsanong, Minh Hoang, Keegan E. Hines, John P. Dickerson, and Chirag Shah. Counterfactual explanations and algorithmic recourses for machine learning: A review. ACM Computing Surveys, 2024. doi: 10.1 10.1145/3677111
Binxu Wang. An analytical theory of power law spectral bias in the learning dynamics of diffusion models. arXiv preprint arXiv:2503.03206. 2025.
Binxu Wang and John J Vastola. The unreasonable effectiveness of gaussian score approximation for diffusion models and its applications. arXiv preprint arXiv:2412.097766 2024.
13

---

Published as a conference paper at ICLR 2026
Stefan Webb, Tom Rainforth, Yee Whye Teh, and M. Pawan Kumar. A statistical approach to assessing neural network robustness. In International Conference on Learning Representations (ICLR), 2019. doi: (0448500aaXii.11111.7722 URL http:s:/aarxv..rrt wttp://g/aa....// 07209
Julia K. Winkler, Christine Fink, Ferdinand Toberer, Alexander Enk, Holger A. Haenssle, et al. As- sociation between surgical skin markings in dermoscopic images and diagnostic performance of deep learning convolutional neural network for melanoma recognition. JAMA Derma- tology, 155(10):
2019. doi: 10. 00010000/ammdemma00000099...... URL https: /7 j/janntworrk com journa 1
Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:31 1172-31143333 2023.
Kevin K. Yang, Alex Lu, and Nicolò Fusi. Convolutions are competitive with transformers for protein sequence pretraining. Cell Systems, 2024. doi: 10.1016/j 2024. 008.. URL https: (24) 00029-2
http:/temmnn_1IIeeeee222000---2 www..cclll. com/ watpm//eellssyseeeeeeee www..pee....oo.eeeeesssssstteeee
Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Minshuo Chen, and Mengdi Wang. Reward- directed conditional diffusion: Provable distribution estimation and reward improve- ment. Advances in Neural Information Processing Systems (NeurIPS), 2023. URL
bttps::/proceesiie www.://s.....rir....e..... cwtparpppppscpppppr.eeeeeeeep be orc33336666558300044433300706060 Sepp://oce-ppppp-e-ococeeeeeeeeo
Giacomo Zanella. Informed proposals for local meme in discrete spaces.
arXiv preprint
arXiv:1711.07424. 2017.
07424
Ruqi Zhang, Xingchao Liu, and Qiang Liu. A langevin-like sampler for discrete distributions. In International Conference Machine Learning, pp. 26375-26396. PMLR, 2022.
Hao Zhu, Daniel M Steinberg, and Piotr Koniusz. Protein fitness landscape: Spectral graph the- ory perspective. In The 28th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 1, 2025.
14

---

Published as a conference paper at ICLR 2026
A Ood d) MIXING FROM CHECKPOINT
Setting.
Let 1) and let
(()
Σ i=1
be sum of degree-1 monomials (linear function) on the Boolean hypercube. We want to sample from the Gibbs measure proportional to exp(f):
1
exp(s(x)) (x) Z
Z II exp()
II (x()), (=1
so is product distribution with one-dimensional marginals ny (2:4) x exp(haa;). Consider random-scan single-site Gibbs: at each step pick I, ~ Unif([d) and resample X1(4) from T (x)4(()))), which for this product target equals the marginal xt.
Claim. For random-scan single-site Gibbs product target,
(=) d(log d + (1/=))
In particular, tmix (1/4) S d(log + log4) = Oddlog d).
Proof.
Let the refresh time
Tee min(t: each coordinate [d] has been selected at least once by time t).
Because is product, whenever coordinate i is selected resample it fresh from m, indepen- dently of everything else. Hence, at time Tier we have resampled every coordinate from its marginal, so
n
n
This makes Tier strong stationary time, which implies
(x,)) π| T||TV I Pr(Tier> 2) for all starting states and times 2 0.
remains to bound the tail Tre- Each step picks a coordinate uniformly from [d], so this the coupon-collector process. For any fixed i,
Pr(i was never chosen in steps) (1 (1-1/d)² s -/dd
A union bound over the d coordinates then gives
Pr(Tui t) Pr(3i i unrefreshed) < de-t/d
(3)
Combining equation and equation 3 and choosing so that t/d < E yields t/d
d(log d log(1/)))
which proves the claim.
B Od d) MIXING FROM 72 CHECKPOINT
At checkpoint 72, let aligned spins y/: 2422 The degree- 2 surrogate be written
f R Sse(v) 2-d
Σ y4
yell
ISigjed
which is the Curie-Weiss Hamiltonian with a positive external field up to scaling.
15

---

Published as a conference paper at ICLR 2026
In the low-temperature regime, existing results show that censored Gibbs dynamics just the degree-2 monomials (ignoring the degree monomials) mixes O(d d) time (Ding et al, 2009). By censoring, mean if proposed update would make the alignment with the pattern negative, we reflect all of the update variables.
Without censoring and with uniform start low temperature and zero field, the chain falls into the + or 7/ basin with probability = 2 each; a positive field (as is the case in our setting) biases toward the +// basin.
After scaling with 3 = 2d, [=) is at least 2d higher than all the other configurations. Thus, under the measure exp 2 fv)(),, the target x* is at least edd more probable than all other configurations. Since there are 2 configurations in total, the lower bound for =*'s likelihood is 24-1+c- 1
edd
e-((d) Consequently, after Ogdlog d) steps the chain is at 2" with high probability (so a constant number parallel chains suffices).
C EMPIRICAL EVIDENCE OF LOWER DEGREES ALIGNING//ROWWIN BEFORE HIGHER DEGREES
Waush L2 mass by degree during training (10-D indicator) For 2-Layer FCNN of width 256
0.6
19999
des 9.9....0 189
des 409
deg
2.9
deg
959
999
deg
0.5
deg
deg
deg
Fourier mass Σ its)?
deg
deg
deg
deg 10
0.1
0.0
0
100
200
300
400
500
Epoch
16

---

Published as a conference paper at ICLR 2026
Fourier-Wale mass by degree during training (10-D indicator) For 3-Layer FCNN of width 2.5
sp
deg
deg
4eg
deg
deg
2.0
deg
its)
deg
15-
deg
Fourier mass
10
1.0
0
50
100
150
200 Epoch
250
300
350
400
Fourier-Wale L2 mass by degree during training (10-D indicator) For 3-Layer FCNN of width 512
30
deg
1
deg.
deg
deg
deg 4
deg
deg
Fourier mass I itsp
deg
deg
deg
deg 10
1.0
0.0
50
100
150
200 Epoch
250
300
350
17

---

Published as a conference paper at ICLR 2026
Fourier-Walsh L2 mass degree during training (10-D indicator) For 3-Layer FCNN of width 512 and activation Silu
2.5
.e0 1.
8922399 28
9
deg 1
deg
deg
deg 4
2.0
deg
deg
deg
Fourier mass E
deg 8
deg
deg
1.0
0.5
0
50
100
150
200
250
300
350
400
Epoch
Fourier-Walsh L2 mass by degree during training (10-D indicator) For 4-Layer FCNN of width 128 2.5
0932922
30
deg
deg
deg
2.0
deg
deg
deg
Fourier mass Σ itsy?
deg
15
deg
deg
deg 10
10-
0.5
20
40
60
80
100
Epoch
18

---

Published as a conference paper at ICLR 2026
Fourier-Waid L2 mass by degree during training (10-D indicator) For 4-Layer FCNN of width
2.5
1.2.9
999 9..92.
deg
deg
deg
2.0
deg
Fourier mass E its)?
deg
deg
10
20
40
60
80
100
Epoch
Fourier-Wake L2 mass by degree during training (10-D indicator) For 5-Layer FCNN of width 128
19399
deg
deg
deg
deg 4
deg
deg
Fourier mass I itsp
deg
deg
deg
deg 10
0.0
20
100
Epoch
19

---

Published as a conference paper at ICLR 2026
Fourier-Was -aalh L2 mass by degree during training (10-D indicator) For 5-Layer FCNN of width
deg
689 4.455
deg
14.
deg
1.2
deg
Fourier mass Σ itsy²
10-
deg
deg
0.8
10
0.6
0,4
0
10
15
20
25
30
Epoch
Walsh L2 mass by degree during training (10-D indicator) For 6-Layer FCNN of width 256
deg
2.5
deg
dl
deg
deg 4
2.0
deg
deg
Fourier mass Σ itsp
deg
deg
deg
deg 10
1.0
0.0
Epoch
20

---

Published as a conference paper at ICLR 2026
Fourier-Walsh L2 mass by degree during training (10-D indicator) 3-Layer FCNN (width 256)
Train 0/1, FWHT on
8 1.
9.28
9 4.440
869 9.9. 9,150
iis? on the (+11) domain
deg 2
699
2.0
deg
deg
deg
deg
Fourier mass Σ
50
100
150
200
250
300
Epoch
Fourier-Walsh L2 mass by degree durining training (5-D indicator)
For 3-Layer CNN with 128 channels and input length (coeffs are averaged over all background and indicator positions)
FWL2 mass from averaged coefts I risy² (EMA)
2.1.0.1 10
deg
deg
deg
0.5
deg
deg
0.4
0.3
0.1
0.0
1000
Epoch
D MULTIPLE NEEDLES
With multiple needles, the degree-2 interaction is equivalent to a Hopfield model with weights set by a Hebbian learning rule; increasing the number of stored needles shrinks attraction basins, poten- tially slowing local MCMC (McEliece et al., 1987 Storkey & Valabrègue 1999). For the standard pairwise Hopfield model with unbiased needles. the capacity scales linearly as Nmax se a.d with Oe = 0.138 (Amit et al. 1985). For pure p-spin Hopfield models, the number of storable patterns scales as Nmax R6 0
with Otp a p-dependent constant (Bovier Niederhauser 2001).
Table shows that across multiple needles, sampling along the training trajectory is still more sample-efficient than running Gibbs-with with-gradiemt the final trajectory.
21

---

Published as a conference paper at ICLR 2026
Table Needle gadget: hit rates across 5 for GWG vs, our method (1 needle hit). Task: 10-D indicator "needle" with 10 additional linear terms; model: 3-layer FCNN (width 128); sampling: particles, 60 total steps per run across checkpoints (5, 25, 50, 75, 100, final] (10 steps per check- point); baseline: 60 steps of GWG on final checkpoint
# Needles
GWG: runs with 21 hit (out of 5)
Ours: runs with 21 hit (out of 5)
5
0/5
5/5
2/5
5/5
3
1/5
5/5
2
0/5
5/5
0/5
5/5
E EMPIRICAL EVIDENCE OF SAMPLING FROM INDICATOR FUNCTION PROJECTED TO TERMS OF ORDER < P
The below graphs contain the median number of steps needed to hit the target when sampling from indicator function that is only non-zero for the target. However, its boolean expan- sion is projected down to terms of degree < P. We consistently observe that the median num- ber of steps increases as P increases. For each value of P, select the best & value across 1,000.000007,0005,0003,,,3
conduct 300 trials, and cap the number of steps at 2d. Sampling done via Gibbs (random-index heat bath).
Best median steps per P (min over B), random-index heat-bath for indicator function of dimension
d=7
Best median steps (min over B gri)
80
60
20
[max degree included)
22

---

Published as a conference paper at ICLR 2026
Best median steps per P (min over B), random-index heat-bath for indicator function dimension
d=88
Best median steps (min over B grid)
175
150
125
75
50
25
[max degree included)
Best median steps per P (min over B), random-index heat-bath for indicator function dimension
d=9
400
Dest median steps (min over B grid)
350
300
230
200
150
100
50
[max degree included)
Best median steps per (min over B), random-index heat-bath for indicator function of dimension 10
d=10
700
Best median steps (min over B grid)
400
300
0
(max degree included)
23

---

Published as a conference paper at ICLR 2026
F NTK-ALIGNED TRAINING MAP
Setup. Let the input domain be the unit sphere $ with the wniform measure. Expand any square-integrable f gd-1 in spherical harmonics Yk.i (degree k E {0,1,.... multiplicity index j):
N₂ Oxj (2),
f(z) -
Cn.j = (f,
k=
For isotropic (zonal) operators the sphere, each degree-k subspace is eigenspace, all coef- ficients {ak. 3=1
;=1 evolve by the same scalar multiplier.
Idealized NTK dynamies. Consider fully connected network in the linearized NTK regime (in- finite width), trained by kernel gradient flow with squared loss and learning rate 77, from zero ini- tialization, data drawn uniformly from gd-1, The analytical NTK K is zonal kernel whose eigenfunctions are the spherical harmonics and whose degree-k eigenvalue we denote by A > 0. Along each degree-k subspace the (prediction) coefficient obeys
d
(mx)(() akks)) -1/ AE (a((aa(()) - okk,), dt
so from (k,(()) we get
akj(1) E -datt) ales
Equivalently, time the entire degree-k block is scaled by
M
M (t) - Aut E [0, 1].
This corresponds to Eq. (7) in Bowman (003)); see that reference for fuller introduction.
x Using 1+z S
for x 2 with x = spl. As. we obtain
nl A
I M S nt AE.
1 + ntt A
Hence for large k (so A + 0),
(F) rtt AR=
i.e., MET x d up to constants depending 1pl.
Activation choice controls the spectrum { and thus the decay of M, NTK across degrees: for ReLU, A = @k (polynomial "spectral bias"); for Tanh, A = 0(k vT (super-
polynomial). For fixed f and large k, .MMTK (t) ss vit du. so high degrees are damped more Murray al. 2022).
Comparison: Gaussian (heat) smoothing. block by
Heat-kernel smoothing on Sd-1 multiplies the degree-
M (t) = exp(-t
#s = k d d 2),
i... an exponential-in--2 decay (stronger high-frequency suppression). Note the time contrast: larger diffusion time 1 means more smoothing, whereas larger NTK training time t means M NNKK(() (t) t and less smoothing (the predictor approaches f*).
Takeaway. Under the NTK idealization, the training trajectory {fe} is family of degree- wise smoothed versions of f*, with the spherical harmonics as eigenfunctions and activation-controll frequency decay. Diffusion performs similar degree-wise smoothing but with heat-kernel multi- pliers.
G BINARY MNIST CHECKPOINT ABLATIONS
We report FID scores for varying numbers of checkpoints used in our method for a fixed budget of 1K total GWG steps. The model was trained for 50,000 epochs. For a given number of checkpoints,
24

---

Published as a conference paper at ICLR 2026
choose them to be evenly spaced along the training trajectory and allocate equal number of sampling steps to each checkpoint. The Temp-GWG baseline samples only from the final checkpoint using temperature annealing.
Table 10: FID (4) on binary MNIST as a function of the number of checkpoints used by our method, with fixed budget of 1K GWG sampling steps. Entries are mean (std) over 10 bootstrap resamples.
# Checkpoints
Mean FID (std)
Temp-GWG (baseline)
29.61 (0.239)
5
16.10 (0.237)
10
14.42 (0.322)
25
12.56 (0.289)
50
11.93 (0.355)
100
12.08 (0.435)
500
11.73 (0.284)
As shown in Table we dramatic gains from using even 5 checkpoints, with additional check- points yielding diminishing marginal returns and performance saturating around 50-500 check- points. The key empirical observation is that the training trajectory tends to evolve from coarse to fine, so any set of evenly spaced epochs can leverage this structure to speed up sampling.
25

---

Published as a conference paper at ICLR 2026
H
SAMPLING FROM BINARY MNIST
Figure 17: First 49 random samples from standard sampling with 1K steps.
$
m th 
038868
068835
8 ****""""
8
#
26

---

Published as a conference paper at ICLR 2026
Figure 18: First 49 random samples from our method with IK steps. These samples are substantially sharper than the above.
S 65
55
2 2877388
 * 3
1713388
8
3 3848208
53//898 8
27

---

Published as a conference paper at ICLR 2026
Figure 19: First 49 random samples from standard sampling with 10K steps.
885568
6
D
1388968
5
1 44
+ #
3 8
8333888
8 3 g
5
28

---

Published as a conference paper at ICLR 2026
Figure 20: First 49 random samples from our method with 10K steps. Again, these samples are sharper.
0
0133772
2
4
Z
6266556
8
5442733
8
875/486
# 4340819
I DISCRETE EXPERIMENTAL DETAILS AND COMPUTE BUDGETS
I.1 CHECKPOINT AND TEMPERATURE SELECTION
For all experiments, we run an equal number of MCMC steps at each stage. We use checkpoints at epoch 25 and temperature 10.0 for most experiments with a few exceptions detailed below.
For the DNA design experiment, we add an additional checkpoint at epoch 50 due to the additional complexity of the task.
For the discrete high-frequency experiment, we use temperature = 100.0 (instead of the 10.0 used in all other settings) because we are sampling from lower-magnitude component.
For the MNIST sampling experiment, we save the model every 100 epochs during training and use these as the checkpoints. We use the default temperature hyperparameters used in the repo; we make no changes.
I.2 KERNEL
All experiments and baselines use the Gibbs-with wire--wwth--hrrre Gradients (GWG) kernel introducee in Grathwohl proposals for categorical flips under a1 encoding, drawing moves from a softmax over approximate first-order energy changes. Each proposal is then corrected by Metropolis-Han Hastings step that accounts for both the approximate energy difference
al. (2021). This kernel uses
29

---

Published as a conference paper at ICLR 2026
and the asymmetric forward/reverse proposal probabilities. The choice of temperature for the kernel is detailed in the subsection above.
For the synthetic experiments
the exact GWG kernel where the y values are the exact y val- ues, not those estimated via gradients. This is to remove the effect of gradient estimation noise. For the subsequent experiments, we use the standard GWG kernel where the y values is approximated from the gradient.
For the MNIST sampling experiment, MH); we make no changes.
the kernel implementation from the repo (also GWG-
For the constrained DNA sampling task, we modify the GWG-MH kernel so that incorporates the Hamming distance constraint. a sampled single-flip would exceed the Hamming cap, we pair with gradient-guided reversion on some already-fipped position the net move stays on the Hamming boundary. Both legs are sampled from the same GWG softmax over their respective candidate sets.
1.3 CHECKPOINT POLICY AND BASELINES
1.3.1 SYNTHETIC
We run 20 steps of our kernel epoch 25, and then another 20 steps at the final epoch. For the baseline, we run 2000 steps of the same kernel at only the final checkpoint. From random starting particles, we measure the fraction of particles that hit the global maxima of the function along their trajectory.
For the synthetic experiments, we only benchmark against GWG. This is because for the indicator function, other methods only manipulate the temperature of the final checkpoint. However, rescaling the temperature does not help- the landscape still a random-walk,
For the High-magnitude high-frequency variation experiment, we additionally benchmark against temperature annealing with linear schedule. Here, the temperature was annealed from 0 to 100 (the temperature value used for this experiment) across 2000 steps, where one MCMC step was taken at each temperature value.
1.3.2 MNIST SAMPLING
We utilize the code, as-is, in GWG repo This trained for 50K epochs, and we save the checkpoint every 100 steps. For sampling with IK steps, we run 2 steps of the kernel at each checkpoint. For sampling with 10K steps, we run 20 steps of the kernel at each checkpoint. We do not adjust any of the kernel hyperparameters for our task we only change the checkpoints. For the baseline, we use the existing sampling code in the repo, which samples from the final checkpoint using linear temperature annealing and takes one kernel step for each temperature value. We control for the number kernel steps across comparisons.
For the FID calculations, pytorch/tree/mmster/ We also use their MNIST LeNet model checkpoint for computing the metrics. We collect 8K random samples from the binary MNIST test set, our sampling method, and the existing sampling method. For the ground-truth FID value, we calculate FID between two random 8K subsets of the binary MNIST test set.
use the rерo https//gituub.com/bbbuuuu mittr
I.3.3 TF-DNA SAMPLING
We start from random length-60 DNA sequence, and run 60 steps of the MCMC kerne.. All methods only use 60 steps. Other than plain GWG, all of the methods use 3 stages with an equal number of steps for each stage.
For our method, we run 20 steps at epoch 25, 20 steps at epoch 50, and 20 steps at the final check- point.
For the GWG baseline, we run 60 steps at only the final checkpoint.
30

---

Published as a conference paper at ICLR 2026
For the parallel tempering baseline, run 3 replicas over geometric 3-ladder from 0.05++10.0, doing 20 local steps per replica with swaps every steps (60 total local updates across replicas). For the replica starting point, the cold chain starts zg. while the mid and hot chains randomly reassign 25% and 50% of positions, respectively, to diversify exploration from the outset.
For annealed importance sampling with automated temperature adjustment, tance sampling with an ESS-targeted adaptive temperature schedule (5 0 > 10) without resam- pling, preserving one-to-one lineages; each stage re-weights particles and picks the next 8 to keep ESS near 0.6 N. After reweighting, it performs K=20 GWG rejuvenation sweeps at the new 3 with stages, this is 60 sweeps per particle.
annealed impor-
For the constrained sampling task, we start from a random DNA sequence and want the best sample within a hamming distance of 7. We use the same setup before except with a modified GWG-MH kernel detailed in Sec. I22
Diversity calculated as the median of pairwise (Hamming) distances within each bootstrap re- sample. Novelty is calculated as median of per-seq min distance to training, within each bootstrap resample.
1.4 ARCHITECTURE AND TRAINING DETAILS
For the synthetic experiment, use 3-layer FCNN with hidden dim 128. It is trained for 30000 epochs with learning rate of 2e-3. Indicator functions have value 10.
For the MNIST experiment, architecture is an EBM with a 3x3 stem conv, two strided residual downsampling blocks (each: Swish - 3x3 conv - 3x3 cony 1x1 projection shortcut), followed by six identity residual blocks (Swish- + 3x3 + 3x3, no projection), then global spatial averaging and a single linear head a scalar energy. In total, has 19 conv layers and fully connected layer. Training is done according to the repo's 's instructions.
use ResNet-EBM architecture specified in the GWG repo The
For the TF-DNA experiment, we use the same architecture from de Almeida et al. (2022). Specif- ically, we apply a Conv 1d4-66, kemnel=11), ReLU, then global max pooling over the sequence dimension for each filter, and then a linear head that outputs a scalar.
1.5 DATASETS
For the synthetic experiments, have 8-10 main variables that are part of the function and 500 spurious variables. We construct synthetic datasets where the spurious variables are randomly sam- pled.
We use the binary MNIST datasets provided GWG repo
For the TF-DNA experiment, the data consists of length-60 DNA sequences. The last 24 DNA letters are always the same; the first 36 are close to random. The TF (MAX, from the mouse species) binds strongly when the motif" "CACGTG" is present. The binding strength increases depending on where the motif is present (upstream leads to stronger binding), the flanking sequences surrounding the motif, the number of times the motif is present, the GC % in the sequence, etc. The dataset taken from Badis et al. 2009).
1.6 CI DETAILS
All reported CIs are 2SD, unless specified otherwise.
For the synthetic experiments, we run the above test on 200 random particles, calculate the hit fraction (whether a particle reaches the global maxima along its trajectory),and report 2SD CIs from these results.
For the MNIST experiment, FID is calculated across sets of 8K samples. The standard deviation is calculated over 10 bootstraps.
For the DNA design experiment, we sample with 300 particles, and calculate 95% bootstrap per- centiles (B=500) from the results.
31

---

Published as a conference paper at ICLR 2026
J SAMPLING ALONG THE TRAINING PATH WITH SMC
Let (2) x (x)()((x)) denote the (unnormalized) target associated with the checkpoint at time (e.g., from the model's energy surrogate negative log-likelihood).
Rather than selecting a single smoothing level, we sample along the NTK training trajectory / E [0, 7] using Sequential Monte Carlo (SMC):
1. Choose schedule 0 to < 4
geometrie).
2. Initialize particles from an easy reference.
3. For l = 1,..., L: compute incremental weights w(x) & nea ( T (x), resample, and apply short MCMC move targeting tt.-
.....,,,,
4. Output particles at t = ' (the desired final target).
This procedure exploits the frequency-selective filtering (L) to traverse from smooth dominated intermediate distribution toward the final target while maintaining particle di- versity.
K
CONTINUOUS EXPERIMENTAL DETAILS AND COMPUTE BUDGETS
K.
SHARED SETTINGS (ALL EXPERIMENTS)
Parallel trajectories. All methods use concurrent trajectories.
Compute parity. Within each experiment, every trajectory performs the same total number of Metropolis-Adjustee Langevin Algorithm (MALA) steps across methods.
Kernels. Ackley and Superconductor use MALA for all five methods.
Checkpoint policy (SMC-Train). Train for 10,000 epochs; checkpoint every 10 epochs. Smooth the training-loss curve and keep the earliest prefix of checkpoints up to (but not beyond) the plateau; exclude later flat checkpoints (each kept checkpoint has strictly lower loss than the previous one).
SMC-Temp schedule. Linear inverse-temperature ladder with the same number of distri- butions SMC-Train for that task.
AIS schedule. Annealed Importance Sampling (AIS) chooses temperatures adaptively each stage to maintain a target conditional effective sample size (cESS); rejuvenation uses the same MALA kernel other methods.
PT schedule. Parallel Tempering (PT) with a fixed temperature ladder across replicas; propose swaps between adjacent replicas every stage; within-replica moves are MALA with the same per-step budget.
K.2 EXPERIMENT-SPECIIIC PARAMETERS
Ackley (10D). Proposal: MALA with step size 10 2, adaptation target acceptance 0.57 (adap- tation off unless stated). SMC-Temp: resample when ESS 0.5N. SMC-Train: resample when ESS < 0.5N. ATS-Aulo: choose temperatures to hit cESS = 0.5N increment (bisection tol. 10-4, max iters). PT-MALA: power temperature ladder (parameter 4.0); per-replica MALA step scales /3100 MCMC-Final: plain MALA with the same step size; burn-in, thinning.
Superconductor. Inputs RT Proposal: MALA with base step size 0.05/ Vd; adap- tation target acceptance 0.57 (off by default). Stabilization in latent 2: per-dimension percentiles [1%, 99%] and radial cap 99.5%. SMC-Temp AlS-Auto: resample choose temperatures to maintain cESS = 0.5N; rejuvenation uses the same MALA step. SMC-Train: same cESS rule; default initialization from Gaussian prior over PT-MALA: geometric temperature ladder up to Smax 1.0, swaps every stage; replicas chosen to evenly factor the parallel budget; report adjacent- swap rates and per-3 MALA acceptance. MCMC-Final: adaptation enabled: target acceptance 0.57 with updates every 10 steps (clip to 10-2, 0.5)).
32

---

Published as a conference paper at ICLR 2026
K.3 COMPUTE BUDGETS
Table 11: Per-trajectory budgets. is the number of intermediate distributions (temperatures for AMCSTI checkpoints for SMC-Train). Total Steps L. K for SMC/AIS/PT and
meme for MCMC.
Experiment
Method
Parallel
Dists
Rejuv./Dist,
Total Steps
Ackley (10D)
MCMC-Fine
N=10,000
50
SMC-Temp
N=10,000
50
SMC-Train
N=10,000
10
5
50
AIS-Auto
N=10,000
10
5
50
PT-MALA
N=10,000
10
5
50
Superconductor
MCMC-Final
-5500
250
SMC-Temp
N=500
50
5
250
SMC-Train
=5500
250
AIS Auto
N-500
50
5
250
PT MALA
-
50
5
250
K.4 DATA MODELS BY TASK
Ackley (10D). Model: MLP with layers [1024, 512, 256], Tanh activations. Training data: synthetic coverage over [-10, 10j with three components: (i) uniform "plateau," (ii) strat- ified radial shells spanning target f-levels, and (iii) small ball near the origin for additional f=0 mass; 3,601,000 total points. Target: regress (z)) Sampling kernels: MALA for all five methods.
Model: MLP with layers [2048, 2048], ReLU activations. Training data: train on the full available dataset; no oracle fine-tuning. Target: regress the provided score f(x) (higher is better). Sampling kernels: MALA for all five methods.
Superconductor (Design-Bench), Dimensions: d=87.
Novelty & Diversity (Superconductorr) Distances are computed standardized feature space using the input scaler fit the training set. Novelty is the per-sample & (Manhattan) distance the nearest training point (scikit-learn Nearest ghbors, Neghoorrs,, met met seetriiimmaatttaa median and IQR across samples. Diversity is the median and IQR of pairwise l distances among generated samples, computed over all unordered pairs (n(n 1)/2 for n samples). All summaries are reported as median [IQR].
we report the
K.5 EXTENDED RESULTS
Ackley (10D, 4). Under matched compute, SMC-Train achieves the best and best-of-set with non-overlapping CIs relative to all baselines (Table 8). Quantitatively, SMC-Train reduces the mean objective vs. MCMC-Final by 17.8% (16.22 + 13.33), vs. SMC-Temp by 18.3% (16.31 13.33), Vs. AIS by 18.3% % (16.31 + 13.33), and vs. PT by 32.9% % (19.87 - 13.33). On best-of-set, SMC-Train improves
3.69). over SMC-Temp by 53.0% (7.86 + 3.69), over AIS by 58.1 % (8.81 + 3.69). and over PT by 72.9% (13.62 + 3.69). The 95% Cls for SMC-Train (13.33 [2.09, 14.58) are disjoint from the tight ranges of the other methods (zz (6.16-16.35) indicating consistent improvement across seeds.
MCMC-Fin
by 56.9% (8.56
Superconductor (*) In
materials design, SMC-Train leads both on best-of- set and mean (Table 8 Mean reward increases by 102.4% vs. MCMC-Final (76.68 + 155.2), by 606.1% % vs. SMC-Temp (21.98 - 155.2), by 535.5% vs. AIS (24.42 - 155.2), and by 497.6% vs. PT (25.97 - 155.2). Relative to the reference score (185.0), the mean reaches 83.9% of the target, while the best-of-ser (318.4) is 172.1 % the reference (i.e., +72.1% % over target). These gains come with wider uncertainty for SMC-Train (95% CI: [105.6, 204.8)), reflecting more aggressive exploration that can land very high-reward candidates.
33

---

Published as a conference paper at ICLR 2026
On novelty and diversity. Table 8 shows that AIS and PT achieve the highest novelty (median ~35.3-35.7) and diversity (median ~25,4-26.
26.1). 1). while SMC-Train moderate on these axes (nov- elty 20.86; diversity 16.60). However, these higher exploration metrics do not translate into better objective quality: both AIS and PT have substantially lower mean rewards (24-26) than SMC- Train (155.2). We observe that the methods with the highest novelty/diversity also retain many low-quality samples, inflating dispersion-based metrics without improving the objective. In con- trast, SMC-Train balances exploration and exploitation: it traverses the space broadly enough to discover strong candidates (best-of-set set 318.4) while concentrating mass to raise the mean reward. Thus, more novelty/dierrsity does not necessarily imply better design quality when significant tail of poor samples is preserved.
Takeaways. (i) Under matched compute, SMC-Train consistently outperforms baselines on Ack- ley and Superconductor by large margins in both mean and best-of-set. (ii) For Superconductor, ap- parent exploration advantages (higher novelty/diversity) from AIS/PT coincide with lower objective quality-Suggesting these methods over-emphasize exploration and retain weak samples. (iii) Re- porting both mean and best-of-set, alongside noveltyddiversity. is essential: together, they show that SMC-Train drives objective gains while maintaining reasonable exploration, rather than chasing dispersion alone.
34

---

Published as a conference paper at ICLR 2026
K.6 FIGURES
K.6.1 ACKLEY
Adkley
MCMC
SMC. TEMP SNC_CKPT
2000
1500
10
14
16
20
Eneryy
Figure 21:
Full energy histogram for Ackley (10D). Samples from MCMC-Final SMC-Temp, and SMC-Train.
K.6.2 SUPERCONDUCTOR
Diveriity across predited TC percenties
20-00k
R0-47%
46.coM
10.0%
68-10k
m5.00
96.100k
Figure 22: Samples are partitioned into Tc deciles (0-10%, 90-100%). Within each bin, we plot the distribution of nearest-neighbor LI distances (in standardized feature space) among the samples in that bin. Broad-and www-sshiniin--wthheee that sample diversity does not collapse as Tc increases.
LI distributions at higher Te percentiles indicate
35

---

Published as a conference paper at ICLR 2026
to training set vs predicted Tc
Min distance to training (standardized space)
30
50
75
100
125
150
175
200
225
Predicted Tc
Figure 23: Each point shows a sample's predicted Tc (x-axis) versus its minimum L1 distance to any training example in standardized space (y-axis). High- Tc proposals do not systematically move closer to the training set; many top-Te samples remain well separated, indicating genuine novelty rather than simple memorization.
Predicted Tr distribution by method
300
250
predicted Tc
100
50
FCMa,Flnal
SMC_Temp
Tran
AIS
PT
Figure 24: Side-by-side violins of the predicted critical temperature (Tc) for the three sampling methods, aggregated over seeds. The plot highlights differences in central tendency and tail behavior across methods.
36

---

Published as a conference paper at ICLR 2026
Percentile summary by method
MCMC Final
SMC, Temp
SMC_Train
PT
Prodicted Tr. value
50
60
100
Percentile of predicted
Figure 25: For each method, we compute selected Te percentiles (50th-100th) per seed and report the mean + standard deviation across seeds. Our method consistently gives higher Te samples than the other two methods.
L USAGE OF LLMS
We utilize LLMS assist with the writing of the paper. We provided GPT-5 outline of our key points for each paragraph, and GPT-5 converted them to paragraph format with latex formatting. We also utilized LLMs to research related work for each of 4 sections.
37


Reviews:
[{"review_id": "8EeLvlBSNn", "reviewer": "Reviewer_zyoC", "summary": "The paper proposes to use intermediate checkpoints from training a predictor model (used to define an energy function) to anneal MCMC from easy to sample distributions to the target distribution, which may have pathological structure that makes it difficult to directly sample from. They provide hierarchical learning analysis to validate the intuition, and provide some experiments to verify that the method works.", "questions": "- How does this method perform when using DLP sampler? Are the gains of using checkpoints / using just the final checkpoint just as large? \n- In the introduction, it is stated that it takes 15 hours to train a diffusion model on CIFAR-10. Is there a citation for this? Are there plots of the loss v.s training time of diffusion methods for the tasks considered?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "The idea of using the training checkpoints to perform some kind of annealing is novel to the best of my knowledge. I agree with their argument on how early stage checkpoints focus on high level features, late checkpoints focus on details, providing a sampling path. \n\nThe use of confidence intervals demonstrates thorough evaluation metrics. \n\nThe hierarchical learning analysis is also quite interesting.", "weaknesses": "**Literature Review**: The paper does not mention [2] or [3], which are focused on accessing modes that are difficult to reach. [3] is more recent so the omission is understandable. \n\n**Base sampler**: Gibbs with gradient changes one coordinate at a time (at least in the default version), which makes it very slow. Are the benefits of annealing via model checkpoints preserved when using samplers proposed in [1, 2, 3]? Or does performance improvement from using this annealing strategy saturate with stronger base samplers? While GWG was chosen due to simplicity, it is important to note that more recent methods may solve this problem without the need for annealing across checkpoints. \n\n**Breadth of Metrics**: It would be nice to include metrics that show sampling accuracy v.s number of sampling steps. For the MNIST experiment, it could take the form of log maximum mean divergence v.s sampling steps. Also, the paper does not include (Effective Sample Size), which is an important metric for evaluating the efficiency [1]. \n\n**Characterization of Diffusion**: They characterize diffusion as requiring training over the entire trajectory instead of single step MLE. However, I am not sure that this is a fair characterization: in practice, the score matching objective is trajectory free [4]. Even in discrete diffusion, the training objective takes the form of corrupting the input and then predicting the clean input [5, 6, 7, 8]. For obtaining SOTA results on CIFAR-10, it might take up to 15 hours. But it is entirely possible that within the first 15 minutes, the model is capable of generating reasonable images. \n\nThis is important because the intuition behind the proposed method is extremely similar to the core idea of diffusion: start from a distribution that is easy to sample, and gradually anneal it to the target distribution. The difference is that diffusion enables this with a single checkpoint, whereas the proposed method requires several checkpoints. Furthermore, diffusion directly supervises learning of the score (which is what GWG requires). \n\nThis is perhaps my largest concern with the submission: it seems to be capturing the intuition of diffusion, but via a more indirect path. If the focus of this paper is small models (as discussed in the introduction), is training a diffusion model really that expensive? And if the focus is on larger models where diffusion training is expensive, then it would also be expensive to store multiple checkpoints of the model and run analysis across all the checkpoints to determine which ones to use for annealing the sampler. \n\nAlso, I do not see a reference to [9], which directly incorporates the logic of diffusion into MCMC to improve mixing via the same intuition presented in this submission. While [9] is focused on the continuous domain, it may be worth considering how to extrapolate their method to the discrete space via gradient based discrete samplers. \n\n[1] A Langevin-like Sampler for Discrete Distributions. Zhang et al. ICML 2022. \n\n[2] Gradient-based Discrete Sampling with Automatic Cyclical Scheduling. Pynadath et al. NeurIPS 2024. \n\n[3] Reheated Gradient-based Discrete Sampling for Combinatorial Optimization. Li, Zhang. TMLR 2025. \n\n[4] Elucidating the Design Space of Diffusion-Based Generative Models. Karras et al. NeurIPS 2022. \n\n[5] Simplified and Generalized Masked Diffusion for Discrete Data. Shi et al. NeurIPS 2024. \n\n[6] Simple and Effective Masked Diffusion Language Models. Sahoo et al. NeurIPS 2024. \n\n[7] Simple Guidance Mechanisms for Discrete Diffusion Models. Schiff et al. Preprint 2024. \n\n[8] The Diffusion Duality. Sahoo et al. ICML 2025. \n\n[9] Diffusive Gibbs Sampling. Chen et al. ICML 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tw4ltnSqlb", "reviewer": "Reviewer_wkfn", "summary": "The paper proposes trajectory annealing: instead of sampling only from the final trained predictor $f^*$, the sampler runs short MCMC updates across saved training checkpoints $\\{f_t\\}$, exploiting the observed coarse‑to‑fine learning dynamics (early checkpoints damp high‑degree/high‑frequency components, later checkpoints restore detail).", "questions": "- With only a few steps at the final checkpoint, how biased are samples relative to the target distribution?\n- How do performance and compute vary with the number and placement of checkpoints (uniform vs geometric; early‑only vs full trajectory)? MNIST uses 500 checkpoints, but discrete synthetic uses just a few.\n- Have you tried saving checkpoints based on a performance metric instead of time? \n- The paper states diffusion “cannot” handle Hamming‑ball constraints (p. 8). Could you qualify this?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Leveraging the existing training trajectory as an annealing schedule is elegant and creative and requires no re‑training or auxiliary generative model. The coarse to fine picture is illustrated empirically and theoretically (Apps A-B).\n- The method works for discrete (GWG) and continuous (SMC) domains. \n- The Hamming‑ball constraint is handled naturally within the MCMC framework and yields large gains on DNA (Table 7).\n- The $O(d\\log d)$ mixing is theoretically supported on degree-1/2 surrogates.", "weaknesses": "- Compute parity (App J) is defined as matching the number of MALA steps, but SMC adds resampling/weight computations and multi‑checkpoint bookkeeping. MNIST also saves hundreds of checkpoints (50k training epochs saving every 100 steps). Reporting the wall‑clock time and memory would substantiate the “no additional compute” claim from the abstract.\n- For the discrete synthetic tasks the only baseline is GWG (authors state this choice explicitly), leaving out stronger informed/non‑local samplers (e.g., locally balanced and discrete‑Langevin families). This narrows the comparison.\n- The method does not apply to transformers (stated in the paper), which constrains scope for many modern applications.\n- Several evaluations measure best‑of‑run (DNA keeps the best of 60 steps per run), which is an optimization metric. Mixing diagnostic or distributional metrics would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "57Iyz34YaY", "reviewer": "Reviewer_45Q3", "summary": "The paper proposes a method called trajectory annealing, which uses intermediate training checkpoints of a predictor to guide sampling. The main idea is that during training, neural networks naturally evolve from coarse to fine representations, early checkpoints smooth high-frequency variations, while later ones add detail. By running MCMC samplers (GWG for discrete and MALA for continuous variables) sequentially across these checkpoints, the method achieves better mixing and sampling efficiency, especially in rugged or synergistic landscapes. Experiments on Boolean, MNIST-EBM, DNA design, and materials datasets show strong empirical gains over standard temperature annealing. The theory part connects this to SGD’s hierarchical learning of low- to high-degree monomials.", "questions": "How sensitive is the performance to checkpoint spacing or number of steps per checkpoint? Could this be integrated with modern optimizer schedules or adaptive checkpoint selection? How does it perform on large-scale, non-convex tasks like ImageNet classifiers? Is there any insight on when the coarse-to-fine property breaks down, e.g., transformers?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method is conceptually simple yet powerful, requires no retraining or architectural change, and leverages an inherent property of neural network learning. The connection to hierarchical degree learning and NTK smoothing is novel. Results are extensive and consistent across very different domains. The Boolean case study is particularly compelling, showing exponential-to-linear mixing improvements.", "weaknesses": "Theoretical results rely on strong assumptions (e.g., degree-wise alignment checkpoints) that may not generalize. Experiments are numerous but some seem cherry-picked to highlight advantages. Limited comparison with recent diffusion-based or amortized samplers. Some derivations could be more formal, and the transition from discrete to continuous domains feels hand-wavy. Also, claims of \"no extra compute\" ignore checkpoint storage and evaluation overhead.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r5dtI3wDEp", "reviewer": "Reviewer_v2vM", "summary": "### Summary\n\nThis paper leverages a well-celebrated theoretical idea, i.e. the coarse to fine, spectral learning dynamics of gradient based learning to help create a series of landscape with different smoothness. Then use these landscapes as annealed landscapes to help sampling from the final complex energy landscape. They conducted theory on simple boolean settings showing the learning order effect in the degree of polynomial, and then validated the idea in various discrete and continuous sampling set up and showed significant improvement upon simply temperature annealed sampling and other MC methods.", "questions": "- C.f. Sec. 3.2, the theory / motivation of the paper also aligns closely with the learning dynamics of score-based diffusion models, i.e. learned score vector fields are simpler, smooth earlier in the training, usually better approximated by a linear vector field. There is a nice spectral ordering for the learning of vector field and distribution. [^1], [^2]\n\n[^1] Wang, & Vastola, (2024). The unreasonable effectiveness of gaussian score approximation for diffusion models and its applications. TMLR\n\n[^2] Wang, & Pehlevan (2025). An analytical theory of power law spectral bias in the learning dynamics of diffusion models. NeurIPS\n\n- For the idea of sampling leveraging a sequence of landscapes from smooth to rugged to help find tricky spiky solutions, the authors could also mention this recent work [^3], which shares a very similar picture, but did not use the learning dynamics to help. I feel [^3] could use many intuition / results from this paper to help learn their sequence of landscapes.\n\n[^3] Du, Y., Mao, J., & Tenenbaum, J. B. (2024). Learning iterative reasoning through energy diffusion. \n\n- In the continuous sampling experiments, is it relevant to compare to or report reference values for evolutionary algorithms? since these are also common functions used to benchmark those, e.g. CMAES.\n\n- Discussing how to best heuristically pick the intermediate checkpoint is quite interesting and useful for the reader of the paper.\n\nMinor\n\n- The Table 2 format is a bit confusing…. \nthe right column (step count) should not be compared on the same table with the mid and left column (success rate). A separatrix should be added or different table should be used to present the median step result. \nFrom the first glance, it’s confusing to see 4.00 and 2.00 in the right column…. and median step is an integer so why do we have 2 decimal points here\n\n- Table 5 and Sec. 4.1.2 are not very clear, which value is the temperature annealed GWG? why does it have two values as authors say they used the final checkpoint.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "### Strength\n\n- The paper is tackling an interesting problem with a very creative solution, connecting ideas from spectral bias, training dynamics and energy based model and sampling. I’m very convinced of the idea.\n- The authors noted significant agreement with the theory in FCNN and MLP setting, and noted the deviation in transformer setting. We commend their honesty about the limitation of the theory.\n- The experimental testing of the idea is very comprehensive and convincing, showing univocal benefit of this idea.", "weaknesses": "### Weakness\n\n- Often it’s not clear which training time check point the authors used in sampling experiments.\n    - e.g. in 4.1.2, it’s not clear from writing which training time checkpoint was used and what 1K and 10K steps denote. is it that 10K steps checkpoint is better intermediate landscape than 1K step?\n    - More generally, I feel there is a lot of heuristics and design space for which checkpoint(s) are best suit for these intermediate landscapes, and how do the authors decide on them?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)