Treinamento Rápido de Transformadores de Difusão com Mascaramento Extremo para Geração de Nuvens de Pontos 3D

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

December 12, 2023

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

cs.AI

Resumo

Os Transformers de Difusão têm demonstrado recentemente uma eficácia notável na geração de nuvens de pontos 3D de alta qualidade. No entanto, o treinamento de modelos de difusão baseados em voxels para voxels 3D de alta resolução permanece proibitivamente caro devido à complexidade cúbica dos operadores de atenção, que surge da dimensão adicional dos voxels. Motivados pela redundância inerente do 3D em comparação com o 2D, propomos o FastDiT-3D, um novo transformer de difusão mascarado projetado para a geração eficiente de nuvens de pontos 3D, que reduz significativamente os custos de treinamento. Especificamente, inspiramo-nos em autoencoders mascarados para operar dinamicamente o processo de remoção de ruído em nuvens de pontos voxelizadas mascaradas. Também propomos uma nova estratégia de mascaramento consciente de voxels para agregar de forma adaptativa informações de fundo/primeiro plano de nuvens de pontos voxelizadas. Nosso método alcança desempenho de ponta com uma taxa de mascaramento extrema de quase 99%. Além disso, para melhorar a geração 3D de múltiplas categorias, introduzimos o Mixture-of-Expert (MoE) no modelo de difusão 3D. Cada categoria pode aprender um caminho de difusão distinto com diferentes especialistas, aliviando conflitos de gradiente. Resultados experimentais no conjunto de dados ShapeNet demonstram que nosso método alcança desempenho de ponta na geração de nuvens de pontos 3D de alta fidelidade e diversidade. Nosso FastDiT-3D melhora as métricas de Acurácia do Vizinho Mais Próximo (1-NN) e Cobertura ao gerar nuvens de pontos voxelizadas de resolução 128, utilizando apenas 6,5% do custo original de treinamento.

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.