Knowing When to Quit: Probabilistic Early Exits for Speech Separation Networks
Authors: Kenny Falkær Olsen
Mads Østergaard
Karl Ulbæk
Søren Føns Nielsen
Rasmus Malik Høegh Lindrup
Bjørn Sand Jensen
Morten Mørup
Venue: ICLR 2026 Poster (2026)
Keywords: speech separation
speech enhancement
deep learning
early exit
dynamic neural networks
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: In recent years, deep learning-based single-channel speech separation has improved
  considerably, in large part driven by increasingly compute- and parameter-efficient
  neural network architectures. Most such architectures are, however, designed with a
  fixed compute and parameter budget and consequently cannot scale to varying compute
  demands or resources, which limits their use in embedded and heterogeneous
  devices such as mobile phones and hearables.
  To enable such use-cases we design a neural network architecture for speech separation
  and enhancement capable of early-exit, and we propose an uncertainty-aware
  probabilistic framework to jointly model the clean speech signal and error variance
  which we use to derive probabilistic early-exit conditions in terms of desired
  signal-to-noise ratios.
  We evaluate our methods on both speech separation and enhancement tasks where we
  demonstrate that early-exit capabilities can be introduced without compromising
  reconstruction, and that when trained on variable-length audio our early-exit
  conditions are well-calibrated and lead to considerable compute savings when used to
  dynamically scale compute at test time while remaining directly interpretable.

Published as a conference paper at ICLR 2026
KNOWING WHEN TO QUIT: PROBABILISTIC EARLY EXITS FOR SPEECH SEPARATION NETWORKS
Kenny Falkær Olsen .22 Mads Ostergaard
Karl Ulbækk Søren Føns Nielsen"
Rasmus Malik Høegh Lindrup
Bjørn Sand Jensen'
Morten Mørup'
'Technical University of Denmark
WS Audiology
ABSTRACT
In recent years, deep Searning-based single-channel speech separation has im- proved considerably, in large part driven by increasingly compute- and parameter- efficient neural network architectures. Most such architectures are, however, de- signed with fixed compute and parameter budget and consequently cannot scale to varying compute demands or resources, which limits their use in embedded and heterogeneous devices such mobile phones and hearables. To enable such use-cases we design a neural network architecture for speech separation and en- hancement capable of early-exit, and we propose an uncertainty-awar probabilis- tic framework jointly model the clean speech signal and error variance which we use to derive probabilistic early-exit conditions terms desired signal-to- noise ratios. We evaluate our methods on both speech separation and enhancement tasks where we demonstrate that carly-exit capabilities can be introduced without compromising reconstruction, and that when trained on variable-length audio our early-exit conditions are well-calibrated and lead to considerable compute savings when used dynamically scale compute test time while remaining directly interpretable.
INTRODUCTION
The cocktail party problem (Cherry, 1953) concerns the separation of (possibly unknown) num- ber of overlapping speakers, potentially corrupted by environmental noise and reverberation. The task is typically divided into speech separation (separating multiple speakers) and speech enhance- ment (removing environmental noise and/or reverberation), which have many applications in, e.g..
lelecommunications and hearables.
Single-channel speech separation has become predominantly deep learning-based with the introduc- tion of the Time-domain Audio Separation Network (TasNet) (Luo & Mesgarani, 2018). Notable follow-ups include Conv-TasNe (Luo & Mesgarani, 2019). which demonstrated competitive speech separation performance could be achieved significantly lower computational cost than previously held, and SepFormer (Subakan et al., 2021), )), which showed large performance improvements using transformer-basse architecture. While the transformer improved modeling performance and training parallelism over prior recurrent and convolutional networks it also tied computational complexity to the length of the attention context window, which may become prohibitive for long sequence in signal-to -o-ltto-
online contexts. For training, these works have primarily relied on the ratio (SI-SNR) loss (Roux et al., 2018) with utterance-level permutation invariant training (uPIT) (Yu et al., 2017; Kolbæk et al., 2017).
Current tath--of- (SOTA) speech separation and enhancement architectures typically cannot vary their compute in response to simplifying conditions in the input such as non-overlapping speech, low environmental noise, or silence. Following the terminology in Montello et al. (2025), we refer to neural networks which can vary compute based either self-estimated difficulty external conditions as dynamic, while networks that cannot are static. One method for introducing dynamism in neural networks is early exit (Scardapane al., 2020b), where the network can be used to make predictions at several depths through the network stack based on an exit condition. Other methods include dynamic ronting where various subsets of the network weights used during inference

---

Published as a conference paper at ICLR 2026
Freq. (Hz)
4000
2000
1000
0
Freq. (Hz)
3000
2000
EXit-SNR (dB)
1000
50
Depth
Exit point
Time (s)
Figure 1: Reconstructed spectrograms of two speakers from the WSJO-2mix test set separated by a PRESS-4-S model with 4 exit points evaluated in segments of T = 2000 samples, showing our proposed exit-SNR exit condition evaluated for each segment with a target level of 22 dB (shown in red). The distributions of each exit-SNR condition is shown shaded by exit point, demonstrating non- trivial improvement for deeper exits. extended version showing our other SNR-like distributions can be seen in appendix B, fig. 7.
based on a routing condition, such as mixture-of-expprts (Mu & Lin, 2025) (MoE) and slimmable networks (Yu et al., 2018).
Early-exit allows direct scaling of both compute and the total number of activated parameters in the network, yielding both energy and prediction latency savings (Jazbec et al., 2024), while allowing fine-grained decision-making using all information and processing available up to the current exit point. End-to-end training of early-exiting neural networks has also been shown to reduce overfitting and improve gradient propagation during training (Scardapane et al., 2020b). In the context of speech separation, a SepReformer (Shin et al., 2024) trained with additional reconstruction losses reminiscent of early exit improved the final performance of the network.
Prior work on early exit typically defines the exit condition implicitly through the loss function by defining a reconstruction loss and a utilization loss and minimizing convex combination of the two (Tan al., 2023; Bralios al., 2023; Elminshawi et al., 2024), which leads an implicit trade-off that is frozen during training, and hence cannot be adapted at infer- ence time. Other work instead defines the exit condition based on similarity between consecutive exit points (Chen et al., 2020), which does not ground the exit condition in any performance metric (which the first kind does through the loss).
performance-compun
Contributions In this work, introduce PRobabilistic Farly-exit for Speech Separation (PRESS) which leverages early exits to reduce inference costs. Importantly, our approach is probabilistic and naturally assign weights to the different exit conditions by their quantified seg-likllihood-baae based reconstruction quality. At inference time, the probabilistic formulation enables PRESS to control when exit computation using the model's confidence that a desired signal-to-noie toiiee ratio (SNR) between the estimated and clean audio signals has been achieved, providing a directly interpretable exit condition grounded in the network's 's performance each exit point.
Specifically, we make the following contributions:
We propose new early-exit framework with uncertainty awareness though a probabilistic model formulation accounting for both the clean speech signal and the variance of the error as well as its associated uncertainty. We demonstrate that our approach provides simple and principled framework to balance optimization of reconstruction quality with early-exit accuracy without the need for careful weighting of multiple objectives. Using the probabilistic model we construct several fully probabilistic SNR-like early-exit conditions which be used early-exit when target SNR level has been reached with a given uncertainty tolerance.
2. To instantiate our framework, we propose a new speech separation architecture based on linear recurrent neural networks (RNNs) designed with the joint goals of (1) achieving SOTA-level reconstruction performance per compute, and (2) being architecturally capable of outputting high-quality reconstructions from its early-exit points.
2

---

Published as a conference paper at ICLR 2026
3. We validate our approach for both speech separation the WSJO-2mix (Garofolo et al., 2007), Libri2Mix (Tzinis et al., 2022), WHAM! (Wichern et al., 2019). WHAMR! (Maciejewski et al., 2020) datasets, and for speech enhancement dataset, demonstrating the viability of training a single, dynamic neural network with early exits achieve performance competitive with SOTA static, single-exit models.
the DNS Challenge 2020 (Reddy et al., 2020)
RELATED WORK
Tasmit-family TasNet (Luo & Mesgarani, 2018) replaced masking in fixed spectral represen- tation with learnable convolutional encoder and decoder layers. The bulk of parameters and com- pute of TasNets are in RNN masker network. which produces masks that are applied to the en- coded mixture before decoding to produce single-source estimates. TasNet inspired follow-up work such computationally efficient variants like the Conv-TasNet (Luo & Mesgarani, 2019) and Su- doRMRF (Tzinis et al., 2020), as well as the DP-RNN (Luo et al., 2020) which introduced dual-path processing, decoupling time-mixing operations into chunks. The dual-path structure was also used with transformer-based networks in SepFormer (Subakan et al., 2021), TF-GridNet (Wang et al., 2023) and MossFormer (Zhao & Ma, 2023; Zhao et al., 2023).
SepReformer SepReformer (Shin et al., 2024) introduces early splitting; instead of projecting a shared encoded representation into separate estimated sources late the network as TasNet, the projection into separate sources occurs early in the network. Further processing happens inde- pendently for each source with only a cross-speaker attention layer enabling information exchange between processing for the sources. SepReformer is constructed as a U-net (Ronneberger et al., 2015) with a transformer-like stack (Vaswani et al., 2023) composed of a convolutional block, a time-mixing attention block, and a cross-speaker attention block, and uses LayerScale (Touvron et al., 2021) l) to train comparatively deep and slim networks. The network maximizes SI-SNR using uPIT with extra loss terms maximizing the SI-SNR of downsampled requenccydomait representa- tions of the estimated sources to encourage early separation of sources.
Diffusion models, SNR estimation and iterative refinement Seplt (Lutati et al., 2023) iteratively refines its estimates and stops processing based on bounding the SNR through mutual information between current estimates and the input mixture-this bound is extended consider generative models in DiftSep (Scheibler et al., 2022). Diffusion/score-basee models like DiftSep can, generally, sample the learned diffusion process with variable compute requirements For instance, DiffWave (Kong et al., 2021) - diffusion-based vocoder also used in separation tasks in Separate and Diffuse Lutati al. (2023)-show shows faster sampling by reducing needed steps through designing an appropriate variance schedule.
trade-off with quality.
Slimmable Networks allow the width the network to adjusted at inference time to trade-off compute and accuracy have been termed slimmable networks (Yu et al., 2018; Yu & Huang, 2019; Li et al., 1). This can be achieved using switchable batch normalization (Yu et al., 2018), which essentially, given predefined set of widths, trains single network with multiple batch normaliza- tion layers, one pr. width. The width specific notch-normalization can be moved to post-training (Yu & Huang, 2019), employing knowledge distillation during training while sampling different widths pr. batch. The idea of dynamic slimmable networks (Li et al., 2021) uses a dynamic gating mecha- nism that routes the signal subset of the next stage in the network to reduce complexity. In the context of speech and audio, Tiim-TaNeet (Elminshawi et al., 2023) combines the TasNet architec- ture for speech separation with width slimming and given input width size. dynamic version of this, dynamic slimmable TasNet (Elminshawi et al., 2024), utilizes predictive scheme, where each subnetwork predicts how much the following subnetwork that should be utilized. Similarly, the context of speech enhancement, dynamic channel pruning (Miccini et al., 2023) is technique for estimating mask applied the channel dimension of convolutional layers, that reduces the runtime computational cost.
Early Exit Various strategies are used to construct early-exit conditions, and to train early-exit models. In classification tasks, such image classification, the entropy at each exit has been used as an exit condition (Teerapittayanon et al., 2017; Scardapane et al., 2020a) using pre-softmax log- its as proxy for model uncertainty. In the context of language modeling the Sparse Universal
3

