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→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≈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θ∖k−vθ∥ ) 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: S3 + AudioSet
L24(0.323)L18(0.293)L17(0.293)
L20(0.180)L21(0.174)L19(0.168)
L1(0.167)L9(0.0695)L5(0.0588)
L1(0.140)L21(0.060)L9(0.0561)
B: A + interleaved BEATs
L23(0.318)L22(0.314)L24(0.309)
L14(0.218)L13(0.212)L15(0.207)
L1(0.192)L2(0.109)L7(0.085)
L1(0.164)L7(0.094)L2(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 ↓
Speech FAD ↓
Audio FAD ↓
Speech MOS ↑
Audio MOS ↑
Base (no layer align.)
5.82
1.84
3.45
3.62 ± .08
3.45 ± .09
REPA @ Layer 4
5.15
1.65
3.12
3.75 ± .07
3.58 ± .08
REPA @ Layer 8
4.93
1.58
3.05
3.79 ± .06
3.64 ± .08
REPA @ Layer 12
5.38
1.72
3.28
3.68 ± .08
3.51 ± .09
REPA @ L4, 8, 12
4.21
1.45
2.88
3.92 ± .06
3.77 ± .07
REPA @ Deep (L20–L22)
5.60
1.79
3.39
3.64 ± .08
3.48 ± .09
Table 2: Comparison of Alignment Strategies. Evaluation of intelligibility (WER ↓ ), generation quality (FAD ↓ ) and perceptual fidelity (MOS ↑ ) 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 ↓
Audio FAD ↓
Rel. Gain
Convergence †
Base (no layer align.)
None
500k
1.84
3.45
0.0%
N/A
Random Control
L5,L14,L19
500k
1.75
3.32
+4.9%
850k
Highest LASP
L22,L23,L24
500k
1.68
3.21
+8.7%
720k
Gradient Norm
L1,L2,L4
500k
1.35
2.71
+26.6%
260k
Highest FoG-A (Ours)
L1,L2,L7
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× larger FAD reduction than aligning representation-rich layers (identified by LASP), and accelerates convergence by 3.3× , 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 ↓
FAD ↓
MOS ↑
Voicebox ( Le et al., 2023 )
2.05
1.20
4.15 ± .06
+ AG-REPA
1.85
0.95
4.28 ± .05
CosyVoice ( Du et al., 2024 )
1.95
0.88
4.25 ± .05
+ AG-REPA
1.78
0.72
4.39 ± .04
F5-TTS ( Chen et al., 2025 )
2.40
1.45
4.05 ± .07
+ AG-REPA
2.12
1.15
4.22 ± .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 ↓ ). 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 S3 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 ( Pstyle ) while utilizing an auxiliary coarse prediction head. (c) Stage 2 (Flow Matching): A Diffusion Transformer (DiT) backbone predicts the velocity field vθ 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
T0
≈0
1.00T0
AG-REPA (Ours)
≈0.30T0
<0.005T0
≈0.305T0
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 ≈3.26× .
Probe Point
Speech Top-3
Audio Top-3
Overlap vs. Warm-up
Warm-up (5k)
L1,L2,L7
L1,L7,L2
3/3, 3/3
Mid-training
L1,L2,L8
L1,L7,L9
2/3, 2/3
Final checkpoint
L1,L2,L7
L1,L7,L3
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 ↓
Audio FAD ↓
Relative Time
Freeze @ Epoch 1 (Ours)
1.29
2.56
1.00×
Refresh set every epoch
1.28
2.55
1.19×
Appendix
Table 8: Freeze vs. Refresh Schedule. Refreshing S every epoch yields marginal gain but inflates wall-clock by 19%.
Alignment Strategy
Routing Complexity
Speech FAD ↓
Audio FAD ↓
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 ↓
Audio FAD ↓
K=2 , FoG-A λ
1.31
2.60
K=3 , equal λ
1.32
2.59
K=3 , FoG-A λ (Ours)
1.29
2.56
K=4 , FoG-A λ
1.30
2.58
Appendix
Table 10: Sensitivity to K and λk . K=3 with FoG-A weighting is the optimal trade-off.
Projection Head
Applied to
Params
Speech FAD ↓
Audio FAD ↓
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 learns probability transports from simple noise to data for fast sampling, density estimation, and controllable generation.
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× and 2.1× 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.
Dohwan Kim, Jung-Woo Choi
School of Electrical Engineering, KAIST, Daejeon, South Korea · Department of AI Systems, KAIST, Daejeon, South Korea
Flow-matching transformers achieve strong audio separation, yet their attention dynamics are opaque. We adapt established causal-intervention principles into a deterministic, inference-time probing protocol for SAM Audio. Orthogonal probing uncovers a dual-pathway text-conditioning mechanism: additive injections control semantic identity, while cross-attention refines acoustic structure. We observe an asynchronous layerwise convergence: stable layers build temporal scaffolds early, whereas fast layers continue resolving artifacts during sampling. The model also attenuates temporal segmentation cues to maintain continuous-flow stability. Using these insights, we propose Layer-Selective Attention Caching (LSAC), a training-free acceleration method that caches attention in stable layers. Across acoustic complexities, LSAC cuts self-attention computation by about ~25% with negligible quality loss and yields up to 6.7x higher quality retention than naive step reduction.
Yuxuan Chen, Haoyuan Yu, Peize He
Jilin University · Hunan University · University of Electronic Science and Technology of China
Generative models, particularly diffusion and score-based approaches, have recently achieved strong performance in speech enhancement, but their iterative sampling process limits real-time deployment. Flow Matching offers an efficient alternative by transporting noisy speech toward clean speech through an ordinary differential equation with few function evaluations. In this work, we propose a skip-free encoder-decoder backbone for flow-matching speech enhancement, guided by Latent Representation Alignment (LRA). Instead of relying on U-Net skip connections, which may transfer noise-correlated low-level features to the decoder, the proposed model aligns its bottleneck and decoder representations with clean latent features extracted from a frozen Descript Audio Codec encoder-decoder without quantization. This codec-aligned supervision promotes compact clean-speech representations while preserving efficient few-step inference. Experiments on WSJ0-CHiME3 and VoiceBank-DEMAND show improved PESQ and perceptual quality, especially on VoiceBank-DEMAND, using only five function evaluations.
Wangyi Pu, Michele Scarpiniti
Department of Information Engineering, Electronics and Telecommunications, Sapienza University of Rome, Italy