AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching
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.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper presents AG-REPA, a framework that enhances audio generation quality by focusing on causal contributions of layers rather than mere representational richness. This innovative approach not only improves training efficiency but also contributes to the interpretability of generative models, marking a meaningful advancement in the field of machine learning.
Methodology
The paper introduces a novel methodology called Attribution-Guided REPresentation Alignment (AG-REPA), which emphasizes causal layer selection for representation alignment in audio flow matching. This approach is grounded in the theoretical concept of Store-Contribute Dissociation (SCD), which reveals that layers rich in semantic information do not necessarily contribute most to the generative process. The methodology includes a forward-only gate ablation (FoG-A) to quantify each layer's causal contribution, allowing for adaptive layer selection and weighting. This is a significant advancement over traditional heuristic methods, providing a more principled basis for layer selection in generative models.
Results & evaluation
The experiments are robust, utilizing well-established datasets such as LibriSpeech and AudioSet for unified speech and general audio training. The results demonstrate that AG-REPA consistently outperforms baseline methods, achieving significant reductions in Fréchet Audio Distance (FAD) and improvements in perceptual quality metrics like Word Error Rate (WER) and Mean Opinion Score (MOS). The comparative analysis against static REPA baselines and other alignment strategies provides strong empirical support for the proposed method.
Reproducibility
The paper outlines a clear methodology and experimental setup, but lacks specific implementation details or code availability, which could hinder reproducibility. The authors mention a "probe-then-intervene" training protocol that separates diagnostic probing from optimization, which is a good practice for ensuring clean experimental conditions.
Limitations
One limitation is the lack of external validation on diverse datasets beyond LibriSpeech and AudioSet, which may limit the generalizability of the findings. Additionally, while the paper discusses potential risks associated with high-fidelity audio generation, it does not provide a detailed framework for mitigating these risks in practical applications.
Broader impact
The work has significant implications for the field of audio generation, particularly in enhancing the intelligibility and quality of synthesized speech and audio. The interpretability toolkit developed in this study could also pave the way for more transparent and controllable generative models in AI, addressing some of the ethical concerns surrounding deepfake technologies. The paper presents AG-REPA, a framework that enhances audio generation quality by focusing on causal contributions of layers rather than mere representational richness. This innovative approach not only improves training efficiency but also contributes to the interpretability of generative models, marking a meaningful advancement in the field of machine learning.
Author of this paper? Send a correction.