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

Pengfei Zhang; Tianxin Xie; Minghao Yang; Li Liu

cs.SDMar 1, 2026

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

Authors: Pengfei Zhang, Tianxin Xie, Minghao Yang, Li Liu

Organizations: AI Thrust, Information Hub, The Hong Kong University of Science and Technology (Guangzhou).

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.

Figures & tables

Figure 1: The Spatiotemporal Anatomy of SCD. We visualize the layer-wise dynamics across diffusion time steps ( $t=0\to 1$ ). (a) Representation: LASP reveals teacher-dependent storage reservoirs: semantic storage is strongest in deep layers, while acoustic/event storage can shift toward middle layers under dense acoustic conditioning. (b) Contribution: In contrast, FoG-A reveals a dynamic functional landscape. Early layers (L1–3) act as strong Causal Drivers , and a Dynamic Transition emerges in the middle layers (L6–12) during the intermediate denoising phase ( $t\approx 0.5$ ). Key Insight: The mismatch between teacher-similarity peaks and velocity-field attribution explains why fixed-depth alignment heuristics can be suboptimal.

Figure 2: Diagnosing Representation Storage. (a) BiT-C establishes a dual-modality supervision baseline using frozen Whisper (semantic) and BEATs (acoustic) teachers to anchor the conditioning interface. (b) LASP probes “what the network knows” by projecting layer-wise representations into a shared teacher space using a frozen head, enabling cross-layer comparison of information storage.

Figure 3: From Causal Attribution to Optimization. (a) FoG-A determines “what the network uses” by measuring the velocity field perturbation ( $\|v_{\theta}^{\setminus k}-v_{\theta}\|$ ) caused by ablating individual layers, creating a functional attribution map. (b) AG-REPA utilizes the causal insights from FoG-A to selectively apply alignment supervision only to critical layers, bridging the gap between representation storage and causal contribution identified as the SCD.

Config	Cos-SEM (top3)	Cos-EVT (top3)	FoG-A Speech (top3)	FoG-A Audio (top3)
A: $S^{3}$ + AudioSet	$L_{24}(0.323)$ $L_{18}(0.293)$ $L_{17}(0.293)$	$L_{20}(0.180)$ $L_{21}(0.174)$ $L_{19}(0.168)$	$L_{1}(0.167)$ $L_{9}(0.0695)$ $L_{5}(0.0588)$	$L_{1}(0.140)$ $L_{21}(0.060)$ $L_{9}(0.0561)$
B: A + interleaved BEATs	$L_{23}(0.318)$ $L_{22}(0.314)$ $L_{24}(0.309)$	$L_{14}(0.218)$ $L_{13}(0.212)$ $L_{15}(0.207)$	$L_{1}(0.192)$ $L_{2}(0.109)$ $L_{7}(0.085)$	$L_{1}(0.164)$ $L_{7}(0.094)$ $L_{2}(0.083)$

Table 1: Store-Contribute Dissociation. We report the top-3 layers for representation storage (Cos-SEM / Cos-EVT from LASP probes, denoting alignment with Whisper semantic and BEATs acoustic event teachers, respectively) and functional contribution (FoG-A sensitivity) on LibriSpeech (Speech) and AudioSet (Audio).

Method	Speech WER $\downarrow$	Speech FAD $\downarrow$	Audio FAD $\downarrow$	Speech MOS $\uparrow$	Audio MOS $\uparrow$
Base (no layer align.)	5.82	1.84	3.45	3.62 $\pm$ .08	3.45 $\pm$ .09
REPA @ Layer 4	5.15	1.65	3.12	3.75 $\pm$ .07	3.58 $\pm$ .08
REPA @ Layer 8	4.93	1.58	3.05	3.79 $\pm$ .06	3.64 $\pm$ .08
REPA @ Layer 12	5.38	1.72	3.28	3.68 $\pm$ .08	3.51 $\pm$ .09
REPA @ L4, 8, 12	4.21	1.45	2.88	3.92 $\pm$ .06	3.77 $\pm$ .07
REPA @ Deep (L20–L22)	5.60	1.79	3.39	3.64 $\pm$ .08	3.48 $\pm$ .09

Table 2: Comparison of Alignment Strategies. Evaluation of intelligibility (WER

\downarrow

), generation quality (FAD

\downarrow

) and perceptual fidelity (MOS

\uparrow

) under Config B. The baseline retains only the input-interface BiT-C anchor and uses no intermediate-layer REPA; AG-REPA achieves the best performance by sparsely selecting and adaptively weighting functionally critical layers.

