Technical Summary: Distilling Drifting Transformers with Representation Autoencoders

Problem Statement

While Representation Autoencoders (RAEs) have significantly improved the performance of flow-based generative models by utilizing semantically rich, label-wise clustered latent spaces derived from pretrained DINO features, distilling these models into efficient one-step or few-step generators remains challenging.

The authors identify a fundamental geometric mismatch: RAE latent spaces exhibit severe anisotropy and large curvatures in their underlying Ordinary Differential Equation (ODE) trajectories compared to traditional Variational Autoencoder (VAE) spaces. Existing trajectory-based distillation methods rely on the assumption of smooth, nearly straight teacher trajectories. The highly curved trajectories in RAE spaces cause these methods to become unstable or inefficient. Furthermore, while the recently proposed Drifting Model offers a distribution-level approach that avoids matching curved trajectories, its original formulation requires an auxiliary, separately trained Masked Autoencoder (MAE) to provide concentrated features, introducing significant computational overhead.

Methodology

The paper proposes Drift-RAE, a distillation framework that adapts the Drifting Model paradigm specifically for RAE latent spaces, eliminating the need for auxiliary feature extractors.

1. Theoretical Analysis of Geometry and Dynamics

The authors first quantitatively analyze the latent spaces, demonstrating that RAE features are substantially more anisotropic (lower Participation Ratio and Spectral Entropy) than SD-VAE features. They theoretically prove (Theorem 1) that in high-dimensional spaces, if positive samples are overly dispersed, the attractive component of the drifting field tends to vanish, leading to a degenerated field.

• Key Insight: While the original Drifting Model requires an MAE to concentrate dispersed features, RAE latent spaces are naturally semantically concentrated. This allows the Drifting Model to operate directly on RAE latents without the auxiliary MAE module.

2. The Drift-RAE Framework

The method distills a pretrained flow model $v_\theta$ into a one-step generator $G_\theta(z) = z - v_\theta(z, 1)$. Instead of matching ODE trajectories, the generator is trained to minimize a drifting objective that pushes generated samples toward real data distributions and away from other generated samples.

• Objective: The loss is computed token-wise on the RAE latent representations, utilizing a normalized drifting field $\tilde{V}$ derived from positive (real) and negative (generated) samples.

3. Practical Modifications for Stability

To align the Drifting Model with the Diffusion-GAN framework and improve training stability, the authors introduce three specific modifications:

• Softmax Dimension: The authors remove the bi-directional softmax trick used in the original Drifting Model. They retain a single softmax over sample indices, which better aligns with the gradient direction of the underlying log-sum-exp potential.
• Input Perturbation: To prevent gradient vanishing caused by non-intersecting manifolds between real and generated data, the generated samples are perturbed with noise sampled from a Laplace distribution before computing the drifting field.
• Partial Detachment of Negative Samples: To increase the number of negative samples used for distribution approximation without exceeding GPU memory budgets, the authors propose using a larger pool of negative samples ($N_{total\_neg}$) where only a subset ($N_{neg}$) carries gradients, while the rest are detached. This improves the empirical approximation of the generated distribution.

Key Contributions

1. Geometric Analysis: A quantitative demonstration that RAE latent spaces possess high curvature and anisotropy, rendering standard trajectory-based distillation unstable.
2. Theoretical Justification: A proof showing that Drifting Models fail on overly dispersed spaces (necessitating an MAE in the original formulation) but thrive in the naturally concentrated semantic spaces of RAEs.
3. Drift-RAE: A novel distillation method that achieves state-of-the-art one-step generation in RAE spaces without requiring an auxiliary MAE feature extractor.
4. Stability Enhancements: Theoretical and empirical insights into modifying the Drifting objective (softmax dimension, input perturbation, and sample detachment) to align with Diffusion-GAN principles and improve convergence.

Experimental Results

The method was evaluated on the ImageNet 256 $\times$ 256 dataset using DiTDH-XL and DiTDH-L backbones.

• Performance: Drift-RAE achieves a FID of 1.77 (with DiTDH-XL) and 2.12 (with DiTDH-L) using a single generation step.
• Efficiency: The model reaches these results in only 10,000 distillation steps (approx. 16 epochs).
• Comparison:
  • It outperforms previous distillation methods in RAE spaces (e.g., MF-RAE).
  • It achieves comparable FID and improved FDDINOv2 scores compared to the original Drifting Model, despite not using the computationally expensive auxiliary MAE.
  • It remains competitive with one-step and few-step generators trained in other latent spaces (e.g., Pixel Space, SD-VAE).

Significance and Claims

The paper claims that Drifting provides an effective and promising framework for distilling generative models in representation spaces. The primary significance lies in resolving the instability of distilling RAE-based models by leveraging their inherent semantic concentration, thereby removing the dependency on auxiliary feature extractors.

The authors modestly note limitations, acknowledging that their theoretical analysis relies on simplified high-dimensional models and that training Drifting Models from scratch (without distillation) in RAE spaces remains an open problem. They also highlight that the requirement for many same-class positive samples may pose scalability challenges for text-to-image generation, suggesting this as a direction for future research.