Source-linked AI summary
Neural-Network Solutions to Real-Space Charge Density and Generalization
Yuxuan Zeng, Taoyuze Lv, Zhicheng Zhong
TL;DR
Learning real-space charge density offers a route to electronic-structure surrogates, while conventional density-functional calculations require iterative self-consistent-field solutions. AIDEN decomposes density into local and environment-dependent terms and reconstructs a continuous field with an equivariant encoder and Gaussian decoder. It reports state-of-the-art performance on periodic-crystal benchmarks, competitive molecular results, and transfer to several out-of-distribution systems.
Problem
Deep-learning surrogates for electronic structure need to predict ground-state charge density, a real-space target that also gives direct access to charge transfer and chemical bonding.
Source: §1 · pp. 1–2Method
Results
AIDEN achieves state-of-the-art performance on periodic-crystal benchmarks, ranks second on QM9 behind the molecularly designed BOA, and shows strong but system-dependent out-of-distribution results.
Source: §5.2 · p. 6, §5.3 · pp. 7–8Takeaways & Limitations
Reusing atomic encodings across spatial query points supports faster inference than probe-based methods and makes density reconstruction practical for larger systems.
Source: §1 · p. 2AIDEN currently focuses on ground-state scalar charge densities; extensions to spin-resolved densities and broader electronic-structure settings remain future work.
Source: §6 · p. 9
Abstract
from arXiv · showhide
The Hohenberg-Kohn theorem establishes that, in principle, the ground state (GS) charge density contains all GS information of a many-electron system, such that all GS observables can be expressed as functionals of the GS charge density. Conventional Kohn-Sham density functional theory requires iterative solution of the self-consistent-field equations at substantial computational cost, motivating the development of deep learning surrogates for electronic structure calculations and, in turn, accelerating computer-aided materials design. Here, we propose AIDEN, an Atomic-Interaction Density Equivariant Network for solving real-space charge density. AIDEN separates the element-dependent one-center density from environment-induced density redistribution and represents the latter through complementary atom- and edge-centered tensor correlations. A continuous low-rank Gaussian decoder then reconstructs the density at arbitrary spatial coordinates while reusing atomic encodings independently of the evaluation grid. AIDEN achieves state-of-the-art accuracy on periodic crystal benchmarks while remaining competitive for molecular systems, and further demonstrates zero-shot transferability across several structurally distinct out-of-distribution case studies. Furthermore, AIDEN provides substantially faster inference than both baseline models and full SCF calculations, enabling efficient charge density reconstruction for large-scale electronic structure calculations.
1 INTRODUCTION
Real-space charge density is a basis-independent and physically interpretable target for learning electronic structure. AIDEN predicts it by separating an element-dependent local baseline from environment-induced density changes.
1 INTRODUCTION
- Deep-learning approaches to DFT commonly target charge density, the Hamiltonian, or the density matrix, which are connected through the self-consistent-field cycle.KS-DFT calculations and molecular-dynamics simulations provide data for this data-driven materials-discovery paradigm.Source: §1 · p. 1
- Real-space charge density is independent of orbital indexing, unlike Hamiltonian and density-matrix representations that change with the atomic-orbital basis.This supports a more unified representation across different basis sets.Source: §1 · p. 1
- Charge density provides direct access to charge transfer and chemical bonding, offering a microscopic basis for interpreting material properties.Its regular three-dimensional grid representation can also serve as a complementary descriptor for downstream representation learning.Source: §1 · p. 2
- AIDEN decomposes charge density into an element-dependent one-center baseline and environment-dependent corrections modeled with higher-order geometric features.Its encoder combines Cartesian ACE and tensor edge cluster expansion, and its Gaussian-type-orbital decoder reconstructs a continuous density field.Source: §1 · p. 2
2 RELATED WORKS
Related work spans equivariant atomistic representations and several approaches to real-space charge-density prediction. AIDEN draws on higher-order many-body and edge-aligned representations within this landscape.
2 RELATED WORKS
- Equivariant atomistic methods preserve geometric transformation laws, while later approaches incorporate higher-order many-body interactions or edge-aligned SO(2) operations.The cited progression includes Tensor Field Networks, e3nn, NequIP, MACE, eSCN, and TECE.Source: §2 · p. 2
- Charge-density methods predict values at spatial queries, on regular grids, or as atom-centered contributions, while NeuralSCF learns the Kohn-Sham density map through neural self-consistent iterations.Examples include DeepDFT, Deep Charge, ChargE3Net, Li et al., and EAC-Net.Source: §2 · p. 2
3 PROBLEM STATEMENT
The problem is to learn a periodic system’s ground-state real-space charge density from its atomic configuration. AIDEN uses a periodic graph and targets a continuous field that respects geometric covariance and lattice periodicity.
3 PROBLEM STATEMENT
- AIDEN represents the periodic atomic system as a graph carrying elemental embeddings, interatomic distances, and edge directions.Its density construction separates an initial local contribution into the two species’ contributions in the Na–Cl illustration.Source: Figure 1 · p. 3
- The learning target is the real-space density on a regular fast Fourier transform grid, with the corresponding electron count denoted Ne.The density is the ground-state quantity associated with the periodic structure.Source: §3 · pp. 2–3
- The model learns a continuous scalar density field with target covariance under Euclidean transformations and periodicity under lattice translations.Training minimizes a grid-averaged absolute deviation.Source: §3 · p. 3
4 MODEL ARCHITECTURE
AIDEN initializes equivariant features from periodic atomic neighborhoods, encodes both atom-centered and edge-centered correlations, then decodes a continuous density as local and environment-dependent terms.
4.1 PERIODIC GEOMETRY AND EQUIVARIANT INITIALIZATION
- Periodic incoming edges are selected within a cutoff, with only the nearest Nnbr images retained when the neighborhood exceeds that limit.Each edge carries a radial vector and an irreducible Cartesian harmonic of its direction.Source: §4.1 · p. 3
- The geometry-informed embedding initializes scalar and higher-order equivariant fields from element features and the atom’s radial environment.Learned radial amplitudes and source-element amplitudes set neighbor contribution strengths, while analytic harmonics determine angular dependence.Source: §4.1 · p. 4
4.2 EQUIVARIANT CLUSTER EXPANSION ENCODER
- The encoder combines atom-centered collective coordination from Cartesian ACE with direction-sensitive source-target correlations from TECE.The two interactions provide complementary geometric inductive biases for density redistribution.Source: §4.2 · p. 4
4.2.1 CARTESIAN ATOMIC CLUSTER EXPANSION
- Cartesian ACE forms higher-order correlations after neighborhood aggregation and applies invariant-sector channel-wise gating to the accumulated atomic field.Its construction uses irreducible Cartesian tensor couplings and correlation polynomials.Source: §4.2 · p. 4, §4.2.1 · p. 4
4.2.2 TENSOR EDGE CLUSTER EXPANSION AND RADIAL ROTARY ATTENTION
- TECE rotates features into edge-aligned local frames, forms edge correlations, and uses radial rotary attention to modulate their aggregation.The radial phase adjusts relative alignment while preserving the SO(2) transformation law.Source: §4.2.2 · p. 5
4.3 CONTINUOUS DENSITY DECODER
- The continuous decoder reconstructs density from element-dependent one-center and environment-dependent terms, both expanded in atom-centered Gaussian-type orbitals.The environment term uses encoded atomic features; the costly equivariant encoding can be reused across query points.Source: §4.3 · pp. 5–6
5 EXPERIMENTS
Experiments benchmark AIDEN on periodic crystals and molecules, then test zero-shot transfer to structurally distinct out-of-distribution systems. AIDEN leads the periodic-crystal benchmark and shows competitive accuracy and faster inference across the reported evaluations.
5.1 DATASET SETTINGS
- The ECD dataset contains 140,646 inorganic structures spanning 94 elements; QM9 contains 133,885 molecules, and OOD tests include water, AlMg, a-Si, and TBG.AIDEN is trained from scratch on the PBE subset and fine-tuned on a smaller HSE subset.Source: §5.1 · p. 6
5.2 EVALUATION ON BENCHMARK DATASETS
- At matched angular order L, AIDEN has substantially faster per-sample inference than ChargE3Net and BOA; its advantage increases with system size.Figure 2 compares inference time, parameter counts, and HSE transfer-learning strategies.Source: §5.2 · p. 6, Figure 2 · p. 7
- AIDEN achieves state-of-the-art performance under both exchange-correlation functionals on periodic crystals and ranks second on QM9, behind molecularly specialized BOA.The benchmark reports NMAE ερ (%); lower values indicate better performance.Source: §5.2 · p. 6, Table 1 · p. 7
- On QM9, BOA's advantage may reflect its molecularly tailored basis, radial corrections, and density-basis message passing; PBE pretraining modestly improves AIDEN's HSE performance.The paper attributes the HSE learning challenge partly to the smaller training set and HSE's screened exact exchange.Source: §5.2 · pp. 6–7
5.3 GENERALIZATION TO OOD SYSTEMS
- On full-grid OOD tests, AIDEN has the lowest density error on AlMg and a-Si, while ChargE3Net is slightly better on water; AIDEN runs in 4.34–23.70 s, 48–76× faster than the evaluated ChargE3Net implementation.Both learned models substantially outperform SAD in density accuracy across all three systems.Source: §5.3 · pp. 7–8, Table 2 · p. 8
- Fixed-density energy and force comparisons differ by system: AIDEN has lower energy errors on water and a-Si and lower or nearly identical force errors there, while ChargE3Net performs better on AlMg.The paper cautions that lower density NMAE does not necessarily imply smaller energy or force errors.Source: §5.3 · p. 8
- Within E_Fermi ± 5 eV, aligned band MAEs are 8.42, 7.41, and 11.47 meV for AIDEN versus 27.35, 10.30, and 12.93 meV for ChargE3Net at 21.79°, 13.17°, and 9.43°, respectively.These aligned errors measure residual bandshape and dispersion after removing one global energy offset per model and cell.Source: §5.3 · pp. 8–9, Figure 4 · p. 9
6 DISCUSSION
AIDEN combines a learnable density decomposition with an equivariant encoder and continuous decoder, with reported accuracy, transfer, and scalability across several settings. The discussion also explains density self-consistency in KS calculations and identifies broader density and functional coverage as future work.
6 DISCUSSION
- AIDEN reports high accuracy on periodic-crystal and molecular benchmarks, PBE-to-HSE fine-tuning, zero-shot transfer to structurally distinct systems, and full-grid inference scaling to millions of points.Fixed-density calculations also assess energies, forces, and TBG band structures; TBG electron counts deviate by 0.0068–0.0095% without explicit normalization.Source: §6 · p. 9
- Future work may extend AIDEN beyond ground-state scalar densities to spin-resolved densities, broader electronic-structure settings, and joint learning of densities with other DFT observables.The proposed broader settings include unified models across exchange-correlation functionals and chemical spaces.Source: §6 · p. 9
A.1 KOHN-SHAM DENSITY FUNCTIONAL THEORY AND CHARGE DENSITY SELF-CONSISTENCY
- The orbital-basis matrix representations of the Hamiltonian and density matrix change with basis choice, whereas the reconstructed real-space density does not.This distinction is identified as one motivation for directly learning the real-space charge density.Source: §A.1 · p. 15
- In KS theory, density both helps determine the effective Hamiltonian and is reconstructed from its eigenstates, making the solution a self-consistency problem.For hybrid functionals such as HSE, the operator also depends on off-diagonal one-particle density-matrix information.Source: §A.1 · pp. 15–16
A.2 DENSITY TARGET AND FIXED-DENSITY EVALUATION
This section defines AIDEN’s density target and fixed-density evaluation, including its atomic-plus-environment decomposition, angular basis, and geometric transformation behavior. It also examines whether the mapping from periodic atomic structures to continuous scalar charge-density fields is E(3)-consistent and lattice-periodic.
A.2 DENSITY TARGET AND FIXED-DENSITY EVALUATION
- AIDEN predicts VASP’s FFT-grid density; OOD property energies, forces, and spectra come from fixed-density VASP calculations using the unrenormalized prediction with ICHARG=11.The predicted grid density integrates to the corresponding electron count; AIDEN and ChargE3Net use identical restart quantities and electronic settings within each structure.Source: §A.2 · p. 16
A.3 THEORY OF SUPERPOSITION OF ATOMIC DENSITIES
- AIDEN generalizes SAD as a learnable decomposition: a pretrained, jointly optimized element-dependent one-center density plus an environment correction for redistribution absent from isolated atoms.Unlike fixed SAD, both contributions are optimized from self-consistent reference densities.Source: §A.3 · pp. 16–17
A.4 SPHERICAL-HARMONIC REPRESENTATION
- AIDEN represents angular structure with analytic symmetric-traceless Cartesian tensors, equivalent to a compact real-spherical basis, with 2ℓ+1 independent components and parity (-1)^ℓ.The Cartesian form supports tensor products and contractions, while the real-spherical form supports edge-aligned operations and density decoding.Source: §A.4 · p. 18
A.5 EQUIVARIANCE OF AIDEN
- AIDEN’s encoder is analytically equivariant to proper rotations and invariant to translations; its predicted density is consistent under joint structure and query-point transforms and periodic under lattice shifts.GaAs tests cover translation, proper rotation, and inversion, including inversion beyond the crystal’s space-group symmetries.Source: §A.5 · pp. 21–22, Figure 7 · p. 22
- Normalized L1 deviations were 7.72×10−6 after translation, 7.27×10−5 after proper rotation, and 9.51×10−7 after inversion.Fig. 7 compares the original prediction with predictions queried at correspondingly transformed spatial coordinates.Source: §A.5 · p. 22
- Inversion of zinc-blende GaAs tests the transformation law because inversion is not a crystal symmetry of that structure.The inverted structure is related to the original by a Euclidean transformation, but not by a space-group operation.Source: §A.5 · p. 22
- Together with exact translation invariance and the analytical SO(3) derivation, the inversion experiment numerically verifies the O(3) extension within measured tolerance.The passage concludes that AIDEN realizes an E(3)-consistent and lattice-periodic mapping from periodic structures to continuous scalar charge-density fields.Source: §A.5 · p. 22
B DETAILED IMPLEMENTATIONS OF AIDEN
The implementation specifies AIDEN’s periodic geometric basis, equivariant interaction blocks, continuous Gaussian decoder, and forward pass. The encoder is evaluated once per structure, after which the decoder can query density at arbitrary positions.
B DETAILED IMPLEMENTATIONS OF AIDEN
- The implementation section retains constructions defining the periodic geometric basis, two equivariant cluster-expansion interactions, and the continuous GTO decoder, omitting numerical bookkeeping that does not affect the model definition.The stated scope is reproducibility of AIDEN’s architecture.Source: B DETAILED IMPLEMENTATIONS OF AIDEN · pp. 22–23
B.1 PERIODIC GRAPH AND GEOMETRIC BASIS
- The elemental input combines a one-hot element vector with 15 descriptors into 133 normalized features, while the periodic graph retains up to 100 neighbors within a 4 Å atomic cutoff.The radial edge basis uses eight components, shared by GIE, Cartesian ACE, and TECE.Source: §B.1 · p. 23
B.2 CARTESIAN ATOMIC CLUSTER EXPANSION
- Cartesian ACE uses explicit correlation order ν=2, building on the preceding GIE representation’s one-hop environmental information.Source: §B.2 · p. 24
B.3 TECE AND RRA
- With L=M=3, TECE retains 16 components; its direct, gated, and product branches preserve local SO(2) equivariance, while RRA uses invariant equal-frequency contractions.The implementation uses four RRA heads and 48 channels.Source: §B.3 · p. 24
B.4 CONTINUOUS GTO DECODER
- The corrected localized GTO basis yields a continuous density independent of the evaluation grid; decoder factors are auxiliary fields, and the jointly optimized split does not impose a unique atomic partition.The same corrected radial basis serves the one-center and environment-dependent branches.Source: §B.4 · pp. 24–25
B.5 FORWARD PROPAGATION
- Forward propagation encodes the structure once before evaluating density at arbitrary query positions; training briefly pretrains ρinit before jointly optimizing the full model.Source: §B.5 · p. 25
C.1 ABLATION STUDIES
AIDEN’s ablations support the full design: density-decoder choices produce the largest validation penalties, and ACE→TECE outperforms the reversed interaction order.
C.1 ABLATION STUDIES
- Removing ρinit raises best validation NMAE by 35.70%, while fixing the GTO basis raises it by 7.38%.These are the largest ablation penalties; the one-center density provides an element-dependent baseline, while the learnable radial correction adapts GTO channel shapes.Source: §C.1 · p. 26, Table 3 · p. 26
- Reversing ACE→TECE to TECE→ACE raises NMAE by 1.26% at the same parameter count.The proposed order first models cross-neighbor correlations, then refines directional information in edge-aligned frames.Source: §C.1 · p. 26
- Replacing RRA with fixed cutoff-normalized weights raises NMAE by 2.62%, and removing GIE raises it by 1.30%.RRA adaptively selects edge messages; GIE supplies higher-order equivariant initial features from one-hop geometry.Source: §C.1 · p. 26
- Across the ablations, ρinit and learnable GTO radial corrections contribute the largest accuracy gains, while ACE→TECE combines atom- and edge-centered interactions.All variants were trained on the same 2,000-structure PBE subset; validation NMAE was evaluated on 256 structures.Source: §C.1 · p. 26, Table 3 · p. 26
C.2 TRAINING DETAILS
AIDEN was trained on PBE crystals and QM9 molecules across angular orders L=0–4, with loss curves and error distributions reported for both datasets.
C.2 TRAINING DETAILS
- Models with maximum angular orders L=0–4 were trained on PBE and QM9, using Adam, batch size 12, and learning rates of 2 × 10^-4 for the encoder and 1 × 10^-3 for the decoder.Training used gradient clipping at 0.5 and a single NVIDIA A100-40GB GPU.Source: §C.2 · p. 26
- Validation and inference used exponential moving averages with decay 0.995.The same settings were applied across the angular-order models.Source: §C.2 · p. 26
- The reported figures show training and validation loss curves and error distributions for PBE and QM9.Violin plots distinguish training from validation errors and mark means, interquartile ranges, and 5%–95% quantiles.Source: Figure 8 · p. 27
C.3 TWISTED BILAYER GRAPHENE
Across twisted bilayer graphene (TBG) cells, AIDEN improves density, fixed-density energy and force estimates, and full-grid inference time relative to ChargE3Net. The accompanying OOD analyses specify matched DFT settings, density evaluation, cross-section comparisons, and timing procedures.
C.3 TWISTED BILAYER GRAPHENE
- AIDEN yields lower ερ at all three TBG angles (0.2892–0.2934% vs 0.3134–0.3160% for ChargE3Net) and smaller electron-count deviations without charge renormalization.Relative charge error decreases from 0.00952% to 0.00679% as cell size grows, despite increasing absolute count deviation.Source: §C.3 · p. 27, Table 5 · p. 28, Figure 10 · p. 29
- AIDEN’s per-atom energy errors are 0.753, 0.747, and 0.741 meV/atom, about 28% of ChargE3Net’s; force-component MAEs are approximately 0.021–0.022 versus 0.040–0.041 eV/Å.AIDEN also reduces component RMSE and maximum per-atom force-vector error in all three cells.Source: §C.3 · p. 28, Table 6 · p. 29, Figure 10 · p. 29
- AIDEN has lower aligned band MAE at every angle, but ChargE3Net has smaller tail-error metrics for some cells.AIDEN’s largest aligned-MAE reduction is at 21.79°, from 27.35 to 8.42 meV; band errors use a rigid shift per model and cell.Source: Figure 10 · p. 29, §C.3 · p. 29, Table 8 · p. 30
C.4.1 STRUCTURES AND DFT SETTINGS
- OOD reference densities and properties were recomputed with PBE-SCF for water, AlMg, and a-Si using matched structures, grids, and electronic settings across methods.The calculations used fixed atomic positions and a 520 eV cutoff.Source: §C.4.1 · p. 30
C.3 TWISTED BILAYER GRAPHENE
- Full-grid inference is 9.71, 10.85, and 15.31× faster for AIDEN than ChargE3Net across the three TBG cells.AIDEN reuses atomic encodings across grid queries; its measured inference time rises from 13.27 to 41.43 s as the grid grows from 2.76 to 15.05 million points.Source: Figure 10 · p. 29, §C.3 · p. 30, Table 9 · p. 31
C.4.2 DENSITY AND PROPERTY EVALUATION
- Candidate densities were evaluated without charge renormalization, then held fixed for VASP energy and force calculations.AIDEN used PBE-pretrained parameters, ChargE3Net its public pretrained checkpoint, and SAD was generated with ICHARG=12.Source: §C.4.2 · p. 31
C.4.3 CROSS SECTIONS
- Cross-section NMAEs for AIDEN/ChargE3Net/SAD are 1.887/1.641/12.103% for water, 1.183/1.648/11.405% for AlMg, and 1.302/1.618/9.772% for a-Si.AIDEN is lowest for AlMg and a-Si; ChargE3Net is lowest for water.Source: §C.4.3 · p. 31
C.4.4 TIMING
- OOD model inference was timed on one NVIDIA A100 after warm-up, excluding synchronization and model-loading overhead; SAD timing includes the complete ICHARG=12 calculation.Model comparisons use full-grid per-sample inference, with ChargE3Net run through its official pipeline.Source: §C.4.4 · p. 31
C.4.5 LOWER NMAE DOES NOT GUARANTEE LOWER DOWNSTREAM ERRORS
Global density errors, spatially localized density errors, and downstream energy or force errors can rank models differently. Density norms are therefore spatial diagnostics, not direct predictors of downstream response.
C.4.5 LOWER NMAE DOES NOT GUARANTEE LOWER DOWNSTREAM ERRORS
- Full-grid NMAE integrates absolute density error uniformly, whereas energy and force errors depend on the error’s location and electronic-structure weighting.Force weights are vector fields, so errors on different sides of an atom can reinforce or cancel; fixed-density observables also include orbital response, while Eqs. 63–64 isolate only the local contribution.Source: §C.4.5 · pp. 31–32
- In water, AIDEN has the higher DR (4.5969 vs 4.1411 e) but the lower total force-component MAE (0.11467 vs 0.12538 eV/Å) than ChargE3Net.AIDEN’s H force-component MAE is 0.08370 eV/Å versus 0.10303 eV/Å for ChargE3Net, outweighing its slightly higher O error.Source: §C.4.5 · p. 32
- For a-Si, AIDEN’s DR is 0.1841 e versus 0.2437 e for ChargE3Net, while force-component MAEs are nearly identical at 0.04087 versus 0.04090 eV/Å.AIDEN also has the smaller energy error, 3.213 versus 4.774 meV/atom.Source: §C.4.5 · p. 32
- For AlMg, ChargE3Net has the lower DR (0.1520 vs 0.2285 e) and maximum pointwise error (0.01001 vs 0.02456 e/Å^3), yet AIDEN has lower global density NMAE.The different density metrics emphasize different spatial characteristics and need not rank downstream energy or force errors the same way.Source: §C.4.5 · p. 32
- For water, AIDEN’s and ChargE3Net’s signed total-energy errors have opposite signs, and energy-component cancellation means a smaller total error does not imply every component is more accurate.Near self-consistency, the stationary Harris-Foulkes construction cancels first-order energy variations for sufficiently small electron-number-conserving perturbations.Source: §C.4.5 · p. 32