FrOGS: Discrete Neural Sampler for Independent Alloy Configurations Across Chemical Conditions

1 Introduction

• MCMC requires separate sequential runs for each chemical condition, followed by an auxiliary calculation to place free energies on a common absolute scale.Near phase boundaries, it may require more steps to produce reasonable configurations [Berg and Neuhaus, 1991, Wolff, 1989].Source: §1 · p. 1
• Neural samplers can miss entire target-distribution modes without increasing their loss, a failure known as mode collapse.The text links this failure to reverse-KL training [Minka, 2005, Soletskyi et al., 2025].Source: §1 · pp. 1–2
• A single FrOGS sampler generates microstates across the chemical-condition region on a common absolute scale.Feature-wise Linear Modulation conditioning on c = (∆µ, T) enables generation at each condition while preserving local equivariance.Source: §1 · pp. 2–3
• FrOGS couples an autoregressive prior to neural CTMC transport through a shared loss, with the prior absorbing first-order-boundary discontinuities and transport refining its proposal.The transport is described as a learned generalization of annealed importance sampling [Neal, 2001, Holderrieth et al., 2025].Source: §1 · p. 3
• FrOGS recovers the AgPd miscibility gap and regions of all three ordered CuAu phases, including near phase boundaries.These results are presented as evidence of accurate free-energy prediction without mode collapse.Source: §1 · p. 3
• FrOGS extends LEAPS to real chemical systems with a locally equivariant multi-scale head for long-range interactions and a rate-matrix reparameterization for low-temperature stability.The reparameterization avoids numerically unstable density-ratio computation at low temperatures.Source: §1 · p. 3

2 Related Work

• Autoregressive models provide exact normalized likelihoods, but reverse-KL training is mode-seeking and learned distributions can miss coexisting phases near transitions.Their biases are usually corrected with an unlearned reweighting scheme [Minka, 2005, Soletskyi et al., 2025].Source: §2 · p. 3
• Metadynamics biases sampling away from visited minima, while its well-tempered variant reduces added bias as it accumulates.The recovered profile is projected onto a collective variable, so states it cannot distinguish can be grouped into one basin; discrete neural extensions train at fixed conditions and recover a profile up to an additive constant [Laio and Parrinello, 2002, Barducci et al., 2008, Bussi and Laio, 2020, Du et al., 2026a].Source: §2 · p. 3
• Flat-histogram methods estimate the density of states and provide global free energies on an absolute scale, but balance only the composition-and-energy marginal.They do not natively provide independent configurations [Berg and Neuhaus, 1991, Wang and Landau, 2001, Neuhaus and Hager, 2003, Dayal et al., 2004].Source: §2 · p. 3
• Learned transport methods move probability mass from a tractable prior toward a complex target, with annealed transport samplers interleaving learned flows and sequential Monte Carlo.The annealed transport samplers inherit asymptotic unbiasedness, while their flows are trained with per-transition KL objectives [Vaikuntanathan and Jarzynski, 2008, Noé et al., 2019, Albergo et al., 2019, Wirnsberger et al., 2020, Arbel et al., 2021, Matthews et al., 2022].Source: §2 · p. 3
• NETS learns a continuous-time drift to minimize weight variance, and LEAPS carries this construction to discrete state spaces.LEAPS is the framework applied here to lattice alloys [Albergo and Vanden-Eijnden, 2025, Holderrieth et al., 2025].Source: §2 · p. 4

3 Background

• LEAPS transports lattice configurations from a tractable prior to the Boltzmann distribution at chemical condition c = (∆µ, T) along an annealing path.The CTMC uses a learnable rate matrix and restricts moves to replacing one site's species.Source: §3 · p. 4
• LEAPS reweights each of its M samples at each path time using a log-weight, and its effective sample size summarizes weight balance.M · ESS approximates the number of independent draws that would give an estimator with the same variance as the M weighted samples.Source: §3 · p. 4
• LEAPS provides a direct unbiased estimate of Z1 and a consistent estimate of observables through sample reweighting.Source: §3 · p. 4
• Bare LEAPS is not conditionable, can transport inefficiently across first-order boundaries from a uniform prior, and had been applied only to square-grid 2D systems.The target alloy systems are 3D, exchange-symmetry-broken, and may use non-simple-cubic lattices such as fcc.Source: §3 · p. 4

4.1 Learned prior

• A conditioned autoregressive prior places mass in the target basin, avoiding the need for local transport to carry probability across first-order boundaries.With a uniform prior, passage through energetically unstable configurations can produce high-variance log weights and collapse ESS.Source: §4.1 · pp. 4–5

