💻 computer science

RePack then Refine: Efficient Diffusion Transformer with Vision Foundation Model

The paper proposes "RePack then Refine," a three-stage framework that enhances Diffusion Transformers by compressing redundant Vision Foundation Model features into a low-dimensional manifold for efficient training and subsequently refining the output to restore high-frequency details, achieving state-of-the-art convergence efficiency and image quality on ImageNet-1K.

Original authors: Guanfang Dong, Luke Schultz, Negar Hassanpour, Chao Gao

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC0 1.0

Original authors: Guanfang Dong, Luke Schultz, Negar Hassanpour, Chao Gao

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to paint beautiful pictures of animals, landscapes, and objects. To do this, the robot needs to understand the "essence" of what it's looking at before it starts painting.

In the world of AI, there are two main tools for this job:

The "Vision Foundation Model" (VFM): Think of this as a super-smart, highly educated art critic who has seen millions of images. It can describe an image in incredible detail, but it speaks in a very long, complex, and redundant language. It might say, "This is a cat," but then spend 768 different words describing the fur texture, the lighting, the background blur, and the exact shade of orange.
The "Diffusion Transformer" (DiT): This is the actual painter. It's talented but gets overwhelmed if the instructions are too long and messy. If you feed it the critic's 768-word description, the painter gets confused, takes a long time to learn, and often produces blurry or weird results.

The Problem:
Previous methods tried to just hand the painter the critic's raw, 768-word description. The paper argues this is like trying to read a novel written in a language where every sentence is repeated three times. It's inefficient, and the painter wastes time sorting through the noise.

The Solution: "RePack then Refine"
The authors propose a clever three-step process to fix this, which they call RePack then Refine.

Step 1: RePack (The "Summarizer")

Imagine the super-smart critic (the VFM) is talking to the painter. Instead of letting the critic ramble on for 768 words, you put a Summarizer in between them.

