Policy-Mask: Budget-Aware Adaptation of Diffusion Features to Discriminative Tasks

AI summary

Overview

Research area: Computer Vision — transfer of frozen diffusion-model representations to discriminative downstream tasks (segmentation, classification, semantic correspondence).

Technical level: Advanced. The method combines diffusion feature extraction with reinforcement-learning-style policy gradients and subset sampling, and assumes familiarity with U-Net, DiT/MMDiT architectures, knowledge distillation, and policy-gradient estimators.

Scope: The paper proposes Policy-Mask, a budget-conditioned policy that jointly selects a denoising timestep, a network layer, and an exact-size channel subset from a frozen diffusion teacher, then distills the selected response into a student that runs alone at inference.

What This Paper Is About

A frozen diffusion model contains useful visual information, but that information is spread across the denoising trajectory, network depth, and channels, so using it for a discriminative task requires deciding which timestep, which block, and which channel subset to extract. Existing approaches normally hand-fix these choices or optimize only a few of them (most often just the timestep) and then reuse the same rule for every input and every budget. The paper's goal is to learn all three choices jointly, per input, under a prescribed channel budget, and to make that joint learning actually work rather than degrade.

Key Contributions

1. A budgeted policy over the triple (t, l, K) — timestep, layer, channel subset — with exact-cardinality sampling via Gumbel top-b over a Plackett–Luce ordering, so the selected subset always has exactly b channels rather than approximately b.
2. A channel-score normalization that divides the channel-head score by the subset size b, keeping the three policy heads on comparable scales; without it, the channel-score contribution grows with the number of selected channels and dominates the timestep and layer updates.
3. Empirical evidence that at 128 of 1280 aligned channels, the method gains 1.2–1.6 mIoU over full-width fixed extraction on four segmentation benchmarks, and beats the reported RepFusion and DDAE compact comparisons on classification and correspondence.
4. A formal action mapping from U-Net features (denoising steps, layers, channels) to DiT, MMDiT, and flow-transformer feature lattices, where the selected coordinate becomes the image-token hidden dimension, instantiated on PixArt-Σ and SD3.5 MMDiT.

Main Findings

• Segmentation at matched budget: With b = 128 channels, Policy-Mask scores 38.9 mIoU on ADE20K, 70.2 on Cityscapes, 79.1 on CelebA-HQ, and 74.8 on Pascal VOC, against 37.4, 68.6, 77.9, and 73.5 for all 1280 channels at a fixed (t, l) — gains of 1.5, 1.6, 1.2, and 1.3 mIoU respectively.
• Comparison to simple selectors: Fixed (t, l) plus fixed K reaches 33.8 / 62.9 / 72.4 / 68.1; top-variance selection reaches 35.6 / 64.8 / 74.6 / 70.3; a Fisher proxy reaches 36.4 / 65.7 / 75.3 / 71.0 on the same four benchmarks.
• Classification at 128 dimensions: 95.2% top-1 on CIFAR-10 and 75.4% on CIFAR-100, versus 94.2% and 74.1% for the reported RepFusion entries, 93.1% and 72.3% for top-variance, and 93.6% and 73.1% for Fisher. Full-width DDAE is listed at 95.4% and 76.8%.
• Semantic correspondence: 76.3 PCK@0.1 on SPair-71k with image normalization and 67.5 with bounding-box normalization; these SPair-71k entries are not reported for DDAE or RepFusion in the table.
• Expanding the action space alone does not help: In the controlled action-axis ablation, a timestep-only policy scores 31.78 mIoU; expanding to the time–layer grid raises this to 33.32, but adding a heuristic top-b subset drops it to 32.82, and the unnormalized joint policy falls to 31.68.
• Normalization matters: With a fixed alignment map, the normalized policy reaches 37.7 ± 0.7 mIoU versus 36.8 ± 0.6 without the 1/b factor, averaged over three seeds and 14 training epochs.
• Budget sweep: Increasing b from 128 to 256 raises the best validation mIoU by 0.72 and the mean over the last five epochs by 0.27 in a sweep that uses top-magnitude selection at a fixed timestep and layer rather than a learned policy. The highest recorded channel-budget score occurs at 256 retained channels.
• The budget sweep is confounded: The masked knowledge-distillation loss is averaged over all 1280 channels, so its mean over epochs 10–19 grows from 0.59 at b = 128 to 0.94 at b = 256 and 1.35 at b = 512, meaning the drop at 512 cannot be attributed to redundant channels.
• Transformer teachers transfer: Under matched capacity fractions between 9.98% and 10.03%, SD-1.5 U-Net at 128/1280 scores 34.1 / 70.8 / 72.8, PixArt-Σ at 115/1152 scores 35.0 / 71.6 / 73.5, and SD3.5 MMDiT at 154/1536 scores 36.4 / 72.9 / 74.6 on ADE20K / CIFAR-100 / SPair-71k. The authors state this is evidence for the interface, not a ranking of teacher architectures.
• Fixed-slot reference runs: On ADE20K over 14 epochs with one seed, all 1280 channels give 38.35 mIoU, a random 128-channel subset gives 37.88, and a top-variance 128-channel subset gives 37.70; a separate three-seed table lists RepFusion at 1280 channels as 33.1 ± 0.3.
• Qualitative behavior: On 2000 ADE20K validation images (seed 0), Policy-Mask has higher image-mIoU than RepFusion on 39.1% of images; class IoU improves for ball, flag, bridge, sand, monitor, and mountain, with losses for sconce, sink, animal, and desk.

