Glonce

Policy-Mask picks diffusion features under a channel budget

researchThe Glonce editorsHugging Face🔥56Glonce rating56Impact45Novelty55Relevance95Surprise25Players10Credibility75
Policy-Mask picks diffusion features under a channel budget

Frozen diffusion models can supply features for discriminative tasks such as segmentation and classification, but choosing which features to read is not simple. The paper says the useful signal is spread across the denoising trajectory, the depth of the network and the channels. In practice that means picking three things: a timestep, a block and a channel subset. According to the authors, existing methods either fix these choices by hand or tune only a few extraction hyperparameters, most often the timestep, and then reuse the same rule across inputs and budgets.

The paper introduces Policy-Mask, a budget-conditioned policy trained on the downstream objective. It selects the timestep, the block and the channel subset jointly, under a prescribed channel budget. One training detail gets attention: without normalization, the channel-score contribution grows with the number of selected channels and dominates the timestep and block updates. The authors normalize it by subset cardinality, which keeps the three policy heads on comparable scales and stabilizes joint training.

At matched budgets, the authors report that Policy-Mask improves over fixed extraction and feature-filter baselines across segmentation, classification and semantic correspondence. The headline results are 38.9 mIoU on ADE20K, 75.4% top-1 accuracy on CIFAR-100 at 128 dimensions, and 76.3 PCK@0.1 on SPair-71k.

Component-wise ablations show that block selection materially affects downstream accuracy. The highest recorded channel-budget score occurs at 256 retained channels, which the authors identify as a practical operating point for compact diffusion representations. Finally, they adapt Policy-Mask to diffusion transformers, including PixArt-Σ and SD3.5 MMDiT, by replacing U-Net channels with image-token hidden dimensions.

Key facts

  • Policy-Mask is a budget-conditioned policy that jointly chooses timestep, block and channel subset when extracting features from a frozen diffusion model.
  • Reported results at matched budgets: 38.9 mIoU on ADE20K, 75.4% top-1 on CIFAR-100 at 128 dimensions, 76.3 PCK@0.1 on SPair-71k.
  • Normalizing the channel score by subset cardinality keeps the three policy heads on comparable scales and stabilizes joint training.
  • Ablations show block selection materially affects accuracy; the highest recorded channel-budget score is at 256 retained channels.
  • The method is adapted to diffusion transformers, including PixArt-Σ and SD3.5 MMDiT, by swapping U-Net channels for image-token hidden dimensions.

Why it matters

Using a frozen diffusion model as a feature extractor means deciding where in the model to look. The authors argue that existing methods fix that choice by hand or tune mostly the timestep, then apply one rule to every input and budget. Policy-Mask instead learns the choice of timestep, block and channels together, and conditions it on a channel budget. That makes compact representations a target in themselves: the paper points to 256 retained channels as a practical operating point.

Who it affects

Researchers and engineers who reuse pretrained diffusion models for segmentation, classification or semantic correspondence, and who need to keep the feature dimension small. The extension to diffusion transformers such as PixArt-Σ and SD3.5 MMDiT widens the group to people working with non-U-Net backbones.

How to use it

The abstract describes the recipe at a high level: train a policy on your downstream objective, give it a channel budget, and let it pick the timestep, block and channel subset. Normalize the channel score by subset cardinality so the three policy heads train stably. For diffusion transformers, replace U-Net channels with image-token hidden dimensions. No code release, compute cost or inference-time figures are mentioned.

How solid is it

This is a preprint summarized from its abstract, and the results are the authors' own. The abstract names no authors or institutions. No baseline numbers are given, so the size of the improvement over fixed extraction or feature-filter baselines is not stated. The backbone model used for the ADE20K, CIFAR-100 and SPair-71k results is not named, and no results for PixArt-Σ or SD3.5 MMDiT are given.

Risks and caveats

Without baseline figures, the three reported scores cannot be judged against anything. The channel budget used for the ADE20K and SPair-71k results is not stated, and the 256-channel finding is described as the highest recorded channel-budget score in the authors' ablations. The diffusion transformer adaptation is described, but no results for it are given.

“the useful signal is distributed across the denoising trajectory, network depth, and channels”

— Policy-Mask paper abstract