Transferable machine learning of excited-state dynamics with extremal pooling

This paper introduces a size-intensive machine learning framework based on extremal pooling of atomic HOMO and LUMO contributions, which enables accurate and scalable simulations of excited-state dynamics in complex periodic systems like solvated electrons in liquid water, overcoming the computational limitations of traditional *ab initio* methods.

The Gist

The Big Problem: Simulating Light-Driven Chemistry

Imagine trying to predict what happens when a drop of water is hit by a flash of UV light. The water molecules get excited, break apart, and create new, reactive particles. This is the kind of chemistry that powers photosynthesis, vision, and solar energy.

To study this on a computer, scientists usually use "ab initio" methods. Think of these as super-accurate, high-definition cameras that take a picture of every single electron in the system. The problem? These cameras are incredibly slow and expensive. They can only take pictures of tiny, short-lived scenes (like a small group of 64 water molecules for a fraction of a second). If you try to simulate a larger pool of water or watch it for longer, the computer crashes because the math is too heavy.

The Old Solution: The "Summing" Mistake

In recent years, scientists have used Machine Learning (AI) to speed things up. Usually, these AI models work like a Lego builder. They look at individual bricks (atoms), calculate the energy of each brick, and then add them all up to get the total energy of the wall.

This works great for normal chemistry (ground state). But for excited states (when light hits the molecule), this "adding up" logic breaks.

• The Analogy: Imagine a spotlight shining on a stage. The brightness of the spotlight (the excitation energy) doesn't get brighter just because you add more empty seats to the audience. It's an intensive property; it depends on the peak of the light, not the total number of people.
• The Failure: Old AI models tried to "add up" the energy of every atom. This meant that if you simulated a bigger box of water, the AI thought the energy of the light hitting it would change, which is physically wrong. This made the models useless for large systems.

The New Solution: The "Peak Finder"

The authors of this paper created a new AI framework that fixes this problem. Instead of adding up every atom, their model acts like a talent scout looking for the best performer.

1. The Scout (The AI): The model looks at every atom in the water and asks, "How much does this atom contribute to the 'Highest Occupied' energy level (HOMO) and the 'Lowest Unoccupied' level (LUMO)?" These are the two energy levels that determine how the molecule reacts to light.
2. The Extremal Pooling (The Rule): Instead of summing the scores, the model uses a special rule called Extremal Pooling.
   • For the HOMO, it finds the highest score (the "SmoothMax").
   • For the LUMO, it finds the lowest score (the "SmoothMin").
   • It then subtracts these two to find the energy gap.
3. Why it Works: Because the model only cares about the extreme values (the best and worst contributors), adding more water molecules to the simulation doesn't change the result. The "spotlight" stays the same brightness regardless of how big the room is. This allows the model to be transferable: you can train it on a small group of water molecules, and it will work perfectly on a huge ocean of them.

The Test: The "Solvated Electron" in Water

To prove their idea worked, the team simulated what happens when liquid water is hit by UV light.

• The Scenario: When water gets excited, it can break apart in two main ways:
  1. HAT (Hydrogen Atom Transfer): A single hydrogen atom gets kicked out like a bullet.
  2. PCET (Proton-Coupled Electron Transfer): The water splits into a proton, a radical, and an electron that gets trapped in a tiny bubble (a "solvated electron").
• The Result: The new AI model successfully predicted both pathways. It didn't just guess the outcome; it learned to "see" where the electron was hiding by looking at which atoms had the lowest energy scores.
• The Scale: While traditional methods could only simulate 64 water molecules for a tiny fraction of a second, this new AI simulated 512 water molecules for much longer.

What They Discovered

By running these larger simulations, they found something interesting about size:

• The Ratio: The mix of the two breaking-apart methods (HAT vs. PCET) stayed mostly the same, which is good news.
• The Timing: However, the speed at which these reactions happened changed with the size of the system. In larger boxes of water, the reactions took slightly longer.
  • Why? In a small box, the broken pieces (like the electron and the proton) bump into the "walls" of the simulation quickly. In a large box, they have more room to drift apart, which changes how long the excited state lasts.

The Bottom Line

The paper presents a new way to teach AI to understand excited chemistry. By changing the math from "adding everything up" to "finding the extremes," they created a model that is accurate, fast, and can handle systems of any size. This allows scientists to study complex photochemical processes (like how water reacts to light) in realistic, large-scale environments that were previously impossible to simulate.

Technical Summary

Technical Summary: Transferable Machine Learning of Excited-State Dynamics with Extremal Pooling

Problem Statement
Accurate simulation of photochemical processes, such as solar energy conversion and atmospheric chemistry, requires modeling excited-state potential energy surfaces (PES) over extended timescales and system sizes. Traditional ab initio methods like CASSCF, CASPT2, and Time-Dependent Density Functional Theory (TDDFT) are computationally prohibitive for condensed-phase systems and long trajectories. While Machine Learning Interatomic Potentials (MLIPs) have revolutionized ground-state simulations by learning total energies as sums of local atomic contributions (assuming energy extensivity), this paradigm fails for excited states. Electronic excitation energy is an intensive property; it does not scale linearly with system size. Standard local descriptors (e.g., Coulomb matrices) either lack size transferability or rely on the extensivity assumption, making them ill-defined for states where the relevant physics involves an electronic degree of freedom not tied to a specific nucleus, such as the solvated electron in liquid water.

