Uni-3DAR: Единая генерация и понимание 3D-данных с помощью авторегрессии на сжатых пространственных токенах

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

March 20, 2025

Авторы: Shuqi Lu, Haowei Lin, Lin Yao, Zhifeng Gao, Xiaohong Ji, Weinan E, Linfeng Zhang, Guolin Ke

cs.AI

Аннотация

Последние достижения в области крупных языковых моделей и их мультимодальных расширений продемонстрировали эффективность объединения генерации и понимания через авторегрессивное предсказание следующего токена. Однако, несмотря на критическую роль генерации и понимания трехмерных структур ({3D GU}) в ИИ для науки, эти задачи в основном развивались независимо, а авторегрессивные методы остаются недостаточно изученными. Чтобы устранить этот разрыв, мы представляем Uni-3DAR — унифицированную структуру, которая бесшовно интегрирует задачи {3D GU} через авторегрессивное предсказание. В основе Uni-3DAR лежит новая иерархическая токенизация, которая сжимает трехмерное пространство с использованием октодерева, используя присущую 3D-структурам разреженность. Затем применяется дополнительная токенизация для детализации структур, захватывая ключевые атрибуты, такие как типы атомов и точные пространственные координаты в микроскопических 3D-структурах. Мы также предлагаем две оптимизации для повышения эффективности. Первая — это стратегия двухуровневого сжатия поддеревьев, которая сокращает последовательность токенов октодерева до 8 раз. Вторая — механизм маскированного предсказания следующего токена, адаптированный для динамически изменяющихся позиций токенов, что значительно повышает производительность модели. Благодаря сочетанию этих стратегий, Uni-3DAR успешно объединяет разнообразные задачи {3D GU} в рамках единой авторегрессивной структуры. Многочисленные эксперименты на различных микроскопических задачах {3D GU}, включая молекулы, белки, полимеры и кристаллы, подтверждают ее эффективность и универсальность. В частности, Uni-3DAR значительно превосходит предыдущие модели диффузии, достигая относительного улучшения до 256\% при скорости вывода до 21,8 раз быстрее. Код доступен по адресу: 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.