Selection Strategy	Top-3 Layers	Training Steps	Speech FAD $\downarrow$	Audio FAD $\downarrow$	Rel. Gain	Convergence †
Base (no layer align.)	None	500k	1.84	3.45	0.0%	N/A
Random Control	$L_{5},L_{14},L_{19}$	500k	1.75	3.32	+4.9%	850k
Highest LASP	$L_{22},L_{23},L_{24}$	500k	1.68	3.21	+8.7%	720k
Gradient Norm	$L_{1},L_{2},L_{4}$	500k	1.35	2.71	+26.6%	260k
Highest FoG-A (Ours)	$L_{1},L_{2},L_{7}$	500k	1.29	2.56	+29.9%	220k

Table 3: The Impact of Alignment Targets: Knowing vs. Doing. Aligning functionally critical layers (identified by FoG-A) yields

3.4\times

larger FAD reduction than aligning representation-rich layers (identified by LASP), and accelerates convergence by

3.3\times

, directly validating the SCD. Audio FAD is shown for all selectors; the measured Gradient Norm vs. FoG-A comparison highlights that FoG-A also dominates on the general-audio domain.

Model	WER $\downarrow$	FAD $\downarrow$	MOS $\uparrow$
Voicebox ( Le et al., 2023 )	2.05	1.20	4.15 $\pm$ .06
+ AG-REPA	1.85	0.95	4.28 $\pm$ .05
CosyVoice ( Du et al., 2024 )	1.95	0.88	4.25 $\pm$ .05
+ AG-REPA	1.78	0.72	4.39 $\pm$ .04
F5-TTS ( Chen et al., 2025 )	2.40	1.45	4.05 $\pm$ .07
+ AG-REPA	2.12	1.15	4.22 $\pm$ .06

Table 4: Generalization of AG-REPA. Comparison of standard training (Baseline) versus our method across different Flow Matching architectures. AG-REPA consistently improves intelligibility (WER), audio quality (FAD), and naturalness (MOS).

Architecture	Base.	Shallow REPA	AG-REPA
Our DiT	1.84	1.36	1.29
F5-TTS ( Chen et al., 2025 )	1.45	1.34	1.15
Voicebox ( Le et al., 2023 )	1.20	1.12	0.95
CosyVoice ( Du et al., 2024 )	0.88	0.85	0.72

Table 5: Shallow REPA vs. AG-REPA Across Architectures (FAD

\downarrow

). AG-REPA transfers better by re-running FoG-A per backbone.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: The Unified Audio Generation Framework. The system utilizes a two-stage cascade architecture decoupling semantic planning from acoustic rendering: (a) Tokenization: Domain-specific pathways process inputs into a unified discrete sequence, utilizing $S^{3}$ tokens for speech and AudioSet (AS) tokens for audio, optionally interleaved with BEATs features to inject dense acoustic priors. (b) Stage 1 (Autoregressive LLM): A causal language model predicts target acoustic tokens, incorporating reference style via a learnable projection and injection module ( $P_{\text{style}}$ ) while utilizing an auxiliary coarse prediction head. (c) Stage 2 (Flow Matching): A Diffusion Transformer (DiT) backbone predicts the velocity field $v_{\theta}$ to transform noise into target mel-spectrograms via Flow Matching, conditioned on projected tokens and reference audio, before final waveform synthesis by a neural vocoder.

Method	Main Training	Diagnostic	End-to-End
LASP-selected REPA	$T_{0}$	$\approx 0$	$1.00\,T_{0}$
AG-REPA (Ours)	$\approx 0.30\,T_{0}$	$<\!0.005\,T_{0}$	$\approx 0.305\,T_{0}$

Appendix

Table 6: End-to-End Efficiency. The FoG-A diagnostic adds

<\!0.5\%

wall-clock; the end-to-end speedup over LASP-selected REPA (the strongest storage-driven baseline reaching FAD = 1.5; Table 3 ) is

\approx\!3.26\times

Probe Point	Speech Top-3	Audio Top-3	Overlap vs. Warm-up
Warm-up (5k)	$L_{1},L_{2},L_{7}$	$L_{1},L_{7},L_{2}$	3/3, 3/3
Mid-training	$L_{1},L_{2},L_{8}$	$L_{1},L_{7},L_{9}$	2/3, 2/3
Final checkpoint	$L_{1},L_{2},L_{7}$	$L_{1},L_{7},L_{3}$	3/3, 2/3

Appendix

Table 7: Cross-Checkpoint Stability of FoG-A Rankings. The top causal layers remain stable after the warm-up diagnostic.

Alignment Schedule	Speech FAD $\downarrow$	Audio FAD $\downarrow$	Relative Time
Freeze @ Epoch 1 (Ours)	1.29	2.56	$1.00\times$
Refresh set every epoch	1.28	2.55	$1.19\times$

Appendix

Table 8: Freeze vs. Refresh Schedule. Refreshing

\mathcal{S}

every epoch yields marginal gain but inflates wall-clock by 19%.

