No Alignment Needed for Generation: Learning Linearly Separable Representations in Diffusion Models
This work addresses the need for more efficient training in large-scale diffusion models for generative AI, offering a novel alternative to alignment-based methods.
The paper tackles the problem of improving discriminative feature representations in diffusion models without relying on computationally expensive external encoders, by proposing a Linear SEParability (LSEP) regularization that achieves an FID of 1.46 on the 256x256 ImageNet dataset.
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 Linear SEParability (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.46 on $256 \times 256$ ImageNet dataset.