Uni-3DAR: Geünificeerde 3D-generatie en -begrip via autoregressie op gecomprimeerde ruimtelijke tokens

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

Samenvatting

Recente vooruitgang in grote taalmodel(len) en hun multimodale uitbreidingen hebben de effectiviteit aangetoond van het verenigen van generatie en begrip via autoregressieve voorspelling van het volgende token. Desondanks, ondanks de kritieke rol van 3D-structuurgeneratie en -begrip ({3D GU}) in AI voor de wetenschap, hebben deze taken zich grotendeels onafhankelijk ontwikkeld, waarbij autoregressieve methoden onderbelicht zijn gebleven. Om deze kloof te overbruggen, introduceren we Uni-3DAR, een geïntegreerd framework dat {3D GU}-taken naadloos integreert via autoregressieve voorspelling. De kern van Uni-3DAR bestaat uit een nieuwe hiërarchische tokenisatie die 3D-ruimte comprimeert met behulp van een octree, waarbij de inherente schaarste van 3D-structuren wordt benut. Vervolgens wordt een aanvullende tokenisatie toegepast voor fijnmazige structurele details, waarbij belangrijke attributen zoals atoomtypen en precieze ruimtelijke coördinaten in microscopische 3D-structuren worden vastgelegd. We stellen verder twee optimalisaties voor om de efficiëntie en effectiviteit te verbeteren. De eerste is een tweeledige compressiestrategie voor subtrees, die de octree-tokenreeks tot 8x verkleint. De tweede is een gemaskeerd voorspellingsmechanisme voor het volgende token, speciaal ontworpen voor dynamisch variërende tokenposities, wat de modelprestaties aanzienlijk verbetert. Door deze strategieën te combineren, slaagt Uni-3DAR erin diverse {3D GU}-taken te verenigen binnen een enkel autoregressief framework. Uitgebreide experimenten over meerdere microscopische {3D GU}-taken, waaronder moleculen, eiwitten, polymeren en kristallen, valideren de effectiviteit en veelzijdigheid ervan. Opmerkelijk is dat Uni-3DAR eerdere state-of-the-art diffusiemodellen met een aanzienlijke marge overtreft, met een relatieve verbetering van tot 256\% en een inferentiesnelheid die tot 21,8x sneller is. De code is publiekelijk beschikbaar op 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.