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

2026-03-01

AI Summary

Existing audio stream matching methods rely on heuristic selection of supervision layers, which struggles to effectively enhance generation quality. This work proposes a causal layer selection strategy, termed AG-REPA, which leverages Forward Gate Ablation (FoG-A) to quantify the causal contribution of each layer to the velocity field, enabling sparse and adaptive representation alignment. The study uncovers a decoupling phenomenon between semantic “storage layers” and generative “driving layers,” and builds a causal contribution evaluation mechanism upon this insight. Evaluated under joint training on LibriSpeech and AudioSet, AG-REPA consistently outperforms the REPA baseline across multiple token-conditioning topologies, demonstrating the efficacy of aligning representations based on causally dominant layers.

Technology Category

• Machine Learning: Causal Learning
• Natural Language Processing: Generation
• Knowledge Representation and Reasoning: Action, Change, and Causality

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• representation alignment
• audio flow matching
• causal layer selection
• token-conditioned generation
• Store-Contribute Dissociation

Innovation

• causal layer selection
• representation alignment
• flow matching
• attribution-guided
• velocity field