Scaling Autoregressive Models for Lattice Thermodynamics

TL;DR

This paper introduces a scalable framework combining any-order autoregressive models and marginalization models to efficiently generate atomic configurations for lattice thermodynamics, capturing phase transitions with reduced computational cost.

Contribution

It develops a flexible, memory-efficient modeling approach that enables training on small lattices and application to larger systems, improving accuracy and scalability in thermodynamic simulations.

Findings

Transformer-based models outperform MLP-based models in free energy accuracy.

The framework scales to larger lattice sizes with less computational cost.

Models faithfully capture phase transitions and critical behavior.

Abstract

Predicting how materials behave under realistic conditions requires understanding the statistical distribution of atomic configurations on crystal lattices, a problem central to alloy design, catalysis, and the study of phase transitions. Traditional Markov-chain Monte Carlo sampling suffers from slow convergence and critical slowing down near phase transitions, motivating the use of generative models that directly learn the thermodynamic distribution. Existing autoregressive models (ARMs), however, generate configurations in a fixed sequential order and incur high memory and training costs, limiting their applicability to realistic systems. Here, we develop a framework combining any-order ARMs, which generate configurations flexibly by conditioning on any known subset of lattice sites, with marginalization models (MAMs), which approximate the probability of any partial configuration in a single forward pass and substantially reduce memory requirements. This combination enables models trained on smaller lattices to be reused for sampling larger systems, while supporting expressive Transformer architectures with lattice-aware positional encodings at manageable computational cost. We demonstrate that Transformer-based any-order MAMs achieve more accurate free energies than multilayer perceptron-based ARMs on both the two-dimensional Ising model and CuAu alloys, faithfully capturing phase transitions and critical behavior. Overall, our framework scales from $10 \times 10$ to $20 \times 20$ Ising systems and from $2 \times 2 \times 4$ to $4 \times 4 \times 8$ CuAu supercells at reduced computational cost compared to conventional sampling methods.