Alignment Strategy	Routing Complexity	Speech FAD $\downarrow$	Audio FAD $\downarrow$
Static AG-REPA (Ours)	Low (fixed heads)	1.29	2.56
Timestep-Adaptive	High (time-variant heads)	1.28	2.55

Appendix

Table 9: Static vs. Timestep-Adaptive Alignment. Static AG-REPA matches timestep-adaptive routing while avoiding time-variant projection heads.

Setting	Speech FAD $\downarrow$	Audio FAD $\downarrow$
$K=2$ , FoG-A $\lambda$	1.31	2.60
$K=3$ , equal $\lambda$	1.32	2.59
$K=3$ , FoG-A $\lambda$ (Ours)	1.29	2.56
$K=4$ , FoG-A $\lambda$	1.30	2.58

Appendix

Table 10: Sensitivity to

K

and

\lambda_{k}

K=3

with FoG-A weighting is the optimal trade-off.

Projection Head	Applied to	Params	Speech FAD $\downarrow$	Audio FAD $\downarrow$
2D Conv (iREPA-style)	Pre-pooled tokens	0.29 M	1.32	2.61
1D Conv ( $K\!=\!3$ )	Pre-pooled tokens	0.24 M	1.31	2.58
2-layer MLP (Ours)	Pooled descriptor	0.18 M	1.29	2.56

Appendix

Table 11: Projection Head Comparison. The pooled MLP attains the best FAD with the fewest parameters.

Supervision Source	Static (L4,8,12)	+ AG-REPA Selection
Base (no layer align.)	1.84 / 3.45	—
Internal self-alignment	1.53 / 2.97	1.40 / 2.74
Whisper / BEATs teachers	1.45 / 2.88	1.29 / 2.56

Appendix

Table 12: Internal Self-Alignment vs. External Teachers. Numbers are Speech FAD / Audio FAD. AG-REPA’s selection rule transfers across supervision sources.

Flow-Based Generative Modeling

Flow-based generative modeling learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.

See all 549 papers

MG-REPA: Multi-Granularity Representation Alignment for Generative Target Speaker Extraction
Oct 23, 2026Dohwan Kim, Jung-Woo ChoiTarget Speaker Extraction Diffusion Model Alignment
School of Electrical Engineering, KAIST, Daejeon, South Korea · Department of AI Systems, KAIST, Daejeon, South Korea
Inside the Latent Flow: Causal Deciphering of Attention Dynamics in Audio Separation Foundation Models
Jun 8, 2026Yuxuan Chen, Haoyuan Yu, Peize HeCross-Attention Audio Diffusion Models
Jilin University · Hunan University · University of Electronic Science and Technology of China
Beyond U-Net: A Latent-Representation-Aligned Skip-Free Backbone for Flow-Matching Speech Enhancement
Jun 23, 2026Wangyi Pu, Michele ScarpinitiFlow Matching Speech Enhancement
Department of Information Engineering, Electronics and Telecommunications, Sapienza University of Rome, Italy

Oct 23, 2026J/K move · Enter open · S save

MG-REPA: Multi-Granularity Representation Alignment for Generative Target Speaker Extraction

Dohwan Kim, Jung-Woo Choi

School of Electrical Engineering, KAIST, Daejeon, South Korea · Department of AI Systems, KAIST, Daejeon, South Korea

Target Speaker Extraction Diffusion Model Alignment

Open paper AskPDF arXiv

Generative target speaker extraction (TSE) models achieve high-quality extraction, but suffer from considerably slow training convergence. Representation Alignment (REPA) accelerates diffusion transformer training by aligning hidden states with a pre-trained Self-Supervised Learning (SSL) teacher. However, existing alignment in the speech and audio domain operates at a single granularity. For TSE, frame-level alignment leaves the global, time-invariant characteristics of the target speech implicit, whereas utterance-level alignment leaves the time-varying content implicit, even though extracting the target requires both. We propose Multi-Granularity Representation Alignment (MG-REPA), aligning the hidden states with an SSL model at the frame level and a speaker verification model at the utterance level. We further apply temporal normalization to the frame-level target, removing its utterance-level statistics so that each branch concentrates on its granularity. On Libri2Mix, MG-REPA reaches the OVRL and PESQ of a MeanFlow-TSE baseline with

3.0\times

and

2.1\times

fewer epochs, respectively. Trained for 750 epochs, it is competitive with the published 2,000-epoch baseline on intrusive metrics while improving non-intrusive metrics. An ablation demonstrates that MG-REPA outperforms alignment at a single granularity.

Abstract

Flow-Based Generative Modeling

Explore similar work

MG-REPA: Multi-Granularity Representation Alignment for Generative Target Speaker Extraction

Inside the Latent Flow: Causal Deciphering of Attention Dynamics in Audio Separation Foundation Models

Beyond U-Net: A Latent-Representation-Aligned Skip-Free Backbone for Flow-Matching Speech Enhancement