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

arXiv: https://arxiv.org/abs/2608.03260 · HTML: https://mlsanity.com/paper/2608.03260/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 32 of its 41 references that are in the map.
- First posted on arXiv 4 Aug 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 32)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling](https://mlsanity.com/paper/2111.14819/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [PointContrast: Unsupervised Pre-training for 3D Point Cloud Understanding](https://mlsanity.com/paper/2007.10985/) (2020)
- [Masked Autoencoders for Point Cloud Self-supervised Learning](https://mlsanity.com/paper/2203.06604/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [Large-Scale Chemical Language Representations Capture Molecular Structure and Properties](https://mlsanity.com/paper/2106.09553/) (2021)
- [Multi-modal Molecule Structure-text Model for Text-based Retrieval and Editing](https://mlsanity.com/paper/2212.10789/) (2022)
- [3D Infomax improves GNNs for Molecular Property Prediction](https://mlsanity.com/paper/2110.04126/) (2021)
- [SODA: Bottleneck Diffusion Models for Representation Learning](https://mlsanity.com/paper/2311.17901/) (2023)
- [Unified 2D and 3D Pre-Training of Molecular Representations](https://mlsanity.com/paper/2207.08806/) (2022)
- [Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids](https://mlsanity.com/paper/2112.00652/) (2021)
- [PointVector: A Vector Representation In Point Cloud Analysis](https://mlsanity.com/paper/2205.10528/) (2022)
- [Self-supervised Learning for Pre-Training 3D Point Clouds: A Survey](https://mlsanity.com/paper/2305.04691/) (2023)
- [UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning](https://mlsanity.com/paper/2405.10343/) (2024)
- [DeepDFT: Neural Message Passing Network for Accurate Charge Density Prediction](https://mlsanity.com/paper/2011.03346/) (2020)
- [X-3D: Explicit 3D Structure Modeling for Point Cloud Recognition](https://mlsanity.com/paper/2404.15010/) (2024)
- [SciReasoner: Laying the Scientific Reasoning Ground Across Disciplines](https://mlsanity.com/paper/2509.21320/) (2025)
- [ELECTRA: A Cartesian Network for 3D Charge Density Prediction with Floating Orbitals](https://mlsanity.com/paper/2503.08305/) (2025)
- [A Recipe for Charge Density Prediction](https://mlsanity.com/paper/2405.19276/) (2024)

## Built on by (most cited first, 1 of 1)

- [Neural-Network Solutions to Real-Space Charge Density and Generalization](https://mlsanity.com/paper/2609.14906/) (2026)

## Abstract (arXiv)

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.