LayerSync: Self-aligning Intermediate Layers

TL;DR

LayerSync presents a self-contained regularizer for diffusion transformers that uses the model's own deeper layers as intrinsic guidance to improve earlier-layer representations, reducing reliance on external pretrained models. The approach, expressed as a simple alignment loss added to the velocity objective, achieves substantial training acceleration and generation quality gains across image, audio, motion, and video domains without extra data or models. Empirical results show improved internal representations, stronger downstream task performance, and robust layer-compatibility across ablations. This method offers a domain-agnostic, low-overhead path toward more efficient and effective diffusion-based generative modeling.

Abstract

We propose LayerSync, a domain-agnostic approach for improving the generation quality and the training efficiency of diffusion models. Prior studies have highlighted the connection between the quality of generation and the representations learned by diffusion models, showing that external guidance on model intermediate representations accelerates training. We reconceptualize this paradigm by regularizing diffusion models with their own intermediate representations. Building on the observation that representation quality varies across diffusion model layers, we show that the most semantically rich representations can act as an intrinsic guidance for weaker ones, reducing the need for external supervision. Our approach, LayerSync, is a self-sufficient, plug-and-play regularizer term with no overhead on diffusion model training and generalizes beyond the visual domain to other modalities. LayerSync requires no pretrained models nor additional data. We extensively evaluate the method on image generation and demonstrate its applicability to other domains such as audio, video, and motion generation. We show that it consistently improves the generation quality and the training efficiency. For example, we speed up the training of flow-based transformer by over 8.75x on ImageNet dataset and improved the generation quality by 23.6%. The code is available at https://github.com/vita-epfl/LayerSync.

Figures (19)

• Figure 1: LayerSync improves training efficiency and generation quality via internal representation alignment. (a) LayerSync aligns deep and shallow layers. (b) LayerSync achieves over 8.75$\times$ training acceleration on the ImageNet 256x256. (c) LayerSync consistently improves generation quality across multiple modalities: by $23.6\%$ on FID for images (ImageNet 256×256), $24\%$ on FAD for audio (MTG-Jamendo), and $7.7\%$ for FID on human motion (HumanML3D).
• Figure 2: LayerSync improves generation quality without relying on external representation. We compare the images generated by SiT-XL/2 when regularized with dispersive and LayerSync. All the models are trained for 400K iterations, share the same noise, sampler, and number of sampling steps, and none of them use classifier-free guidance.
• Figure 3: Assessing the quality of intermediate features shows that LayerSync improves average validation accuracy across layers (shown with dashed lines in the figures) for both classification and segmentation, and enhances alignment with DINOv2. In this experiment, layer 8 is aligned with layer 16.
• Figure 4: Qualitative comparison of generated samples from SiT XL/2 and SiT XL/2 with LayerSync when layers 7 and 16 are synced. After dropping blocks [9, 11, 13, 15], we observe that LayerSync helps preserve visual quality despite block removal.
• Figure 5: Qualitative results when dropping blocks from SiT XL/2 with LayerSync, where layers 7 and 16 are synced. Dropping blocks outside the synced layers leads to a more noticeable degradation in sample quality.