Uni-3DAR: Vereinheitlichte 3D-Generierung und -Verständnis durch Autoregression auf komprimierten räumlichen Tokens

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

March 20, 2025

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

cs.AI

Zusammenfassung

Jüngste Fortschritte bei großen Sprachmodellen und ihren multimodalen Erweiterungen haben die Effektivität der Vereinheitlichung von Generierung und Verständnis durch autoregressive Vorhersage des nächsten Tokens demonstriert. Trotz der entscheidenden Rolle der 3D-Strukturgenerierung und des 3D-Verständnisses ({3D GU}) in der KI für die Wissenschaft haben sich diese Aufgaben weitgehend unabhängig voneinander entwickelt, wobei autoregressive Methoden weitgehend unerforscht blieben. Um diese Lücke zu schließen, stellen wir Uni-3DAR vor, ein einheitliches Framework, das {3D GU}-Aufgaben nahtlos durch autoregressive Vorhersage integriert. Im Kern von Uni-3DAR wird eine neuartige hierarchische Tokenisierung verwendet, die den 3D-Raum mithilfe eines Oktobaums komprimiert und dabei die inhärente Sparsität von 3D-Strukturen nutzt. Anschließend wird eine zusätzliche Tokenisierung für feinkörnige Strukturdetails angewendet, die Schlüsselattribute wie Atomtypen und präzise räumliche Koordinaten in mikroskopischen 3D-Strukturen erfasst. Wir schlagen außerdem zwei Optimierungen vor, um Effizienz und Wirksamkeit zu steigern. Die erste ist eine zweistufige Subbaum-Komprimierungsstrategie, die die Oktobaum-Token-Sequenz um bis zu 8x reduziert. Die zweite ist ein maskierter Mechanismus zur Vorhersage des nächsten Tokens, der für dynamisch variierende Token-Positionen maßgeschneidert ist und die Modellleistung erheblich steigert. Durch die Kombination dieser Strategien vereinheitlicht Uni-3DAR erfolgreich diverse {3D GU}-Aufgaben innerhalb eines einzigen autoregressiven Frameworks. Umfangreiche Experimente über mehrere mikroskopische {3D GU}-Aufgaben, einschließlich Moleküle, Proteine, Polymere und Kristalle, validieren seine Effektivität und Vielseitigkeit. Bemerkenswerterweise übertrifft Uni-3DAR frühere state-of-the-art Diffusionsmodelle deutlich und erreicht eine relative Verbesserung von bis zu 256\% bei gleichzeitiger Beschleunigung der Inferenzgeschwindigkeit um bis zu 21,8x. Der Code ist öffentlich verfügbar unter 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.