Fourier Transformers for Latent Crystallographic Diffusion and Generative Modeling

The discovery of new crystalline materials calls for generative models that handle periodic boundary conditions, crystallographic symmetries, and physical constraints, while scaling to large and structurally diverse unit cells. We propose a reciprocal-space generative pipeline that represents crystals through a truncated Fourier transform of the species-resolved unit-cell density, rather than modeling atomic coordinates directly. This representation is periodicity-native, admits simple algebraic actions of space-group symmetries, and naturally supports variable atomic multiplicities during generation, addressing a common limitation of particle-based approaches. Using only nine Fourier basis functions per spatial dimension, our approach reconstructs unit cells containing up to 108 atoms per chemical species. We instantiate this pipeline with a transformer variational autoencoder over complex-valued Fourier coefficients, and a latent diffusion model that generates in the compressed latent space. We evaluate reconstruction and latent diffusion on the LeMaterial benchmark and compare unconditional generation against coordinate-based baselines in the small-cell regime ($\leq 16$ atoms per unit cell).

💡 Research Summary

The paper introduces a novel generative pipeline for crystalline materials that operates entirely in reciprocal space rather than directly on atomic coordinates. The authors represent each crystal by a species‑resolved unit‑cell density and compute a truncated Fourier transform of this density using a fixed set of wave vectors (|j|∞ ≤ j_max = 9). This yields a compact matrix of complex Fourier coefficients, one column per chemical species. The zero‑frequency (DC) coefficient encodes the atom count for each species, allowing the representation to naturally support variable multiplicities without changing tensor shape.

Crystallographic space‑group symmetries become simple algebraic operations in this Fourier domain: a symmetry consisting of a lattice‑preserving rotation M and a fractional translation δ maps coefficient α_j to exp(‑2πi j·δ) α_{M^T j}. Consequently, symmetry constraints are embedded in the token ordering and phase relationships rather than enforced through geometric post‑processing. To preserve these relationships during learning, the authors design a complex‑valued rotary positional encoding tied directly to the wave vectors.

The learning architecture consists of two stages. First, a complex‑valued transformer variational auto‑encoder (VAE) encodes the full set of Fourier tokens together with a global lattice token and a small auxiliary token block. After each transformer layer, the auxiliary tokens are extracted and stacked into a depth‑aligned “auxiliary ladder,” which serves as the latent representation passed through a variational bottleneck. The decoder mirrors the encoder, re‑injecting the appropriate ladder slice at each depth to condition reconstruction. This auxiliary‑ladder mechanism enables strong compression while retaining global regularities such as symmetry‑induced correlations.

Second, a latent diffusion model (U‑Net style) is trained to generate the compressed latent codes. During sampling, the diffusion process denoises a random latent vector, which is then decoded by the VAE decoder to produce lattice parameters and the truncated Fourier coefficients. An inverse Fourier transform, guided by the DC coefficients, reconstructs fractional atomic coordinates and species labels.

Empirical evaluation on the LeMaterial benchmark demonstrates that with only nine basis functions per spatial dimension the method can exactly reconstruct unit cells containing up to 108 atoms per species. In the small‑cell regime (≤ 16 atoms per unit cell), the approach matches or exceeds coordinate‑based transformer and GAN baselines in terms of structural similarity and physical plausibility, while also offering the unique ability to generate crystals with variable atom counts. The authors report stable training, high reconstruction fidelity, and effective unconditional sampling that respects space‑group symmetries.

Overall, the work presents a “representation‑first” solution to the long‑standing challenges of periodicity, symmetry, and compositional flexibility in crystal generation. By moving the generative problem to reciprocal space, the method eliminates the need for explicit geometric handling of periodic boundaries and fixed‑size point clouds, opening avenues for conditional generation, high‑resolution electron density modeling, and extensions to more complex or disordered materials.