ADT: All-Atom Diffusion Tokenizer for 3D Molecular Structures

Xuantuo Huang; Encheng Su; Aoran Wang; Pan Tan; Shengyong Xu; Yizhou Wang; Chen Tang; Shixiang Tang

cs.LGOct 16, 2026

ADT: All-Atom Diffusion Tokenizer for 3D Molecular Structures

Authors: Xuantuo Huang, Encheng Su, Aoran Wang, Pan Tan, Shengyong Xu, Yizhou Wang, Chen Tang, Shixiang Tang,

Organizations: Peking University · University of Science and Technology of China · Shanghai AI Laboratory · The Chinese University of Hong Kong

Abstract

Discrete structural tokenization provides a promising route to autoregressive 3D molecular generation. Yet, existing tokenizers face a trade-off between compact representation and all-atom reconstruction: atom-level tokens preserve atomic details but produce long sequences, while compact residue-level protein tokens typically omit side-chain geometry. To overcome this limitation, we propose All-atom Diffusion Tokenizer (ADT), which decouples token granularity from reconstruction resolution. ADT uses one structural token per atom for small molecules and one token per protein residue that jointly represents backbone and side-chain geometry. To mitigate geometric distortion introduced by quantization, we leverage the generative capacity of diffusion models to recover fine-grained atomic geometry beyond deterministic decoding from discrete tokens. Extensive experiments across small molecules and proteins demonstrate that ADT achieves strong reconstruction from compact discrete representations while enabling high-quality autoregressive generation of both small molecules and all-atom protein structures. These results establish compact structural tokens as an effective discrete interface for all-atom molecular reconstruction and generation.

Figures & tables

Figure 1: Illustration of tokenizer granularity. (a) Atom-level tokenization assigns one token to each atom. (b) Backbone-only tokenization omits side-chain geometry. (c) ADT jointly encodes backbone and side-chain geometry into one token per residue and reconstructs the all-atom structure.

Figure 2: Overview of ADT. (a) All-atom structures are represented at chemical-unit granularity. (b) Each unit is mapped to a discrete structural token via FSQ. (c) Conditional diffusion reconstructs all-atom geometry, with residue-to-atom expansion for proteins. (d) ADT enables autoregressive generation of all-atom molecular structures.

Type	Method	Match $\uparrow$	RMSD $\downarrow$
Continuous	GeoLDM ( Xu et al., 2023 )	87.2	0.1820
	ADiT ( Joshi et al., 2025 )	94.6	0.0890
	SinAE ( Ren et al., 2026 )	100.0	0.0007
Discrete	Mol-StrucTok ( Gao et al., 2026 )	90.7	0.1280
	CoordToken Geikyan et al. (2026)	99.4	0.0478
	ADT (Ours)	99.7	0.0343

Table 1: Reconstruction results on QM9.

Method	Conformer		Structure		Scaffolds
Method	RMSD $\downarrow$	Valid $\uparrow$	RMSD $\downarrow$	Valid $\uparrow$	RMSD $\downarrow$	Valid $\uparrow$
mol2token ( Liu et al., 2025 )	0.2010	47.1	0.2010	47.1	0.2010	47.1
CoordToken Geikyan et al. (2026)	0.0671	100.0	0.0673	100.0	0.0671	100.0
ADT (Ours)	0.0499	100.0	0.0512	100.0	0.0511	100.0

Table 2: Reconstruction performance on the

\nabla^{2}

DFT benchmark.

	Method	C $\alpha$ -lDDT $\uparrow$	BB RMSD $\downarrow$	AA lDDT $\uparrow$	AA RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
CASP14	DPLM-2 ( Wang et al., 2025 )	0.8820	2.8576	–	–	0.8860	0.7718	0.5724
	+ Rosetta fixbb	0.8820	2.8576	0.7532	3.4918	0.8859	0.7717	0.5723
	+ PDBFixer	0.8820	2.8576	0.6632	3.9657	0.8859	0.7717	0.5723
	+ cg2all (nofix)	0.8820	2.8093	0.7710	3.4120	0.8859	0.7718	0.5722
	+ cg2all (fixed)	0.8819	2.8576	0.7611	3.4358	0.8859	0.7717	0.5723
	ADT (Ours)	0.8970	0.9628	0.8223	1.3157	0.9695	0.9398	0.7686

Table 3: Protein reconstruction performance on CASP14 and CASP15. DPLM-2 variants recover all-atom structures using different backbone-to-all-atom reconstruction methods.

Type	Method	Valid $\uparrow$	Unique $\uparrow$	Div. $\uparrow$	PB Valid $\uparrow$
Continuous	EQGAT-diff ( Le et al., 2024 )	0.99	1.00	0.89	0.94
	SemlaFlow ( Irwin et al., 2024 )	0.99	0.95	0.89	0.94
	FlowMol ( Dunn and Koes, 2024 )	0.97	0.97	0.92	0.94
	ADiT ( Joshi et al., 2025 )	0.97	0.97	0.90	0.92
	SinAE ( Ren et al., 2026 )	0.97	0.97	0.91	0.94
