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

Policy-Mask. Budget-Aware Adaptation of Diffusion Features to Discriminative Tasks. By Ivan Listopadov and two others, from HSE University, Moscow State University, and two other institutions. Published on the twentieth of October. No code repository is linked to it.

A frozen diffusion model holds features that can be put to work on discriminative tasks, but pulling the right ones out is harder than it looks. The useful signal is spread across three places at once: along the denoising trajectory, through the depth of the network, and over its channels. Using it means picking a timestep, a block, and a subset of channels. Existing methods either fix those choices by hand, or tune only a few extraction hyperparameters, most often just the timestep. Then they apply the same rule to every input and every budget.

Policy-Mask replaces the fixed rule with a learned one. It is a policy, conditioned on a budget and trained on the downstream objective, that selects the timestep, the block and the channel subset together, under a prescribed limit on how many channels it may keep. The budget is an input to the policy rather than an afterthought, so the selection can change when the budget does.

Training all three choices jointly raises a balance problem. Left unnormalized, the contribution from the channel scores grows with the number of channels selected, and it ends up dominating the updates for timestep and block. The fix is to normalize that contribution by the size of the chosen subset. That keeps the three policy heads on comparable scales and stabilizes joint training. The approach also carries over to diffusion transformers, including PixArt and the SD3.5 MMDiT, by swapping U-Net channels for the hidden dimensions of image tokens.

At matched budgets, Policy-Mask improves on fixed extraction and on feature-filter baselines across segmentation, classification, and semantic correspondence. On ADE20K it reaches thirty-eight point nine mIoU. On CIFAR-100, at one hundred and twenty-eight dimensions, it reaches seventy-five point four percent top-one accuracy.

On SPair-71k it reaches seventy-six point three PCK. Ablations of the separate components show that the choice of block materially affects downstream accuracy. The best channel-budget score recorded comes at two hundred and fifty-six retained channels, which the authors identify as a practical operating point for compact diffusion representations.

The thread running through this is that feature extraction from a diffusion model is a selection problem with three coupled parts, and treating it as one learned decision under a budget works better than freezing two parts and tuning the third. The normalization step is small, but it is what lets the three parts be learned together without the channel choice drowning out the rest.