4.2 Coupling and conditioning the autoregressive and transport networks

• A shared loss jointly trains the prior and transport networks, conditioned on chemical conditions c = (∆µ, T), and avoids a mode-seeking objective.The objective for the autoregressive prior is a generalized form of VarGrad.Source: §4.2 · p. 5

4.3 Extending the transport network

• The model retains local equivariance by factoring probability flux so its output at site i does not depend on the current token σi.Swapping σi and τ then flips only the projector difference.Source: §4.3 · p. 5
• A U-net-style locally equivariant architecture gives the transport network a larger receptive field and better performance at equal scale than stacked locally equivariant convolution layers.Its decoder combines upsampled coarse features with encoder features through skip connections.Source: §4.3 · p. 5
• The rate reparametrization uses the equilibrium Metropolis acceptance factor to yield a bounded corrector without the density ratio that can diverge at low temperatures.The reparametrized rates span the same family of CTMCs, preserving the universal-representation property.Source: §4.3 · p. 6

5 Experiments

• The reported alloy thermodynamics are those of the chosen cluster expansion, which can differ considerably from the real phase diagram.For CuAu, the cluster expansion places Cu3Au above 700 K versus 663 K experimentally.Source: §5 · p. 6

5.1 Ising Model

• FrOGS reproduces the 20 × 20 Ising model’s finite-size thermodynamics with high accuracy, and its learned transport tracks references more closely than the autoregressive prior alone.The evaluation uses 5000 samples per condition; results across 10 × 10, 15 × 15, and 20 × 20 systems match Kaufman/Wolff references.Source: §5.1 · p. 6

5.2 AgPd

• For 125-site AgPd, FrOGS reproduces the miscibility gap; ESS is ≥0.1 at 98.38% of the chemical-condition grid, with median 0.903.The binodal shape agrees with the metadynamics reference; the sweep uses 5000 samples at each condition.Source: §5.2 · p. 7

5.3 CuAu

• For 128-site CuAu, FrOGS reliably detects Cu3Au, CuAu, and CuAu3, with phase-boundary shapes agreeing with the reference.ESS is ≥0.1 at 98.58% of the chemical-condition grid, with median 0.907, supporting reliable free-energy estimation.Source: §5.3 · pp. 7–8, Figure 3 · p. 8
• FrOGS distributes raw samples across degenerate ordered-phase variants without mode collapse, while importance-weighted counts are close to uniform.The reported χ2/dof test compares counts with the flat 1/g expectation.Source: §5.3 · pp. 7–8, Table 1 · p. 7

5.4 Comparison with previous work

• Only FrOGS reliably captures the CuAu3 stability range, and its phase-boundary shapes track metadynamics more closely than Damewood et al. [2022].Metadynamics indicates CuAu3 is thermodynamically unstable above 400K, and the comparison remains consistent across tested phase-diagram thresholds and FrOGS initialization seeds.Source: §5.4 · pp. 8–9, Figure 3 · p. 8
• FrOGS better agrees with Du et al. [2026b] on Ising heat capacity and has higher ESS across chemical conditions than prior samplers.The authors report that higher ESS leads to more accurate predicted phase-boundary curves.Source: Figure 2 · p. 7, §5.4 · p. 9, Figure 4 · p. 9

6 Conclusion

• FrOGS matches established 2D Ising thermodynamics and reconstructs AgPd and CuAu phase diagrams more faithfully than previous autoregressive approaches.It provides i.i.d. samples and directly computable free energies at any chemical condition, without an external reference.Source: §6 · p. 9
• FrOGS generates i.i.d. samples from the cluster expansion and computes free energy directly at any chemical condition, unlike conventional MCMC's separate chains and thermodynamic integration.The cited MCMC approach requires phase-specific chains initialized from an ordering identified beforehand.Source: §6 · p. 9
• 4.5 GPU-hours train FrOGS, and its CuAu grid sweep takes 31 GPU-hours at 75 s per condition for M = 5000 samples.The comparison is to a few core-hours for anchored Metropolis on the same system, with an ordering supplied beforehand; FrOGS training is paid once and sweeps only requested conditions.Source: §6 · p. 9
• FrOGS computes free-energy gradients from the same samples via standard thermodynamic identities, opening a way to gradient-based inverse design.The conclusion presents multicomponent alloys (K ≥3) as a natural next application.Source: §6 · p. 9

A Technical appendices and supplementary material

