Flux en quelques étapes pour la génération 3D via la distillation par transport de données marginales

Few-step Flow for 3D Generation via Marginal-Data Transport Distillation

September 4, 2025

papers.authors: Zanwei Zhou, Taoran Yi, Jiemin Fang, Chen Yang, Lingxi Xie, Xinggang Wang, Wei Shen, Qi Tian

cs.AI

papers.abstract

Les modèles de génération 3D basés sur des flux nécessitent généralement des dizaines d'étapes d'échantillonnage lors de l'inférence. Bien que les méthodes de distillation en quelques étapes, en particulier les Modèles de Consistance (CMs), aient réalisé des avancées significatives pour accélérer les modèles de diffusion 2D, elles restent peu explorées pour les tâches de génération 3D plus complexes. Dans cette étude, nous proposons un nouveau cadre, MDT-dist, pour la distillation de flux 3D en quelques étapes. Notre approche repose sur un objectif principal : distiller le modèle pré-entraîné pour apprendre le Transport Marginal des Données. L'apprentissage direct de cet objectif nécessite d'intégrer les champs de vitesse, mais cette intégrale est difficile à mettre en œuvre. Par conséquent, nous proposons deux objectifs optimisables, l'Appariement de Vitesse (VM) et la Distillation de Vitesse (VD), pour convertir de manière équivalente la cible d'optimisation du niveau de transport vers le niveau de vitesse et de distribution respectivement. L'Appariement de Vitesse (VM) apprend à faire correspondre de manière stable les champs de vitesse entre l'élève et le professeur, mais fournit inévitablement des estimations de gradient biaisées. La Distillation de Vitesse (VD) améliore encore le processus d'optimisation en exploitant les champs de vitesse appris pour effectuer une distillation de densité de probabilité. Lorsqu'elle est évaluée sur le cadre de génération 3D pionnier TRELLIS, notre méthode réduit les étapes d'échantillonnage de chaque transformateur de flux de 25 à 1 ou 2, atteignant des latences de 0,68s (1 étape x 2) et 0,94s (2 étapes x 2) avec des accélérations de 9,0x et 6,5x sur A800, tout en préservant une haute fidélité visuelle et géométrique. Des expériences approfondies démontrent que notre méthode surpasse significativement les méthodes de distillation CM existantes et permet à TRELLIS d'atteindre des performances supérieures dans la génération 3D en quelques étapes.

English

Flow-based 3D generation models typically require dozens of sampling steps during inference. Though few-step distillation methods, particularly Consistency Models (CMs), have achieved substantial advancements in accelerating 2D diffusion models, they remain under-explored for more complex 3D generation tasks. In this study, we propose a novel framework, MDT-dist, for few-step 3D flow distillation. Our approach is built upon a primary objective: distilling the pretrained model to learn the Marginal-Data Transport. Directly learning this objective needs to integrate the velocity fields, while this integral is intractable to be implemented. Therefore, we propose two optimizable objectives, Velocity Matching (VM) and Velocity Distillation (VD), to equivalently convert the optimization target from the transport level to the velocity and the distribution level respectively. Velocity Matching (VM) learns to stably match the velocity fields between the student and the teacher, but inevitably provides biased gradient estimates. Velocity Distillation (VD) further enhances the optimization process by leveraging the learned velocity fields to perform probability density distillation. When evaluated on the pioneer 3D generation framework TRELLIS, our method reduces sampling steps of each flow transformer from 25 to 1 or 2, achieving 0.68s (1 step x 2) and 0.94s (2 steps x 2) latency with 9.0x and 6.5x speedup on A800, while preserving high visual and geometric fidelity. Extensive experiments demonstrate that our method significantly outperforms existing CM distillation methods, and enables TRELLIS to achieve superior performance in few-step 3D generation.