What it does: The Summarizer listens to the critic and instantly compresses that long, redundant speech into a short, 32-word "cheat sheet."
The Magic: It throws away the repetitive fluff (like saying "orange" three times) but keeps the most important structural information (that it's a cat, it's sitting, and it's orange).
The Result: The painter now receives a clean, compact instruction manual. Because the instructions are short and clear, the painter learns much faster. In the paper's experiments, this allowed the AI to reach a high level of skill in just 64 training sessions, whereas other methods needed hundreds or even thousands.

Step 2: The Painter (The DiT)

Now, the painter (the Diffusion Transformer) works on this clean, 32-word cheat sheet.

Because the instructions are so clear, the painter quickly figures out the big picture: the shape of the cat, where the eyes go, and the general pose.
However, because the instructions were so short (compressed), the painter misses the tiny, fine details. The cat might look perfect in shape, but its fur might look a bit smooth or plastic-like, lacking the "crunchy" texture of real fur.

Step 3: Refine (The "Detail Artist")

This is where the second part of the magic happens. The authors introduce a Refiner.

The Analogy: Think of the painter as a master sculptor who gets the shape right, and the Refiner as a master texturizer who adds the final touches.
What it does: The Refiner looks at the painter's smooth, basic cat and says, "I see the shape is perfect. Now, let me add the real fur, the whiskers, and the pores on the nose."
How it works: It uses the painter's "cheat sheet" as a guide to know where to put the details, but it works directly on the final image to add those high-frequency textures that were lost during the compression.

The Result

By combining these steps, the team created an AI that:

Learns incredibly fast: It reaches top-tier quality in record time (64 epochs) because it isn't bogged down by redundant data.
Produces stunning images: The final pictures are not only structurally perfect but also have realistic, high-definition textures.

In a Nutshell:
Instead of forcing the painter to read a 768-page manual, the authors gave them a 32-page summary (RePack) to learn the basics quickly, and then hired a specialist (Refine) to add the fine details at the end. This "pack it down, then build it up" strategy makes the whole process faster and the final result better.

Technical Summary: RePack then Refine

Problem Statement

Recent advancements in Latent Diffusion Models (LDMs) have increasingly leveraged Vision Foundation Models (VFMs), such as DINOv2/v3 and CLIP, as feature encoders to enhance semantic richness. However, the direct utilization of raw VFM features presents a critical inefficiency. As noted in the paper, raw features from standard ViT encoders (e.g., DINOv3-B/16) are high-dimensional (768 channels) and often redundant. The total number of elements in these feature maps frequently equals or exceeds the pixel count of the original image, contradicting the core LDM principle of operating on a compact, compressed latent space.

The authors argue that performing diffusion in such a high-dimensional, redundant space forces the Diffusion Transformer (DiT) to learn distributions dominated by noise and less critical information, significantly slowing convergence and reducing training efficiency. Empirical observations show that methods using raw features (e.g., RAE) converge slower than those using compressed spaces, despite the former's richer semantic input.

Methodology: RePack then Refine

To address the trade-off between semantic richness and training efficiency, the authors propose RePack then Refine, a three-stage framework designed to compress VFM features into a compact manifold while restoring high-frequency details.

1. Representation Packing (RePack)

The first stage introduces a lightweight, bias-free linear projection layer ( $P_\theta$ ) that maps high-dimensional raw VFM features ( $z_{raw} \in \mathbb{R}^{D \times h \times w}$ ) onto a compact, low-dimensional manifold ( $z \in \mathbb{R}^{d \times h \times w}$ , where $d \ll D$ , e.g., $d=32$ ).

Mechanism: The frozen VFM extracts features, which are then projected by $P_\theta$ . A decoder ( $D_\psi$ ) reconstructs the image from this compressed latent.
Objective: This stage filters out redundancy while preserving essential structural information. The authors hypothesize that semantic information lies on a low-dimensional manifold, and projecting onto this manifold simplifies the distribution modeling task for the subsequent DiT.
Training: Optimized end-to-end using a composite loss including $L_1$ , LPIPS, adversarial loss, Watson perceptual loss, and Focal Frequency Loss.

2. Generative Modeling (DiT)

In the second stage, a standard Diffusion Transformer (specifically LightningDiT) is trained on the compact latent space generated by the RePack encoder.

Process: The DiT learns the velocity field to transform Gaussian noise into the data distribution within the compressed 32-channel manifold.
Efficiency: Because the latent space is clean and semantically coherent, the DiT converges significantly faster than when trained on raw, high-dimensional features.

3. Latent-Guided Refinement

To compensate for the high-frequency texture details lost during the compression in the RePack stage, a Latent-Guided Refiner is introduced.

Architecture: A U-Net based model that takes the base image (decoded from the latent) and the upsampled latent features as input.
Function: It explicitly models high-frequency details in the pixel domain, guided by the structural semantics of the latent space. This decouples semantic generation (handled by DiT) from texture refinement (handled by the Refiner).
Training: Trained offline in a third stage using ground-truth images to learn the mapping from coarse semantics to fine details.

Key Contributions

RePack Framework: The proposal of a semantic compression framework that identifies redundancy in raw VFM features and employs a compact manifold bottleneck to filter noise while preserving core semantics.
Efficient Generative Modeling: Demonstration that generative modeling in this compact space yields state-of-the-art convergence speeds. The RePack-DiT-XL/1 achieves an FID of 1.82 in only 64 training epochs.
Latent-Guided Refiner: The introduction of a module that decouples structure generation from texture refinement, recovering high-frequency details lost during packing and improving the final FID to 1.65.

Experimental Results

Evaluated on ImageNet-1K (256×256), the framework demonstrates superior efficiency and quality:

Convergence Speed: RePack-DiT-XL/1 reaches an FID of 1.82 in 64 epochs. In comparison, standard DiT requires 1,400 epochs to reach an FID of 2.27, and RAE requires 800 epochs to reach 1.41 (though with a much larger parameter count and different training dynamics).
Quality Improvement: With the addition of the Refiner, the FID improves to 1.65 at 64 epochs, surpassing recent Pixel Diffusion Models (PDMs) and other LDMs in terms of convergence efficiency.
Ablation Studies:
- Projection Design: Bias-free linear projection outperforms biased variants and static PCA projections, confirming the necessity of learnable projection for capturing the semantic manifold.
- Latent Dimension: Reducing dimensions to $d=16$ offers the fastest convergence, but $d=32$ is chosen as a default to balance reconstruction fidelity and fair comparison with baselines like VA-VAE.
- Refiner Generalization: The Refiner significantly improves RePack but degrades performance on other VAEs (like VA-VAE), suggesting the Refiner is specifically tuned to the semantic-pure, texture-sparse nature of RePack's latent space.

Significance and Claims

The paper claims that RePack then Refine represents a highly effective strategy for balancing generative fidelity with training efficiency. By acknowledging the redundancy in raw VFM features, the authors argue that "packing" these features into a compact manifold is superior to direct usage. The framework achieves the fastest-converging results among recent methods, reaching high-quality generation (FID 1.65) in a fraction of the training epochs required by state-of-the-art baselines.

The authors emphasize that their approach does not require complex, parameter-heavy encoders (unlike FAE) or the fine-tuning of foundation models (unlike PS-VAE), thereby preserving the universality of the VFM while drastically reducing the computational burden of the generative model. The work suggests that the future of efficient diffusion lies in decoupling semantic compression from detail restoration, leveraging the complementary strengths of latent and pixel spaces.

Note: The authors acknowledge limitations regarding the current scope of experiments (limited to 80 epochs and ImageNet-1K) and note that scaling to text-to-image generation and exploring extended training schedules are directions for future work.