Uni-3DAR: Generación y Comprensión Unificada en 3D mediante Autoregresión sobre Tokens Espaciales 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

Resumen

Los recientes avances en los modelos de lenguaje de gran escala y sus extensiones multimodales han demostrado la efectividad de unificar la generación y la comprensión mediante la predicción autoregresiva del siguiente token. Sin embargo, a pesar del papel crítico de la generación y comprensión de estructuras 3D ({3D GU}) en la IA para la ciencia, estas tareas han evolucionado en gran medida de manera independiente, con los métodos autoregresivos aún poco explorados. Para cerrar esta brecha, presentamos Uni-3DAR, un marco unificado que integra de manera fluida las tareas de {3D GU} mediante predicción autoregresiva. En su núcleo, Uni-3DAR emplea una novedosa tokenización jerárquica que comprime el espacio 3D utilizando un octree, aprovechando la inherente dispersión de las estructuras 3D. Luego, aplica una tokenización adicional para capturar detalles estructurales finos, incluyendo atributos clave como tipos de átomos y coordenadas espaciales precisas en estructuras 3D microscópicas. Además, proponemos dos optimizaciones para mejorar la eficiencia y efectividad. La primera es una estrategia de compresión de subárboles de dos niveles, que reduce la secuencia de tokens del octree hasta 8 veces. La segunda es un mecanismo de predicción del siguiente token enmascarado, diseñado para posiciones de tokens que varían dinámicamente, lo que mejora significativamente el rendimiento del modelo. Al combinar estas estrategias, Uni-3DAR logra unificar diversas tareas de {3D GU} dentro de un único marco autoregresivo. Experimentos exhaustivos en múltiples tareas de {3D GU} microscópicas, incluyendo moléculas, proteínas, polímeros y cristales, validan su efectividad y versatilidad. Notablemente, Uni-3DAR supera a los modelos de difusión anteriores de última generación por un margen considerable, logrando una mejora relativa de hasta el 256\% y velocidades de inferencia hasta 21.8 veces más rápidas. El código está disponible públicamente en 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.