aiqbench analysis
This research introduces AG-REPA, a causal layer selection method that optimizes representation alignment in audio flow matching by identifying layers that actively drive the generation velocity field rather than just storing semantic information. By resolving the "Store-Contribute Dissociation" phenomenon through forward-only gate ablation, the technique achieves superior performance across unified speech and general-audio tasks compared to heuristic depth-based baselines. The work advances generative modeling efficiency by ensuring alignment strategies target functionally dominant layers critical for token-conditioned audio synthesis.
AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching
arXiv:2603.01006v3 Announce Type: replace-cross Abstract: 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.
Abstract: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.
Accepted to ICML 2026. 17 pages, 4 figures, 12 tables
Sound (cs.SD); Artificial Intelligence (cs.AI); Machine Learning (cs.LG); Multimedia (cs.MM)
I.2.6; I.5.4; G.3
arXiv:2603.01006 [cs.SD]
(or arXiv:2603.01006v3 [cs.SD] for this version)
Full text under the publisher's open licence · 222 words
Related coverage
J++ Lens: Jacobian Filtering Enables More Faithful Workspace Lenses
LessWrong · Oct 8, 2026
MIDI-REPA: Improving Text-to-Music Generation with Symbolic Representation Alignment
arXiv · Oct 16, 2026
Represented but Ignored: A Causal Account of Prosodic Underuse in Audio-Language Models
arXiv · Aug 21, 2026
Emotion Across Speech and Faces: Shared Affective Mechanisms in Multimodal Foundation Models
arXiv · Aug 19, 2026