---

Published as a conference paper at ICLR 2026
Transformer (Tan al., 2023) used stick- stick-breaking construction to define monotonically increas- ing halting probabilities which down-weighted intermediate activations in later layers to scale down their influence on the final prediction, which was learned through single loss. In the context of speech separation, the Euclidean norm difference between successive blocks has been used as an exit condition in transformer-basee used based on a learned gating function in an iterative model (Bralios et al., 2023). In speech en- hancement, PDRE (Nakatani al., 2025) iteratively applies deterministic U-net network to the noisy signal and predicts the parameters of Gaussian mixture model (GMM), to get a distribution over the clean signal at each iteration, and the network is trained to maximize the weighted sum of log-likelihoods of each step's GMM. However, no stopping criteria or exit-conditions were explored to determine when stop the enhancement process.
architecture (Chen et al., 2020) and early-exit has also been
METHODS
The goal of single-channel speech separation neural networks (and enhancement as a special case) to separate input mixture signal ã E R. into set of estimated sources क; RT which approximate the target sources 2; E RT the input mixture, where is the total number of sources and j.j [] To incorporate early-exit into such networks, we require (1) a set of early-exit conditions which can be used to decide when an estimate is acceptable enough to exit, (2) an objective to learn both the source reconstruction and the exit conditions, and (3) neural network architecture which can support early exit without compromising reconstruction performance.
3.1 PROBABILISTIC SPEECH MODELLING
The performance of speech separation and enhancement systems is often reported as signal-to-noke ratios (SNRs) SNR improvements (SNRis) (i.e. the relative gain in SNR by using 2; over the input signal T)), computed as the ratio of the power of the target signal æy to be estimated and the power of the error signal x) Z
|
SNR(2,, 20) ANRi(æ,,,,,.32 SNN((2,; (
| क 2 ||. 3.
SNR(æ,, 2.)
(I)
||.. 24
Using SNR and SNRi as early-exit conditions for speech separation is desirable because they con- ceptually measure the loudness the error relative to the loudness of the target (or the improvement over the target), ensuring that the system can optimistically exit when acceptable speech-to-noise balance is obtained.
Both SNR and SNRi however require access to the target 2y,, which is unknown in predictive setting. We therefore propose to probabilistically model both the target æ, and the error of the prediction using a simple Bayesian objective where the target signal æj is modelled by jointly pre- dicting an estimated signal Z; and a variance parameter of E R, where we assume Gaussian distribution on the signal error and conjugate -aarse-gamma prior on the variance. Marginalizing out the variance, we obtain multivariate Student t-likelihood:
C. = / N(0x) a, If.ooi)nnGGamcc aa (x).3:)dd do =St 2j
B
R,
20
(2)
T In x Inl lnnnnr + T 2T) -InI(a:))) T In 3 (ac + 2
ar.
In 1
(3)
23g
where r'(-) is the gamma function and a; and 3, are the shape and scale in the inverse-gamma parameterization, which are to predicted by the model along with the estimated signal š. This objective strikes a simple balance between reducing the ratio of the signal error and variance scale through the last term, while also being penalized for underestimating the variance by the second-to- last term. See appendix F for scale-invariant version of this objective and relationship to SI-SNR. When optimizing this objective we use uPIT to assign targets 2; to estimated sources and parameters L, ais, Br by taking the maximum likelihood permutation. When training with multiple exits, jointly permute all exits together such that speakers cannot swap between consecutive exits, and optimize the total likelihood by summing over all exits and speakers without any weighting.
4

---

Published as a conference paper at ICLR 2026
Using the distributional assumptions the target and error signals allows us to construct early exit conditions expressed directly as predictive signal-to momal-to-noise ratios, providing interpretable thresholding mechanism for early exit. Based on our model as- sumptions, we have
Predictive Signal-to-Noise Ratios
2j of
||ا| oî
(4)
of
a2 2
Ø. 2:112 ~ x,x
| - m|| o,
(5)
o
where (A) is non-central chi-square distribution with non-centrality parameter and T degrees of freedom. The SNR and SNRi can now be written as ratios of (non-central) chi-square distribu- tions,
||æ,,
|
SNNR
SNR(
ØSNR XT
Xi:
(6)
||
6
o?
||0,
1:112
ŌSNRI
|
SNRi(æy, B,,
PSNRI
XT
||00
1:
of
These chi-square distributions are not independent, but the ratio even dependent chi-square vari- ables with equal degrees of freedom quickly concentrates around its mean for large 'T' (see ap- pendix A for details), and in the limit of large means and see that the expressions takes the form of shifted gamma distributions,
approximate the ratios with their conditional
@ 2
1
SNR(Ø3, ..))
T-000
Tof
+ ZSNR
ZSNR
Gam
3,T
(8)
SNRi(a,, R. x) T-000 [ - 3112
|
2
To?
ZSNRi
ZSNRi Gam
Olis
B,T
(9)
Unified Early-Exit SNR Both SNR and SNRi can be overly pessimistic as exit conditions when used in isolation, silence in the target such that | O will cause the local SNR approach zero, and a lack of interfering sources such that ||ا,, - . 2 z 0 similarly causes the SNRi to vanish. Additionally, in the case of total silence in all sources both SNR and SNRi vanish. To mitigate these silence issues, we propose to use a third auxiliary loudness condition measuring the SNR between a fixed reference signal Fret R.T with average power P2 R and the predicted noise signal,
|
P2
P B
SNRret(az,))
T
=2
Gam
Oi,
(10)
| 2.12
This third condition be used to exit whenever the predicted noise is quieter than the reference signal by some level difference, where p should set based on some tolerable lower bound on the loudness of the noise.
We can now integrate all three exit conditions into a single unified condition, where we consider the maximum of the individual complementary CDFs for target SNR level t R.
(PSNN j. ..) 2 f),
p(SNRexid Exxin (I,, t) max
(SNNi [zj, Cr T 2 1), (SNR (x:,) 1)
(11)
which corresponds optimistically exiting when least one of the conditions met with sufficient confidence. We note that p P(NN 2 t) is itself valid complementary CDF whose full distribution
can be given terms of its individual component PDFs and CDFs.
To obtain our final combined exit condition across all speakers we take the minimum over the com- plementary exit-SNR CDFs for all speakers and exit when the target level is exceeded with confi- dence [0, 1].
min PSSNPe (x), D., x)
(12)
5

---

Published as a conference paper at ICLR 2026
encoder
decoder
Toxporal mixing
Tempord wbxing
Spe attention
Temporal mixing
Brodde keaa
Speaker splt
eNi??
InuGon
exit
AdM
Figure 2: Detailed architecture of PRESS-Net consists of three parts: an encoder, an early split module and a reconstruction decoder with the ability to reconstruct early.
which corresponds to pessimistically exiting only when all speaker meet at least one condition with sufficient confidence. We show the predicted distributions of this condition in fig. 1, where P is set to -35 decibel (dB) relative to full scale (dBFS) such that a target level of / = 22 dB corresponds estimated noise loudness below -57 dBFS, and full version in appendix B, fig. 7 where the individual component distributions are also shown. We leave the extension of early-exiting each speaker individually important future work.
Calibration Using to our probabilistic modelling setup we can readily investigate whether our predicted o² distributions are calibrated - that is, they model the actual distribution of the observed
mean eтror power o? æ - well. We consider two methods for evaluating this: (1) the probability integral transform (PIT) (Dheur & Taieb, 2023) evaluates the predicted CDF on the observed mean error oe i.e. computing Z = (ā?) which, given perfect calibration, should itself be uniformly distributed - we use this property to construct calibration curves in fig. and (2) the continuous ranked probability score (CRPS) (Matheson & Winkler, 1976) which is proper scoring rule generalizing the mean absolute error distributional regression.
Block-wise Likelihood Modelling a single, global of assumes stationary error signal which not realistic for real-world audio. We also experiment with a blocked variant of the likelihood in eq. (2), where model blocks of T consecutive samples by independent of, for every block b This redefines T to fixed block size rather than the full length of the audio signal. Note that this likelihood does not require the underlying neural network to be evaluated blocks, and only serves control which samples in the estimated sources modelled by the corresponding predicted block-wise inverse-gamma parameters.
3.2 MODEL ARCHITECTURE
We design fasNet-family model suitable for early exit using PRESS, and denote the architecture PRESS-NV -et as illustrated in fig. 2. We base the initial design of the PRESS-NG Net architecture building blocks from the SepReformer (Shin et al., 2024) model due to its strong performance in speech separation. Following recent trends (Shin et al., 2024; Wang et al., 2023), opt for design where a shallow encoder/decoder pair (referred to as encoder/de- coder heads in our models) down-/upsamples the audio signals whereas separator module directly maps the encoder output to the decoder, rather than masking the encoder output as in TasNet. The encoder/decoder heads and aaverse-gamma parametrization blocks are described in appendix C. We do not perform any down- upsampling in the separator network in order to be able to reconstruct estimated sources early the network without excessive artifacts from upsampling. We use the GELU (Hendrycks & Gimpel, 2023) activation function.
encoder-separrato--eeeoae
We construct our separator module as a deep transformer-lik stack with pre-norm, skip connections, and LayerScale (Touvron al., 2021; Wortsman al., 2023). Specifically, for every layer f with input æ in the stack we compute its output as x +- x + 7 (normm(x)) (x)) where norm(.) is RMSNorm (Zhang & Sennrich, 2019) layer with per-channel scales initialized to
Separator

---

Published as a conference paper at ICLR 2026
Ablation
S1-SNRi
SDRi
Params
Ablation
SI-SNRi
SDRi
Receptive field
(a) ST-SNR loss
22.95
23.1
3.55M
T 8000
22.82
22.98
1000ms
(b) Normal likelihood loss
22.42
22.58
3.55M
4000
22.81
22.99
500ms
t-likelihood per-exit uPTT
21.1
20.97
3.55M
T 2000
22.79
22,96
250ms
(d) t-likelihood + exits
22.89
23.01
3.57M
T 1000
22.69
22.91
125ms
(c) t-likelihood 12 exits
22.9
22.99
3.66M
T 500
22.69
22.86
62ms
(1) t-likelihood 200K finetune
22.9
23.11
3.55M
(b) Block size ablations.
(a) Training setup ablations.
Table I: Ablation results the WSJO-2mix test set.
and = 1002 as in Parker et al. (2024), and E R is the per-channel LayerScale scaling factors, which we initialize to 10-5 This allows very deep networks to be trained stably the separator is approximately skip connection at initialization and only later learns to integrate non-linearity as y
grows.
Due to not downsampling in the separator module the time-resolution of the intermediate activations much higher than SepReformer, precluding the direct application of self-attention over time due to its prohibitively high cost from the quadratic compute scaling. Instead, we use linear RNNs with self- self-gating as our primary building block, as well as speaker attention layers from SepReformer throughout the separator. These blocks are described in more detail in appendix C, and illustrated in fig. 8.
We use early split as in SepReformer, where the first Nnn layers (consisting only of linear RNN blocks) in the stack process the encoded speech mixture after which a speaker split mod- ule SpeakerSplit RTTD + TxXD projects the refined speech mixture into S separate groups of channels, where the speaker dimension is considered batch dimension in later processing.
After splitting, we repeat decoder block NDec times which consists of linear RNN blocks and speaker attention blocks repeated in a 5:1 ratio. After each block, we may place an exit point E; < Nor which reconstructs the current latent representation with separate decoder head. Apart from reconstruction, each exit point parametrizes the inverse gamma distribution which we can use to evaluate whether should exit not.
RESULTS AND DISCUSSION
PRESS is evaluated on four speech separation corpora: WSJO-2mix (Garofolo et al., 2007), Libri2Mix (Cosentino et al., 2020), WHAM! (Wichern et al., 2019), and WHAMR! (Maciejewski al., 2020), as well one denoising corpora: DNS2020 (Reddy et al., 2020). Complete descrip- tions of the corpora can be found in appendix D.
We train two main model configurations: (I) PRESS-4 (S), smaller model with model dimension D = 64, P = 4, Dune = 256, NEne = 8, Nee = 12 and 4 exit points placed at every 3 decoder blocks, and (2) PRESS-12 (M). larger model with model dimension increased to D 128. NEnc 4, Nec = 24 and having 12 exits, placed at every second decoder block. Further details of model training can be found in appendix E.
We show the performance of our models in table 2 compared with other SOTA methods, and in fig. 3 we plot performance in SI-SNR improvement (SI-SNRi) as a function of giga-multiply (GMAC) per second (GMAC/s), showing how PRESS models can scale their compute dynamically while remaining competitive on WSJO-2Mix.
In fig. 5 we plot calibration curves for the WSJ0-2Mix training and test sets, showing that models are not well-calibrated after training on 4-second audio clips, and experiment with finetuning these with just 200K extra training steps (ca. 3% of base training time) on full-length audio clips from the training set, after which we see our models become well-calibrated, and performance also increases substantially which we show in table in the bottom rows.
We train several ablation variants of PRESS-4 (S) to determine the effectiveness of (a) using SI-SNR as loss instead of t-likelihood, (b) using normal likelihood with single predicted variance
7

---

Published as a conference paper at ICLR 2026
WSJ0-2mix
Libri2Mix
WHAM!
WHAMR!
SI-SNRi
SDRI
SI-SNRi
SDRI
SI-SNRI
SDRI
SI-SNRi
SDRi
Params
GMACK
Model
(dB)
(dB)
(dR)
(JB)
(dB)
(dB)
(dR)
(dB)
(M)
(G/s)
Conv-TasNet'
15.3
15.6
12.2
12.7
12.7
8.3
5.1
10.5
DualPathRNN
18.8
19.0
16.1
16.6
13.7
14.1
10.3
2.6
42.5
SepFormer
20.4
20.5
19.2
19.4
14.7
16.8
14.0
26.0
86.9
SepFormer + DM
22.3
22.5
16.4
16.7
14.0
13
26.0
86.9
MossFormer (S)
20.9
10.8
_2
MossFormer (M) DM
22.5
25.3
MossFormer (L) DM
22.8
17.3
16.3
42.1
70.4
MossFormer2 DM
24.1
21.7
18.1
17.0
55.7
84.2
TF-GridNet (S)
20.6
8.2
19.2
TF-GridNet (M)
22.2
36.2
TF-GridNet (L)
23.5
17.3
15.8
14.4
SepMamba DM
21.2
21.4
7.2
12.5
SepMamba (M) +
22.7
22.9
22.0
37.0
SepReformer
22.4
22.6
20.2
17.2
17.5
3.7
10.4
SepReformer (S)
23.0
23.1
20.6
21.0
17.3
17.7
21.3
SepReformer (M)
24.2
24.4
22.0
22.2
17.8
18.1
17.3
81.3
SepReformer (I.) + DM
25.1
25.2
18.4
18.7
17.2
16.0
55.3
155.5
PRESS-4 @ 4 (S)
22.91
23.08
20.04
20.41
16.49
16.91
14.54
13.37
3.4
11.3
PRESS-12 @ (M)
22.64
22.93
19.76
19.71
16.43
16.71
14.24
13.09
8.7
29.1
PRESS-12 @ (M)
23.47
24
20.42
20.86
16.57
17.03
14.67
13.45
15.6
54.4
PRESS-12 (M)
24.28
24.46
20.88
16.65
17.12
14.69
13.47
22.4
79.7
PRESS-4 @ (S) FT
23.41
23.56
21.01
21.36
17.25
17.58
15.13
13.92
3.4
11.3
PRESS-12 (M)
23.27
23.43
20.31
20.72
16.47
16.81
14.91
13.73
29.1
PRESS-12 (M)+ FT
24.18
24.40
20.92
21.33
17.21
17.34
15.61
14.38
15.6
54.4
PRESS-12 @ 12 (M) + FT
24.36
24.55
21.29
21.68
17.49
17.89
15.67
14.43
22.4
79.7
Table 2: Speech separation performance the WSJO-2Mix, Libri2Mix. WHAM! and WHAMR! 2-speaker test sets with our PRESS variants highlighted with gray background. We evaluate our PRESS-4 (S) model its final exit point, and also show results for our larger PRESS-12 (M) model at exit points. +DM: models trained with dynamic mixing. +FT: models finetuned on full-length training data. : some values from Shin al. (2024).
26
SepReformer DL
Refermees
SepRef
ossFormer + DM
TF-GridNet (L)
SepRaformer
Sep Mamba (M
ossForrer (L)
3ISNRI (db)
22
RepRefor
DM
F F-GridNet (M)
SepFormer DM
TF- TF-GridNet
SeaFormer
PRESS-4
SepReformer
PRESS 4 (5) (dynamic)
TF GridNet
PRESS-4 (S) (single)
Sepformer
SepMamba
PRESS-12 (M) (single)
Conv-TasNet
Conv Tachlet
Mossfomer
DR-ANN
10
40
50
100
150
200
Compute (GMAC/S)
Figure 3: Source separation performance on WSJO-2mix in terms of SI-SNR improvement (SI- (GMAC) per second (GMAC/s)), with the area of points corresponding to parameter count of models. The static performance of every exit point is shown for PRESS models, as well as the dynamic performance the PRESS-4 (S) model using probabilistic exit condition for varying target levels, beating the static performance curve in efficiency. We also include the performance of single-exit models, which underperform the jointly trained model at deeper exits.
SNRi) per compute wwwpmmmllllllcccullaaeee
8

---

Published as a conference paper at ICLR 2026
20
30 dB
Exit SNR Regret, (JB)
5.5
Dynemic
0.45
Static
0.40
2.0
5.0
Oracle
Uniform
0.35
1.5
4.5
0.30
4,0
1.0
10
10 12
8
12
Com putp (GMAC/S)
Compute (GMAC/S)
Compute (GMAC/S)
Figure 4: One-sided exit-SNR regret the WSJO-2mix test set for PRESS-4 (S) model trained with a block size of 2000 samples using different early-exit strategies with target levels = 20, 25, 30 dB: (dynamic) our probabilistic exit strategy eq. (11) evaluated for varying confidence thresholds (static) using single exit for all blocks, (oracle) a best-case strategy that always exits when the target is achieved using the ground-truth exit-SNR, (uniform) uninformed strategy that selects exit points uniformly at random.
1.0
1.0
1.0
1.0
Emplrical cDr (PIT)
Empirical cDF (PlT)
Empiriaal CDF (PTT)
Empirical cDF (P1))
0.6
0.6
0.6
0.6
0.4
0.4
0.4
0.4
MT
PIT
0.2
0.2
0.2
0.2
Ideal
0.8
0.2
0.2
1.0
Nominal quantile
Nominal quantile
Nominal quantile
Nominal quantile
(a) Full-length training
(b) Full-length test data. CRPS 2.96
(c) Full-length training data finetuned.
(d) Full-length test data finetuned.
data.
CRPS 1.61
CRPS 1.43
CRPS 2.80
Figure 5: Calibration curves for the predicted of mean error distributions on the WSJO-2mix test set for a PRESS-4 (S) model with a block size of 2000 samples. In (a) and (b) we see that the distributions are uncalibrated when the model is trained on 4-second clips and evaluated on full- length sequences on both training and test data. In (c) and (d) we see that the model predictions become well-calibrated on both training and test data after finetuning full-length training data.
instead of normal-inverse-gammet likelihood (c) joint (i.e. all per-exit losses permuted together) vs. per-exit permutation-invariaane where speakers can swap between exits, (d) and (e) the number of exits we train 4- 6-, and 12-exit variants of our small model configuration with exit exit-points placed uniformly over depth, and (f) where we ablate that simply finetuning on more 4-second audio does not lead to the same performance improvements training on full-length data.
Ablation results can be found in table 1. For ablation (a) we see that our Student t-likelihood can be used in place of SI-SNR without loss of performance, even though our t-likelihood is scale- invariant. Ablation (b) shows that simpler normal likelihood results in worse reconstruction per- formance in terms of SI-SNR, possibly due
log-scaling the error as both SI-SNR and the t-likelihood do. Ablation (c) reveals that permuting consecutive early exits together is crucial for stable joint training of early exits, likely because per-exit commutation-nvariaan network to swap sources through the speaker attention layers, defeating much of the point of the early split architecture design. ablations (d) and (e) we found that increasing the number of exits from 4 to 6 12 did not worsen performance at any of the exit points, which motivated us train larger models at the 12-exit level. We also see in ablation (f) that additional finetuning with 4-second audio does not increase performance.
would allow the
We further investigate the use of the PRESS architecture the DNS2020 dataset by treating the speech enhancement task source separation, where predict both the clean speech and noise signals as separate sources. Surprisingly, as seen in table 3, this leads to very competitive performance after accounting for total GMAC/s even though model also explicitly recovers the noise signal

