Быстрое обучение диффузионного трансформера с экстремальным маскированием для генерации 3D-точечных облаков

Fast Training of Diffusion Transformer with Extreme Masking for 3D Point Clouds Generation

December 12, 2023

Авторы: Shentong Mo, Enze Xie, Yue Wu, Junsong Chen, Matthias Nießner, Zhenguo Li

cs.AI

Аннотация

Диффузионные трансформеры недавно продемонстрировали выдающуюся эффективность в генерации высококачественных 3D-облаков точек. Однако обучение воксельных диффузионных моделей для высокоразрешающих 3D-вокселей остается чрезмерно дорогостоящим из-за кубической сложности операторов внимания, обусловленной дополнительным измерением вокселей. Вдохновленные внутренней избыточностью 3D по сравнению с 2D, мы предлагаем FastDiT-3D — новый маскированный диффузионный трансформер, оптимизированный для эффективной генерации 3D-облаков точек, что значительно снижает затраты на обучение. В частности, мы черпаем идеи из маскированных автоэнкодеров для динамического выполнения процесса удаления шума на маскированных вокселизованных облаках точек. Мы также предлагаем новую стратегию маскирования, учитывающую воксели, для адаптивного агрегирования информации о фоне и переднем плане из вокселизованных облаков точек. Наш метод достигает наилучших результатов с экстремальным коэффициентом маскирования, близким к 99%. Более того, для улучшения генерации 3D-объектов нескольких категорий мы внедряем подход Mixture-of-Expert (MoE) в 3D-диффузионную модель. Каждая категория может изучать отдельный путь диффузии с разными экспертами, что снижает конфликт градиентов. Экспериментальные результаты на наборе данных ShapeNet демонстрируют, что наш метод достигает наилучших показателей по точности и разнообразию генерации 3D-облаков точек. FastDiT-3D улучшает метрики 1-Nearest Neighbor Accuracy и Coverage при генерации воксельных облаков точек с разрешением 128, используя всего 6,5% от исходных затрат на обучение.

English

Diffusion Transformers have recently shown remarkable effectiveness in generating high-quality 3D point clouds. However, training voxel-based diffusion models for high-resolution 3D voxels remains prohibitively expensive due to the cubic complexity of attention operators, which arises from the additional dimension of voxels. Motivated by the inherent redundancy of 3D compared to 2D, we propose FastDiT-3D, a novel masked diffusion transformer tailored for efficient 3D point cloud generation, which greatly reduces training costs. Specifically, we draw inspiration from masked autoencoders to dynamically operate the denoising process on masked voxelized point clouds. We also propose a novel voxel-aware masking strategy to adaptively aggregate background/foreground information from voxelized point clouds. Our method achieves state-of-the-art performance with an extreme masking ratio of nearly 99%. Moreover, to improve multi-category 3D generation, we introduce Mixture-of-Expert (MoE) in 3D diffusion model. Each category can learn a distinct diffusion path with different experts, relieving gradient conflict. Experimental results on the ShapeNet dataset demonstrate that our method achieves state-of-the-art high-fidelity and diverse 3D point cloud generation performance. Our FastDiT-3D improves 1-Nearest Neighbor Accuracy and Coverage metrics when generating 128-resolution voxel point clouds, using only 6.5% of the original training cost.