Entraînement rapide de Transformateur de Diffusion avec masquage extrême pour la génération de nuages de points 3D

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

December 12, 2023

Auteurs: Shentong Mo, Enze Xie, Yue Wu, Junsong Chen, Matthias Nießner, Zhenguo Li

cs.AI

Résumé

Les Transformers de Diffusion ont récemment démontré une efficacité remarquable dans la génération de nuages de points 3D de haute qualité. Cependant, l'entraînement de modèles de diffusion basés sur des voxels pour des voxels 3D haute résolution reste prohibitivement coûteux en raison de la complexité cubique des opérateurs d'attention, qui découle de la dimension supplémentaire des voxels. Motivés par la redondance inhérente de la 3D par rapport à la 2D, nous proposons FastDiT-3D, un nouveau transformer de diffusion masqué conçu pour une génération efficace de nuages de points 3D, réduisant considérablement les coûts d'entraînement. Plus précisément, nous nous inspirons des autoencodeurs masqués pour opérer dynamiquement le processus de débruitage sur des nuages de points voxélisés masqués. Nous proposons également une nouvelle stratégie de masquage adaptée aux voxels pour agréger de manière adaptative les informations de premier plan/arrière-plan des nuages de points voxélisés. Notre méthode atteint des performances de pointe avec un taux de masquage extrême de près de 99%. De plus, pour améliorer la génération 3D multi-catégories, nous introduisons le Mixture-of-Expert (MoE) dans le modèle de diffusion 3D. Chaque catégorie peut apprendre un chemin de diffusion distinct avec différents experts, atténuant les conflits de gradient. Les résultats expérimentaux sur le jeu de données ShapeNet démontrent que notre méthode atteint des performances de pointe en termes de fidélité et de diversité dans la génération de nuages de points 3D. Notre FastDiT-3D améliore les métriques de précision du 1-Nearest Neighbor et de couverture lors de la génération de nuages de points voxélisés à 128 résolutions, en utilisant seulement 6,5% du coût d'entraînement original.

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.