Self-Conditioned Denoising for Atomistic Representation Learning

Tynan Perez; Rafael Gomez-Bombarelli

cs.LGMar 17, 2026

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

Abstract

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%
$U_{0}$ (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%
$U_{0}$ (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 $\downarrow$	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 $\downarrow$	Pear. $\uparrow$	Spear. $\uparrow$	RMSE $\downarrow$	Pear. $\uparrow$	Spear. $\uparrow$
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	$U_{0}$	avg % impv.
		(meV)	(meV)	over baseline
baseline		20.3	6.15	—
FE-Fwd	Standard	15.9	12.9	-44.0%
FE-Fwd	Reset Head, $\sigma_{\text{reg}}=0.005$	15.0	10.2	-19.9%
FE-Fwd	Reset Head, $\sigma_{\text{reg}}=0$	13.6	3.40	38.9%
FE-Fwd-long	Reset Head, $\sigma_{\text{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
$\beta_{1}$	0.9
Weight decay	0.05
DropPath	0.1
Corrupt noise $\sigma$	0.04

Appendix

Table 17: Base SCD pretraining hyperparameters: applied in all SCD training runs

Materials augmentations
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

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
$\beta_{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
$\beta_{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
$\beta_{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.

Self-Supervised Learning

Self-supervised learning uses structure within unlabeled data to train representations and models for downstream tasks.

See all 258 papers

Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
Oct 6, 2026Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet +2Joint-Embedding Predictive Architecture Materials Property Prediction
Department of Applied Mathematics and Computer Science Technical University of Denmark
Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains
Oct 8, 2026Miruna Cretu, Alex Abrudan, Antonia Panescu +11Generative Modeling Molecular Generation
University of Cambridge · Yale University · MIT +3
TriForces: Augmenting Atomistic GNNs for Transferable Representations
May 20, 2026Ali Ramlaoui, Alexandre Duval, Hannah Bull +4Disentangled Representation Learning Machine Learning Interatomic Potentials
Entalpic, Paris, France · Universit´e Paris-Saclay, CentraleSup´elec, Inria, Gif-sur-Yvette, France

Oct 6, 2026J/K move · Enter open · S save

Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

Kasper Helverskov Petersen, Rasmus Hannibal Tirsgaard, François R J Cornet, Mikkel Jordahn+1

Department of Applied Mathematics and Computer Science Technical University of Denmark

Joint-Embedding Predictive Architecture Materials Property Prediction 3D Representation Learning Molecular Property Prediction Molecular Representation Learning

Open paper AskPDF arXiv

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

Abstract

Self-Supervised Learning

Explore similar work

Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems

Zatom-2: Multitask Pretraining on Atomistic Data for Generative Modeling across Domains

TriForces: Augmenting Atomistic GNNs for Transferable Representations