---

Published as a conference paper at ICLR 2026
ST-SDR
STOI
WR-PESQ
# Params
GMACk
MFNet (Liu et al,, 2023)
20.31
97.98
3.43
6.1
MP-SENet al., 2023)
21.03
98.16
3.62
2.26M
40.7
TF-Locoformer (Saijo et al,, 2024)
23.30
98.
3.72
14.97M
248.6
ZipEnhancer (Wang & Tian, 2025)
22.22
98.65
3.81
11,34M
133.5
PRESS-4 @ 4 (S)
20.53
96.34
2.69
3.55M
11.6
PRESS-12 (M)
20.97
96.52
2.92
8.57M
29.1
PRESS-12 (M)
21.98
96.97
3.10
14.95M
53.7
PRESS-12 @ 12 (M)
22.15
97.13
3.10
18.14M
78.3
Table Speech enhancement performance on the DNS2020 non-blind test set without reverbera- tion.
while other methods do not, with our smaller PRESS- 12 (M) model configuration matching the ZipEnhancer (Wang & Tian, 2025) method in terms of SI-SNRi using substantially less compute.
We also experiment with evaluating our exit-SNR condition in fig. 4, at 3 different target levels by varying the confidence threshold p, and measuring the one-sided regret of exit condition, i.e. the difference between the achieved exit-SNR and the target level, or 0 if the exit-SNR exceeded the target. We that using our early exit strategy closely matches the oracle strategy the appropriate confidence level.
5 CONCLUSION & FUTURE WORK
We introduced the PRESS method and PRESS-Net model, achieving competitive performance on WSJO-2Mix, Libri2Mix. WHAM!, WHAMR! and DNS Challenge 2020 while allowing flexible compute scaling based on probabilistic early-exit conditions.
Our probabilistic approach allows well-defined early exit conditions to be formulated with integrated uncertainty quantification using the CDF of the exit-SNR distributions, and can be used in place of conventional training objectives such as SI-SNR no apparent cost to the reconstruction objective.
Our predictive SNR-like distributions proved to be very well-calibrated after finetuning full- length data, which also provided significant performance improvements to reconstruction.
An extension of our work would be to consider iterative models, i.e. special cases of our model with a single shared block in the decoder stack repeated for each exit point. This would allow theoretically infinite scaling with compute, but if done naively couples the total parameter count to the size of the iterative block requiring more careful network design, possibly using width-scaling neural networks.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work is partly funded by WS Audiology and the Innovation Fund Denmark (IFD) under File No. 3129-00075B.
REFERENCES
Libri 2005. URL https::://l/brrivo
Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs, Febru- ary 2025. URL https::///////c arxiv. wtg/ abaaabbs/2222 15938v1.
Guy E. Blelloch. Prefix sums and their applications. 1990. URL http:://sheeeeeee library. pdf.
cmu. www.eeec//eeeeeeee University Pittsburg,, PA, USA.
Tecch 4 54461 pdf.. Publisher: School of Computer Science, Carnegie Mellon
Dimitrios Bralios, Efthymios Tzinis, Gordon Wichern, Paris Smaragdis, and Jonathan Le Roux. Latent Iterative Refinement for Modular Source Separation. In ICASSP 2023 2023 IEEE In- ternational Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5, June 2023. doi: 10.1 N0.1109//CASSP499335 202. URL /tL.ttp://arravv.oo//aaaas 11917. arXiv:2211.1191 [es].
George Casella and Roger Berger. Statistical inference. Chapman and Hal//CRC, 2024. URL https://www www.ttaylorrrrano http:///ssocooooooooooooo/////// 1tt.://111111111111111110000
Sanyuan Chen, Yu Wu, Zhuo Chen, Takuya Yoshioka, Shujie Liu, and Jinyu Li. Don't 't shoot butterfly with rifles: Multi-channel Continuous Speech Separation with Early Exit Transformer, October 2020. URL http:/// /arxiv.. org alagg//ba// 22100.2000.. 12180. arXiv:2010.12180.
Edward Collin Cherry. Some experiments on the recognition of speech, with one and with cars. Journal of the aconstical society of America, 25:975-979, 1953. URL https:/ /pure mmp. be/rest/itee ettemttmmmmm 2309493 99999999_9550 http///oppeet/tceeeeeeeeeeeocett /file 2acilee2309493 2309492 cono9492/cooeett...
Joris Cosentino, Manuel Pariente, Samuele Cornell, Antoine Deleforge, and Emmanuel Vincent. LibriMix: An Open-Source Dataset for Generalizable Speech Separation, May 2020. http:
// //arxiiv.org wtt..//g/.b./..........// 11262. arXiv:2005. 1262 [eess].
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Des- jardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre. Griffin: Mixing Gated Linear Recurrences with Local Attention for Ef- ficient Language Models, February 2024. URL httttp:/arxiv..og//Baa. arXiv:2402. 19427 [es] version: 1.
Victor Dheur and Souhaib Ben Taieb, A Large-Scale Study of Probabilistic Calibration in Neural Network Regression, June 2023. URL http:// htttp:///axii....t wttp://..rg///////////////////// 02738. arXiv:2306.02733 [es].
Mohamed Elminshawi, Srikanth Raj Chetupalli, and Emanuël A. P. Habets. Slim-Tasnet A Slimmable Neural Network for Speech Separation. In 2023 IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (WASPAA), pp. 1-5, October 2023. doi: 10.1 URL http::/Leeexpoors Leee wot///occeet...../ooooooo
WASPAA58266 200.A886666.22..1000 202., 102.41441..... 10248143. ISSN: 1947-1629.
Mohamed Elminshawi, Srikanth Raj Chetupalli, and Emanuël A. P. Habets. Dynamic Slimmable Network for Speech Separation. IEEE Signal Processing Letters, 11:2225-22001 2024. ISSN 1558-2361. doi: 10. 00.10999LPP22224444434444 URL http://///eeexxeeeeeeeeeee
11

