PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ivan Listopadov, Maksim Nakhodnov, Aibek Alanov

arXiv:2610.28889v1cs.CV

TL;DR

Existing methods fix extraction choices or optimize only some hyperparameters, often the timestep, rather than jointly selecting timestep, layer, and channel subset. Policy-Mask trains a budget-conditioned policy to jointly select timestep, layer, and an exact-size channel subset, with channel-score normalization controlling its relative contribution. With 128 of 1280 aligned channels, POLICY-MASK exceeds full-width fixed extraction by 1.2–1.6 mIoU on four segmentation benchmarks and reaches 75.4% top-1 on CIFAR-100, 1.3 points above RepFusion.

  • Problem

    Existing methods fix extraction choices or optimize only some hyperparameters, often the timestep, rather than jointly selecting timestep, layer, and channel subset.

    Source: §1 · p. 1

  • Method

    Policy-Mask trains a budget-conditioned policy to jointly select timestep, layer, and an exact-size channel subset, with channel-score normalization controlling its relative contribution.

    Source: §1 · p. 1, §3.3 · p. 4

  • Results

    With 128 of 1280 aligned channels, POLICY-MASK exceeds full-width fixed extraction by 1.2–1.6 mIoU on four segmentation benchmarks and reaches 75.4% top-1 on CIFAR-100, 1.3 points above RepFusion.

    Source: §1 · p. 2, §4.2 · pp. 4–5

  • Takeaways & Limitations

    The experiments support joint selection at a fixed budget, and show that channel-score normalization improves the tested joint policy.

    Source: §5 · p. 7

    The budget sweep changes both selected features and distillation strength, so it does not isolate the effect of channel count.

    Source: §4.3 · p. 5, §4.5 · p. 7

Abstract

from arXiv · show

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-$\Sigma$ and SD3.5 MMDiT, by replacing U-Net channels with image-token hidden dimensions.

1 INTRODUCTION

Diffusion features for downstream tasks require choices across denoising timesteps, network layers, and channels. Policy-Mask jointly selects these under a channel budget, extending timestep selection and using channel-score normalization to balance policy updates.

1 INTRODUCTION

  • Existing approaches fix extraction rules, aggregate across timesteps and layers, or select timesteps, leaving layer and retained feature dimensions as additional choices.Those choices also determine which activations must be stored and used as distillation targets.Source: §1 · p. 1
  • Policy-Mask selects a timestep, layer, and exact-size channel subset under a prescribed budget, using task-decoder reward and selected features to supervise a student.At inference, the student processes the image without evaluating the diffusion teacher.Source: §1 · p. 1
  • A budgeted (t, l, K) policy uses exact-cardinality sampling, while channel-score normalization controls the relative contributions of the action heads.The method extends timestep selection to a joint choice of extraction point and feature subset.Source: §1 · p. 1
  • At 128 of 1280 aligned channels, gains are 1.2–1.6 mIoU over full-width fixed extraction on four segmentation benchmarks.The paper also formalizes mapping U-Net features to DiT, MMDiT, and flow-transformer feature lattices.Source: §1 · p. 2

2 RELATED WORK

Prior work studies diffusion representations for varied downstream tasks and feature-selection methods under budgets. Policy-Mask shares RepFusion’s teacher–student setup while expanding the selection to timestep, layer, and channel subset.

Diffusion representations

  • RepFusion distills a recognizer from a frozen diffusion teacher after selecting a useful noise step; Policy-Mask expands selection to (t, l, K).Both retain the teacher–student division.Source: §2 · p. 2

Budgeted subset selection

  • Feature-selection approaches include filters, wrappers, continuous relaxations, and greedy methods; here the budget constrains feature coordinates used for distillation.The policy uses a Plackett–Luce ordering and Gumbel top-b sampler.Source: §2 · p. 2

3 METHOD

Policy-Mask treats diffusion feature extraction as a budgeted decision and trains a selector from task-decoder reward. Its policy samples timestep, layer, and channel choices, while channel-score normalization controls the channel head’s relative contribution.

3.1 PROBLEM SETUP

  • A task decoder evaluates the selected teacher response and supplies the policy reward, while that response also supervises the student.The policy receives one scalar reward per sampled action; the teacher is not differentiated through.Source: §3 · p. 2, §3.1 · pp. 2–3
  • For budget b, feasible actions choose a timestep, layer, and channel subset K with |K| = b; alignment maps candidate layers into a common feature space.The selected response supervises the student, which alone consumes the image at inference.Source: §3.1 · p. 2

