A diffusion model’s useful visual signal is not concentrated in one layer or denoising step, so choosing features once and reusing them can waste a tight representation budget. This work treats feature extraction as a budget-conditioned decision: which timestep, block, and channels to keep can all depend on the downstream task and available dimensionality.
Key Findings
- Jointly selecting timesteps, blocks, and channels outperforms fixed extraction and feature-filter baselines at matched budgets across segmentation, classification, and semantic correspondence.
- Normalizing channel scores by the number of selected channels prevents that policy head from overwhelming timestep and block updates, making joint training more stable.
- Reported results include 38.9 mIoU on ADE20K, 75.4% top-1 accuracy on CIFAR-100 with 128 dimensions, and 76.3 [email protected] on SPair-71k. The best recorded channel-budget score is at 256 retained channels.
- The approach also extends to diffusion transformers such as PixArt-Σ and SD3.5 MMDiT by selecting image-token hidden dimensions instead of U-Net channels.
Method and Trade-offs
The policy learns feature-selection decisions against the downstream objective rather than relying on a single hand-picked extraction rule. Its practical value is strongest when a frozen diffusion model must supply compact representations for varied discriminative tasks. It adds a policy-training step and does not establish that one selected feature set will transfer equally well to every task or budget.
Who It’s For
Great fit for researchers evaluating frozen diffusion backbones as reusable visual feature extractors, especially when representation size is constrained. Look elsewhere if you need a general-purpose diffusion model, a generative method, or a plug-and-play feature extractor with no task-specific policy training.