---

Published as a conference paper at ICLR 2026
Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadegh. Were RNNs All We Needed?, October 2024. URL URRhttt://aaxtt org abs/ /2100.410... 01201. arXiv:2410.01201 [es].
Frederic Font, Gerard Roma, and Xavier Serra. Freesound technical demo. In Proceedings of the 21st ACM international conference on Multimedia, MM '13, pp. 411-412, New York. NY, USA.
Octoher 2013. Association for Computing Machinery. ISBN 978-1-4503-2444-5 doi: 10.1145/
2550225.. URL https: / /dl. acm. weccm.rrg//ae hwt./oogg/ooooooeoooooooooo 1145 1ttp:///....111552222222222.0000
John S. Garofolo, David Graff, and David Pallett. (SR-I(WSJO) Complete LDC93S6A, 2007. URL attps://catall
1dc. upenn wwt...cddd/......
Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces, December 2023, URL http://arxiv org/ abs/ 22122.1. 00752. arXiv:2312.00752 [cs].
Jens Heitkaemper, Darius Jakobeit, Christoph Boeddeker, Lukas Drude, and Reinhold Haeb- Umbach. Demystifying TasNet: A dissecting approach. In ICASSP 2020-2020 IEEE Interna- tional Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3363. IEEE, 99052981.
2020. URL htttttt//eeeepppiieeee
Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), June 2023. URL http: ///aaxiv..rg/ab///////aaa
08415. arXiv:1606.08415 [cs].
Sukjun Hwang, Aakash Lahoti, Tri Dao, and Albert Gu. Hydra: Bidirectional State Space Mod- /2407.
els Through Generalized Matrix Mixers, July 2024. URL 09941. arXiv:2407.0994 [cs].
Metod Jazbec, Alexander Timans, Tin Hadi Veljkovi, Kaspar Sakmann, Dan Zhang, Christian A. Naesseth, and Eric Nalisnick. Fast yet Safe: Early-Exiting with Risk Control, November 2024.
URL hww.hhtp///aaaaa //axiv. org/ abs//455.225 20915. arXiv:2405.20915 [cs].
Anwar H. Joarder. Moments of the product and ratio of two correlated chi chi-square vari- ables. Statistical Papers, 50(3):581-50) June 2009, ISSN 0932-5026, 1613-9798. doi: 10.1007760002440-000055,,, 50C362-007-0105-1
URL http:// springer. com/10.
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization, January 2017.
URL httpt:tt///aiiee org/ abs//4/2/111./ 6980. arXiv:1412.660 [cs].
Morten Kolbæk, Dong Yu, Zheng-Hua Tan, and Jesper Jensen. Multi-talker Speech Separation with Utterance-level Permutation Invariant Training of Deep Recurrent Neural Networks, July 2017.
URL
org/ abs/ 110.. 0 06284, 4. arXiv:1703.06284 [cs].
Zhifeng Kong. Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DifiWave: A Versatile Diffusion Model for Audio Synthesis, March 2021. URL 09761. trXivv2009.0972 [eess].
2209.
Changlin Li, Guangrun Wang, Bing Wang, Xiaodan Liang, Zhihui Li, and Xiaojun Chang. Dy- namic Slimmable Network, March 2021. URL arXiv:2103. 13255 [cs].
22252113..33. 13258.
Liang Liu, Haixin Guan, Jinlong Ma, Wei Dai, Guangyong Wang, and Shaowei Ding. A Mask Free Neural Network for Monaural Speech Enhancement, June 2023. URL abs/2300 04286. trXiv:2306.04285 [cs].
Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization, January 2019. URL http: org wwwg//b///// /1711. 0510111111111111111 01. 1rXiv:1711.051 1.05111 [es, math].
Ye-Xin Lu, Yang Ai, and Zhen-Hua Ling. MP-SENet: A Speech Enhancement Model with Parallel Denoising of Magnitude and Phase Spectra. August 2023. doi: 00221437/ntrrppecch2223 1441.
URL http: / /tt///aaiiii rxiv. org/ aas//bb//3355 13686. arXiv:2305.1 13666 [eess).
12

