Uni-3DAR: Geração e Compreensão 3D Unificadas via Autoregressão em Tokens Espaciais Comprimidos

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

March 20, 2025

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

cs.AI

Resumo

Os avanços recentes em modelos de linguagem de grande escala e suas extensões multimodais demonstraram a eficácia de unificar geração e compreensão por meio da previsão autoregressiva do próximo token. No entanto, apesar do papel crítico da geração e compreensão de estruturas 3D ({3D GU}) na IA para ciência, essas tarefas evoluíram em grande parte de forma independente, com métodos autoregressivos permanecendo pouco explorados. Para preencher essa lacuna, introduzimos o Uni-3DAR, um framework unificado que integra de forma contínua tarefas de {3D GU} por meio de previsão autoregressiva. No seu núcleo, o Uni-3DAR emprega uma nova tokenização hierárquica que comprime o espaço 3D usando uma octree, aproveitando a esparsidade inerente das estruturas 3D. Em seguida, aplica uma tokenização adicional para detalhes estruturais refinados, capturando atributos-chave como tipos de átomos e coordenadas espaciais precisas em estruturas 3D microscópicas. Propomos ainda duas otimizações para melhorar a eficiência e a eficácia. A primeira é uma estratégia de compressão de subárvore em dois níveis, que reduz a sequência de tokens da octree em até 8x. A segunda é um mecanismo de previsão mascarada do próximo token, adaptado para posições de tokens que variam dinamicamente, aumentando significativamente o desempenho do modelo. Ao combinar essas estratégias, o Uni-3DAR unifica com sucesso diversas tarefas de {3D GU} em um único framework autoregressivo. Experimentos extensos em múltiplas tarefas de {3D GU} microscópicas, incluindo moléculas, proteínas, polímeros e cristais, validam sua eficácia e versatilidade. Notavelmente, o Uni-3DAR supera os modelos de difusão anteriores de última geração por uma margem substancial, alcançando uma melhoria relativa de até 256\% enquanto oferece velocidades de inferência até 21,8x mais rápidas. O código está disponível publicamente em 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.