RePack then Refine: Efficient Diffusion Transformer with Vision Foundation Model

TL;DR

The paper introduces RePack then Refine, a three-stage framework that efficiently leverages Vision Foundation Model features for diffusion transformers, significantly improving training speed and generative quality on ImageNet-1K.

Contribution

It proposes a novel feature compression and refinement strategy that enhances diffusion transformer training efficiency and generative performance.

Findings

RePack module reduces feature dimensionality while preserving structure.

RePack-DiT-XL/1 achieves an FID of 1.82 in 64 epochs.

Adding the Refiner improves FID to 1.65, surpassing recent LDMs.

Abstract

Semantic-rich features from Vision Foundation Models (VFMs) have been leveraged to enhance Latent Diffusion Models (LDMs). However, raw VFM features are typically high-dimensional and redundant, increasing the difficulty of learning and reducing training efficiency for Diffusion Transformers (DiTs). In this paper, we propose Repack then Refine, a three-stage framework that brings the semantic-rich VFM features to DiT while further accelerating learning efficiency. Specifically, the RePack module projects the high-dimensional features onto a compact, low-dimensional manifold. This filters out the redundancy while preserving essential structural information. A standard DiT is then trained for generative modeling on this highly compressed latent space. Finally, to restore the high-frequency details lost due to the compression in RePack, we propose a Latent-Guided Refiner, which is trained lastly for enhancing the image details. On ImageNet-1K, RePack-DiT-XL/1 achieves an FID of 1.82 in only 64 training epochs. With the Refiner module, performance further improves to an FID of 1.65, significantly surpassing latest LDMs in terms of convergence efficiency. Our results demonstrate that packing VFM features, followed by targeted refinement, is a highly effective strategy for balancing generative fidelity with training efficiency. Source code is publicly available at https://github.com/guanfangdong/RePack-then-Refine.