3.2 POLICY OVER BUDGETED SUBSETS

  • The policy factorizes timestep, layer, and ordered channel sampling; Gumbel top-b draws b distinct channels without replacement, and a score-function estimator updates the policy.A stopped moving baseline reduces variance; selector operations are small relative to frozen-teacher queries and layer alignment.Source: §3.2 · pp. 3–4

3.3 CHANNEL-SCORE NORMALIZATION

  • Dividing the channel score by b replaces its sum of b terms with a mean, controlling its relative weight against the single-term timestep and layer scores.The factor is a preconditioner, not a claim that gradient norms grow linearly with b or that the estimator remains unbiased.Source: §3.3 · p. 4
  • With fixed alignment, normalization raises mean best-validation mIoU from 36.8 to 37.7 over three seeds.In the action-axis ablation, the unnormalized joint policy reaches 31.68 mIoU versus 31.78 for the timestep-only policy.Source: §3.3 · p. 4

4 EXPERIMENTS

At a 128-channel budget, POLICY-MASK improves segmentation over full-width fixed extraction across four benchmarks and also reports gains on classification and correspondence. Ablations examine action choices, channel-score normalization, and budget; transformer experiments transfer the selection interface to image-token dimensions.

4.1 EXPERIMENTAL SETUP

  • The main selector comparisons use a frozen SD-1.5 U-Net and a common width of 1280 aligned channels, with compact selectors retaining 128 channels.Segmentation is evaluated on four datasets using mIoU; classification uses CIFAR-10 and CIFAR-100 top-1 accuracy, and correspondence uses SPair-71k PCK@0.1.Source: §4.1 · p. 4

4.2 MAIN RESULTS

  • At 128 selected channels, POLICY-MASK exceeds full-width extraction at a fixed (t, l) by 1.5, 1.6, 1.2, and 1.3 mIoU on ADE20K, Cityscapes, CelebA-HQ, and Pascal VOC, respectively.These results compare feature-selection rules at 10% of the full aligned width.Source: §4.2 · p. 4, Table 1 · p. 5
  • At 128 dimensions, POLICY-MASK reaches 95.2% top-1 on CIFAR-10 and 75.4% on CIFAR-100, 1.0 and 1.3 points above the reported RepFusion entries; on SPair-71k it reaches 76.3 PCK@0.1 with image normalization and 67.5 with bounding-box normalization.The comparisons are reported in Table 2.Source: Table 2 · p. 5, §4.2 · p. 5

4.3 ABLATION STUDY

  • In the controlled action-axis ablation, expanding a timestep-only policy to the time–layer grid raises mIoU from 31.78 to 33.32; the heuristic top-b subset reaches 32.82, and the unnormalized joint policy 31.68.With Al fixed, POLICY-MASK reaches 37.7 ± 0.7 mIoU versus 36.8 ± 0.6 without the 1/b factor, using the same three seeds and training settings. In the top-magnitude budget sweep, increasing b from 128 to 256 raises best validation mIoU by 0.72 and the last-five-epoch mean by 0.27, but the sweep also changes selected features and KD-loss scale.Source: §4.3 · p. 5

4.4 TRANSFORMER TEACHERS

  • For transformer teachers, the selected feature coordinate changes from U-Net channels to image-token hidden dimensions, with the selected target restored to a spatial grid.PixArt-Σ and SD3.5 MMDiT are evaluated at matched fractions of 9.98%–10.03% of available feature width; the evaluation supports the transferred interface, not teacher rankings.Source: §4.4 · pp. 5–6, Table 3 · p. 7, §4.5 · p. 7

4.5 LIMITATIONS

  • The budget sweep changes both selected features and distillation strength, and the action-axis ablation, budget sweep, and qualitative analysis use a single seed.Transformer interface evaluations use one evaluation per teacher and do not support ranking architectures.Source: §4.5 · p. 7

5 CONCLUSION

POLICY-MASK jointly selects a timestep, layer, and exact-size channel subset to distill a student that runs without the diffusion teacher at inference. The experiments report gains over fixed extraction and RepFusion, benefits from channel-score normalization, and transfer to transformer teachers.