Methodology
The authors propose a size-intensive machine learning framework that decouples the prediction of the excited-state energy gap from the extensive ground-state energy. The core innovation is the use of extremal pooling on predicted atomic contributions to frontier molecular orbitals.

1. Theoretical Framework:

   • Inspired by frontier molecular orbital theory, the model predicts the Highest Occupied Molecular Orbital (HOMO) and Lowest Unoccupied Molecular Orbital (LUMO) energies separately.
   • A Graph Neural Network (GNN), specifically the Point Edge Transformer (PET) architecture, maps atomic configurations to per-atom contributions ($h_i^{\text{HOMO}}$ and $h_i^{\text{LUMO}}$).
   • Instead of summing these contributions (which would yield an extensive quantity), the system-level orbital energies are obtained via extremal pooling:
     • $E_{\text{HOMO}} = \text{SmoothMax}(\{h_i^{\text{HOMO}}\})$
     • $E_{\text{LUMO}} = \text{SmoothMin}(\{h_i^{\text{LUMO}}\})$
   • The excitation energy (energy gap) is calculated as $E_{\text{gap}} = E_{\text{LUMO}} - E_{\text{HOMO}}$. Because the maximum and minimum operators select specific local environments rather than summing over all atoms, $E_{\text{gap}}$ is naturally size-intensive.
   • The total excited-state energy ($E_{S1}$) is constructed by adding this intensive gap to a separate, size-extensive ground-state MLIP ($E_{S0}$): $E_{S1} = E_{S0} + E_{\text{gap}}$. Forces are derived via automatic differentiation.

2. Training and Architecture:

   • The model uses the PET architecture with two readout heads for HOMO and LUMO contributions.
   • Training data is generated via an iterative active learning strategy involving Restricted Open-Shell Kohn-Sham (ROKS) calculations.
   • The HOMO-LUMO model is trained exclusively on excitation energies ($E_{\text{gap}}$) and gap forces ($F_{\text{gap}}$), while the ground-state model is fine-tuned separately on $S_0$ energies and forces.
   • The framework was applied to photoexcited liquid water, a system exhibiting competing Hydrogen Atom Transfer (HAT) and Proton-Coupled Electron Transfer (PCET) pathways.

Key Results

• Accuracy and Transferability: The model achieves a Mean Absolute Error (MAE) of 0.137 eV for energy gaps and 53 meV/Å for forces. Crucially, it demonstrates seamless transferability across system sizes (32, 64, 128, and 512 water molecules) without retraining, a feat unattainable by standard local-descriptor MLIPs.
• Interpretability: Despite being trained only on total gap energies, the model learns physically meaningful atomic decompositions. The per-atom LUMO contributions ($h_i^{\text{LUMO}}$) successfully identify the location of the excess electron and distinguish between chemical species (e.g., OH• radicals, bulk protons) without explicit spin density training.
• Dynamics Validation: The framework reproduces the branching ratios of HAT (58%) and PCET (42%) and excited-state lifetimes in quantitative agreement with reference ROKS simulations for 64-molecule systems.
• Finite-Size Effects: By enabling simulations of 512-molecule systems (accessible only via ML), the study reveals systematic size dependencies:
  • HAT lifetimes increase from ~33 fs (64 molecules) to ~52 fs (512 molecules).
  • PCET lifetimes increase from ~415 fs to ~486 fs, with the tail of the distribution extending significantly (up to 3.4 ps in 512-molecule systems).
  • These findings suggest that while the initial branching is governed by local solvation shells, the stabilization of charge-separated states (PCET) and the timescales of decay are sensitive to long-range solvent reorganization and box size.

Significance and Claims
The paper establishes a general strategy for machine learning-driven excited-state dynamics that overcomes the fundamental extensivity assumption of standard MLIPs. By utilizing extremal pooling of frontier orbital contributions, the authors demonstrate that it is possible to predict intensive excitation energies while maintaining the linear scaling and transferability of local models.

The authors claim this work enables:

1. Quantitative agreement with high-level electronic structure methods (ROKS) for photochemical branching and dynamics.
2. Access to length and time scales (picosecond dynamics in 512-molecule systems) that are inaccessible to reference ab initio methods.
3. Explicit study of finite-size effects in photochemistry, revealing how system size influences reaction lifetimes and charge separation.
4. General applicability to diverse photochemical systems, from molecular chromophores to extended condensed-phase systems, provided the physics is dominated by frontier orbital behavior.

The work does not claim to solve all excited-state problems but specifically addresses the challenge of size-transferability in systems where the excitation is an intensive property, offering a pathway to simulate complex photochemical dynamics in condensed phases.