---

Published as a conference paper at ICLR 2026
Yi Luo and Nima Mesgarani. TaSNet: Sine-Domain Audio Separation Network for Real-Time, Single-Chamnn Speech Separation. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). pp. 696-700, April 2018. doi: 10.1109/CAASPP2211....M 8.88.. URL https: / / eeeepllreeeee....e.e. leee. org/ document.//////o////mocomm 84 6211 6. ISSN: 2379-190X.
Yi Luo and Nima Mesgarani. Conv-TasNet: Surpassing Ideal Time-Frequency Magnitude Masking for Speech Separation. JEEE/ACM Transactions on Audio, Speech, and Language Processing, 27 (8):1256-1266, August 2019. ISSN 2329-9290, 2329-9304.
doi: 10.1099TTAS.P220999955567
arXiv: 1809.07454 [es, eess].
Yi Luo, Zhuo Chen, and Takuya Yoshioka. Dual-path RNN: efficient long sequence modeling for time-domain single-channel speech separation, March 2020. URL htttttt//axxvv.ooo/
abs 11t//190... 10.00.0633 06379. arXiv:1910.0637 [cs, eess].
Shahar Lutati, Eliya Nachmani, and Lior Wolf. Separate And Diffuse: Using a Pretrained Diffusion Model for Improving Source Separation, June 2023. URL http:://taxtttt org/ wwg/abs///./.oooo.oeeo..oooe.e abs/2301. 10752. arXiv:2301.10755 [eess].
Matthew Maciejewski, Gordon Wichern, Emmett McQuinn, and Jonathan Le Roux. WHAMR!: Noisy and Reverberant Single-Channe Speech Separation, February 2020. URL http:// arxiv. org/ abs / 91 0 0279. arXiv: 10279 [es].
Eric Martin and Chris Cundy. Parallelizing Linear Recurrent Neural Nets Over Sequence Length, February 2018. URL http: / aawai.iiiiv.///.i///aaaaaaaaaava /tapg///b///ia 04057.
[es].
James E. Matheson and Robert L.. Winkler. Scoring Rules for Continuous Probability Distri- butions. Management Science, 22(10)) 1976. ISSN 0025-1909. URL https:
// //wwww www.jstorr sestorr.ogg//talleee/eee b2tpsbb/lee22222.... Publisher: INFORMS.
Riccardo Miccini, Alaa Zniber, Clément Laroche, Tobias Piechowiak, Martin Schoeberl, Luca Pezzarossa, Ouassim Karrakchou, Jens Sparsó, and Mounir Ghogho. Dynamic nsNet2: Efficient Deep Noise Suppression with Early Exiting. In 2023 IEEE 33rd International Workshop Machine Learning for Signal Processing (MLSP), pp. 1-6, Septemher 2023. www.xiiv.org/aaoo http://aab/33381666663
doi: 10. 11109//LLS555444202...12222 arXiv:2308. 1667788888..888888888 [cs].
URL
Fabio Montello, Ronja Güldenring, Simone Scardapane, and Lazaros Nalpantidis. A Survey Dynamic Neural Networks: from Computer Vision to Multi-modal Sensor Fusion, January 2025.
URL h axxiv org/ /251 07451. arXiv:2501.0743 [es].
Siyuan Mu and Sen Lin. A Comprehensive Survey of Mixture-of-Experts: Algorithms, Theory, and Applications, April 2025. URL
arXiv:2503.07137 [cs].
Tomohiro Nakatani, Naoyuki Kamo, Marc Delcroix, and Shoko Araki. A Hybrid Probabilistic- Deterministic Model Recursively Enhancing Speech. In ICASSP 2025 2025 IEEE Interna- tional Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5, April 2025.
doi: 10.1109/CCASPP496600 2tw.:60..222.1.87799999 10887979. URL https:://eeepploree ISSN: 2379-190X.
document/108877
Julian D. Parker, Anton Smirnov, Jordi Pons, C. J. Carr. Zack Zukowski, Zach Evans, and Xubo Liu. Scaling Transformers for Low-Bitrate High-Quality Quality Speech Coding, November 2024. URL hwtp://b//22 /2.bs/24111 19842. arXiv:241 1.9942 [eess].
http:// /ttp://axxii..v//
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, and Luca Antiga. Pytorch: An impera- tive style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. URL https: /ttp:////////eeeeeee ngs neurips ccpppapp/20111111 9 /hahh/ bdbca288fee 0atoc2288fee7992223222 www.22Ea9E000222211711. ott. Abstract html.
Serge B. Provost and Edmund M. Rudiuk. The exact density function the ratio two dependent linear combinations of chi-square variables. Annals of the Institute of Statistical Mathematics, 46(3):557-571, 1994. doi: 10.1007/BF007/3517. URL https: /tt////inkk..///////////// springer. com/ article/10. 10t.01000000000000000 7. Publisher: Springer.
13

---

Published as a conference paper at ICLR 2026
Chandan K. A. Reddy, Ebrahim Beyrami, Harishchandra Dubey, Vishak Gopal, Roger Cheng, Ross Cutler, Sergiy Matusevych, Robert Aichner, Ashkan Aazami, Sebastian Braun, Puneet Rana, Sriram Srinivasan, and Johannes Gehrke. The INTERSPEECH 2020 Deep Noise Suppression Challenge: Datasets, Subjective Speech Quality and Testing Framework, January 2020. URL /lbs/ ht.ps:://rxiis org/ alss / 2001 08662v2.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomed- ical Image Segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejan- dro F. Frangi (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015, pp. 234-241, Cham, 2015. Springer International Publishing. ISBN 978-3-319-245744. doi: 010009978-3-311-24544-4444
Jonathan Le Roux, Scott Wisdom, Hakan Erdogan, and John R. Hershey, SDR half-baked well done?, November 2018. URL http:://attt org/ abs 11111111111111111111111111111111 02508. arXiv:1811.02500 [es].
Kohei Saijo, Gordon Wichern, François G. Germain, Zexu Pan, and Jonathan Le Roux. TF-Locoformer: Transformer with Local Modeling by Convolution for Speech Separa- tion and Enhancement, August 2024. URL http: http///axvvvv.o// wwwp:ii.org/ab//2ee/ 03440. arXiv:2408.034m. 404404000444 [eess].
Simone Scardapane, Danilo Comminiello, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini. Differentiable Branching In Deep Networks for Fast Inference. In ICASSP 2020 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4167- 4171, May 2020a. doi: I0.1009/CAASP44076660000 200.099400776222000904422.. URL httppa://ieeexpporr feee. wwg. beeeee.rg//occmment///e
ISSN: 2379-190X.
Simone Scardapane, Michele Scarpiniti, Enzo Baccarelli, and Aurelio Uncini. Why should we add early exits to neural networks? Cognitive Computation, 12(5):94--966 September 2020b. ISSN 1866-9956, 1866-9964, doi: 10. 1007//1225990020-0977444 URL hRttht///aaxixi http:://wxww.....aaabba 2004. 12814. arXiv:2004.128: [cs, stat].
Robin Scheibler, Youna Ji, Soo-Whan Chung, Jaeuk Byun, Soyeon Choe, and Min-Seok Choi. Diffusion-based Generative Speech Source Separation, November 2022. URL htth: //axxiv. httg/ab/2210..1
arXiv:2210.17327 [eess].
Ui-Hyeop Shin, Sangyoun Lee, Taehan Kim, and Hyung-Min Park. Separate and Reconstruct: Asymmetric Encoder-Decodee for Speech Separation, October 2024. URL http:://rxix: org/abs/208 2406 05983. MrXiv:2406.05983.
Cem Subakan, Mirco Ravanelli, Samuele Cornell, Mirko Bronzi, and Jianyuan Zhong. Attention is All You Need in Speech Separation, March 2021. arXiv: 2010.13154 [es, eess] Number:
arXiv:2010.13154
Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse Universal Transformer, October 2023. URL http: http://wwwww/////1///00000/ arXiv:2310.07096 [cs].
Surat Teerapittayanon, Bradley McDanel, and H. T. Kung. BranchyNer: Fast Inference via Early Ex- iting from Deep Neural Networks, September 2017. URL htt.:tpppp/.......gg//bb//////// 01 01686.2 100i::1709006866 [cs].
Joachim Thiemann, Nobutaka Ito, and Emmanuel Vincent. The Diverse Environments Multi- channel Acoustic Noise Database (DEMAND): A database of multichannel environmental noise recordings. Proceedings of Meetings on Acoustics, 19(1):03081,, May 2013. ISSN 1939-800X. doi: 10.1121/1 URL htwpt:ttps.ocmmm /ww.///oii//////////// /10. 1121/1. 4799597.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Go- ing deeper with Image Transformers, April 2021. URL http: //arriv....i..// wwg//brraooooooooooooo 2abs/20 /2003. 17239. arXiv:2103.172 version: 2.
14

---

