Entrenamiento Rápido de Transformadores de Difusión con Enmascaramiento Extremo para la Generación de Nubes de Puntos 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

Resumen

Los Transformadores de Difusión han demostrado recientemente una notable eficacia en la generación de nubes de puntos 3D de alta calidad. Sin embargo, entrenar modelos de difusión basados en vóxeles para vóxeles 3D de alta resolución sigue siendo prohibitivamente costoso debido a la complejidad cúbica de los operadores de atención, que surge de la dimensión adicional de los vóxeles. Motivados por la redundancia inherente del 3D en comparación con el 2D, proponemos FastDiT-3D, un novedoso transformador de difusión enmascarado diseñado para la generación eficiente de nubes de puntos 3D, lo que reduce significativamente los costos de entrenamiento. En concreto, nos inspiramos en los autoencodificadores enmascarados para operar dinámicamente el proceso de eliminación de ruido en nubes de puntos voxelizadas enmascaradas. También proponemos una nueva estrategia de enmascaramiento consciente de vóxeles para agregar de manera adaptativa la información de fondo/primer plano de las nubes de puntos voxelizadas. Nuestro método alcanza un rendimiento de vanguardia con una proporción de enmascaramiento extrema de casi el 99%. Además, para mejorar la generación 3D de múltiples categorías, introducimos el enfoque de Mezcla de Expertos (MoE) en el modelo de difusión 3D. Cada categoría puede aprender una ruta de difusión distinta con diferentes expertos, aliviando los conflictos de gradiente. Los resultados experimentales en el conjunto de datos ShapeNet demuestran que nuestro método logra un rendimiento de vanguardia en la generación de nubes de puntos 3D de alta fidelidad y diversidad. Nuestro FastDiT-3D mejora las métricas de Precisión del Vecino Más Cercano (1-NN) y Cobertura al generar nubes de puntos voxelizadas de resolución 128, utilizando solo el 6.5% del costo original de entrenamiento.

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.