Self-Conditioned Denoising for Atomistic Representation Learning
Authors: Tynan Perez, Rafael Gomez-Bombarelli
Organizations: Department of Chemistry, Massachusetts Institute of Technology, 77 Massachusetts Avenue, Cambridge, 02139, MA, USA · Department of Materials Science and Engineering, Massachusetts Institute of Technology, 77 Massachusetts Avenue, Cambridge, 02139, MA, USA
The success of large-scale pretraining in NLP and computer vision has catalyzed growing efforts to develop analogous foundation models for the physical sciences. However, pretraining strategies using atomistic data remain underexplored. To date, large-scale supervised pretraining on DFT force-energy labels has provided the strongest performance gains to downstream property prediction, out-performing existing methods of self-supervised learning (SSL) which remain limited to ground-state geometries, and/or single domains of atomistic data. We address these shortcomings with Self-Conditioned Denoising (SCD), a backbone-agnostic reconstruction objective that utilizes self-embeddings for conditional denoising across any domain of atomistic data, including small molecules, proteins, periodic materials, and 'non-equilibrium' geometries. When controlled for backbone architecture and pretraining dataset, SCD significantly outperforms previous SSL methods on downstream benchmarks and matches or exceeds the performance of supervised force-energy pretraining. We show that a small, fast GNN pretrained by SCD can achieve competitive or superior performance to larger models pretrained on significantly larger labeled or unlabeled datasets, across tasks in multiple domains. Our code is available at: https://github.com/TyJPerez/SelfConditionedDenoisingAtoms
Figures & tables
Figure 1: Self-Conditioned Denoising Pretraining
Pretraining Dataset
Domain
Unique Structures
Total Geometries
Origin/functional
PCQ (PCQM4Mv2)
Small Molecules
3.378M
3.378M
B3LYP/6-31G* DFT
GEOM10
Small Molecules
304,339
2.7M
GFN2-xTB , conformers
AMP20 (Alex-MP-20)
Periodic Materials
607,673
607,673
PBEsol DFT
SAIR
Protein-Ligand
888k
4.4M
Boltz-1x, conformers
ALL (AllAtoms)
Above, mixed
5.2M
11.3M
PCQ, GEOM10, AMP20, SAIR
OMol25 (4M split)
mixed organic molecules
—
4M
MD, Non-equilibrium
Table 1: Summary of Pretraining and Benchmarking Datasets
Method
Baseline
Coord
SCD
% imp.
Backbone
ET
ET
CT
vs. Coord
HOMO (meV)
20.3
17.7
12.7
28.2%
LUMO (meV)
17.5
14.3
11.5
19.6%
Gap (meV)
36.1
31.8
24.5
23.0%
ZPVE (meV)
1.84
1.71
1.18
21.0%
U0 (meV)
6.15
6.57
3.58
45.5%
Table 2: SCD vs. Standard Node Denoising . Given the same backbone and pretraining datasets (PCQ), SCD significantly outperforms standard denoising ( Coord ). Best results are bolded . All values reported as MAE on the test set.
Method
Frad
SliDe
SCD
SCD
% imp.
Backbone
GET
GET
CT
CGT
vs. SliDe
HOMO (meV)
15.2
13.6
12.7
9.65
29.0%
LUMO (meV)
13.7
12.3
11.5
9.05
26.4%
Gap (meV)
27.8
26.2
24.5
19.7
24.8%
ZPVE (meV)
1.42
1.52
1.18
1.22
19.8%
U0 (meV)
5.33
4.28
3.58
3.97
8.1%
Table 3: SCD vs. Variations on Node Denoising . All models are pretrained on PCQ. Percent improvement (% imp.) is calculated between SliDe and SCD pretraining with the most similar backbone, CGT. Bold denotes best; underline denotes second best. All values reported as MAE on the test set.
Pretraining Dataset
Source / Quality
Unique Mol.
Total Geom.
QM9 HOMO
(MAE, meV)
baseline
20.3
PCQ
DFT (B3LYP/6-31G*)
3,378,606
3,378,606
12.7
PCQ-300k
DFT (B3LYP/6-31G*)
300,000
300,000
13.0
GEOM10
xTB (Conformers)
304,339
2,791,929
13.2
GEOM1
xTB (Lowest Energy)
304,339
304,339
13.7
Table 4: SCD Pretraining on Diverse Data Sources . SCD enables effective pretraining across ground-state and non-equilibrium conformers from varied sources. Every pretraining configuration tested shows improvement over the non-pretrained baseline.
Figure 2: UMAP of molecule embeddings colored by atom count: SCD results in smoother, less extensive, semantically rich representation. ’Coord’ refers to standard node denoising pretraining. See figure 4 for additional plots and details.
Metric
EquiformerV2
Baseline (ET)
CT-SCD
CGT-SCD
JMP-S
JMP-L
Model Parameters
11.2M
7M
10M
17M
27M
235M
Pretraining Structures
0
0
3.4M (U)
3.4M (U)
120M (L)
120M (L)
Pretraining GPU-hours
0
0
298
1,308
5,700
34,400
Finetuning GPU-h/task
137
—
46
96
—
—
Total hours (10 tasks)
1,370
—
758
2,268
>5,700
>34,400
QM9 Task (MAE)
Table 5: Comparing Computational Costs and Performance . SCD pretraining enables competitive performance at a lower cost than more advanced Clebsch-Gordan tensor product based architectures, and large-scale force-energy (FE) pretrained models. Bold denotes best; underline denotes second best.
Model
MAE (eV)
Params
Data
No Pretraining
coGN ( Ruff et al., 2023 )
0.156
—
n/a
CT (Baseline)
0.186 †
10M
n/a
JMP-S ( Shoghi et al., 2023 )
0.235 †
25.2M
n/a
JMP-L ( Shoghi et al., 2023 )
0.228 †
235M
n/a
MODNet ( De Breuck et al., 2021 )
0.220
—
n/a
Table 6: Matbench Bandgap Prediction . SCD models achieve competitive performance with significantly fewer parameters and unlabeled structures. Within each category best results are bolded ; second best are underlined .
Model
RMSE: id30 / id60 ↓
Params
No Pretraining
ProtNet ( Wang et al., 2023 )
1.463 / 1.343
—
EPT-Scratch ( Jiao et al., 2025 )
1.378 / 1.277
30M
CT (Baseline)
1.510 / 1.386
10M
Other Pretrained
Uni-Mol ( Zhou et al., 2023 )
1.520 / 1.619
47.6M
Table 7: Ligand binding affinity: RMSE on the 30% (id30) and 60% (id60) sequence identity overlap splits. Best results are bolded ; second best are underlined .
QM9 (Small Mol.)
Matbench (Crystals)
LBA (Proteins)
Model
HOMO
% imp.
MP Gap
% imp.
RMSE
% imp.
Avg %
(Pretraining Data)
(meV)
(eV)
(id60)
imp.
ET/CT (Baseline)
20.3
0.0%
0.186
0.0%
1.386
0.0%
0.0%
CT-SCD-PCQ
12.7
37.4%
0.174
6.5%
1.226
11.5%
18.5%
CT-SCD-GEOM10
13.2
35.0%
0.177
4.8%
1.211
12.6%
17.5%
CT-SCD-AMP20
14.6
28.1%
0.122
34.4%
1.219
12.0%
24.8%
Table 8: Comparing the Impact of Different Pretraining Data . Pretraining on any data source benefits every task, though downstream performance gains are greatest when pretraining on in-domain data. Datasets with high elemental diversity yield the highest average improvements. CT-FE models denote supervised force-energy pretraining; CT-SCD denotes self-supervised pretraining. Best results are bold ; second best are underlined .
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Table 9: Architecture Efficiency Comparison . Benchmarked using QM9 molecules, with batch size 128, averaged over 32 batches, on a NVIDIA RTX A5000 GPU. All models share identical embedding dimensions, layer counts, and attention heads. Parameter counts in this table encompass only backbone architecture, excluding the embedding and prediction heads. All metrics shown are computed from singular forward passes.
Figure 3: Adaptive layer normalization for SCD: On the left we illustrate a simplified diagram of the original residual layer block in TorchMD-Net (ET). Here, ’X’ represents invariant (L0) embeddings while ’V’ represents equivariant vector embeddings (L1), and ’C’ represents the conditioning embedding. On the right, we show how we modify the original residual block with a conditional scale, shift, and gate of the invariant embeddings. In the gate block the gamma vector is attenuated by a tanh function and multiplied to the output of the ET Attention block. For this work, we use a two layer MLP for conditioning in each layer.
Hyperparameter
CT
CGT
CT-small
Layers
8
8
6
Attention heads
8
8
8
Embedding dimension
256
256
128
Layer pooling
sum
sum
sum
Graph cutoff (Å)
5
5
5
L0 head pooling
sum
sum
sum
Appendix
Table 10: Backbone architecture hyperparameters used in this work. CT-small is only used in the MD17 benchmark
Figure 4: Pretraining Effects on the Embedding Space. UMAP visualizations of embeddings for 10k QM9 molecules, colored by atom count ( N ). (a) Untrained GNNs exhibit a strong bias towards ”extensive” representations, clustering molecules primarily by size ( N ). (b) Standard node denoising smooths the latent manifold but embeddings remain highly extensive. (c–d) SCD pretraining yields a smoother, less extensive (more semantic) embedding space. We observe that downstream performance improves as representations becomes smoother and less extensive (Performance: a<b<c<d ). ’Molecule embeddings’ are created by passing sum pooled atom embeddings through a two-layer MLP ’embedding head’. In the case of ’ET-Coord’ the embedding head is untrained. In SCD ’Molecule embeddings’ are used for self-embeddings. In our ’coord’ model, the molecule ’embedding head’ remains untrained.
Method
Model
Seq. overlap 30%
Seq. overlap 60%
Params
RMSE ↓
Pear. ↑
Spear. ↑
RMSE ↓
Pear. ↑
Spear. ↑
No Pretraining
HoloProt-Full Surface ( Somnath et al., 2021 )
1.464
0.509
0.500
1.365
0.749
0.742
1.4M
ProtNet-All-Atom ( Wang et al., 2023 )
1.463
0.551
0.551
1.343
0.765
0.761
—
ATOM3D-3DCNN ( Townshend et al., 2022 )
1.416
0.550
0.553
1.621
0.608
0.615
—
ATOM3D-ENN ( Townshend et al., 2022 )
1.568
0.389
0.408
1.620
0.623
0.633
—
ATOM3D-GNN ( Townshend et al., 2022 )
1.601
0.545
0.533
1.408
0.743
0.743
—
Appendix
Table 11: SCD on Proteins : Ligand Binding Affinity (LBA) predictions over two sequence identity splits (30% and 60%). SCD models pretrained on synthetic or diverse data (e.g., SAIR, OMol25) can outperform larger models pretrained on PDB-based datasets. Best results within each category are bolded ; second best are underlined .
Category
Model
MP Gap (eV)
# Params
Pretraining
Fold 0
Mean (0–4)
Data Size
No Pretraining
MODNet ( De Breuck et al., 2021 )
0.215
0.220
—
n/a
coGN ( Ruff et al., 2023 )
0.153
0.156
—
n/a
JMP-S ( Shoghi et al., 2023 )
0.235
—
25.2M
n/a
JMP-L ( Shoghi et al., 2023 )
0.228
—
235M
n/a
CT (Baseline)
0.186
—
10M
n/a
Appendix
Table 12: SCD on Periodic Materials : Matbench band gap prediction performance for SCD and supervised force-energy (FE) pretrained models. Best results within each category are bolded ; second best are underlined . SCD models achieve competitive performance while utilizing significantly fewer parameters and unlabeled structures.
Model Name
Initialization
Force Method
Val Force MAE
Val Energy MAE
num training steps
(meV/Å)
(meV)
FE-Fwd
From Scratch
Forward
44.8
648
800k
FE-Fwd-long
From Scratch
Forward
42.6
606
1600k
FE-Bwd
From Scratch
Backward
50.0
618
800k
SCD-FE-Fwd
CT-SCD-OMOL25
Forward
42.9 (+4.2%)
617 (+4.8%)
800k (+800k SCD)
SCD-FE-Bwd
CT-SCD-OMOL25
Backward
41.4 (+17.2%)
665 (-7.6%)
800k (+800k SCD)
Appendix
Table 13: OMol25 4M Force-Energy Pretraining . Comparison of supervised force-energy training baselines (FE) against fine-tuning from SCD pretrained weights (SCD-FE). We compare direct force prediction (Forward) against gradient-based force calculation (Backward). CT-SCD-OMOL25 is pretrained by SCD using the same Omol25 4 geometries.
Model
Fine-tuning Protocol
HOMO
U0
avg % impv.
(meV)
(meV)
over baseline
baseline
20.3
6.15
—
FE-Fwd
Standard
15.9
12.9
-44.0%
FE-Fwd
Reset Head, σreg=0.005
15.0
10.2
-19.9%
FE-Fwd
Reset Head, σreg=0
13.6
3.40
38.9%
FE-Fwd-long
Reset Head, σreg=0
13.7
3.47
38.0%
Appendix
Table 14: Fine-tuning Force-Energy Pretrained Models . Force-energy pretrained models performed best when the energy head was reset and regularization noise was turned off during finetuning. Forward (non-conservative) force pretraining, conveyed the best average results on downstream property prediction and is reported as ’CT-FE-OMOL25’ elsewhere.
QM9 Task
CT
CGT
EPT-10
Uni-Corn
(Ours)
(Ours)
( Jiao et al., 2025 )
( Feng et al., 2024 )
Params
10M
17M
30M
—
Pretrain Data
3.4M
3.4M
5.9M
15M
HOMO (meV)
12.7
9.65
15.2
13.0
LUMO (meV)
11.5
9.05
13.6
11.9
Gap (meV)
24.5
19.7
29.0
24.9
Appendix
Table 15: SCD vs. other SSL Methods using different backbones and pretraining data. Models CT and CGT below are both pretrained on the PCQ dataset. Best results are bolded ; second best are underlined .
Method
Baseline
Coord
SCD
SCD (force-denoise)
Frad (force-denoise)
SliDe
Architecture
ET-small
ET-small
CT-small
CT-small
GET-small
GET-small
Pretraining dataset
none
PCQ
PCQ
PCQ
PCQ
PCQ
Forces MAE (kcal/mol/Å)
Aspirin
0.253
0.211
0.210
0.200
0.209
0.174
Benzene
0.196
0.169
0.171
0.175
0.199
0.169
Ethanol
0.109
0.096
0.120
0.116
0.091
0.088
Appendix
Table 16: SCD achieves comparable performance on MD17 force–energy prediction to prior methods while using a smaller and faster architecture. Best results are in bold , second best are underlined .
Hyperparameter
Value
Warmup steps
10,000
Learning rate
0.005
β1
0.9
Weight decay
0.05
DropPath
0.1
Corrupt noise σ
0.04
Appendix
Table 17: Base SCD pretraining hyperparameters: applied in all SCD training runs
Backbone
CT
CGT
CT
CT
CT
CT
CT
Dataset
PCQ
PCQ
GEOM10
AMP20
SAIR
ALL
Omol25
Total steps
800k
800k
800k
800k
800k
1.2M
800k
Batch size / GPU
128
32
48
142
12
54
48
# GPUs (A100)
4
16
16
4
8
16
16
Effective batch size
512
512
768
568
96
864
768
Materials augmentations
Appendix
Table 18: SCD Pretraining: dataset and model specific hyperparameters.
Hyperparameter
HOMO, LUMO, Gap
U0, U, H, G
Alpha, ZPVE, CV
Batch size
128
32
128
Learning rate
5e-4
3e-4
3e-4
Total steps
300k
500k
400k
Warmup steps
10k
10k
10k
β1
0.995
0.995
0.95
Weight decay
0.01
0.01
0.01
Appendix
Table 19: QM9 finetuning hyperparameters for different target groups.
Hyperparameter
Value
Batch size
8
Learning rate (from scratch)
2e-4
Learning rate (finetuning)
8e-5
Total steps
300k
Warmup steps
10k
β1
0.995
Appendix
Table 20: Matbench finetuning hyperparameters for the mp_gap task.
Hyperparameter
Value
Batch size
8
Learning rate (from scratch)
2e-4
Learning rate (finetuning)
1e-4
Total steps
30k
Warmup steps
500
β1
0.995
Appendix
Table 21: Finetuning hyperparameters for the Ligand Binding Affinity (LBA) benchmark
Hyperparameter
Value
Backbone
CT-small
Training set size
1,000
Batch size
8
Energy loss EMA
0.05
Energy loss weight
0.8
Force loss weight
0.2
Appendix
Table 22: Training hyperparameters for the MD17 benchmark.
Large-scale self-supervised pretraining has reshaped modern machine learning, substantially advancing the ability of language and vision models to generalize across downstream tasks. While deep learning has driven considerable progress in modeling atomistic systems in recent years, self-supervised pretraining in this domain has not yet achieved comparable downstream generalization. To address this, we introduce Atom-JEPA, a self-supervised pretraining framework that learns latent representations from unlabeled 3D structures through complementary atom-level and substructure-level objectives inspired by joint-embedding predictive architectures. We pretrain Atom-JEPA on large-scale molecular and crystalline datasets and evaluate its transfer performance by fine-tuning on a diverse set of downstream property prediction tasks. Atom-JEPA achieves state-of-the-art performance on molecular ADMET and quantum-chemical property prediction tasks, and is highly competitive in predicting the physical properties of crystalline materials. These results demonstrate the potential of latent-space predictive pretraining to support broad downstream generalization from structural data alone. Code and pretrained model checkpoints are publicly available at https://github.com/khelverskovp/atom-jepa
Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet +2
Department of Applied Mathematics and Computer Science Technical University of Denmark
Unified atomistic modeling has the potential to accelerate discovery in chemistry, materials science, and biology by bridging data-rich chemical domains and data-scarce biological contexts. However, existing generative approaches to atomistic modeling remain highly specialized to scientific disciplines (chemistry vs. biology) or do not leverage both high-volume organic (molecule) and inorganic (material) data for general-purpose pretraining. To this end, we introduce Zatom-2, an atomistic generative model pretrained on approximately five million structures from the OMol25 and OMat24 electronic structure datasets. Zatom-2 features a multiscale Transformer architecture coupled with conditional flow matching that supports force conditioning and foundational pretraining tasks such as generation, structure prediction, and prediction of molecular and material energies and forces. Empirically, Zatom-2 achieves better molecular distribution fidelity than Zatom-1 and achieves strong performance on existing molecule and material generation benchmarks. Zatom-2 demonstrates the ability to control sample generation across low- and high-force regimes, and enhances protein generation in a low-data setting through joint generative-predictive pretraining and transfer learning, increasing protein backbone designability in a length extrapolation setting from 67.8% without pretraining to 74.8% after finetuning on 2,000 protein domains.
Miruna Cretu, Alex Abrudan, Antonia Panescu +11
University of Cambridge · Yale University · MIT +3
Machine learning interatomic potentials (MLIPs) achieve excellent accuracy when trained on large Density Functional Theory (DFT) data. To be useful in practice, they must often be adapted to target chemistries using small and expensive task-specific datasets. However, MLIPs transfer inconsistently across domains, with representations that often loose accessible composition and structure information. To address this, we present TriForces, a model-agnostic three-stream framework that separates composition and structure information, combined with self-supervised learning to preserve transferable representations. TriForces improves performance on MatBench and QM9 over baselines without needing DFT labels and enables efficient similar structure retrieval through its learned latent space. On OMat24, in limited-data training regime, TriForces reduces energy MAE by 57% at 20K samples only and improves force MAE across sample sizes. We release pretrained TriForces variants across multiple MLIP architectures with code at https://github.com/Ramlaoui/triforces.
Ali Ramlaoui, Alexandre Duval, Hannah Bull +4
Entalpic, Paris, France · Universit´e Paris-Saclay, CentraleSup´elec, Inria, Gif-sur-Yvette, France