Uni-3DAR: 압축된 공간 토큰에 대한 자기회귀를 통한 통합 3D 생성 및 이해

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

March 20, 2025

저자: Shuqi Lu, Haowei Lin, Lin Yao, Zhifeng Gao, Xiaohong Ji, Weinan E, Linfeng Zhang, Guolin Ke

cs.AI

초록

최근 대규모 언어 모델과 그 다중모달 확장의 발전은 자동회귀적 다음 토큰 예측을 통해 생성과 이해를 통합하는 접근법의 효과를 입증했습니다. 그러나 과학 분야의 AI에서 3D 구조 생성 및 이해({3D GU})가 중요한 역할을 함에도 불구하고, 이러한 작업들은 대체로 독립적으로 발전해 왔으며, 자동회귀적 방법은 여전히 충분히 탐구되지 못했습니다. 이러한 격차를 해소하기 위해, 우리는 자동회귀 예측을 통해 {3D GU} 작업을 원활하게 통합하는 통합 프레임워크인 Uni-3DAR를 소개합니다. Uni-3DAR의 핵심은 3D 공간의 고유한 희소성을 활용하여 옥트리를 사용해 3D 공간을 압축하는 새로운 계층적 토큰화 방식을 채택합니다. 또한 미세한 구조적 세부 사항을 포착하기 위해 추가적인 토큰화를 적용하여, 원자 유형 및 정확한 공간 좌표와 같은 핵심 속성을 미시적 3D 구조에서 포착합니다. 우리는 효율성과 효과성을 향상시키기 위해 두 가지 최적화 방안을 추가로 제안합니다. 첫 번째는 옥트리 토큰 시퀀스를 최대 8배까지 줄이는 두 단계 서브트리 압축 전략입니다. 두 번째는 동적으로 변화하는 토큰 위치에 맞춤화된 마스크된 다음 토큰 예측 메커니즘으로, 모델 성능을 크게 향상시킵니다. 이러한 전략들을 결합함으로써, Uni-3DAR는 단일 자동회귀 프레임워크 내에서 다양한 {3D GU} 작업을 성공적으로 통합합니다. 분자, 단백질, 고분자, 결정 등 여러 미시적 {3D GU} 작업에 걸친 광범위한 실험을 통해 그 효과성과 다용도성을 검증했습니다. 특히, Uni-3DAR는 이전의 최첨단 확산 모델을 상당한 차이로 능가하며, 최대 256%의 상대적 개선을 달성하면서 최대 21.8배 빠른 추론 속도를 제공합니다. 코드는 https://github.com/dptech-corp/Uni-3DAR에서 공개되어 있습니다.

English

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. The code is publicly available at https://github.com/dptech-corp/Uni-3DAR.