Differentiable SO(3) Power Spectrum Objective

• The paper introduces a differentiable SO(3) power spectrum loss that quantifies deviations in atomic environments via rotation-invariant descriptors.
• It employs GPU-accelerated, batched PyTorch operations to enable end-to-end gradient optimization in both direct and latent spaces.
• Quantitative benchmarks show significant improvements in unique and low-energy structure yields, highlighting its potential in computational materials discovery.

A differentiable SO(3) power spectrum objective is a loss function formulated on crystal structures that quantifies the deviation of atomic local environments from a prescribed reference, in a way that is both fully rotationally invariant and end-to-end differentiable. This objective enables the steering of structure generation processes—particularly those based on deep generative models such as VAEs—toward configurations exhibiting targeted local symmetries, such as the sp²-bonding motif characteristic of graphite. Recent implementations leverage GPU-accelerated tensor operations and full support for automatic differentiation, facilitating high-throughput, gradient-based optimization in both direct (atomic configuration) and latent (model embedding) spaces (Ridwan et al., 8 Jan 2026).

1. Mathematical Foundation of the SO(3) Power Spectrum Objective

At the core is a rotation-invariant descriptor, constructed for each Wyckoff site $j$ in a candidate crystal. The local atomic neighborhood within a cutoff $r_\mathrm{cut}$ is represented by a continuous, Gaussian-smeared neighbor density:

$$\rho_j(\vec{r}) = \sum_{k} e^{-\alpha \lVert \vec{r}_{jk} \rVert^2 } \delta(\vec{r} - \vec{r}_{jk})$$

with $\vec{r}_{jk} = \vec{r}_k - \vec{r}_j$ and $\alpha$ the smearing width. This density is expanded in a combined radial–spherical-harmonic basis with orthonormal radial functions $R_n(r)$ and spherical harmonics $Y_{\ell m}(\hat{r})$, yielding coefficients:

$$c^{(j)}_{n\ell m} = \int_0^{r_\mathrm{cut}} dr\, r^2 R_n(r) \sum_k e^{-\alpha r_{jk}^2} Y_{\ell m}(\hat{r}_{jk})$$

The rotation-invariant SO(3) power spectrum is evaluated as:

$$p^{(j)}_{nn'\ell} = \sum_{m=-\ell}^{\ell} c^{(j)}_{n\ell m} \left[c^{(j)}_{n'\ell m}\right]^*$$

which is flattened to a descriptor vector $P_j \in \mathbb{R}^L$ for each Wyckoff site.

Given a reference descriptor $P_\mathrm{ref}$ (e.g., extracted from graphite), a squared-error loss is constructed:

$$\ell_j = \sum_{k=1}^L (P_{j,k} - P_{\mathrm{ref},k})^2$$

and the structure-level mean-squared loss is

$$\ell_\mathrm{struct} = \frac{1}{WL} \sum_{j=1}^W \sum_{k=1}^L (P_{j,k} - P_{\mathrm{ref},k})^2$$

where $W$ is the number of Wyckoff sites in the unit cell.

2. Fully Differentiable Implementation

Every step in the computation of the SO(3) power spectrum objective is implemented as a fully differentiable, batched GPU pipeline in PyTorch, distinguishing it from prior CPU-based workflows reliant on numerical gradients. Key elements include:

• Symmetry mapping: Precomputed tensors reconstruct Cartesian coordinates from the reduced, learnable crystal representation tensor $R \in \mathbb{R}^{B \times N_f}$, where $B$ is batch size and $N_f$ the number of free parameters.
• Neighbor listing: Constructed as differentiable tensor operations, vectorized over crystal batches.
• Basis function computation: Radial basis $R_n(r)$ and spherical harmonics $Y_{\ell m}(\theta, \phi)$ are constructed from elementary PyTorch operations, preserving compatibility with reverse-mode autodiff.
• Loss and gradient calculation: Mean-squared error in the descriptor space is minimized using exact gradients derived from PyTorch autograd throughout the computational graph. No finite-difference approximations are used, in contrast to legacy CPU/SciPy implementations (Ridwan et al., 8 Jan 2026).

3. Integration with Symmetry-Conditioned Variational Autoencoders