5 CONCLUSION

  • With 128 of 1280 channels, POLICY-MASK exceeds full-width fixed extraction on all four segmentation benchmarks and RepFusion on CIFAR-10 and CIFAR-100.With fixed alignment, normalization raises ADE20K mIoU from 36.8 to 37.7 over three seeds.Source: §5 · p. 7
  • With fixed alignment, normalization raises ADE20K mIoU from 36.8 to 37.7 over three seeds.Without the 1/b factor, the joint policy falls below the timestep-only policy in the action-axis ablation.Source: §5 · p. 7
  • The action transfers to DiT and MMDiT teachers by replacing U-Net channels with image-token hidden dimensions while retaining the student loss and exact-cardinality constraint.Source: §5 · p. 7
  • Separating channel count from distillation strength remains the main open question for choosing the operating budget.Source: §5 · p. 7

A.1 CHANNEL-BUDGET SWEEP

The channel-budget sweep uses top-magnitude selection at a fixed timestep and layer, training a separate student for each budget. Its highest validation mIoU occurs at 256 channels, but the experiment does not isolate channel count from distillation-loss scale.

A.1 CHANNEL-BUDGET SWEEP

  • The highest validation mIoU in the sweep occurs at b = 256.Source: §A.1 · p. 11
  • The sweep uses top-magnitude selection at a fixed timestep and layer, with a separate student trained for each budget.Source: §A.1 · p. 11
  • Because the masked distillation loss is averaged over all 1280 channels, changing the budget also changes its scale and does not isolate the effect of channel count.Source: §A.1 · p. 11

A.2 FIXED ALIGNMENT

In diagnostic runs with trainable alignment, the KD loss fell below 0.005 within two epochs. The authors attribute this behavior as consistent with projection shrinkage and fixed alignment for two comparisons.

A.2 FIXED ALIGNMENT

  • With trainable Al, the KD loss fell below 0.005 within two epochs in each diagnostic run.Source: §A.2 · p. 11
  • The observed loss behavior is consistent with projection shrinkage toward zero.Source: §A.2 · p. 11
  • The authors fixed Al for the channel-score normalization ablation and the RepFusion comparison.Source: §A.2 · p. 11

B ADDITIONAL RESULTS

The additional results include fixed-slot ADE20K reference runs and comparisons of channel-score normalization with RepFusion. The supplied passages describe their protocols and metrics, but do not provide the table values.

B ADDITIONAL RESULTS

  • Table 4 reports best validation mIoU for 14-epoch, one-seed fixed-slot reference runs on ADE20K.The full-width baseline uses 1280 channels; random and top-variance selection retain 128.Source: Table 4 · p. 11, §B.1 · p. 11

B.1 FIXED-SLOT REFERENCE RUNS

  • Table 5 reports mean and sample standard deviation of best validation mIoU over three seeds after 14 training epochs with fixed Al.Source: Table 5 · p. 11

B.2 CHANNEL-SCORE NORMALIZATION AND REPFUSION COMPARISON

  • The normalization and RepFusion comparisons use fixed Al and the same seeds and training settings; they differ only in the 1/b channel-score factor.RepFusion uses all 1280 dimensions, while POLICY-MASK retains 128.Source: §B.2 · p. 11

C QUALITATIVE AND CLASS-WISE ANALYSIS

The seed-0 comparison shows POLICY-MASK has higher image-mIoU than RepFusion on 39.1% of ADE20K validation images, alongside class-wise gains and losses. The figures highlight the largest class differences and outcome-selected image losses.

C QUALITATIVE AND CLASS-WISE ANALYSIS

  • 39.1% of 2000 ADE20K validation images have higher image-mIoU with POLICY-MASK than RepFusion in the seed-0 comparison.This image-level comparison is separate from dataset mIoU reported across three seeds.Source: C QUALITATIVE AND CLASS-WISE ANALYSIS · p. 11
  • Class IoU gains include ball, flag, bridge, sand, monitor, and mountain, while losses include sconce, sink, animal, and desk.Figure 5 shows the 12 largest gains and losses among classes present in at least 20 validation images; positive values favor POLICY-MASK.Source: C QUALITATIVE AND CLASS-WISE ANALYSIS · p. 11, Figure 5 · p. 12
  • Figure 6 presents the 12 ADE20K validation images with the largest image-mIoU losses relative to RepFusion, selected by outcome.Its columns follow Figure 3.Source: Figure 6 · p. 12
Loading 2610.28889v1…