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

arXiv: https://arxiv.org/abs/2603.01006 · HTML: https://mlsanity.com/paper/2603.01006/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Audio, speech and music, one of 14 directions in Diffusion models.
- Builds on 25 of its 33 references that are in the map.
- First posted on arXiv 1 Mar 2026.

## Placed in

- Diffusion models / [Audio, speech and music](https://mlsanity.com/diffusion/audio-speech-and-music/)

## Builds on (most cited first, 25 of 25)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Robust Speech Recognition via Large-Scale Weak Supervision](https://mlsanity.com/paper/2212.04356/) (2022)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [BEATs: Audio Pre-Training with Acoustic Tokenizers](https://mlsanity.com/paper/2212.09058/) (2022)
- [Voicebox: Text-Guided Multilingual Universal Speech Generation at Scale](https://mlsanity.com/paper/2306.15687/) (2023)
- [F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching](https://mlsanity.com/paper/2410.06885/) (2024)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [CosyVoice: A Scalable Multilingual Zero-shot Text-to-speech Synthesizer based on Supervised Semantic Tokens](https://mlsanity.com/paper/2407.05407/) (2024)
- [Does Localization Inform Editing? Surprising Differences in Causality-Based Localization vs. Knowledge Editing in Language Models](https://mlsanity.com/paper/2301.04213/) (2023)
- [Vocos: Closing the gap between time-domain and Fourier-based neural vocoders for high-quality audio synthesis](https://mlsanity.com/paper/2306.00814/) (2023)
- [Matcha-TTS: A fast TTS architecture with conditional flow matching](https://mlsanity.com/paper/2309.03199/) (2023)
- [Similarity of Neural Network Models: A Survey of Functional and Representational Measures](https://mlsanity.com/paper/2305.06329/) (2023)
- [Audiobox: Unified Audio Generation with Natural Language Prompts](https://mlsanity.com/paper/2312.15821/) (2023)
- [What matters for Representation Alignment: Global Information or Spatial Structure?](https://mlsanity.com/paper/2512.10794/) (2025)
- [REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training](https://mlsanity.com/paper/2505.16792/) (2025)
- [RepCodec: A Speech Representation Codec for Speech Tokenization](https://mlsanity.com/paper/2309.00169/) (2023)
- [LAFMA: A Latent Flow Matching Model for Text-to-Audio Generation](https://mlsanity.com/paper/2406.08203/) (2024)
- [TARO: Timestep-Adaptive Representation Alignment with Onset-Aware Conditioning for Synchronized Video-to-Audio Synthesis](https://mlsanity.com/paper/2504.05684/) (2025)
- [LayerSync: Self-aligning Intermediate Layers](https://mlsanity.com/paper/2510.12581/) (2025)

## Built on by (most cited first, 1 of 1)

- [What Visual Generators Need from Teachers: Rethinking Representation Alignment](https://mlsanity.com/paper/2609.34732/) (2026)

## Abstract (arXiv)

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.