• The appendices cover semi-grand canonical sampling, LEAPS, and derivation of the reparameterized corrector.They also include phase-diagram construction and joint optimization of the autoregressive network.Source: A Technical appendices and… · p. 14
• Supplementary sections analyze prior-only and transport-only ablations, model and training details, and multiscale 2D Ising comparisons.The contents also list a sensitivity analysis of FrOGS phase count.Source: A Technical appendices and… · p. 14
• The appendices include proofs and implementation material on U-net local equivariance, FiLM conditioning, and fcc-lattice re-indexing.They also compare FrOGS with conventional Monte Carlo.Source: A Technical appendices and… · p. 14

A.1 Sampling from the SGC ensemble

• A K-component alloy configuration assigns one of K species to each of N fixed lattice sites, yielding a configuration space of size K^N.The site occupations are denoted σ_i and the full configuration σ.Source: §A.1 · p. 14
• In the semi-grand canonical ensemble, N and T are fixed while composition fluctuates through exchange with a chemical-potential reservoir; the condition is c = (∆µ, T).Only chemical-potential differences are thermodynamically meaningful because the species counts sum to N.Source: §A.1 · p. 14
• The partition function sums over K^N configurations and is intractable to evaluate directly, motivating a sampler that draws from the target distribution and estimates it across conditions.The free energy is defined as F_SG(c) = −k_BT ln Z_SG(c).Source: §A.1 · p. 15

A.2 LEAPS

• Exact transport of ρ_0 along the path requires the network transition rates to satisfy the Kolmogorov forward equation.LEAPS records the mismatch between the forward-equation terms and the change in free energy in an operator.Source: §A.2 · p. 15
• Each sample integrates the transport residual along its trajectory into a log-weight, analogous to accumulated work in the Jarzynski equality.At perfect transport, the residual equals the negative time derivative of the free energy.Source: §A.2 · p. 15
• FrOGS estimates the partition function directly from the weights, and reweighting recovers observables consistently.The expected exponential weight equals the ratio of final to initial partition functions.Source: §A.2 · p. 15
• LEAPS trains the transport by minimizing weight variance, which is zero exactly when the transported samples follow the target distribution.The cited objective is the variance of the log-weight A_t.Source: §A.2 · p. 15
• Locally equivariant rates let LEAPS evaluate reverse rates for all single-site neighbors in one forward pass.The free-energy network tracks a scalar function of time and is a training-time device, not part of the free-energy readout; the reference measure may draw states from a replay buffer.Source: §A.2 · p. 15

• The corrector needs one flux-network evaluation and no standalone density ratio; its Metropolis factor avoids overflow from strongly downhill moves.For transport parameters θ, Mt is fixed, so the flux gradient is scaled by Mt as well.Source: §A.3 · p. 16
• The nonnegative reparametrized off-diagonal rates, paired with the matching corrector, preserve the continuous-time LEAPS identity and its unbiasedness result.Source: §A.3 · p. 16
• Every one-way rate matrix can be represented as [G]+Mt for an antisymmetric flux G, preserving the one-way rate family.A one-way matrix has at most one nonzero direction for each configuration pair.Source: §A.3 · p. 16
• The Metropolis factor explicitly damps uphill moves, and the reparametrized rates show a faster early decrease in PINN loss.Source: §A.3 · p. 16

• The constructor uses each sample’s log-weight and composition from one trained FrOGS model, with no phases supplied manually.Source: §A.4 · p. 16
• At fixed composition, a sample-average estimator yields the partition function on the integer composition lattice after dividing out the chemical-potential factor.The resulting Zc curve does not depend on chemical potential, so conditions at a given temperature independently estimate the same curve.Source: §A.4 · p. 17
• Precision-weighted merging combines independent unbiased estimates, with conditions contributing most strongly near their equilibrium compositions and not to bins they never reach.Source: §A.4 · p. 17
• The cross-check statistic assesses the accuracy of the estimates, and its one-sided error-bar correction can only widen them.The leading Jensen-bias term stays below the reported standard error when the mean effective sample count over contributing conditions exceeds kn/4.Source: §A.4 · p. 17
• A hull edge is accepted as a tie-line only if it skips at least four compositions, meets an intermediate free-energy or unvisited-composition gap condition, and the gap recurs at the lower temperature.Source: §A.4 · p. 17
• AgPd yields a miscibility gap that closes at a critical point, while CuAu yields Cu3Au, CuAu, and CuAu3 ordered phases at their tangent compositions.Tie-line closure temperatures are not, by themselves, precise order–disorder transition temperatures.Source: §A.4 · p. 17

