Purify Before You Align: Improving Representation Alignment for Diffusion Models

Yingheng Wang, Yaoqiang Li, Yaqin Wu, Junwen Bai, Jiatao Gu +2

Representation alignment improves and accelerates diffusion training by encouraging intermediate features to match a pretrained visual encoder. However, standard alignment matches pretrained representations without distinguishing what information may be more useful for generation. We introduce local appearance purification (LAP), which improves generation quality and makes encoders such as MAE more effective alignment targets by removing locally predictable content and aligning to the residual. LAP modifies only the alignment target, and requires no changes to diffusion training or inference. Through controlled interventions within a fixed encoder, we show that effective alignment depends on what information the target emphasizes, not simply how much image information it contains. Our theory gives a variational interpretation of standard alignment and, for fixed local predictors, a sum of token-wise information bounds conditioned on local appearance. To understand why purification benefits some encoders more than others, we derive an exact constant-readout test that measures how well the target can be matched without using the diffusion model's hidden state. LAP improves…