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

Ivan Listopadov; Maksim Nakhodnov; Aibek Alanov

cs.CVOct 20, 2026

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

Authors: Ivan Listopadov, Maksim Nakhodnov, Aibek Alanov

Organizations: HSE University · MSU · Constructor University · FusionBrain Lab, AXXX

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 [email protected] 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.

Figures & tables

Figure 1: Policy-Mask training and inference. Given an image $x$ and a budget $b$ , the policy $\pi_{\theta}$ outputs distributions over timesteps and layers and channel logits $z$ . A timestep $t$ and a layer $l$ are sampled. An ordered sequence $\sigma\sim\pi_{\sigma}$ of $b$ channels is drawn without replacement by Gumbel top- $b$ , giving the mask $m_{K}$ . The frozen diffusion model $D$ returns $G_{t,l}(x)$ after the alignment $A_{l}$ , and masking gives $\widetilde{F}_{a}(x)$ . This feature is the distillation target for the student $S_{\phi}$ and, through the task decoder $q_{\omega}$ , gives the reward $R=-\ell_{\mathrm{dec}}$ . The policy is updated by REINFORCE with baseline $c$ and the channel score scaled by $1/b$ (Equation 8 ). At inference, only $S_{\phi}$ is evaluated. Images and predictions are ADE20K validation examples (Figure 3 ).

Selector	Width	ADE20K	Cityscapes	CelebA-HQ	Pascal VOC
All channels at fixed $(t,l)$	1280	$37.4$	$68.6$	$77.9$	$73.5$
Fixed $(t,l)$ + fixed $K$	128	$33.8$	$62.9$	$72.4$	$68.1$
Top-variance $K$	128	$35.6$	$64.8$	$74.6$	$70.3$
Fisher proxy $K$	128	$36.4$	$65.7$	$75.3$	$71.0$
Policy-Mask $(t,l,K)$	128	$\mathbf{38.9}$	$\mathbf{70.2}$	$\mathbf{79.1}$	$\mathbf{74.8}$

Table 1: Segmentation results with different feature selection methods, measured in mIoU. All methods use the same SD-1.5 feature lattice with

b=128

, except for the “All channels” baseline, which retains all 1280 channels. Bold values indicate the best result among 128-channel selectors.

Method	Teacher	Width	CIFAR-10	CIFAR-100	SPair-71k [email protected]
					img	bbox
DDAE	DDPM	full	$95.4$	$76.8$	—	—
RepFusion	SD-1.5	full	$94.2$	$74.1$	—	—
Top-variance $K$	SD-1.5	128	$93.1$	$72.3$	—	—
Fisher $K$	SD-1.5	128	$93.6$	$73.1$	—	—
Policy-Mask	SD-1.5	128	$\mathbf{95.2}$	$\mathbf{75.4}$	$\mathbf{76.3}$	$\mathbf{67.5}$

Table 2: Classification and correspondence comparison with DDAE ( Xiang et al., 2023 ) and RepFusion ( Yang and Wang, 2023 ) . Full-width diffusion methods and compact selectors are listed separately. A dash denotes a method–task pair not reported under the corresponding source protocol. Bold values indicate the best result among 128-dimensional selectors.

Figure 2: Effect of the channel budget on ADE20K segmentation. Channels are selected by activation magnitude at a fixed timestep and layer. A separate student is trained for each budget. Predictions are shown for one validation image, while mIoU is computed over the validation set. No learned policy is used.

Figure 3: ADE20K validation, seed 0. The 12 images with the largest gains in image-mIoU over RepFusion. Selected by outcome. See the loss panel (Figure 6 ). Columns in each block show input, ground truth, RepFusion, Random 128, and Policy-Mask .

Figure 4: Controlled action-axis ablation on ADE20K, mIoU. The time–depth grid exceeds the timestep-only policy. Appending a heuristic subset does not retain that gain, while the unnormalized joint action is lower. The values come from one ablation run and are not compared numerically with Table 1 .

Teacher	Selected coordinates	$b/D$	ADE20K	CIFAR-100	SPair-71k
SD-1.5 U-Net	$(t,l,K_{\rm ch})$	$128/1280$	$34.1$	$70.8$	$72.8$
PixArt- $\Sigma$ DiT	$(\lambda,d,K_{\rm hid})$	$115/1152$	$35.0$	$71.6$	$73.5$
SD3.5 MMDiT	$(\lambda,d,K_{\rm img})$	$154/1536$	$36.4$	$72.9$	$74.6$

Table 3: Separate matched-fraction interface evaluation across teacher architectures. The selected coordinate changes from U-Net channels to image-token hidden dimensions. These rows are not a continuation of Table 1 : each uses its teacher-specific feature interface.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Table 8

Figure 5: Class IoU differences between Policy-Mask and RepFusion on ADE20K, seed 0. The plot shows the 12 largest gains and losses among classes present in at least 20 validation images. Positive values favor Policy-Mask .

Figure 6: ADE20K validation, seed 0. The 12 images with the largest losses in image-mIoU relative to RepFusion. Selected by outcome. Columns follow Figure 3 .

Diffusion Model Adaptation

See all 34 papers

MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Jun 7, 2026Lianyu Pang, Tianlin Pan, Cheng Da +5Diffusion Model Alignment Diffusion Transformer
The Hong Kong University of Science and Technology · Kuaishou Technology · University of Chinese Academy of Sciences
Adaptive Order Policies for Masked Diffusion
May 29, 2026Jama Hussein Mohamud, Mohsin Hasan, Mirco Ravanelli +1Masked Diffusion Models Diffusion Model Sampling
Université de Montréal · Mila · Concordia University +1
What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
Oct 22, 2026Jinya Sakurai, Xun Xu, Jaehong YoonDiffusion Language Model Decoding Diffusion Language Models
NTU Singapore · Institute for Advanced Intelligence and Computing (IAIC), A*STAR · Centre for Frontier AI Research (CFAR), A*STAR

Jun 7, 2026J/K move · Enter open · S save

MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

Lianyu Pang, Tianlin Pan, Cheng Da, Changqian Yu+4

The Hong Kong University of Science and Technology · Kuaishou Technology · University of Chinese Academy of Sciences

Diffusion Model Alignment Diffusion Transformer Diffusion Model Acceleration Representation Alignment

Open paper AskPDF arXiv

Representation alignment with pretrained vision models has recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations from self-supervised vision encoders, existing methods improve convergence and generation quality. However, such alignment also introduces a non-trivial constraint: diffusion models operate on noisy inputs whose usable information varies across timesteps, while the reference features are extracted from clean images. In this paper, we revisit this mismatch from a token-level perspective. We find that, under full-token representation alignment, tokens with large alignment-gradient norms exhibit a stable spatial preference, suggesting that the alignment objective does not affect all tokens uniformly and may encourage the model to rely on the complete set of clean-image tokens. To address this issue, we propose MaskAlign, a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations, MaskAlign reduces the dependence of representation alignment on the complete token set and encourages alignment behavior that is more stable under token-subset perturbations. To mitigate the information loss caused by directly dropping tokens, we further introduce a lightweight pre-mask token mixing block that shares information across tokens before masking.

Abstract

Diffusion Model Adaptation

Explore similar work

MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

Adaptive Order Policies for Masked Diffusion

What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models