Skip to content
Preprint

ED-DiT: Physics-Guided Diffusion Pretraining for Transferable Molecular Representations from Electron Density

Aug 2026 · 0 citations · 41 references
Computer Science

TL;DR

ED-DiT, a physics-guided Diffusion Transformer for self-supervised pretraining on electron-density point clouds, demonstrates the effectiveness of physics-guided electron-density pretraining for learning transferable molecular representations.

Abstract

Pretraining has shown strong potential for learning transferable representations, yet it remains underexplored for electron-density-based molecular learning. Electron density provides a continuous three-dimensional description of molecular electronic structure, capturing both local spatial patterns and global physical quantities. This raises a key question: can electron-density fields be used for self-supervised pretraining to learn a shared representation that transfers across diverse electronic-structure-related tasks? We propose ED-DiT, a physics-guided Diffusion Transformer for self-supervised pretraining on electron-density point clouds. ED-DiT learns reusable representations by reconstructing corrupted and partially masked log-density fields across diffusion noise levels. An electron-number consistency constraint is further introduced to preserve the total electronic mass. The pretrained encoder can be adapted to property prediction, open-/closed-shell classification, molecule-electron-density retrieval, and molecule-conditioned electron-density prediction. Experiments on six EDBench tasks show that ED-DiT consistently outperforms the same architecture trained from scratch, especially under limited supervision. For molecule-conditioned electron-density prediction, it reduces RMSE from 2.2474 to 1.3753 and surpasses the available baseline. With only 10% labels, it improves orbital energy prediction RMSE from 0.0293 to 0.0138. These results demonstrate the effectiveness of physics-guided electron-density pretraining for learning transferable molecular representations.

View source

Similar papers

Preprint Sep 2026

SALTED: a symmetry-adapted machine-learning program for predicting electron-densities in molecules and materials

SALTED provides an open-source Python package for machine learning the quantum-mechanical electron density, $n(\mathbf{r})$, in molecular and condensed-phase systems based on input atomic coordinates and species. The program adopts a linear atom-centered decomposition of the electron density, which makes it highly transferable across diverse atomistic configurations sharing similar chemical environments. Because of this representation choice, SALTED is naturally interfaced with state-of-the-art electronic-structure programs based on atomic orbitals, namely CP2K, FHI-aims, and PySCF, from which reference electron-density data can be generated and used to train a model. The learning algorithm is based on a symmetry-adapted extension of Gaussian process regression, making SALTED especially efficient in small-data regimes. Thanks to the implementation of vector-field kernel functions, SALTED can also learn the first-order response of the electron density to applied electric fields, $\partial n(\mathbf{r})/\partial \mathbf{E}$. The application of SALTED within computational workflows has already shown its utility in a wide variety of contexts, including the calculation of polarization vectors and polarizability tensors, the accurate evaluation of Coulomb forces in QM/MM molecular-dynamics simulations, and electronic-structure studies of large-scale 2D materials.

Zekun Lou, Alan M. Lewis, Théophane Bernhard et al. · 0 citations
Preprint Jul 2026

Machine Learning Materials Properties by Encoding Orbital-Projected Density of States

Graph neural networks have become the dominant machine-learning architecture for predicting materials properties from crystal structures. Yet the initialization of atomic node features has received comparatively little attention, and conventional approaches rely on static elemental descriptors that carry no information about the quantum-mechanical electronic environment of each atom in its crystalline host. Here we show that augmenting atomic node representations with site-projected orbital density of states (pDOS) fingerprints, computed directly from density functional theory calculations, yields systematic and substantial improvements in predictive performance.These representations are fused with Pettifor elemental embeddings at each atomic site before message passing. For the superconducting critical temperature $T_c$ and the optical dielectric constant $\epsilon_{\infty}$,the pDOS augmentation reduces prediction errors by 22.9% and 27.9%, respectively, relative to the elemental-descriptor baseline. These improvements are comparable to those achieved by doubling the training-set size. The gains are, however, contingent on training-set size. For the magnetic exchange energies of Heusler compounds, a substantially smaller dataset, the improvement is reduced,indicating that pDOS augmentation is most effective when the training data exceeds the length of the pDOS feature vector. We introduce an interpretable spectral attention-gating mechanism that reveals that the model autonomously learns to prioritize the orbital channels and energy windows most physically relevant to each target property. These results establish pDOS-augmented graph nodes as a broadly applicable strategy for infusing first-principles electronic-structure knowledge into graph networks, opening a practical route to high-accuracy property prediction in data-scarce regimes.

Paulo R. Pires, Pierre-Paul De Breuck, Mauro Fava et al. · 0 citations
Preprint Jul 2026

Rem3Di: Learning smooth, chiral 3D molecular descriptors from atomistic foundation models

Rem3Di is introduced, a representation-learning framework that repurposes latent features from atomistic foundation models as transferable molecular descriptors for property prediction and virtual screening and provides a route from simulation-trained atomistic representations to transferable, chirality-aware molecular representations for chemical machine learning.

Steffen Wedig, Felix Burton, Rokas Elijošius et al. · 0 citations
Jul 2026

Learning Diffusion from Sparse Data: A Machine-Learning Bridge between Molecular Motion and Macroscopic Transport

Predicting molecular self-diffusion coefficients (D*) across chemical space remains challenging due to sparse experimental data and the high computational cost of molecular simulations. We present a data-centric machine learning framework that integrates experimental diffusion measurements with molecular dynamics simulations through a semisupervised distillation strategy. Unlike conventional approaches that treat limited experiments or simulations as direct ground truth, our method selectively incorporates simulation-derived D* only when they align with the uncertainty bounds of a Random Forest model initially trained on experimental D*. This enables controlled data set expansion, from 130 unique experimentally measured molecules to over 1,000 unique substances, while preserving label reliability. We further employ pretrained large language model embeddings to encode transferable chemical context beyond conventional descriptors, reducing predictive variance and improving generalization across chemically diverse systems. The resulting model achieves improved accuracy, with a held-out test set R2 of 0.87, an overall R2 of 0.92, and a test set mean absolute error of 0.13 × 10-9 m2 s-1 after iterative distillation across more than 1,000 chemically diverse molecules under near-ambient conditions (295-300 K). We further demonstrate that learned D* act as transferable, physically interpretable descriptors in skin permeability modeling, where their inclusion reduces test set error and improves correlation relative to models relying solely on conventional physicochemical descriptors. This work establishes a scalable framework for bridging sparse experimental measurements with broadly generalizable predictions, enabling interpretable and transferable modeling across distinct molecular transport phenomena.

Hojin Jung, Su-min Song, Sabari Kumar et al. · 0 citations
Jul 2026

GTAttn-XC: Physically constrained attention for nonlocal density functionals.

A novel exchange-correlation functional approximation model-GTAttn-XC is proposed, which introduces a learnable attention mechanism to enable unsupervised modeling of long-range electronic responses and establishes a new technical pathway toward high-accuracy nonlocal exchange-correlation approximations.

Xin Wang, Jianzhou Feng, Hao Zhang et al. · 0 citations
Open access Feb 2026

Machine learning of electronic structure and atomistic properties from the external potential.

This work proposes an operator-centric framework in which the external (nuclear) potential, expressed in an AO basis, serves as the model input and builds hierarchical, body-ordered representations of atomic configurations that closely mirror the principles underlying several popular atom-centered descriptors.

Jigyasa Nigam, T. Smidt, G. Dusson · 2 citations

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.