Crystal Generation using the Fully Differentiable Pipeline and Latent Space Optimization

TL;DR

This work tackles rapid, targeted crystal structure generation under crystallographic constraints by coupling a symmetry-conditioned CVAE with a differentiable SO(3) descriptor. The method advances by (i) conditioning generation on discrete symmetry to enable controlled exploration, (ii) implementing a fully differentiable, GPU-accelerated pipeline for batchwise optimization in both direct and latent spaces, and (iii) introducing an iterative dual-level refinement that uses latent-space perturbations to escape local minima and then re-optimizes in the representation space. Key findings show improved valid and unique structure yields, faster descriptor-based refinements (approximately 5x speedup), and substantial yield gains from latent-space refinements, all while maintaining comparable energy distributions to CPU-based baselines. These advances enhance scalable, AI-driven crystal-gen design and pave the way for discovering complex, targeted motifs in multi-component materials; the authors also provide public code to enable broader adoption.

Abstract

We present a materials generation framework that couples a symmetry-conditioned variational autoencoder (CVAE) with a differentiable SO(3) power spectrum objective to steer candidates toward a specified local environment under the crystallographic constraints. In particular, we implement a fully differentiable pipeline that performs batch-wise optimization on both direct and latent crystallographic representations. Using the GPU acceleration, the implementation achieves about fivefold speed compared to our previous CPU workflow, while yielding comparable outcomes. In addition, we introduce the optimization strategy that alternatively performs optimization on the direct and latent crystal representations. This dual-level relaxation approach can effectively overcome local barrier defined by different objective gradients, thus increasing the success rate of generating complex structures satisfying the targe local environments. This framework can be extended to systems consisting of multi-components and multi-environments, providing a scalable route to generate material structures with the target local environment.

Figures (6)

• Figure 1: Architecture of the Conditional VAE with DiffGMM data transformation. The encoder processes GMM-transformed continuous features to produce latent variables $\mu$ and $\log\sigma^2$, while discrete conditions (space group and Wyckoff positions) are embedded separately and concatenated with the sampled latent $z$ before decoding.
• Figure 2: SO(3) power spectrum comparison for CVAE-generated carbon structures before and after refinement. Shown are the initial structures, results after energy-based relaxation using MACE-FF, and results after SO(3) descriptor-based optimization, together with the graphite sp$^2$ reference (gray dashed line).
• Figure 3: Iterative dual-level refinement workflow. Structures satisfying the target $\mathrm{CN}=3$ motif are collected in $\mathcal{D}_{\rm valid}$; remaining samples ($\mathcal{W}$) undergo latent refinement followed by representation-space re-optimization for up to $I$ rounds.
• Figure 4: Baseline generation comparison between a VAE and a CVAE. Both models generate 100k candidates and are evaluated with the same post-processing pipeline. Left: Summary counts of valid structures satisfying the target environment ($N_{\rm valid\_env}$), unique structures ($N_{\rm unique}$), and low-energy candidates ($N_{\rm low\_E}$). Right: MACE-relaxed energy distributions (eV/atom) for the screened unique structures, shown as counts with identical binning.
• Figure 5: GPU acceleration of direct representation optimization.Left: CPU--SciPy (prior LEGOxtal implementation) versus GPU--PyTorch (this work), including device, optimizer, batching strategy, throughput (time per 1k candidates), and post-screening counts. Right: MACE-relaxed energy histograms (eV/atom) for the screened unique structures, shown as counts with identical binning.