Discrete	Mol-StrucTok ( Gao et al., 2026 )	0.97	0.85	0.87	0.88

Table 4: Unconditional molecular generation results on QM9.

Method	Length Match	C $\alpha$ -lDDT	BB-lDDT	TM-score	AA-lDDT
DPLM-2	1.0	0.5607	0.5683	0.633	–
+ Rosetta fixbb	1.0	0.5607	0.5683	0.633	0.482
+ PDBFixer	1.0	0.5607	0.5683	0.633	0.372
+ cg2all nofix	1.0	0.5607	0.5683	0.633	0.503
+ cg2all fixed	1.0	0.5607	0.5683	0.633	0.496
ADT (Ours)	1.0	0.5740	0.5801	0.665	0.543

Table 5: Sequence-conditioned protein structure generation on CAMEO22.

	Method	C $\alpha$ -lDDT $\uparrow$	BB RMSD $\downarrow$	AA lDDT $\uparrow$	AA RMSD $\downarrow$	SC RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
CASP14	Backbone only	0.8802	1.2804	0.7699	1.9269	2.3858	0.9555	0.9038	0.7241
	Separate BB+SC	0.8762	1.3105	0.8262	1.5075	1.6763	0.9549	0.9106	0.7238
	Joint All-atom (ADT)	0.8970	0.9628	0.8223	1.3157	1.5938	0.9695	0.9398	0.7686
CASP15	Backbone only	0.8864	1.6517	0.7866	2.2386	2.6672	0.9365	0.8875	0.7150
	Separate BB+SC	0.8827	1.4884	0.8306	1.7711	2.0341	0.9416	0.8896	0.7065
	Joint All-atom (ADT)	0.8972	1.5191	0.8258	1.8261	2.0779	0.9470	0.9001	0.7285

Table 6: Ablation of protein structural encoding strategies on CASP14 and CASP15.

Decoder	Layers	Match $\uparrow$	RMSD $\downarrow$	PB Valid $\uparrow$	RMSD Gain $\uparrow$
Transformer	4	97.5	0.0695	0.88	27.19%
Diffusion	4	98.3	0.0506	0.89	27.19%
Transformer	8	98.4	0.0636	0.90	33.18%
Diffusion	8	99.1	0.0425	0.91	33.18%
Transformer	12	98.9	0.0584	0.92	41.27%
Diffusion	12	99.7	0.0343	0.93	41.27%

Table 7: Effect of decoder depth on QM9 reconstruction.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

	Model	CASP14		CASP15
	Model	AA RMSD $\downarrow$	TM-score $\uparrow$	AA RMSD $\downarrow$	TM-score $\uparrow$
Atom-level	Bio2Token ( Liu et al., 2025 )	0.644	0.983	0.693	0.981
Atom-level	BiHiTo ( Zheng et al., 2026 )	0.515	0.987	0.585	0.986
Residue-level	DPLM-2 ( Wang et al., 2025 )	–	0.8860	–	0.8644
	+ cg2all (no fix)	3.4120	0.8859	4.5037	0.8644
	+ PDBFixer	3.9657	0.8859	5.0213	0.8644

Appendix

Table 8: All-atom protein reconstruction across token granularities on CASP14 and CASP15.

	CASP14					CASP15
Method	C $\alpha$ -lDDT $\uparrow$	BB RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$	C $\alpha$ -lDDT $\uparrow$	BB RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
ESM3	0.9653	1.9934	0.9552	0.9041	0.8068	0.9672	2.0685	0.9455	0.8693	0.7470
DPLM-2	0.8820	2.8576	0.8860	0.7718	0.5724	0.8933	4.0223	0.8644	0.7032	0.5104
ADT (Ours)	0.8970	0.9628	0.9695	0.9398	0.7686	0.8972	1.5191	0.9470	0.9001	0.7285

Appendix

Table 9: Protein structure reconstruction performance on CASP14 and CASP15.

Method	Des. $\uparrow$	Div. $\uparrow$	DPT $\downarrow$	Nov. $\uparrow$
ESM3	0.61	127	0.37	0.84
DPLM-2	0.63	130	0.37	0.72
ADT (Ours)	0.68	139	0.35	0.76

Appendix

Table 10: Unconditional protein structure generation.

Figure 5: Qualitative results of samples generated by ADT for small molecules from QM9 and proteins.

	Method	C $\alpha$ lDDT $\uparrow$	BB RMSD $\downarrow$	AA lDDT $\uparrow$	AA RMSD $\downarrow$	SC RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
CASP14	Shuffled	0.3100	21.166	0.2978	21.709	22.278	0.2253	0.0373	0.0058
CASP14	ADT	0.8970	0.9628	0.8223	1.3157	1.5938	0.9695	0.9398	0.7686
CASP15	Shuffled	0.3233	22.245	0.3103	22.733	23.274	0.2250	0.0354	0.0057
CASP15	ADT	0.8972	1.5191	0.8258	1.8261	2.0779	0.9470	0.9001	0.7285

Appendix

Table 11: Ablation of structural tokens on protein structure reconstruction.

