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 πθ outputs distributions over timesteps and layers and channel logits z . A timestep t and a layer l are sampled. An ordered sequence σ∼πσ of b channels is drawn without replacement by Gumbel top- b , giving the mask mK . The frozen diffusion model D returns Gt,l(x) after the alignment Al , and masking gives Fa(x) . This feature is the distillation target for the student Sϕ and, through the task decoder qω , gives the reward R=−ℓdec . The policy is updated by REINFORCE with baseline c and the channel score scaled by 1/b (Equation 8 ). At inference, only Sϕ 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
38.9
70.2
79.1
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.
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,Kch)
128/1280
34.1
70.8
72.8
PixArt- Σ DiT
(λ,d,Khid)
115/1152
35.0
71.6
73.5
SD3.5 MMDiT
(λ,d,Kimg)
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 .
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.
Lianyu Pang, Tianlin Pan, Cheng Da +5
The Hong Kong University of Science and Technology · Kuaishou Technology · University of Chinese Academy of Sciences
Masked diffusion models have seen great success in capturing data distributions over discrete sequences in domains such as text and proteins. These models generate data by iteratively unmasking tokens starting from a fully masked sequence, with the unmasking order typically chosen at random or using a heuristic based on denoiser probabilities. In this work, we propose a scheme for learning the unmasking order using an additional lightweight policy network on top of a diffusion model. Our proposed loss reweights terms in the masked diffusion loss according to policy probabilities, and results in a policy that prefers positions where the denoiser is more likely to be correct. We study this loss in two settings: (i) training solely the policy while using a frozen pre-trained denoiser, and (ii) training the policy and denoiser jointly with the weighted loss to allow for mutual adaptation. We demonstrate that our approach outperforms common heuristics on problems that are sensitive to token ordering, such as combinatorial tasks and proteins.
Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajectory-contraction targets, which accept only commitments staying on a reference trajectory, using confidence and hidden states. On LLaDA-8B-Instruct, the resulting 856K-parameter policy achieves the best accuracy--NFE frontier and highest throughput on GSM8K, is competitive at moderate targets on MATH-500 and HumanEval, and largely transfers to Dream-7B-Instruct.
Jinya Sakurai, Xun Xu, Jaehong Yoon
NTU Singapore · Institute for Advanced Intelligence and Computing (IAIC), A*STAR · Centre for Frontier AI Research (CFAR), A*STAR