Technical Summary: Teacher-Feature Drifting (TFD)

Problem Statement

Pretrained diffusion and flow-matching models achieve high-fidelity image generation but suffer from expensive inference costs, often requiring tens or hundreds of forward passes to generate a single image. While distillation methods exist to compress these multi-step teachers into efficient one-step or few-step generators, current approaches face significant complexity. Methods like Distribution Matching Distillation (DMD) and its variants rely on auxiliary networks (e.g., fake-score networks, adversarial discriminators) or multi-stage training schedules. Other approaches, such as Consistency Models, may suffer from unstable training or require complex Jacobian-vector products. Furthermore, recent "Drifting Model" objectives, which offer a direct attraction-repulsion loss for one-step generation, typically depend on an external pretrained feature extractor (e.g., self-supervised models or latent-MAE) to define the representation space. This introduces additional training overhead and architectural complexity.

The core problem addressed is how to achieve efficient, high-quality one-step diffusion distillation using a single objective without relying on auxiliary networks, complex optimization pipelines, or external feature extractors.

Methodology: Teacher-Feature Drifting (TFD)

The authors propose Teacher-Feature Drifting (TFD), a framework that simplifies one-step distillation by leveraging the internal representations of the pretrained teacher model itself.

1. Teacher-Induced Representation Space

Unlike the original Drifting Model, which requires an external feature encoder, TFD uses the intermediate hidden states of the frozen pretrained teacher as the feature space for computing the drifting loss.

• Mechanism: For a generated sample $x$ and a real sample $x^+$, features are extracted from specific layers $\ell$ of the teacher network $T$: $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$.
• Drifting Loss: The drifting field is computed directly in this teacher-induced space. The generator $f_\theta$ is optimized to minimize the distance between its output and a "frozen transported target" defined by the attraction-repulsion field of the teacher's feature space. This aligns the distillation objective with the model being distilled, eliminating the need for a separate representation network.

2. Noised Feature Extraction

A key empirical finding is that extracting features from moderately noised inputs yields better results than using clean samples.

• Before feeding a sample into the teacher for feature extraction, it is perturbed with the teacher's forward noising process: $x_{tf} = x + \sigma_{tf} \xi$.
• The authors hypothesize that this perturbation smooths the teacher's feature geometry, suppressing high-frequency discrepancies and making the representation space more suitable for distribution-level matching. Experiments indicate a noise level of $\sigma_{tf} = 0.1$ is optimal.

3. Anchor-Margin Regularization

To address the issue of mode collapse (where the generator fails to cover diverse regions supported by the teacher), the authors introduce a lightweight anchor-margin loss.

• Anchors: A set of anchor features $\{a_i\}$ is extracted from real or teacher-generated samples.
• Mechanism: The loss encourages generated samples to occupy regions around these anchors. It adaptively adjusts the required "support" for each anchor based on the local density of the anchor distribution itself.
• Goal: This provides a complementary coverage signal to the local attraction-repulsion drifting objective, ensuring the student generator covers diverse teacher-supported regions without introducing new networks.

The final objective function combines the Teacher-Feature Drifting loss ($L_{TFD}$) and the anchor-margin loss ($L_{anchor}$):
$$L = L_{TFD} + \lambda_{anchor} L_{anchor}$$

Key Contributions

1. Simplified Framework: TFD proposes a one-step diffusion distillation framework that computes the drifting objective directly in the hidden representation space of the pretrained teacher, removing the need for an additional representation network.
2. Noised Feature Extraction: The work demonstrates that extracting teacher features from moderately noised inputs improves the feature geometry for distillation, leading to better FID scores.
3. Mode Collapse Mitigation: A lightweight anchor-margin loss is introduced to encourage sample diversity and coverage of teacher-supported feature regions, addressing a common failure mode in distillation.
4. Performance: The method achieves competitive image quality and diversity with a single forward pass, substantially simplifying the overall distillation pipeline compared to state-of-the-art baselines.

Experimental Results

The authors evaluated TFD on ImageNet-64×64 (class-conditional) and SDXL (text-to-image) generation.

• ImageNet-64×64:

  • TFD achieved an FID of 1.58 with a single forward pass after 200k training iterations.
  • This result is competitive with DMD2 (FID 1.51) and significantly outperforms other one-step baselines like DMD (2.62), EMD (2.20), and CTM (1.92).
  • Convergence: TFD reached FID $\le$ 10 in 1.0k updates, whereas the DMD2 baseline required 5.5k updates (5.5x faster). It reached FID $\le$ 3 in 12.5k updates compared to DMD2's 24.0k.

• SDXL Text-to-Image:

  • On the COCO 2014 dataset, TFD achieved an FID of 18.39 with one forward pass, outperforming other distilled SDXL variants (e.g., DMD2 at 19.01, SDXL-Turbo at 24.57).
  • It maintained a competitive CLIP score of 0.332, indicating strong text-image alignment.
  • Qualitative results showed that TFD preserves coherent structure, visual detail, and prompt alignment comparable to 4-step distilled baselines and the 50-step teacher, despite using only one step.

Significance and Claims

The paper claims that TFD offers a simpler and more efficient alternative to existing distillation frameworks. By utilizing the teacher model itself as the feature extractor, TFD avoids the complexity of training auxiliary networks or designing multi-stage pipelines. The introduction of noised feature extraction and anchor-margin regularization addresses specific challenges in representation geometry and mode coverage, respectively.

The authors position TFD as a method that achieves efficient one-step generation with competitive image quality and diversity, demonstrating that a single objective in the teacher's native representation space is sufficient for high-fidelity distillation. The work suggests that the "drifting" concept can be effectively realized without external feature encoders, provided the teacher's internal states are utilized correctly.

Limitations: The authors note that the current work focuses on U-Net-based teachers and has not yet explored Diffusion Transformer (DiT) architectures. Additionally, the selection of feature layers and noise levels currently lacks theoretical guidance, and the method has not been tested on larger-scale datasets or higher-resolution settings beyond the current experiments.