GeoOpt-Net: One-Shot DFT Geometry Optimization

• GeoOpt-Net is a SE(3)-equivariant neural network that refines molecular structures to DFT-level quality in a single-shot optimization.
• It integrates multi-branch message passing of 2-, 3-, and 4-body features with fidelity-aware feature modulation for scalable accuracy.
• Benchmarks show sub-milli-Å RMSD and near-zero energy deviations, dramatically reducing DFT optimization steps and wall-clock time.

GeoOpt-Net is a multi-branch SE(3)-equivariant neural architecture designed for one-shot refinement of molecular geometries to density functional theory (DFT) quality. It predicts B3LYP/TZVP-level structures directly from initial conformers produced with inexpensive force-field methods. Trained through a two-stage protocol—consisting of low-fidelity pre-training and high-fidelity fine-tuning with a fidelity-aware feature modulation mechanism—GeoOpt-Net facilitates rapid, accurate, and physically consistent geometry optimizations. The network achieves sub-milli-Å all-atom root mean square deviation (RMSD), near-zero single-point energy errors, and demonstrably reduces wall-clock time and DFT optimization steps when integrated into quantum-chemical screening workflows (Liu et al., 30 Jan 2026).

1. Network Design and SE(3)-Equivariance

GeoOpt-Net consumes an initial 3D molecular conformer and its associated graph representation $G=(V,E)$. Its architecture comprises three SE(3)-equivariant message-passing streams, each encoding:

• 2-body features (bond lengths $r_{ij}$),
• 3-body features (bond angles $\theta_{ijk}$),
• 4-body features (dihedral angles $\phi_{ijkl}$).

Scalar features ($\ell=0$) are expanded via radial basis functions, while directional components ($\ell\geq1$) involve real spherical harmonics $Y^{(\ell)}(\hat{r}_{ij})$. Message formation leverages Clebsch–Gordan tensor products to rigorously ensure equivariance under rotation and translation: $$m_{ij}^{(\ell)} = \sum_{\ell_1, \ell_2} \left[ h_i^{(\ell_1)} \otimes Y^{(\ell_2)}(\hat{r}_{ij}) \right]_\text{CG} \cdot \phi(r_{ij})$$ Nonlinear activations (GELU) and LayerNorm are used only on scalar channels. The outputs of all streams are fused with a lightweight Transformer decoder, producing a latent $F_\theta(G, R_\text{initial}, d)$. The SE(3)-equivariant coordinate update is

$$R_\text{refined} = R_\text{initial} + F_\theta(G, R_\text{initial}, d)$$

guaranteeing that the network output transforms consistently with the input geometry under arbitrary rigid motions.

2. Fidelity-Aware Feature Modulation (FAFM) and Training Protocol

GeoOpt-Net introduces “fidelity-aware feature modulation” to inject theory and basis-set dependence into the model without full retraining. Each hidden feature vector $r_{ij}$0 in the message-passing layers is modulated as: $r_{ij}$1 where $r_{ij}$2 is a one-hot domain embedding (e.g., “6-31G(2df,p)” vs. “TZVP”), and $r_{ij}$3, $r_{ij}$4 are small, learnable vectors. The training consists of two stages:

• Stage 1 (Pre-training): On 290k QM9+QM40 molecules at B3LYP/6-31G(2df,p), with $r_{ij}$5 (no modulation), optimizing a composite loss:

  $r_{ij}$6

• Stage 2 (Fine-tuning): FAFM is enabled ($r_{ij}$7="TZVP"), optimizing $r_{ij}$8, $r_{ij}$9 and new output layers on 180k QMe14S molecules at B3LYP/TZVP, using the same loss but referencing high-fidelity targets.

All optimization is performed with AdamW ($\theta_{ijk}$0, weight decay $\theta_{ijk}$1, batch size 64, epoch-based learning rate decay and gradient clipping).

3. Benchmarking and Accuracy

On the ZINC20 benchmark (N=1,000), GeoOpt-Net surpasses classical and ML baselines—UMA, xTB, Auto3D, RDKit—by a wide margin. Key results include:

