X-MACE: DL for Excited-State PES

• X-MACE is a deep learning architecture that uses message passing atomic cluster expansion, an autoencoder, and a Hermitian decoder to accurately model non-smooth excited-state potential energy surfaces.
• It employs a MACE backbone to capture many-body atomic interactions and a DeepSets autoencoder to enforce permutation invariance, ensuring smooth latent representations.
• Experiments demonstrate that X-MACE significantly reduces energy and force prediction errors near conical intersections, outperforming previous models by up to 50%.

X-MACE (eXcited-state Message Passing Atomic Cluster Expansion) is a deep learning architecture designed to model excited-state potential energy surfaces (PES) with high fidelity, transferability, and efficiency, particularly in the vicinity of conical intersections where standard neural regressors fail due to the non-smooth and degenerate character of the adiabatic surfaces. By integrating a message passing atomic cluster expansion (MACE) backbone with a DeepSets autoencoder and a Hermitian companion-matrix decoder, X-MACE learns smooth, permutation-invariant representations of the entire set of adiabatic energies. This allows accurate reconstruction of non-smooth energy landscapes and improved generalization across chemical space and electronic state manifolds (Barrett et al., 18 Feb 2025).

1. Motivation and Problem Context

Conical intersections—a ubiquitous feature of photochemistry—are points or seams in nuclear coordinate space where two (or more) Born–Oppenheimer electronic states become degenerate. These configurations underpin ultrafast, nonradiative transitions governing primary events in vision, photosynthesis, and photostability. Quantum chemical computation of PES near conical intersections is expensive; learning to model these features with standard ML architectures is impeded by the sharp, cusp-like non-smoothness and by permutation indeterminacy in state labels. Most neural potentials, including previous excited-state extensions of ground-state models, either smear out the degeneracy or introduce spurious avoided crossings, failing to obtain physical crossing topologies.

X-MACE was developed specifically to resolve these deficiencies by projecting the learning problem into a latent, permutation-invariant space—mirroring the smoothness of underlying diabatic or quasi-diabatic surfaces—such that the complex intersection topology is faithfully reconstructed from learned invariants via a mathematically controlled decoder (Barrett et al., 18 Feb 2025).

2. Model Architecture

The architecture of X-MACE is fundamentally modular and composed of three primary components: (i) a MACE graph neural network backbone, (ii) a DeepSets autoencoder to encode and regularize energy invariants, and (iii) a Hermitian companion-matrix decoder to invert the latent representation back to physical energy levels.

• MACE Backbone: Molecular geometries are represented as undirected graphs $G=(V,E)$ with nodes (atoms) characterized by atomic number $Z_i$, position $\mathbf R_i$, and hidden feature vectors. Two message-passing layers iteratively aggregate environment information using equivariant multi-layer perceptrons (MLP), with angular dependencies resolved up to $l=3$ through spherical harmonics and radial functions within a 5 Å cutoff. The message update is:

$$m_i^{(t)} = \sum_{\nu=1}^\nu \sum_{j_1,\dots,j_\nu\in\mathcal N_\nu(i)} u_\nu\bigl(\sigma_i^{(t)},\sigma_{j_1}^{(t)},\dots,\sigma_{j_\nu}^{(t)}\bigr)$$

Sums over many-body terms capture higher-order correlations. The hidden state is updated by $$\mathbf h^{(t+1)}_i = U_t(\mathbf h^{(t)}_i, m^{(t)}_i)$$. Summing per-atom outputs yields a global descriptor $\mathcal I^{(\mathrm{MACE})} \in \mathbb R^k$.

• DeepSets Autoencoder: To enforce permutation invariance, an encoder of the form

$$\mathcal I = \rho\left(\sum_{i=1}^n \phi(E_i)\right)$$

maps the set of energies to a latent vector. Both $\phi$ and $\rho$ are 2-layer perceptrons with ReLU activation (hidden dimension 128), and $k = 16$.

• Companion-Matrix Decoder: A two-layer MLP with SiLU activations outputs a Hermitian matrix $M$, whose real eigenvalues reconstruct the adiabatic energies:

$\{\tilde E_i\} = \mathrm{Eig}(M)$

Training enforces alignment between the MACE descriptor and the autoencoder latent vector, so that only MACE plus the decoder is required at inference.

3. Mathematical Formalism and Training Objective

The X-MACE loss combines multiple objectives:

• Energy reconstruction:

$$\mathcal L_{\mathrm{recon}} = \frac{1}{n} \sum_{i=1}^n (E_i - \tilde E_i)^2$$

• Latent invariant matching:

$$\mathcal L_{\mathrm{match}} = \frac{1}{k}\sum_{j=1}^k \left[ \mathcal I^{(\mathrm{MACE})}_j - \mathcal I^{(\mathrm{AE})}_j \right]^2$$

• Auxiliary terms: force matching and phase-free nonadiabatic couplings, weighted by $a_F = a_C = 100$.

The total loss is:

$$\mathcal L_\mathrm{total} = \mathcal L_\mathrm{recon} + \mathcal L_\mathrm{match} + a_F \sum_i \| F_i^{\rm ref} - F_i \|^2 + a_C \sum_{m<n}\min \| C^{\rm ref}_{mn} - C_{mn} \|^2$$

This enforces both physically accurate observables and representation regularization.

4. Dataset Composition and Training Protocols

Datasets span both prototypical intersections and broad chemical space:

• SO₂ slices from linear-vibronic-coupling models.
• Methylenimmonium cation (CH₂NH₂⁺), Fulvene, Ethene, Propene, and Butene using multireference CASSCF or MR-CISD electronic structure references (2–3 lowest singlet states).
• 368 chromophore SMILES with conformers generated via metadynamics, recomputed at ADC(2) for the five lowest singlets (12,183 points).

Standard 80/10/10 train/validation/test splits are used, with metadynamics for sampling near intersection regions but no further augmentation. The model employs Adam (learning rate $10^{-3}$), with the MACE backbone using two interaction layers, 5 Å cutoff, $l_\mathrm{max} = 3$, and hidden channels $H \in \{16, 32, 64, 128\}$.

5. Transferability and Generalization

X-MACE incorporates two principal transferability mechanisms:

• Ground-to-Excited Fine-Tuning: A ground-state MACE model (“MACE-OFF”) is pre-trained, after which the composite X-MACE is fine-tuned on only 1–10% of the excited-state dataset. Readout functions are re-initialized, enabling X-MACE to achieve mean absolute errors (MAEs) in energies and forces significantly below those attainable when training from scratch on equivalent data.
• Chemical-Space and Size Transfer: After training on the chromophore dataset, X-MACE is evaluated on five unseen chromophores. The transfer model realizes energy MAE ≈ 0.07–0.09 eV and force MAE ≈ 0.06–0.08 eV/Å, representing $\sim$30% and $\sim$15% improvements, respectively, over non-transfer baselines. This suggests the learned invariants enable broad chemical generalization.

6. Performance Comparison and Physical Accuracy

Evaluation across diverse systems demonstrates that X-MACE with a DeepSets autoencoder achieves superior accuracy relative to prior excited-state machine learning models, including SchNarc, SPaiNN, and E-MACE. Error statistics for energies, forces, and nonadiabatic couplings on chromophores, butene, ethene, and CH₂NH₂⁺ are summarized below.

Model	Chromophores (E, eV)	Butene (E, eV)	Ethene (E, eV)	CH₂NH₂⁺ (E, eV)
SchNarc	1.648	0.0759	0.0599	0.1190
SPaiNN	0.5762	0.0475	0.0279	0.1018
E-MACE	0.1326	0.0235	0.0048	0.1295
X-MACE (AE)	0.1031	0.0151	0.0041	0.0944

X-MACE consistently yields the lowest energy and force errors, often reducing these by factors of 2–10 relative to other models. It recovers conical intersections sharply—where standard MACE introduces avoided crossings—and reduces the energy-gap RMSE by nearly 50% (from 0.50 eV to 0.28 eV). Histograms of error versus true energy gap confirm systematic advantages near degeneracies.

7. Limitations and Prospects

Current challenges for X-MACE include residual instabilities precisely at degeneracies (eigenvalue sensitivity), sign ambiguities in couplings, and the polynomial computational cost of the companion-matrix decoder for large numbers of states. Future research directions include adaptive, uncertainty-driven sampling of conical seams, extension to higher-lying electronic manifolds (including spin-orbit and triplet states), incorporation of explicit diabatization or symmetry constraints, and scalable decoder architectures for handling >10 states.

X-MACE demonstrates that integrating advanced equivariant message passing with learned permutational invariants enables accurate excited-state predictions in domains previously inaccessible to machine learning, including large-scale photodynamics simulations and the study of complex molecular materials (Barrett et al., 18 Feb 2025).