Methodology in Plain English

The setup keeps the diffusion model frozen and puts a small student network in its place at inference time. The problem is framed as a choice: for a given image and a given channel budget b, pick a denoising timestep, pick a layer, and pick exactly b channels from that layer's feature map.

Because raw layers differ in width and resolution, a per-layer alignment map projects any candidate layer into one shared feature space of 1280 channels, so "128 channels" means the same thing everywhere. The selected channels are masked out of the aligned response, and that masked feature serves two purposes: it is the distillation target for the student, and it is fed to a small task decoder that produces a reward signal.

The policy is split into three heads — a categorical distribution over timesteps, a categorical distribution over layers, and channel logits. Channels are drawn using Gumbel noise added to the logits followed by a top-b selection, which yields an ordered draw without replacement and guarantees exactly b distinct channels. The probability of that draw is scored exactly with a Plackett–Luce expression. The policy is updated by REINFORCE with a baseline: the reward is the negative decoder loss, and the gradient weights the score of the sampled action.

The key practical fix is scaling the channel-head score by 1/b. Since the channel score is a sum of b terms while the timestep and layer scores each contain one term, the channel head otherwise dominates training as b grows. The authors describe this factor as a preconditioner rather than an unbiased gradient estimator, and note it is equivalent to shrinking only the channel-head learning rate for disjoint heads under plain SGD. The policy and task decoder are trained alongside the student, whose loss combines a task loss with a distillation loss to the masked frozen feature; the teacher is never differentiated through, and at inference only the student is evaluated.

Why This Matters

Impact on research. The paper reframes diffusion feature extraction as a resource-allocation decision rather than a hand-tuned constant, and shows that a broader action space is not automatically better — the unnormalized joint policy underperforms a timestep-only policy. That negative result, plus the 1/b normalization, is a useful lesson for anyone composing multiple policy heads with different score lengths. The interface mapping to DiT and MMDiT teachers extends the idea beyond U-Net diffusion models at a time when transformer-based generators dominate.

Real-world applications:

• Label-efficient dense prediction (segmentation, depth, correspondence) where a compact student must run without the diffusion teacher at inference.
• Edge and on-device vision, where retaining 128 of 1280 channels lowers the storage and compute needed for distillation targets and for the deployed student.
• Annotation-scarce domains such as medical or industrial inspection, where a pretrained generator's internal features can be repurposed with a small task decoder.
• Asset pipelines that must store teacher activations as distillation targets, where the budget directly controls how much must be cached.

Industry relevance. Any team already holding a frozen Stable Diffusion checkpoint can apply the selection policy as a preprocessing decision layer over existing distillation infrastructure, without retraining the teacher. The explicit channel budget gives a tunable knob trading accuracy against memory, which is the kind of control a deployment team needs.

Future Directions

• Decoupling budget from distillation strength: The masked KD loss is averaged over all 1280 channels, so changing b changes both the selected features and the loss scale. The authors call for a follow-up sweep with KD normalized over the selected channels to identify an operating budget independently of distillation strength; they identify this as the main open question.
• More repetitions: The normalization ablation and the RepFusion comparison use three seeds, but the action-axis ablation, the budget sweep, and the qualitative analysis use a single seed, leaving variance in those results unquantified.
• Broader transformer evaluation: The PixArt-Σ and SD3.5 MMDiT results use one evaluation per teacher, so they support the transferred interface but not a ranking across architectures; more teachers and repeated runs would test generality.
• Extending the budget dimension: The current action fixes cardinality exactly; whether the budget itself should be learned or conditioned more richly per input, rather than swept at fixed values, is left open.

Target Audience

Researchers and engineers working on diffusion-based representation transfer, knowledge distillation, and efficient dense prediction will get the most from this paper. It is also relevant to practitioners building compact students from large generative teachers, and to readers interested in budget-conditioned policy learning with exact-cardinality subset sampling. Readers without background in policy gradients, diffusion sampling, or transformer versus U-Net feature layouts will find the method section demanding.

Authors’ abstract

Selecting intermediate features from a frozen diffusion model for discriminative tasks is non-trivial because the useful signal is distributed across the denoising trajectory, network depth, and channels. In practice, this requires choosing a timestep, a block, and a channel subset. Existing methods either fix these choices by hand or optimize only a few extraction hyperparameters, most often the timestep, then reuse the same rule across inputs and budgets. We introduce Policy-Mask, a budget-conditioned policy trained on the downstream objective that jointly selects all three under a prescribed channel budget. Without normalization, the channel-score contribution grows with the number of selected channels and dominates the timestep and block updates. We normalize it by subset cardinality, keeping the three policy heads on comparable scales and stabilizing joint training. At matched budgets, Policy-Mask improves over fixed extraction and feature-filter baselines across segmentation, classification, and semantic correspondence, reaching 38.9 mIoU on ADE20K, 75.4% top-1 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, and the highest recorded channel-budget score occurs at 256 retained channels, identifying a practical operating point for compact diffusion representations. We further adapt Policy-Mask to diffusion transformers, including PixArt-$Σ$ and SD3.5 MMDiT, by replacing U-Net channels with image-token hidden dimensions.