Published as a conference paper at ICLR 2026
Efthymios Tzinis, Zhepei Wang, and Paris Smaragdis. Sudo mm -Rf: Efficient Networks for Uni- versal Audio Source Separation. In 2020 IEEE 30th International Workshop on Machine Learn- ing for Signal Processing (MLSP), pp. 1-6, September 2020. doi: 10.1099/LLSP4996220220 9231900. arXiv: 2007.06833 [cs, eess, stat].
Efthymios Tzinis, Gordon Wichern, Aswin Subramanian, Paris Smaragdis, and Jonathan Le Roux. Heterogeneous Target Speech Separation. In Interspeech 2022, pp. 1-5, September 2022. doi: 1022143/mmerspecch.0222
arXiv: 2204.03594 [cs, eess].
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, August 2023. URL http:
//arxiv.... org/ aabs/1056 03762. arXiv:1706.0331 [cs].
Haoxu Wang and Biao Tian. ZipEnhancer: Dual-Path Down-Up Sampling- Zipformer for http:/iw...gg//////00111111
Monaural Speech Enhancement, January 2025. URL
05183v1.
Shong-Qiu Wang, Samuele Cornell, Shukjae Choi, Younglo Lee, Byeong-Yeou Yeol Kim, and Shinji Watanabe. TF-GridNet: Making Time-Frequency Domain Models Great Again for Monau- ral Speaker Separation, March 2023. URL http://arxiv. www.:/iwwwwrgg..///oaeeooaa http://as//22....... 03952.
arXiv:2209.03952 [cs, eess].
Gordon Wichern, Joe Antognini, Michael Flynn, Licheng Richard Zhu, Emmett McQuinn, Dwight Crow, Ethan Manilow, and Jonathan Le Roux. WHAM!: Extending Speech Sepa- ration Noisy Environments, July 2019. URL http::/arxiv..rrrr arXiv:1907.01 [cs].
/997 01160.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari S. Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and low- www-preciiion training for large-scale scale vasion-language models. November 2023. URL https:t////. //te://oopppee openreviev. wwt..eee//t.... www.eeettoooooee http://rrmmmdssssssspppppppp d==ssisSsss...........
Greg Yang. James B. Simon, and Jeremy Bernstein. A Spectral Condition for Feature Learning, May 2024. URL
abs /2100.2/31.. 17813. arXiv:2310.17813.
Dong Yu, Morten Kolbæk, Zheng-Hua Tan, and Jesper Jensen. Permutation Invariant Training of Deep Models for Speakcer-Indepeedent Multi-talker Speech Separation, January 2017. URL http:/ www....org/bb//117...11 00325. arXiv: 160.00035 [cs, eess].
Jiahui Yu and Thomas Huang. Universally Slimmable Networks and Improved Training Techniques, 903. 051 34. 1.XXvv1900305. 134 [cs).
October 2019. URL htt::/// /ttp///...tttt////////..oor/////
Jiahui Yu, Linjie Yang, Ning Xu, Jianchao Yang, and Thomas Huang. Slimmable Neural Networks, 1812.1111 08828. arXiv:1812.08929 [cs].
December 2018. URL http: /ttp:://axvvvvo.voee hatp///b///tt /ttas/1111
Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization, October 2019. URL 07467. arXiv: [cs].
http:/ /ttpxvaxvvvv/////////////////av www./org/////oooooooo...ooooe///
Shengkui Zhao and Bin Ma. MossFormer: Pushing the Performance Limit of Monaural Speech Sep- aration using Gated Single-Head Transformer with Convolution-Auggenn February 2023. URL
Joint Self-AAtentionss, 1824 arXiv:2302. 1824 [cs].
Shengkui Zhao, Yukun Ma, Chongjia Ni, Chong Zhang, Hao Wang. Trung Hieu Nguyen, Kun Zhou, Jiaqi Yip, Dianwen Ng, and Bin Ma. MossFormer2:: Combining Transformer and RNN-Free Recurrent Network for Enhanced Time-Domain Monaural Speech Separation, December 2023.
URL attp:///a /awpiv..ivvv//v.////v//..//./.iv httg/bbs/23122 11825. arXiv:23 12.1 1825 [cs, eess].
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A X. -RATIO APPROXIMATION
In section 3.1 we claim that our model assumptions allow us to express a probabilistic SNRi in terms of the estimated error variance o². From our likelihood we have (omitting speaker indices and j).
x (2 11). I).
o² InvGam(a, 3).
(13)
which we can equivalently write as,
x= 2 + o2:
~ (O..).
(14)
We then obtain
||0
2 = | oz az -ūl12 0²
2 2
XT
1
2
(15)
0
c= x a
0² ||
XT-
(16)
Φ
SNRi
(17)
These distributions are dependent on the same draw of and consequently the distribution of does nof take the form of non-central F-distribution, nor can we rely on existing dependence results that assume shared sub-summations in the two associated x² distributions (Provost & Rudiuk, 1994) or
the correlation to be known (Joarder, 2009).
Notably, the o² scaling factors cancel the ratio, so that the only sources of randomness are (1) the draw of x² in the non-centrality parameter of e and (2) the draw of z. We now instead consider the limiting behavior of the ratio as T - oo to see how it may be approximated for large T. Since
|| - 2 is a constant which scales linearly with T, we can rewrite the non-centrality parameter of #
é to TX, where we have absorbed the constant scale into a new random variable T
Gam a,
- We then have to show that for TB
Xr X2 (TA),
c= Yu X2:
(18)
the ratio Xx Yy
converges in distribution to
X7 YT
1
as + 00.
(19)
We start by noting that the normalized variables
Xx T
and
Yr have decreasing variance T -
Do,
Xxr
E
X]
Xx
Var (rr A
2
4A
E
A
1+
Var
A
0
T
T
T
72
T
(20)
Yr
EY T
Yr
Var Yrl
2
1 T
E
1,
Var
=
0
(21)
T
T
T
T
which implies that both normalized variables converge in probability their expected values as T -> co. By Slutsky's theorem (Casella & Berger, 2024), the ratio Xy then also converges in distribution as
Yr
X7
X T
Xx { E
=1
as T 7
(22)
Yr
[
YF
T
This convergence holds even when XT and Yr are independent, the normalized variables converge to the expected values separately. The SNR case is analogous, substituting | - x for ||ق||.
Xx. To get a feeling for the rate of convergence, we simulate the true distribution of + for finite T with
~ Gam a, +) and perform one-sample Kolmogorov-Smirovv tests comparing the empirical
16

---

Published as a conference paper at ICLR 2026
distribution of the ratio for 1 000 000 samples to the limiting distribution 1 + À in fig. 6. We set the parameters a = 60, 3 = 0.1, and c = 0.1 based on the average model predictions seen during training.
In table we show the average Kolmogorov-Smin Serirovv (KS) distance between approximate and simu- lated SNRi distributions the WSJO-2mix test set for Hocked-likelihode PRESS-4 (S) models, demonstrating that in practice even relatively small values T are still approximated well by our asymptotic distribution.
T (seconds, 8 kHz)
0.00
0.25
0.50
0.75
1.00
0.10
KS Distance
0.08
0.06
0.04
0.02
0.00
2000
4000
6000
8000
T
Figure 6: Simulated KS distances between the simulated and approximate SNRi distributions for varying values of T. The KS statistic represents the maximum deviation between the empirical and limiting CDFs, which is below 1% after T x 2000.
Ablation
KS distance
Receptive field
Block size T 8000
0.007
1000ms
Block size T 4000
0.008
500ms
Block size T =2000
0,012
250ms
Block size T =1000
0.019
125ms
Block size T -500
0.049
62ms
Table 4: Test KS distances on the WSJO-2mix test set.
B FULL EARLY EXIT DISTRIBUTIONS
complete version of fig. can be seen in fig.
C MODEL ARCHITECTURE DETAILS
Audio Encoder and Decoder Heads We base our encoder and decoder heads on the design from SepReformer, with the audio encoder head processing the time-domain audio signal as RT - Ponc X by passing it through 1-D convolution with kernel size K 16, stride P and encoding dimension Denc 256 with bias, which used to avoid bivision-bb-zero in later RM- SNorm layers that would otherwise occur if an all-zero input signal passed into the model. The representation is then passed through a GELU, RMSNorm and finally linear layer which projects down the model dimension Dmodel 64 128. See fig, 8d.
The decoder head consists of gated linear unit (GLU) layer followed by a transposed convolution 4 that maps the latent representation back the original
with kernel size K 16, and stride sampling as D.Dmdeet//S R Every exit point has its own decoder head; see fig. 8e.
17

---

Published as a conference paper at ICLR 2026
4000
Freq. (Hz) Freq. (Hz)
3000
20
2000
1000
4000
20
3000
2000
40
1000
SNR (dB)
60
30
20
Depth
10
SNRi (d8)
30
20
10
Ref- Ref-SNA SNR (dB)
30
20
Exit SNR (d8)
30
10
0
1
2
3
4
5
6
7
8
Time (s)
Figure 7: Full version of fig. including separate exit conditions.
Linear RNNs
To capture rich, long-range temporal relationships without the quadratic cost of attention (Vaswani et al., 2023), we use linear RNN based on minGRU (Feng et al., 2024) and RG-LRU (De et al., 2024). The recurrence R(-,)) takes input 2, R and re E RD to produce an output sequence h4 E R as follows:
h4 gt h4-1 +1 9t) gt) Ft.
9: ())(())
(23)
where learnable diagonal matrix, and (+) is the sigmoid function. This parametrization ensures that 9. is bounded between O and 1, making the recurrence stable. X is initialized such that (A)) uniformly distributed between 0.9 and 0.999 as De et (2024).
Notice that À being a diagonal matrix, makes the computation of Sr elementwise. Moreover, the recurrence itself operates entirely an clementwise fashion. Since the recurrence is linear in terms of fu, and since gr does not depend on ht, the recurrence can be parallelized along the time dimension using a parallel associative scan (Blelloch, 1990; Martin & Cundy, 2018) for efficient training.
We the quasiseparable matrix framework from Hwang et al. (2024), which refer to just Hydra bidirectionality. order to construct a bidirectional variant our recurrence:
Hydran (Fr, .x) = shift (R(r,, æx)) + flip(shift [ip(shiit(R((ip(r) ((ipp())),
(24)
where (lip(-) reverses sequence along the time dimension and shift(-) shifts the sequence one time index to the right, discarding the final time step and prepending the sequence with a vector 0 E RD, Notice that and Ft are shared in both terms of the bidirectionality, such that the Hydra approach does not require additional learned parameters. This bidirectionality has been shown to give better performance than simpler alternatives like additive or concatenative approaches (Hwang al., 2024).
We use linear RNNs in a self-gating block that projects the input into x; and for use in eq. (23), applies the recurrence, multiplies the output with a multiplicative branch gated by a GELU activation, before finally projecting the output with another linear layer. This block closely resembles the block used Mamba (Gu & Dao, 2023), but without convolutions which we found did not affect performance. See fig. 8b.
Inverse Gamma Parametrization We model the parameters a and 3 of the rnverse-gammn dis- tribution in eq. (2) with the InvGam block shown in fig. 8f. The block is composed of a GLU layer followed by GELU activation, which is then linearly projected to two scalars followed by a final softplus activation to the parameters to be non-negative.
18

---

Published as a conference paper at ICLR 2026
Linear
scale
MHSA
Lingar
block
recuirence
xb
r
Linear
Linear
Linear
rms
(a)
(b)
(c)
softplus
softplus
linear
linear
linear
Linear
Linear
rmS nor'm
linear
Linear
Linear
bias
deconv
(d)
(e)
(I)
Figure 8: Detailed architecture of the blocks used. (a) shows the overall block architecture where the block which may be either (b) the linear RNN or (c) speaker attention blocks. (d) and (e) are the audio encoder and decoder heads respectively and (f) InyGamma parametrization block.
19

---

Published as a conference paper at ICLR 2026
Block-wise Exrly-Exits Early-exiting with our block-wise likelihood requires managing RNN state and reconstruction from overlapping transposed convolutions in the decoder heads. We handle RNN states by initializing all states to zero whenever an RNN layer is "restarted" after an earlier exit, and we combine contributions from different exit points across block boundaries by overlap-adding the transposed convolutions from separate layers, corresponding smoothly interpolating between partial results from both contributions.
D DATASET DESCRIPTIONS
WSJO-2mix The WSJO-2mix (Garofolo et al., 2007) dataset is generated according to Python replica' of the original Matlab implementation. 2-speaker mixtures are generated by randomly se- lecting utterances from different speakers in the Wall Street Journal (WSJO) corpus. The training set contains 20,000 mixtures (30 hours), the validation set has 5,000 mixtures (8 hours), and the test set comprises 3,000 mixtures (5 hours). Speech mixtures are created by mixing pairs of utterances from different speakers at random SNRs drawn uniformly between 0 and 5 dB.
Libri2Mix The Libri2Mix (Cosentino et al.. 2020) dataset is generated according to the official script', which creates 2-speaker mixtures by randomly selecting utterances from different speakers in the LibriSpeech corpus. The loudness of each individual utterance is uniformly sampled between 25 and -33 LUFS. This mixing approach results in signal- signal-to-noise ratios (SNRs) that follow a normal distribution with mean of dB and standard deviation of 4.1 dB. Only the 100se, which has //hourr/9300utteannees of data, is used for training and results are reported on the test set which contains 3000 samples.
WHAM! The WHAM! (Wichern al., 2019) dataset extends the WSJO-2mix mixtures with ad- ditive environmental noise recorded in various urban environments. The mixtures are generated fol- lowing the same procedure for WSJO-2mix, but with noise added at SNR sampled uniformly in the -6 and 3 dB range between the loudest speaker and noise.
WHAMR!
The WHAMR! (Maciejewski al., 2020) dataset in turn extends WHAM! by introduc- ing reverberation into the WSJ0-2mix speech mixtures before adding the WHAM! noise by convolv- ing the speech signals with artificially generated room impulse responses designed to approximate domestic and class-room environments.
DNS2020 The Deep Noise Suppression (DNS) Challenge 2020 (Reddy et al., 2020) dataset in- cludes two main parts: 441 hours of clean speech extracted from LibriVor (Noa, 2005) audiobooks and noise library of about 195 hours. The noise library combines 60,000 clips from AudioSet with 10,000 clips from Freesound (Font et al., 2013) and DEMAND (Thiemann et al., 2013) datasets. DNS2020 provides mixing script but sets no guidelines for how much data should be mixed. There- fore, speech-noise mixtures were synthesized on the fly during training at SNRs between 0-20 dB. This synthesis followed the procedure provided in the original script". Results are reported on the non-blind test set without reverberation. The test set contains 150 premixed samples.
Sampling Rate and Downsampling The speech separation datasets operate at 8kHz sampling rate, while DNS2020 uses a 16kHz sampling rate. This creates the only difference between our models. For 8kHz speech separation, we use downsampling factor of 4 in the encoder, while for the 16kHz DNS2020 we use factor of 8, resulting in the same amount of compute per second.
E TRAINING DETAILS: DATASETS, RESOURCES AND SOFTWARE
We use the AdamW Kingma& & Ba (2017); Loshchilov & Hutter (2019) optimizer with 31 = 0.9 and B3 = 0.99, and weight decay of 0.01 which we apply to only linear and convolutional layers. We use base learning rate of 5 10-1, which was found for D 64 model with 4 exit points,
htttts:://gittub.com/ httpr////////tttooooc/////e/eeer
http://aaa.eeeeeeeeeeeeeeeeeeeea /maater
20

---

Published as a conference paper at ICLR 2026
and we transfer this learning rate to wider models (e.g. D 128) by per-layer factor of Dotd. as Dnew described in Yang et al. (2024). We use a linear straight-to-zero learning rate schedule as described in Bergsma et al. (2025) with 5t0p-step linear warmup period. During training we clip the total gradient if its L2 norm exceeds 1. All weight matrices were initialized from normal distribution truncated at 30 with standard deviations set according to Yang et al. (2024).
For all datasets the models train 4-second segments while at evaluation time the model process samples of varying length. The models train for up to million steps with batch size of 1, amount- ing 6,666 hours of training data exposure.
We trained all models on NVIDIA GPUs with Ampere architecture or higher, using any H100, A100, A40, A10, RTX 4090, and RTX 4070 Ti at either university HPC cluster commercially available GPUS with PyTorch (Paszke et al., 2019) 2.7 using torch. compile, The PRESS-4 (S) configurations took around 2-3 days train, while the PRESS-12 (M) configurations took around days to train.
F SCALE-INVARIANT STUDENT T-LIKELIHOOD
Similar to SI-SNR (Roux et al., 2018), we can construct a scale-invariant version of our likelihood by introducing a scaling parameter on the estimated signal & which yields a modified likelihood.
L= St y 20,
(25)
Any conjugate prior can be imposed on 7, but its effect would quickly be overwhelmed by the posterior for non-trivial signal length. The maximum-likelihod estimate for y yields,
(26)
which closely resembles the scale-invariance parameter in the conventional SI-SNR, but normalized by the energy in the prediction rather than the target. This scale-invariance formulation has also been used with SI-SNR to train TF-GridNet (Wang et al., 2023).
Empirical studies show that TasNets trained with log- -og-RMS error perform comparably to those trained with SI-SDR (Heitkaemper et al., 2020). While SI-SDR losses yield better models than RMS error, the key factor may be the logarithmic error scaling rather than scale invariance. Notably, the multivariate Student t-distribution likelihood in eq. (2) also measures errors on log scale.
21


Reviews:
[{"review_id": "cQPWuXFXSP", "reviewer": "Reviewer_eFq7", "summary": "PRESS introduces an early-exit architecture for speech separation that enables dynamic compute scaling by allowing the network to stop processing at different depths. \nThis is achieved by a probabilistic framework that models the predicted SNR distributions, exiting when the predicted SNR , SNRi or a loudness condition exceeds a target threshold with specified confidence. \nThe proposed method offers competitive results with state-of-the-art models on WSJ0-2mix, Libri2Mix, and DNS2020 while enabling compute reduction.", "questions": "I ve put questions in the weaknesses section.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper is the first to apply uncertainty-aware probabilistic modeling specifically for early-exit decisions in speech separation. \nOne great advantage of the proposed method is some interpretability in the chosen early exit strategy. \nThe paper has strong empirical results with respect to the SotA non early exiting models on 3 different datasets including DNS challenge which features speech enhancement.", "weaknesses": "This work misses a crucial comparison with another work which also explored early exit for speech separation already in 2023: \n\nBralios, D., Tzinis, E., Wichern, G., Smaragdis, P., & Le Roux, J. (2023, June). Latent iterative refinement for modular source separation. In ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 1-5). IEEE.\n\nThis work goes into the direction outlined by the authors in the future possible works where instead of using different blocks one uses iterative models. \nAnother key different between the two works is that the early exiting mechanism is fundamentally different as PRESS relies on estimated SNR. \nHowever I feel that this work should be discussed and maybe compared against e.g. by having a baseline system adopting shared parameters and gating-based stopping. \n\nAnother work is mentioned which explored early exit for speech separation: \n\nChen, S., Wu, Y., Chen, Z., Yoshioka, T., Liu, S., Li, J., & Yu, X. (2021, June). Don’t shoot butterfly with rifles: Multi-channel continuous speech separation with early exit transformer. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 6139-6143). IEEE.\n\nI think another drawback of this paper is that the proposed early exit strategy is not compared with the strategy adopted in this latter paper. Does proposed probabilistic SNR work better than Euclidean norm difference ?\nWhich exit criterion transfers better across domains? is PRESS estimation robust ? \nRight now these questions are unsolved and it is unclear how significant is the proposed method compared to the previously proposed two. \n\nAlso the paper lacks other ablations to justify the proposed probabilistic early exiting framework. \nFor example what if instead of the probabilistic framework I use simply a supervised predictive head to estimate the SNR ? \nOr what if only p(SNR > t) is used as an early exit condition ? Figures 1and 6 compare the different strategies but there are no tables. \n\nUnfortunately due to the lack of these experiments the paper as it is right now is not suitable for publication. While the proposed strategy seems very strong it is unclear if it is more effective than others (and possibly simpler) strategies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iilmuAkAhG", "reviewer": "Reviewer_2GXJ", "summary": "The paper proposes PRESS, a technique for performing speech separation and enhancement that can dynamically scale to various compute budgets. They achieve this by designing network architectures equipped with “early exits”. These exits are driven by probabilistic SNR (&SNRis) based criterion, and the network stops at the first depth where the probability of SNR > a set target. Results on WSJ0-2mix, Libri2Mix, and DNS Challenge datasets were presented.", "questions": "Questions:\nI would like the authors to comment on these:\n1. Eqs. 8 and 9 approximate SNR by 1 + $\\lambda/T$ where $\\lambda$ is the non-centrality parameter when T tends to infinity. How does this behave for shorter time windows and how does this hold in real scenarios?\n2. Other than simplifications, what is the authors’ rationale behind using the inverse-gamma as a conjugate prior? Anything specific about the variance that motivates the prior? Does restricting the noise model to the Normal-Inverse-Gamma family limit performance? Because once the classes are restricted, the performance may be bounded? Is this the case? \n3. A stretch, but can we see this depth-quality as analogous to the sampling time-quality in diffusion models? i.e, can diffusion models naturally achieve this dynamism simply by varying the number of sampling steps? Can the authors draw any parallels between the two?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The problem is well motivated. Having the ability to adaptively exit from the computational chain is especially useful on compute and energy-constrained devices such as VR headsets, hearbles, and smart speakers. So studying this seems a useful task for practical deployments.\n\n2. Basing the exit criterion on SNR-type variants seems intuitive and logical, given that they have easy-to-tune parameters.\n\n3. Experimental results at the deepest exit are competitive with shown baselines. Abaltions are helpful.", "weaknesses": "1. Novelty: While the paper tries to solve a practical problem, it seemed like an incremental improvement to previous works (e.g., SepReformer) and did not come across as novel enough.\nThe architecture largely follows SepReformer (early split, transformer stack, CLA and other layers), with the main additions being extra exit heads and a standard likelihood derived from conjugate priors. Also, early exit itself is a known idea, and turning the SNR threshold into an exit confidence feels like a natural application of that likelihood.\n\n2. Empirical results: While the depth-wise exits show incremental improvements (Table 2), they are not surprising. Moreover, why not simply use SepReformer Tiny/Small, which has fewer parameters yet achieves similar (and in some cases better) performance? Early exits don’t really reduce the on-device model size anyway; the full network (with the exit heads) must still be stored, so gains are mainly runtime latency or energy saving on easy inputs, not footprint. But even the #GMAC/s reported don’t clearly beat other baselines. So it’s unclear to me what the concrete gains of the proposed approach are relative to baselines. \n\n3. Unified early exit: The reasoning behind using a combined SNR exit criterion (Eq. 11)  when at least one of them exceeds the set threshold seems heuristic and is not quite convincing. What if you make all three exceed the threshold? Also, I wonder if the thresholds are always triggered mostly by a particular exit. Ablations corresponding to these are missing.\n\n4. Conjugate prior:  The choice of an inverse-Gamma prior on $\\sigma^2$ seems to be made mainly for conjugacy and convenience, and there’s little justification for why that is the case. It would help to compare results against a simpler case, such as a Gaussian with both mean and sigma learned as point estimates, removing the need for marginalization.\n\n5. SNR-centric: The paper relies heavily on the SNR-style variants for both the exit criterion and evaluations. It would help the reader to see whether these SNR-based exit choices are a key design choice. Also, reporting with other perceptual metrics, such as PESQ, STOI would help make better judgment of the method.\n\n6. Writing: The paper is somewhat of a hard read; Section 3.1, in particular, feels under-motivated and confusing. For instance $x$ is not defined in Eq. 1, and similarly for T in Eq. 3 (the signal length). At line 200, in Eq. (3), doesn’t the second-to-last term actually penalize over-estimating the variance?\n\nMinor comments:\nL234: correction at p(SNR(x, \\hat{x}) \\geq t) \\geq p to Pr(.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A5ll8JGyx9", "reviewer": "Reviewer_DGuZ", "summary": "This paper proposes PRESS, a probabilistic early-exit framework for speech separation and enhancement networks that enables dynamic computation scaling. The method introduces a neural architecture with multiple exit points, each predicting a clean speech signal and an error variance parameterized by an inverse-gamma distribution to model uncertainty. Evaluated on standard benchmarks, PRESS achieves performance competitive with state-of-the-art static models while providing interpretable, SNR-based exit conditions that can be calibrated on validation data for efficient inference.", "questions": "The paper shows the performance available at each exit point, but not the performance achieved when the probabilistic exit condition is actively used. What is the average computational cost (e.g., in GMAC/s) and the corresponding average SI-SNRi when processing the entire WSJ0-2mix test set with your proposed exit condition (e.g., with a target of t=20 dB and p=0.9)?\n\n In Equation 7, the unified exit condition is defined as the maximum of the individual complementary CDFs. Could you provide more intuition for why the maximum is the right operator here? Was there an ablation comparing this to other operators (e.g., a weighted average or a product of probabilities)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper demonstrates high originality. While early-exit mechanisms and probabilistic deep learning are known concepts, their application to speech separation is novel. The formulation of a probabilistic model that jointly estimates the clean speech signal and its error variance, parameterized by an inverse-gamma distribution, allows the authors to derive a principled early-exit condition. This is different from common heuristics like entropy or difference between successive blocks. The proposed \"unified exit-SNR\" condition, which combines SNR, SNR improvement, and a loudness condition, is a clever solution to the problem of silent segments, further showcasing the depth of the original contribution.\n- The proposed PRESS is technically sound, evidenced by rigorous experimental design and compelling results. The paper includes thorough ablations that validate core design choices (e.g., the Student-t likelihood vs. SI-SNR, joint permutation invariance). The models are evaluated on three standard benchmarks (WSJ0-2mix, Libri2Mix, DNS Challenge), demonstrating the generality of the approach for both separation and enhancement tasks. The results show that PRESS achieves state-of-the-art performance at its deepest exit points.\n- The paper is very well-written and clearly structured. The motivation -- enabling dynamic computation for heterogeneous devices, is immediately established. The probabilistic framework is explained step-by-step, with a clear derivation of the likelihood and the subsequent SNR distributions, making the theoretical contribution accessible. The results are presented effectively with tables and informative plots, such as the performance-vs-compute graph (Figure 3) and the calibration curves (Figure 4), which directly illustrate the core advantages of the method. The appendix provides valuable supplementary material, including convergence analysis and implementation details.", "weaknesses": "- The ablation study in Table 1 shows that the proposed Student-t likelihood performs on par with the standard SI-SNR loss. While the probabilistic framework is elegant and enables the exit condition, the paper does not demonstrate that this formulation provides a training advantage. The core contribution of the probabilistic model is its utility for uncertainty-aware exiting, but if a model trained with standard SI-SNR and a heuristic exit condition (e.g., based on the norm of the change between layers) performed similarly, the necessity of the more complex probabilistic training objective would be questioned.\n- The paper's most significant limitation is the modelling of a single, global variance parameter $\\sigma^2$ per exit for the entire audio segment. It must process the entire segment (e.g., 4 seconds at training, variable at test) to compute a global variance before any exit decision can be made. This undermines the goal of \"fine-grained dynamic compute-scaling\" for live applications, as the compute savings are only realized at the segment level, not within a segment.\n  - To truly meet its stated goal, the authors should explore a time-varying variance model, predicting a frame-level or chunk-level variance $\\sigma_t^2$ at each exit point. The early-exit condition could then be evaluated on a rolling window (e.g., the past 500ms). This is a non-trivial but crucial extension that would make the system causal and demonstrate true fine-grained, adaptive computation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eAzMNoYdDB", "reviewer": "Reviewer_9ZZr", "summary": "The paper propose a neural network for single channel speech separation and enhancement.\nThe proposed PRESS algorithm is capable of early-exit with the ability to output uncertainty-aware probabilistic early-exit parameter. \nThey test it on both speech separation and enhancement tasks and shows very good results without compromising reconstruction (state of the art results).", "questions": "1. Can you please provide more results on speech separation tasks, i.e. on other dataset (not only WSJ-mix)?\n\n2. Can you evaluate the proposed method on more than two speakers (3 or more)? \n\n3. Could you consider more perceptual metric for evaluation such as STOI and PESQ?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The idea of early exit is novel is interesting, and the SNR based criterion is also important to the field\n\n2. When adding more exits it save the results of the latest exit performance, and the last exit get closer results to the state of the art results.\n\n3. The proposed method can be use in both separation task and enhancement task which is very useful. \n\n4. The proposed method introduce the Student t-likelihood to use as likelihood measure.", "weaknesses": "1. The results of the proposed architecture on early exits doesn't achieves state of the art results \n\n2. The calculation of the inverse-gamma variance at the head is global for the whole reconstruction, in real audio case one scalar is not enough to express the uncertainty\n\n3. There may be a problem if each exit peak its own permutation since it can choose different speakers in the following exits which can lead to degradation in the training", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WD62fWT3A1", "reviewer": "Reviewer_fRD5", "summary": "This paper deals with the problem of single-channel speech separation, a crucial task in telecommunication, hearables, and other speech processing applications. Recent deep neural network models have achieved strong performance in this area based on commonly used metrics and loss functions such as SI-SNR and uPIT. However, neural networks with fixed computation paths may be inefficient in real-world scenarios where input signals often contain silence or non-overlapping speech. To address this, the authors propose a novel metric and loss function based on Bayesian and GMM formulations, combined with an early-exit mechanism. They extend the SepReformer architecture integrating these components to support resource-adaptive computation. Experiments on widely used benchmarks like LibriMix and WSJ0Mix demonstrate that the model achieves competitive performance comparable to SepReformer, and a favorable accuracy-efficiency trade-off with adaptive computation.", "questions": "For better readability among non-experts in Bayesian modeling, it would be beneficial to include a brief introduction to Bayesian Gaussian models, particularly the Normal–Inverse-Gamma (NIG) family, which appears to underpin the proposed probabilistic speech modeling framework. A concise explanation of how this family models both mean and variance uncertainty would make the paper more accessible and improve conceptual clarity.\n\nThe post-calibration procedure also seems to play a crucial role in the method’s generalizability. However, it is described as a form of training on validation data, effectively learning the variance distribution from the validation set. This raises concerns about potential data leakage or overfitting. If the model’s performance relies heavily on this post-calibration step, it would be important to clarify whether the approach generalizes well without it.\n\nThe paper introduces several architectural modifications on top of SepReformer, such as replacing the Transformer with a Linear RNN and using the Snake activation function. It would strengthen the paper to include ablation studies quantifying the individual contributions of these design choices to the overall performance.\n\nThe 2-speaker separation benchmarks on datasets such as WSJ0Mix and LibriMix are largely saturated, with SI-SNR values above 20 dB already reaching perceptual saturation. Therefore, it would be valuable to discuss whether the reported improvements, though numerically measurable, are practically significant in real-world speech separation scenarios.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The proposed model demonstrates strong speech separation performance comparable to the state-of-the-art SepReformer, while offering support for adaptive computation and maintaining relatively low computational and memory costs.\n\n\nThe work provides a mathematically grounded formulation of speech separation objectives and metrics using GMM theory, offering potential alternatives to existing dominant approaches.\n\n\nThe evaluation includes multiple speech separation models and benchmarks, adding rigor and relevance to the empirical results.", "weaknesses": "The motivation and evaluation are not well aligned. The paper emphasizes potential computational savings in speech separation scenarios involving non-overlapping speech or silence; however, the evaluation only measures overall end-to-end computational efficiency. It does not quantitatively assess the specific advantages in those targeted scenarios.\n\nThe learning objective restricts the generality of the proposed framework. While metrics such as uPIT or SI-SNR can be readily applied to most speech separation architectures that output raw signals, the proposed early-exit SNR objective additionally requires the network to produce a variance output. This constraint limits applicability, as it necessitates architectural components explicitly designed for variance computation.\n\nThe reported performance improvements over existing state-of-the-art methods are marginal, and the situations where the proposed approach would provide a clear advantage are not convincingly demonstrated. In particular, the paper lacks experiments on realistic deployment settings, such as mobile or edge devices, where adaptive computation would be most beneficial. Moreover, under comparable computational budgets, the proposed model achieves slightly lower accuracy than SepReformer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)