In a crystal generation pipeline, the differentiable SO(3) loss $\ell_\mathrm{struct}$ is coupled to a symmetry-conditioned variational autoencoder (CVAE):

• Training: The CVAE is trained using the standard reconstruction and KL-divergence losses, optimizing for a latent space $Z$ conditioned on specified discrete variables (space group and Wyckoff positions).
• Inference: New candidates are sampled by decoding latent vectors $(Z, C) \rightarrow X$ (raw crystal parameters). Prior to first-principles relaxation, the differentiable power spectrum loss refines candidate structures via gradient-based updates, enforcing the target local environment downstream of the CVAE.

This environmental objective provides a descriptor-guided refinement layer, bridging the latent generative model and geometric local environment specification under space-group symmetry constraints (Ridwan et al., 8 Jan 2026).

4. Dual-Level Gradient-Based Optimization

Optimization is performed at two hierarchical levels:

Algorithm 1: Representation-space (direct geometry) optimization

• Treat the reduced free-parameter tensor $R$ as learnable.
• Compute $\ell_\mathrm{struct}$ and its gradient with respect to $R$.
• Update $R$ with AdamW optimizer, clamp coordinates, and repeat for $T_\mathrm{rep}$ iterations.

Algorithm 2: Latent-space optimization

• Fix the discrete condition $C$, treat the latent vector $Z$ as learnable.
• Decode $(Z, C) \rightarrow R$, reconstruct geometry, compute $\ell_\mathrm{struct}$.
• Update $Z$ by backpropagation and AdamW optimization over $T_\mathrm{lat}$ steps.

Gradient flows from descriptor-space MSE all the way back to $R$ or $Z$. Adaptive techniques such as per-sample gradient clipping and learning-rate scheduling are employed to regularize optimization for diverse and heterogeneous samples (Ridwan et al., 8 Jan 2026).

5. High-Throughput Batched Processing and Computational Efficiency

Efficient, high-throughput implementation is achieved by sorting candidates into batches (e.g., $B \approx 1,000$), and fusing all mapping, neighbor listing, projection, and averaging operations into single batched PyTorch kernels. This architecture yields:

• Approximately $5\times$ speedup over previous CPU/SciPy routines (1 minute per $1,000$ samples on NVIDIA H100 GPUs vs 5 minutes on 96-core CPUs).
• Stable convergence for large and statistically diverse sample batches, facilitated by the AdamW optimizer and robust gradient regularization.
• Optional post-processing with MACE-FF relaxation and duplicate removal for database finalization (Ridwan et al., 8 Jan 2026).

6. Quantitative Performance Benchmarks

Empirical performance, as benchmarked on $100,000$ generated candidates, demonstrates the efficacy and throughput of the differentiable SO(3) power spectrum objective:

Scenario	Valid Env†	Unique	Low-Energy
VAE, 1-pass refinement	18,248	4,878	508
CVAE, 1-pass refinement	23,372	5,695	571
CPU/SciPy batch (L-BFGS-B, num. grad)	18,248	4,878	-
GPU/PyTorch batch (AdamW, autograd)	12,874	4,817	-
Dual-level, 4 rounds (CVAE, GPU, full)	48,308	11,546	1,014

Sequential dual-level optimization (alternating between representation- and latent-space) more than doubles the number of unique and low-energy sp² carbon nets relative to single-pass representation-space optimization, in under $9$ hours of GPU runtime. CVAE-augmented pipelines yield approximately $28\%$ more target-environment structures and $12\%$ more low-energy hits than baseline VAEs (Ridwan et al., 8 Jan 2026).

7. Significance and Extension

The differentiable SO(3) power spectrum objective functions as an effective, scalable molecular descriptor and environmental constraint, simultaneously enabling rigorous symmetry preservation and high-yield, high-throughput generation of complex materials. Its extension to multi-component, multi-environmental systems is plausible given the generality of the underlying descriptor and the flexibility of the differentiable implementation. Current evidence substantiates its impact in increasing both the rate and diversity of successful crystal structure generation, underlining its role as a central methodological advancement in computational materials discovery (Ridwan et al., 8 Jan 2026).