• All-atom RMSD: GeoOpt-Net distribution sharply peaks at $\theta_{ijk}$2 Å, while baselines typically range $\theta_{ijk}$3–$\theta_{ijk}$4 Å.
• Single-point energy deviations ($\theta_{ijk}$5): GeoOpt-Net errors are tightly clustered around 0 kcal/mol ($\theta_{ijk}$6), where baselines reach several kcal/mol.
• Error decomposition: Bonds: $\theta_{ijk}$7 Å (GeoOpt-Net) vs. $\theta_{ijk}$8–$\theta_{ijk}$9 Å (baselines); Angles: $\phi_{ijkl}$0 vs. $\phi_{ijkl}$1–$\phi_{ijkl}$2; Dihedrals: $\phi_{ijkl}$3 vs. $\phi_{ijkl}$4–$\phi_{ijkl}$5.

Molecule	Method	RMSD (Å)	$\phi_{ijkl}$6 (kcal/mol)
Ex1	GeoOpt-Net	0.0001	0.002
	UMA	0.5718	0.625
	xTB	1.1529	5.894
	Auto3D	0.9740	2.461
	RDKit	0.9830	2.466

Performance is robust on drug-like molecules with up to 20 rotatable bonds and 40 heavy atoms, maintaining $\phi_{ijkl}$7 kcal/mol while baseline errors increase significantly (Liu et al., 30 Jan 2026).

4. DFT Convergence and Workflow Acceleration

GeoOpt-Net output structures intrinsically match DFT (B3LYP/TZVP) optimization convergence thresholds:

• Convergence metric satisfaction: ~40–58% of geometries meet each DFT convergence criterion (baselines: ≈0%).
• “All-YES” rate: 65.0% under loose, 33.4% under default criteria (baselines: 0% in both).
• DFT re-optimization: Average optimization steps are halved and wall-clock time is reduced by ~60% starting from GeoOpt-Net versus UMA/xTB/Auto3D/RDKit initial guesses.
• High-throughput compatibility: Single-shot refinement eliminates the need for pre-optimization or force-field minimization loops in screening settings.

5. Electronic Property Preservation

Sub-milli-Å accuracy leads to near-exact reproduction of electronic observables:

• Dipole moments ($\phi_{ijkl}$8) at B3LYP/TZVP: Reference, 3.165 D; GeoOpt-Net, 3.167 D ($\phi_{ijkl}$9 D); baselines, errors range from –0.369 to –0.498 D.
• Energy scaling: Energy deviation and RMSD remain nearly invariant under increased molecular complexity, while baseline errors escalate.
• Conclusion: High geometric fidelity directly translates to reliable electronic-structure descriptors.

6. Implementation and Limitations

• Implementation: Developed in PyTorch with e3nn for equivariant operations and a custom Transformer decoder. Training set sizes: 290k (pre-training), 180k (fine-tuning). Validation uses 1,000 ZINC20 molecules. Training duration: ~48 hours on multi-GPU A100 cluster.
• Model configuration: Scalar channel dim 256, vector channel dim 64, radial basis size 64, $\ell=0$0.
• Known limitations: Current model is restricted to neutral, closed-shell organic molecules; FAFM supports two fidelities (“6-31G(2df,p)”, “TZVP”) and is not immediately extensible to more advanced levels (e.g., MP2, CCSD) without further modifications. Memory and computational cost scales with $\ell=0$1, potentially impacting performance on molecules exceeding 50 heavy atoms.

7. Assessment and Outlook

GeoOpt-Net enables the replacement of traditional, iterative DFT geometry optimization with a single-shot, SE(3)-equivariant neural approach that is robust, accurate, and scalable. The paradigm provides direct DFT-ready geometries, preserves electronic features, and accelerates quantum-chemical screening without the need for iterative force-field or machine learning pre-optimizations. Achieving sub-milli-Å accuracy and near-zero energy deviation across a range of chemical complexity, GeoOpt-Net significantly increases computational throughput for both method developers and practitioners in molecular design (Liu et al., 30 Jan 2026). Extensions to support a broader class of molecules and theoretical fidelities represent immediate future directions.