Decoder	Match $\uparrow$	RMSD $\downarrow$	PB Valid $\uparrow$	Bond-length $\downarrow$	Bond-angle $\downarrow$	Dihedral $\downarrow$
Transformer	98.9	0.0584	0.92	0.02219	1.4971	4.5234
Diffusion	99.7	0.0343	0.93	0.01557	1.0929	3.7622

Appendix

Table 12: Fine-grained Geometric Reconstruction on QM9.

	Coarse:Fine	C $\alpha$ lDDT $\uparrow$	BB RMSD $\downarrow$	AA lDDT $\uparrow$	AA RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
CASP14	12:0	0.8568	1.1196	0.7934	1.5559	0.9528	0.9027	0.7089
	8:4	0.8892	0.9876	0.8161	1.3574	0.9668	0.9324	0.7571
	10:2	0.8970	0.9628	0.8223	1.3157	0.9695	0.9398	0.7686
	0:12	0.8624	1.0738	0.8067	1.5126	0.9587	0.9068	0.7115
CASP15	12:0	0.8574	1.7668	0.7967	2.1594	0.9307	0.8630	0.6688
	8:4	0.8895	1.5583	0.8196	1.8840	0.9444	0.8930	0.7170

Appendix

Table 13: Ablation of coarse-to-fine diffusion decoding on protein structure reconstruction.

	Codebook	C $\alpha$ lDDT $\uparrow$	BB RMSD $\downarrow$	AA lDDT $\uparrow$	AA RMSD $\downarrow$	TM-score $\uparrow$	GDT-TS $\uparrow$	GDT-HA $\uparrow$
CASP14	2048	0.8634	1.0342	0.7864	1.4734	0.9563	0.9172	0.7243
	4096	0.8721	0.9818	0.8013	1.3869	0.9634	0.9195	0.7431
	8192	0.8970	0.9628	0.8223	1.3157	0.9695	0.9398	0.7686
	16384	0.9056	0.9467	0.8315	1.2768	0.9711	0.9584	0.7869
CASP15	2048	0.8628	1.6324	0.7899	2.0413	0.9338	0.8812	0.6853
	4096	0.8706	1.5587	0.8049	1.9284	0.9409	0.8876	0.7041

Appendix

Table 14: Effect of discrete codebook size on protein structure reconstruction.

Strategy	Valid $\uparrow$	Unique $\uparrow$	Div. $\uparrow$	PB Valid $\uparrow$
Interleaved AR	0.94	0.93	0.90	0.88
G2C-AR	0.92	0.89	0.91	0.85
C2G-AR	1.00	0.96	0.91	0.94

Appendix

Table 15: Effect of autoregressive factorization on QM9 generation.

Protein Structure Generation

See all 20 papers

BiMol-Diff: A Unified Diffusion Framework for Molecular Generation and Captioning
Apr 27, 2026Aditya Hemant Shahane, Anuj Kumar Sirohi, Devansh Arora +3Diffusion Sampling Molecular Generation
Indian Institute of Technology Delhi · Indian Institute of Technology Ropar · Indian Institute of Science Bengaluru
Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation
Mar 10, 2026Junyi An, Chao Qu, Yun-Fei Shi +33D Molecular Generation Molecular Generation
Shanghai Academy of Artificial Intelligence for Science, SAIS · Artificial Intelligence Innovation and Incubation (AI3) Institute, Fudan University
ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization
Oct 13, 2026Romain Hardy, Xiaoman Zhang, Pranav RajpurkarProtein Representation Learning Protein Language Models
Department of Biomedical Informatics, Harvard Medical School

Apr 27, 2026J/K move · Enter open · S save

BiMol-Diff: A Unified Diffusion Framework for Molecular Generation and Captioning

Aditya Hemant Shahane, Anuj Kumar Sirohi, Devansh Arora, Nitin Kumar+2

Indian Institute of Technology Delhi · Indian Institute of Technology Ropar · Indian Institute of Science Bengaluru

Diffusion Sampling Molecular Generation Diffusion Language Models Conditional Diffusion Models

Open paper AskPDF arXiv

Bridging molecular structures and natural language is essential for controllable design. Autoregressive models struggle with long-range dependencies, while standard diffusion processes apply uniform corruption across positions, which can distort structurally informative tokens. We present BiMol-Diff, a unified diffusion framework for the paired tasks of text-conditioned molecule generation and molecule captioning. Our key component is a token-aware noise schedule that assigns position-dependent corruption based on token recovery difficulty, preserving harder-to-recover substructures during the forward process. On ChEBI-20 and M3-20M, BiMol-Diff improves molecule reconstruction with a 15.4% relative gain in Exact Match and achieves strong captioning results, attaining best BLEU and BERTScore among compared baselines. These results indicate token-aware noising improves fidelity in molecular structure-language modelling.

Abstract

Protein Structure Generation

Explore similar work

BiMol-Diff: A Unified Diffusion Framework for Molecular Generation and Captioning

Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation

ProTok: Unified Protein Tokenization via Modular Multi-Codebook Quantization