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

1. INTRODUCTION

• Existing speech and audio alignment approaches supervise either frame-level content or utterance-level characteristics, although TSE needs both dimensions.Prior work includes speech and audio SSL teachers and utterance-level pooling.Source: §1 · p. 1
• MG-REPA aligns generator hidden states at both frame and utterance levels, using a speech SSL teacher for phonetic content and a speaker verification model for global characteristics.Separate adapters map the shared hidden states for the two branches.Source: §1 · p. 1
• 3.0× and 2.1× fewer epochs let MG-REPA reach MeanFlow-TSE’s OVRL and PESQ, respectively, on Libri2Mix.At 750 epochs, it is also comparable to the published 2,000-epoch baseline on intrusive metrics and improves non-intrusive metrics.Source: §1 · p. 2
• Temporal normalization removes utterance-level statistics from frame-level targets so that each alignment branch focuses on its own granularity.REPA was introduced to accelerate diffusion transformer training by aligning intermediate hidden states with representations from a frozen SSL encoder.Source: §1 · pp. 1–2

2.1. Flow-based TSE

• Flow-based TSE estimates target speech from a mixture and an identifying enrollment utterance, with a network conditioned on the enrollment trained to minimize flow-matching loss.Flow matching prescribes a path between prior and data samples, with optimal transport using linear interpolation.Source: §2.1 · p. 2
• AD-FlowTSE places the mixture on the path between interference and target, allowing inference to start from the mixture and integrate only over the remaining interval.The mixing-ratio predictor supplies the interval’s starting point.Source: §2.1 · p. 2

2.2. MeanFlow-based TSE

• MeanFlow models average velocity over an interval so one evaluation can move the state between its endpoints; MeanFlow-TSE applies this to the AD-FlowTSE path for one-step extraction.Its endpoints correspond to the mixture-side mixing ratio and the target.Source: §2.2 · p. 2
• MeanFlow-TSE trains toward a hybrid target interpolating ground-truth direction and self-consistency, with an adaptive error weight and α annealed from 1 toward 0.Despite one-step inference, it still requires 2,000 epochs of training.Source: §2.2 · p. 2

3. MG-REPA

• MG-REPA extends REPA by aligning one intermediate generator layer at frame and utterance levels through separate branches with distinct teachers and adapters.Both branches use the same hidden states as input.Source: §3 · p. 2

3.1. Frame-level REPA

• Temporal normalization removes each channel’s across-frame mean and standard deviation, suppressing global statistics so frame-level alignment focuses on local, time-varying structure.The frame-level teacher is frozen WavLM, whose upper layers are associated with phonetic content.Source: §3.1 · pp. 2–3

3.2. Utterance-level REPA

• The utterance-level branch uses frozen CAM++ to produce a time-invariant speaker embedding and aggregates hidden states into a single prediction for alignment.Its loss compares one pair of vectors per utterance rather than averaging similarity over frames.Source: §3.2 · p. 3

3.3. Training Objective

• The generator and both alignment branches are trained jointly with the generative objective, while the teachers remain frozen and the generative loss can vary.This work applies MG-REPA to MeanFlow-TSE with Lgen = Lα; the alignment branches and teachers are discarded at inference, leaving backbone inference cost unchanged.Source: §3.3 · p. 3

4.1. Datasets

• Libri2Mix train-100 and train-360 are combined for training, while dev and test serve validation and evaluation, using six-second 16 kHz segments.Each input combines a three-second enrollment and a three-second mixture.Source: §4.1 · p. 3

4.2. Training Configuration

• MG-REPA is instantiated on the uDiT velocity estimator shared by AD-FlowTSE and MeanFlow-TSE, using the MeanFlow-TSE recipe with Lgen = Lα.The model has eight encoder blocks, a bottleneck, and eight decoder blocks; optimization uses AdamW and distributed mixed-precision training.Source: §4.2 · p. 3
• Both alignment branches attach at bottleneck block l = 9, using WavLM Large layer 24 and CAM++ with βf = βu = 0.5.Comparisons include generative TSE baselines; evaluation covers fidelity, intelligibility, non-intrusive quality, and speaker similarity, while the granularity ablation also reports LPS and UTMOS.Source: §4.2 · p. 3
• Table 1 compares generative TSE model performance on the Libri2Mix Clean and Noisy sets.The supplied caption does not provide table values.Source: Table 1 · p. 4

5.1. Overall Performance

• With 2.7× fewer training epochs, MG-REPA matches or exceeds the MeanFlow-TSE backbone on every metric except noisy SI-SDR.MG-REPA is trained for 750 epochs versus 2,000 for the backbone.Source: §5.1 · p. 4
• Clean PESQ (3.29) and SI-SDR (18.88 dB) are best among reporting methods; OVRL improves by +0.03 on both clean and noisy sets.Noisy PESQ is 2.22, and noisy SI-SDR (−0.47 dB) is the sole metric below the backbone; MG-REPA also reaches backbone ESTOI on both sets (0.93 clean, 0.82 noisy).Source: §5.1 · p. 4

5.2. Model Convergence Speed

• 3.0× fewer epochs for OVRL and 2.1× fewer for PESQ let MG-REPA reach the retrained MeanFlow-TSE baseline’s 750-epoch level at epochs 250 and 350, respectively.The comparison is on Libri2Mix Noisy; the gap to the baseline persists throughout training.Source: §5.2 · p. 4

5.3. Ablation on Alignment Granularity

• Frame-only WavLM and utterance-only CAM++ alignment each improve every metric over the MeanFlow-TSE baseline, with comparable gains.Table 2 reports alignment-granularity ablations on Libri2Mix Noisy; all variants are trained for 150 epochs.Source: §5.3 · p. 4, Table 2 · p. 4
• With teachers and training held fixed, assigning WavLM to frame-level and CAM++ to utterance-level alignment matches or exceeds both two-teacher single-granularity variants on every metric.Aligning both teachers at the frame level improves over frame-only WavLM except on ESTOI; aligning both at the utterance level raises ESTOI over CAM++ alone but lowers several other metrics.Source: §5.3 · p. 4
• Temporal normalization leaves ESTOI and LPS unchanged and improves every other metric over the multi-granularity variant, producing the best ablation performance.Compared with aligning both teachers at frame level, it improves PESQ by 0.02, SI-SDR by 0.09 dB, ESTOI by 0.01, OVRL by 0.03, and UTMOS by 0.04.Source: §5.3 · p. 4

• MG-REPA aligns an intermediate generative-TSE layer with temporally normalized WavLM features at frame level and a CAM++ speaker embedding at utterance level.Source: §6 · p. 4
• On Libri2Mix, MG-REPA reaches the MeanFlow-TSE backbone’s OVRL and PESQ with 3.0× and 2.1× fewer epochs, respectively.Source: §6 · p. 4
• Ablations favor supervising both granularities over alignment at either granularity alone; future work will test generality on other backbones such as AD-FlowTSE.Source: §6 · p. 4