Uni-3DAR : Génération et compréhension 3D unifiées via autorégression sur des tokens spatiaux compressés

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

March 20, 2025

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

cs.AI

Résumé

Les récentes avancées dans les modèles de langage de grande taille et leurs extensions multimodales ont démontré l'efficacité de l'unification de la génération et de la compréhension grâce à la prédiction autoregressive du prochain jeton. Cependant, malgré le rôle crucial de la génération et de la compréhension des structures 3D ({3D GU}) dans l'IA pour la science, ces tâches ont largement évolué de manière indépendante, les méthodes autoregressives restant peu explorées. Pour combler cette lacune, nous introduisons Uni-3DAR, un cadre unifié qui intègre de manière transparente les tâches de {3D GU} via la prédiction autoregressive. Au cœur d'Uni-3DAR se trouve une nouvelle tokenisation hiérarchique qui compresse l'espace 3D en utilisant un octree, exploitant ainsi la sparsité inhérente des structures 3D. Une tokenisation supplémentaire est ensuite appliquée pour capturer les détails structurels fins, tels que les types d'atomes et les coordonnées spatiales précises dans les structures microscopiques 3D. Nous proposons également deux optimisations pour améliorer l'efficacité et l'efficience. La première est une stratégie de compression à deux niveaux des sous-arbres, qui réduit la séquence de jetons de l'octree jusqu'à 8 fois. La seconde est un mécanisme de prédiction masquée du prochain jeton, conçu pour les positions de jetons variant dynamiquement, ce qui améliore significativement les performances du modèle. En combinant ces stratégies, Uni-3DAR réussit à unifier diverses tâches de {3D GU} dans un seul cadre autoregressif. Des expériences approfondies sur plusieurs tâches microscopiques de {3D GU}, incluant les molécules, les protéines, les polymères et les cristaux, valident son efficacité et sa polyvalence. Notamment, Uni-3DAR surpasse les modèles de diffusion précédents de manière substantielle, atteignant une amélioration relative allant jusqu'à 256\% tout en offrant des vitesses d'inférence jusqu'à 21,8 fois plus rapides. Le code est disponible publiquement à l'adresse suivante : 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.