• The annealing path linearly interpolates U0 and U1, with U0 defined as the negative log density of the autoregressive prior.Source: §A.5 · p. 18
• Both ingredients of the corrector carry explicit dependence on the autoregressive weights ψ along the annealing path.Source: §A.5 · p. 18
• A single backward pass through the joint objective returns gradients for θ, ϕ, and ψ.Source: §A.5 · p. 18

A.6 The prior objective as the VarGrad generalized

• With transport off, FrOGS’s prior gradient is VarGrad [Richter et al., 2020]; with transport on, independent reference states keep mode collapse from becoming self-reinforcing.The log-variance loss constrains the prior only on the reference distribution’s support, so preventing collapse still requires that the reference overlap the relevant modes.Source: §A.6 · pp. 18, 20

A.7 Proof of the local equivariance of the U-net pyramid

• The U-net head is locally equivariant when every axis at the coarsest level has length at least two, because its features exclude the site’s own cell.The axis-length condition matters: circular convolution on a length-one axis can wrap a noncentral offset back to the excluded cell.Source: §A.7 · p. 21

A.8 The FiLM setup

• FiLM conditioning rescales each condition axis to [-1, 1] before an MLP produces feature-wise scales and shifts, initialized to the identity.The normalization prevents temperature from dominating the first conditioning layer when axes differ by three orders of magnitude.Source: §A.8 · p. 22

A.9 Re-indexing the fcc lattice onto an integer grid

• Representing fcc sites by primitive-basis integer coordinates makes the lattice a dense 3D grid while preserving physical-geometry energy evaluation.The coordinate map is a shear from the fcc lattice to simple-cubic index space.Source: §A.9 · p. 22

A.10 Ablation study of prior-only and transport-only networks

• Both prior-only and transport-only ablations sample less effectively than FrOGS; the uniform-prior model’s ESS collapses to 1/M in the ordered region.The uniform-prior model’s raw samples populate degenerate ordering variants, but importance weighting concentrates on effectively one variant.Source: §A.10 · pp. 22–23, Figure 6 · p. 23, Figure 7 · p. 23

A.11 Model and training details

• Training experiments selected a transport pyramid with 3 encoders and 2 decoders, prior width 32, and nmcmc = 10.The choices followed comparisons of training statistics over the final 10k optimizer steps on CuAu and the L = 20 2D Ising model.Source: §A.11 · p. 24, Table 3 · p. 24, Table 4 · p. 25

A.13 Sensitivity analysis of FrOGS phase count

• Across 2,268 settings per sampler, FrOGS resolved three ordered compounds and SEGAL two under the stated phase-counting scheme, with the difference occurring at xAu = 3.The sweep varied (wmin, nσ), three persistence rules, measured-absence routing, and χ2 bin-floor values.Source: §A.13 · p. 25

A.12 Multiscale 2D Ising result comparison

• Across L = 10, 15, and 20, FrOGS’s partition-function error per site stayed within 2.5 × 10^-4 of the Kaufman value, with mass balance near 0.5 and critical G(r) mostly within 2σ of Wolff references.These results used 5,000 walkers per temperature without MCMC during evaluation; Var[A]/N at Tc was 0.014, 0.032, and 0.016 across the three sizes.Source: §A.12 · p. 25, Figure 9 · p. 26

A.14 FrOGS vs. conventional Monte Carlo

• At 200 K, seven of 36 chemical potentials produced Monte Carlo chains frozen in different basins, with seed-to-seed x differences up to 0.14 despite reported error bars below 10^-5.Anchoring one chain per phase removes this issue but requires supplying the phases in advance.Source: §A.14 · p. 26

A.13 Sensitivity analysis of FrOGS phase count

• All four FrOGS training seeds resolved three ordered CuAu compounds and closely matched metadynamics phase boundaries.The four-seed phase diagrams show the agreement was not a single lucky event.Source: §A.13 · p. 25, Figure 10 · p. 27

A.14 FrOGS vs. conventional Monte Carlo

• At ten non-ergodic Monte Carlo conditions, FrOGS ESS ranged from 0.55 to 0.97 (median 0.85), while Monte Carlo remained ergodic at FrOGS’s eight lowest-ESS conditions.The costliest of those eight conditions required 210 core-seconds; pooling precision-weighted free-energy estimates from neighboring conditions recovers free energies even in FrOGS’s low-ESS regions.Source: §A.14 · p. 26, Table 5 · p. 27