AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching

REPresentation Alignment (REPA) improves the training of generative flow models by aligning intermediate hidden states with pretrained teacher features, but its effectiveness in token-conditioned audio Flow Matching critically depends on the choice of supervised layers, which is typically made heuristically based on the depth. In this work, we introduce Attribution-Guided REPresentation Alignment (AG-REPA), a novel causal layer selection strategy for representation alignment in audio Flow Matching. Firstly, we find that layers that best store semantic/acoustic information (high teacher-space similarity) are not necessarily the layers that contribute most to the velocity field that drives generation, and we call it Store-Contribute Dissociation (SCD). To turn this insight into an actionable training guidance, we propose a forward-only gate ablation (FoG-A) that quantifies each layer's causal contribution via the induced change in the predicted velocity field, enabling sparse layer selection and adaptive weighting for alignment. Across unified speech and general-audio training (LibriSpeech + AudioSet) under different token-conditioning topologies, AG-REPA consistently outperforms REPA baselines. Overall, our results show that alignment is most effective when applied to the causally dominant layers that drive the velocity field, rather than to layers that are representationally rich but functionally passive.

When training AI models to generate high-quality speech and sound effects, researchers often use a shortcut called "representation alignment"—speeding up learning by forcing the model's internal layers to mimic a smart "teacher" model. However, choosing which layers to align has traditionally been a game of guesswork. In this work, we discovered a surprising paradox we call "Store-Compute Dissociation": the internal layers that are best at storing rich acoustic knowledge are not necessarily the ones that actively compute and drive the audio generation process. It is like a team where the wisest member stays passive, while a quieter member does all the heavy lifting. To solve this, we introduce AG-REPA, a smart selection strategy that automatically identifies and targets only the functionally critical layers. By putting the training supervision exactly where it matters most, our method significantly improves audio quality and training efficiency across both speech and general sound generation, making AI audio synthesis faster and more realistic.