Uni-3DAR: Generazione e Comprensione 3D Unificate tramite Autoregressione su Token Spaziali Compressi

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

March 20, 2025

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

cs.AI

Abstract

I recenti progressi nei modelli linguistici di grandi dimensioni e nelle loro estensioni multimodali hanno dimostrato l'efficacia dell'unificazione della generazione e della comprensione attraverso la previsione autoregressiva del token successivo. Tuttavia, nonostante il ruolo cruciale della generazione e comprensione delle strutture 3D ({3D GU}) nell'IA per la scienza, questi compiti si sono evoluti in gran parte in modo indipendente, con i metodi autoregressivi che rimangono poco esplorati. Per colmare questa lacuna, introduciamo Uni-3DAR, un framework unificato che integra senza soluzione di continuità i compiti di {3D GU} attraverso la previsione autoregressiva. Al suo nucleo, Uni-3DAR utilizza una nuova tokenizzazione gerarchica che comprime lo spazio 3D utilizzando un octree, sfruttando l'intrinseca sparsità delle strutture 3D. Successivamente, applica un'ulteriore tokenizzazione per i dettagli strutturali di livello fine, catturando attributi chiave come i tipi di atomo e le coordinate spaziali precise nelle strutture 3D microscopiche. Proponiamo inoltre due ottimizzazioni per migliorare l'efficienza e l'efficacia. La prima è una strategia di compressione a due livelli dei sottoalberi, che riduce la sequenza di token dell'octree fino a 8 volte. La seconda è un meccanismo di previsione mascherata del token successivo, progettato per posizioni di token che variano dinamicamente, migliorando significativamente le prestazioni del modello. Combinando queste strategie, Uni-3DAR unifica con successo diversi compiti di {3D GU} all'interno di un unico framework autoregressivo. Esperimenti estensivi su molteplici compiti di {3D GU} microscopici, tra cui molecole, proteine, polimeri e cristalli, ne validano l'efficacia e la versatilità. In particolare, Uni-3DAR supera i precedenti modelli di diffusione all'avanguardia con un margine sostanziale, raggiungendo un miglioramento relativo fino al 256\% e offrendo velocità di inferenza fino a 21,8 volte più veloci. Il codice è disponibile pubblicamente all'indirizzo 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.