Uni-3DAR: Unified 3D Generation and Understanding via Autoregression on Compressed Spatial Tokens

Shuqi Lu1†, Haowei Lin3†, Lin Yao1†, Zhifeng Gao1, Xiaohong Ji1, Weinan E2,3*, Linfeng Zhang1,2* Guolin Ke1,2*
1DP Technology, 2AISI, 3Peking University

† These authors contributed equally to this work.

* Corresponding authors.

Paper arXiv Code
Schematic illustration of the Uni-3DAR framework

Schematic illustration of the Uni-3DAR framework.

Abstract

Recent advancements in large language models and their multi-modal extensions have demonstrated the effectiveness of unifying generation and understanding through autoregressive next-token prediction. However, despite the critical role of 3D structural generation and understanding (3D GU) in AI for science, these tasks have largely evolved independently, with autoregressive methods remaining underexplored. To bridge this gap, we introduce Uni-3DAR, a unified framework that seamlessly integrates 3D GU tasks via autoregressive prediction. At its core, Uni-3DAR employs a novel hierarchical tokenization that compresses 3D space using an octree, leveraging the inherent sparsity of 3D structures. It then applies an additional tokenization for fine-grained structural details, capturing key attributes such as atom types and precise spatial coordinates in microscopic 3D structures. We further propose two optimizations to enhance efficiency and effectiveness. The first is a two-level subtree compression strategy, which reduces the octree token sequence by up to 8x. The second is a masked next-token prediction mechanism tailored for dynamically varying token positions, significantly boosting model performance. By combining these strategies, Uni-3DAR successfully unifies diverse 3D GU tasks within a single autoregressive framework. Extensive experiments across multiple microscopic 3D GU tasks, including molecules, proteins, polymers, and crystals, validate its effectiveness and versatility. Notably, Uni-3DAR surpasses previous state-of-the-art diffusion models by a substantial margin, achieving up to 256% relative improvement while delivering inference speeds up to 21.8x faster.

Overview of Uni-3DAR

Adaptive coarse‑to‑fine subdivision of grid cells; darker nodes indicate non‑empty cells eligible for further partitioning
(a) Adaptive coarse‑to‑fine subdivision of grid cells; darker nodes indicate non‑empty cells eligible for further partitioning.
Partitioning builds an octree, offering lossless compression of the full‑size 3D grid
(b) Partitioning builds an octree, offering lossless compression of the full‑size 3D grid.
Uni‑3DAR’s tokenization couples hierarchical spatial compression with fine‑grained structural tokens
(c) Uni‑3DAR’s tokenization couples hierarchical spatial compression (octree) with fine‑grained structural tokens, each located by tree level + cell center.
2‑level subtree compression further reduces octree tokens up to 8× (4× in a 2D quadtree example)
(d) 2‑level subtree compression further reduces octree tokens up to 8× (4× in this 2D quadtree example).
Masked next‑token prediction handles dynamically shifting token positions
(e) Masked next‑token prediction handles dynamically shifting token positions;
Masked next‑token prediction handles dynamically shifting token positions
(f) Unified framework integrates generation and understanding in one model.

Results

Unconditional Molecular Generation Results
Unconditional Molecular Generation Results
De Novo Crystal Generation Results
De Novo Crystal Generation Results
Crystal Structure Prediction Results
Crystal Structure Prediction Results
Binding Site Prediction Results
Binding Site Prediction Results
Docking Results
Docking Results
Molecular Property Prediction Results
Molecular Property Prediction Results
Polymer Property Prediction Results
Polymer Property Prediction Results

BibTeX

@article{lu2025uni3dar,
  author    = {Shuqi Lu and Haowei Lin and Lin Yao and Zhifeng Gao and Xiaohong Ji and Weinan E and Linfeng Zhang and Guolin Ke},
  title     = {Uni-3DAR: Unified 3D Generation and Understanding via Autoregression on Compressed Spatial Tokens},
  journal   = {Arxiv},
  year      = {2025},
}