Distilling Drifting Transformers with Representation Autoencoders

2026-06-13

AI Summary

This work addresses the instability and convergence difficulties in existing representation autoencoder (RAE)-based distillation methods, which stem from high curvature and strong anisotropy in the latent space. To mitigate these geometric distortions, we propose Drift-RAE, the first approach to integrate a drift model into the RAE distillation framework. By constructing a theoretically aligned drift field in the latent space, Drift-RAE substantially enhances training stability without requiring an auxiliary MAE feature extractor. On ImageNet 256, our method achieves a state-of-the-art FID of 1.77 using only 10k distillation steps, outperforming current RAE-based distillation techniques. Furthermore, our analysis reveals distinct curvature and isotropy characteristics across different autoencoder latent spaces.

Abstract

Representation Autoencoders (RAEs) have improved diffusion and flow models by semantically richer latent space owing to the strongly label-wise clustered DINO features in the pretrained encoders. Yet in the distillation stage, the severe anisotropy and large curvatures caused by the rich semantic representations would hinder the convergence and performance, making the trajectory-based distillation unstable. In this work, we argue that the RAE latent space is compatible with distillation via the newly proposed Drifting Models. We first quantitatively study the curvatures and isotropy statistics across different autoencoders, and theoretically reveal that Drifting Model itself is highly likely to fail on extremely scattered spaces like reconstruction-based VAEs. These motivate us to apply the drifting paradigm directly to representation autoencoders. Our proposed method, Drift-RAE, distills pretrained flow models in RAE latent spaces using Drifting, together with insightful modifications that improve training stability by thereotically aligning drifting fields with other frameworks. Regarding the experimental evidences, we achieve 1.77 FID on ImageNet 256 dataset using only 10k distillation steps, surpassing state-of-the-art RAE distillation methods and appearing comparative with the original Drifting Model without requiring an auxiliary MAE feature extractor. The code will be made publicly available.

Problem

• Representation Autoencoders
• distillation
• latent space anisotropy
• trajectory-based distillation
• curvature

Innovation

• Drifting Models
• Representation Autoencoders
• Latent Space Distillation
• Anisotropy
• Flow Model Distillation