No Alignment Needed for Generation: Learning Linearly Separable Representations in Diffusion Models

Junno Yun, Yasar Utku Alcalar, Mehmet Akcakaya

Keywords: Diffusion models, Representation Learning, Linear Separability

Abstract: Efficient training strategies for large-scale diffusion models have recently emphasized the importance of improving discriminative feature representations in these models. A central line of work in this direction is representation alignment with features obtained from powerful external encoders, which improves the representation quality as assessed through *linear probing*. Alignment-based approaches show promise but depend on large pretrained encoders, which are computationally expensive to obtain. In this work, we propose an alternative regularization for training, based on promoting the **L**inear **SEP**arability (LSEP) of intermediate layer representations. LSEP eliminates the need for an auxiliary encoder and representation alignment, while incorporating linear probing directly into the network’s learning dynamics rather than treating it as a simple post-hoc evaluation tool. Our results demonstrate substantial improvements in both training efficiency and generation quality on flow-based transformer architectures such as SiTs, achieving an FID of 1.44 on $256 \times 256$ ImageNet and FID of 1.66 on $512 \times 512$ ImageNet dataset.

Primary Area: generative models

Submission Number: 20427