PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Model as Representation Learner

Xingyi Yang, Xinchao Wang

arXiv:2308.10916v1cs.CVcs.AI

TL;DR

The paper investigates whether pretrained DPM representations, whose usefulness for recognition remains unclear, can be reused despite architectural and computational challenges. It introduces RepFusion to transfer timestep-specific DPM representations to student networks, with reinforcement learning selecting the timestep. RepFusion consistently improves performance across several recognition benchmark types, highlighting DPMs’ potential as representation learners.

  • Problem

    The usefulness of pretrained DPM features for recognition remains unclear, and direct reuse is complicated by model architecture, computational cost, and timestep selection.

    Source: §1 · pp. 1–2

  • Method

    RepFusion distills intermediate representations from pretrained DPMs into student networks as auxiliary supervision, using reinforcement learning to select the timestep.

    Source: §1 · p. 2

  • Results

    RepFusion consistently improves performance across image classification, semantic segmentation, and landmark detection benchmarks.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The findings highlight DPMs’ potential as representation learners for recognition tasks beyond sample generation.

    Source: §1 · p. 2

    The analysis assumes a linear autoencoder with a skip connection to mimic the U-Net used in DPMs.

    Source: §3.2 · p. 3

Abstract

from arXiv · show

Diffusion Probabilistic Models (DPMs) have recently demonstrated impressive results on various generative tasks.Despite its promises, the learned representations of pre-trained DPMs, however, have not been fully understood. In this paper, we conduct an in-depth investigation of the representation power of DPMs, and propose a novel knowledge transfer method that leverages the knowledge acquired by generative DPMs for recognition tasks. Our study begins by examining the feature space of DPMs, revealing that DPMs are inherently denoising autoencoders that balance the representation learning with regularizing model capacity. To this end, we introduce a novel knowledge transfer paradigm named RepFusion. Our paradigm extracts representations at different time steps from off-the-shelf DPMs and dynamically employs them as supervision for student networks, in which the optimal time is determined through reinforcement learning. We evaluate our approach on several image classification, semantic segmentation, and landmark detection benchmarks, and demonstrate that it outperforms state-of-the-art methods. Our results uncover the potential of DPMs as a powerful tool for representation learning and provide insights into the usefulness of generative models beyond sample generation. The code is available at \url{https://github.com/Adamdad/Repfusion}.

1. Introduction

The paper asks whether DPM representations can be reused for recognition, while noting that direct reuse is challenged by model architecture, computational cost, and timestep selection. It frames DPMs as denoising autoencoders and introduces RepFusion, which transfers their representations to student networks and improves performance across recognition benchmarks.

1. Introduction

  • The paper addresses the underexplored question of whether DPM features can be effectively reused for recognition tasks.Prior work had explored related uses, sometimes through complicated model modifications, but the usefulness of DPM feature spaces remained unclear.Source: §1 · pp. 1–2
  • DPMs can be viewed as denoising autoencoders with varying denoising scales, connecting their features to established representation-learning methods.The paper relates this view to score matching and argues that DPMs produce meaningful encodings.Source: §1 · pp. 1–2
  • Direct DPM reuse for discriminative tasks is challenging because time-conditioned U-Nets are specialized, computationally heavy, and indexed by nontrivial timesteps.These constraints make the representations difficult to use directly for non-generative tasks.Source: §1 · p. 2
  • RepFusion transfers intermediate DPM representations to student networks as auxiliary supervision, using reinforcement learning to select the timestep.The method is motivated by the view that DPMs balance meaningful feature learning with regularization.Source: §1 · p. 2
  • RepFusion consistently improves performance across image classification, semantic segmentation, and landmark detection benchmarks.The paper presents these results as evidence of DPMs’ representation-learning capability.Source: §1 · p. 2

2. Related Work

Related work connects denoising autoencoders and diffusion models to representation learning, while prior generative-model reuse approaches often require modifications. This paper positions distillation from pretrained DPMs as an approach to general representation learning.

Related work

  • Denoising autoencoders learn representations through reconstruction from corrupted inputs [18] [36], and score matching connects them to generative modeling [62] [7].The paper describes DPMs as score-based models operating across multiple scales [24] [63] [62].Source: §2 · p. 2
  • Prior generative-model work explored semantic features and tasks including segmentation [76], correspondence [43] [64], and image editing [61], while some approaches require complex modifications [5].Other work learns from model-generated data [4] [74].Source: §2 · p. 2
  • The paper presents distilling pretrained DPMs for general representation learning as a novel approach to reusing generative models for recognition tasks.It situates this contribution among work on pretrained-model reuse [28] [77] [79] [53] and knowledge distillation [22] [80] [65].Source: §2 · pp. 2–3

3. Diffusion Models are Auto-Encoders

The section connects diffusion models to denoising autoencoders and examines how noise scale shapes their representations. Its analysis and visualizations show that feature quality varies with timestep, with moderate timesteps yielding stronger structure than very small or large ones.

3.1. Recap on Diffusion Probabilistic Models

  • A diffusion model progressively adds Gaussian noise to clean data, then learns a time-conditioned denoiser to reverse the process.Gaussian parameterization permits forming an arbitrary noisy timestep through a closed-form expression.Source: §3.1 · p. 3

3.2. Probing the Representations in DPMs

  • A linear DPM with a skip connection has an objective combining representation learning with regularization, which constrains the latent space.The representation-learning term favors a meaningful latent space, while spectral regularization drives the combined mapping toward a noise-dependent scale.Source: §3.2 · p. 3
  • As timestep increases, regularization takes priority and learned feature singular values and effective rank rapidly decline.The section identifies an ideal representation as having uniform singular values near 1 and larger effective rank.Source: §3.2 · pp. 3–4, Figure 2 · p. 4
  • T-SNE embeddings cluster classes most distinctly at t = 300 on MNIST and t = 100 on CIFAR-10, while very small or large timesteps blur class boundaries.These visualizations use unconditional DDPM features.Source: §3.2 · p. 4, Figure 3 · p. 4
  • On CIFAR-10, higher noise levels correspond to attention over more image tokens and long-range dependencies, whereas smaller timesteps capture local, fine-grained patterns.The attention maps are averaged across 128 samples and attention heads.Source: §3.2 · p. 4, Figure 4 · p. 4

4. RepFusion

RepFusion transfers intermediate representations from a pretrained diffusion model to a recognition student, using reinforcement learning to select a timestep for each sample. The student is distilled and then initialized for task-specific training.

4.1. Distill Knowledge from Diffusion Teacher

  • The student matches intermediate DPM and student features using a knowledge-distillation loss, with L2, attention-transfer, and relational losses as implementations.After distillation, the student is used as a feature extractor and initialized for task-label training.Source: §4.1 · pp. 4–5

4.2. Reinforced Time Selection for Distillation

  • RepFusion addresses the mismatch between a set of timestep-indexed DPM teachers and a single student by selecting the most predictive timestep for each sample.The authors describe exhaustive training across timestep and input combinations as generally impractical and use REINFORCED for selection.Source: §4 · p. 4, §4.2 · p. 5
  • RepFusion uses a policy network and REINFORCED to select a timestep per input, with negative task loss as reward, then distills the corresponding DPM feature into the student.The policy and teacher decoder are updated alongside student distillation; entropy encourages action diversity.Source: §4.2 · p. 5

5. Experiment

RepFusion is evaluated on image classification, semantic segmentation, and landmark detection, using pretrained DPMs as teachers for selected tasks and transferring a segmentation student to landmark detection. Results show improvements across tasks, while timestep-selection ablations favor the reinforced per-sample strategy.

5. Experiment

  • The experiments assess whether representations distilled from generative DPMs transfer to segmentation, keypoint detection, and classification.The authors report competitive performance across these downstream applications.Source: §5 · p. 6

5.1. Experimental Setting

  • Experiments cover CelebAMask-HQ segmentation, WFLW landmark detection, and CIFAR-10 and Tiny-ImageNet classification, with task-specific pretrained DPM teachers.Distillation is not applied on WFLW; instead, the CelebA-HQ segmentation weights initialize landmark detection.Source: §5.1 · p. 6

5.2.1 Semantic Segmentation

  • A 6.15 mIOU improvement over raw ResNet-18 and a 4.25 improvement over the ResNet-50 distilled student are reported for CelebAMask-HQ segmentation.RepFusion also matches the performance of up-to-date self-supervised representation learners; Figure 6 highlights rare patterns and minority categories.Source: Table 1 · p. 6, §5.2.1 · p. 7, Figure 6 · p. 7

5.2.2 Keypoint Detection

  • RepFusion achieves the smallest mean errors for both ResNet-50 and HRNetv2w18 in WFLW keypoint detection, while self-supervised methods also improve performance.Reported improvements in challenging subsets include 0.39/0.18 on NMEpose, 0.25/0.15 on NMEocclusion, and 0.24/0.09 on NMEillumination.Source: §5.2.2 · p. 7, Table 3 · p. 8

5.2.3 Image Classification

  • On CIFAR-10, accuracy gains average +0.42/+0.38 versus the student; on Tiny-ImageNet, gains are +1.83 over ResNet-18 and +0.42 on MBNv2.All three distillation losses boost student accuracy, with hint-based distillation showing the best improvement; RepFusion also outperforms students distilled from supervised teachers.Source: §5.2.3 · p. 8, Table 4 · p. 8, Table 5 · p. 8

5.3. Ablation Study and Analysis on Time Selection

  • Random or large fixed timesteps significantly degrade performance, even below no distillation; fixed small timesteps can help slightly, but are not universally effective.The reinforced per-sample strategy outperforms fixed-timestep selection, and selected timesteps tend to converge to a narrow range of 0–200 during training.Source: §5.3 · p. 9, Figure 8 · p. 9

6. Conclusion

The paper presents RepFusion as a way to transfer DPM representations to recognition tasks, with consistent improvements across recognition benchmarks. It also connects DPMs to denoising auto-encoders and uses reinforcement learning to choose when to distill representations.

6. Conclusion

  • RepFusion demonstrates consistent improvements across various recognition benchmarks.The conclusion describes DPMs as effective visual representations that facilitate transfer to downstream tasks.Source: §6 · p. 9
  • RepFusion uses reinforcement learning to determine the optimal timing for representation distillation.The approach transfers representation learning from trained generative models to discriminative tasks.Source: §6 · p. 9
  • The paper connects DPMs with denoising auto-encoders and empirically validates statistical properties of DPM-extracted features.Source: §6 · p. 9

Related papers

Loading 2308.10916v1…