Policy-Mask picks frozen diffusion features under a channel budget

Pulling perception features out of a frozen diffusion model usually means picking a timestep, a block, and a slice of channels by hand. In a paper on arXiv, Ivan Listopadov, Maksim Nakhodnov and Aibek Alanov train what they call "a budget-conditioned policy ... that jointly selects all three under a prescribed channel budget," and name the method Policy-Mask.

They report 38.9 mIoU on ADE20K segmentation, 75.4% top-1 on CIFAR-100 at 128 dimensions, and 76.3 [email protected] on SPair-71k semantic correspondence, and single out one operating point: "the highest recorded channel-budget score occurs at 256 retained channels, identifying a practical operating point for compact diffusion representations."

A normalization trick holds the training together. "Without normalization, the channel-score contribution grows with the number of selected channels and dominates the timestep and block updates," the abstract says; dividing by subset cardinality keeps the three policy heads "on comparable scales." The authors adapt the setup to diffusion transformers, "including PixArt-Σ and SD3.5 MMDiT, by replacing U-Net channels with image-token hidden dimensions." The abstract reports no wall-clock, memory, or per-backbone accuracy numbers for those transformer variants.