L1 Signal-to-Noise Ratio (SNR) loss functions for audio source separation in PyTorch. This package provides four loss functions that combine implementations from recent academic research with novel extensions, designed to integrate easily into any audio separation or enhancement training pipeline.
The core L1SNRLoss is based on the loss function described in [1]. L1SNRDBLoss adds adaptive level-matching regularization proposed in [2]. STFTL1SNRDBLoss provides a spectrogram-domain L1SNR-style loss (real/imag STFT components as in [1] / [3]). MultiL1SNRDBLoss combines time-domain and spectrogram-domain losses into a single loss function for convenience and flexibility. Optional novel algorithmic extensions have also been included (such as multi-resolution STFT averaging, spectrogram-domain adaptation of the level-matching regularizer from [2], and blending of standard L1 loss) with the goal of increasing flexibility for improved performance depending on the specific task.
import torch
from torch_l1_snr import MultiL1SNRDBLoss
# Create combined time + spectrogram domain loss function with adaptive regularization
loss_fn = MultiL1SNRDBLoss(name="multi_l1_snr_db_loss")
# Calculate loss between model output and target
estimates = torch.randn(4, 32000, requires_grad=True) # (batch, samples)
targets = torch.randn(4, 32000)
loss = loss_fn(estimates, targets)
loss.backward()- Time-Domain L1SNR Loss: A basic, time-domain L1-SNR loss, based on [1].
- Regularized Time-Domain L1SNRDBLoss: An extension of the L1SNR loss with adaptive level-matching regularization from [2], plus an optional L1 loss component.
- Multi-Resolution STFT L1SNRDBLoss: A spectrogram-domain L1SNR-style loss (real/imag STFT components as in [1] / [3]), computed over multiple STFT resolutions, with optional spectrogram-domain level-matching regularization inspired by its time-domain counterpart in [2].
- Combined Multi-Domain Loss:
MultiL1SNRDBLosscombines time-domain and spectrogram-domain losses into a single, weighted objective function.
- L1 Loss Blending: The
l1_weightparameter allows mixing between L1SNR and standard L1 loss, softening the "all-or-nothing" behavior of pure SNR losses for more nuanced separation. - Multi-Resolution STFT Averaging - Extending an STFT-based loss to multiple resolutions is common in recent literature.
- Spectrogram-Domain Adaptation of Level-Matching Regularizer [2] - Options to extend adaptive level-matching regularization to spectrogram-domain. Experimental and not used by default.
- Time vs. Spectrogram Loss Balancing - Allows fine-tuning the relative contribution of time-domain and spectrogram-domain losses in
MultiL1SNRDBLossvia thespec_weightparameter. Not a novel extension: the authors' ownbanditexposes equivalenttime_weight/freq_weightcontrols, and the single-knobspec_weightis a convenience over the same idea. - Numerical Stability: Robust handling of
NaNandinfvalues during training inSTFTL1SNRDBLoss(and inMultiL1SNRDBLossthrough it), controlled bycheck_finite(defaultTrue; set itFalseto skip the scan and let non-finite values propagate visibly, which also removes four host-device synchronizations per call on CUDA). The time-domain lossesL1SNRLossandL1SNRDBLossdo not sanitize non-finite input: aNaNestimate propagates to aNaNloss, which is visible rather than silent. - Short Audio Fallback: Graceful fallback to time-domain loss when audio is too short for STFT processing.
- Gradient Scaling: These losses produce gradients a few hundred times larger than a plain L1 loss. That is fine on its own, but a run using the common default gradient-clip threshold of
1.0would clip away most of a useful gradient.grad_scale(default1.0, no effect) scales the gradient without changing the loss value, so you can size it for clipping without touching what you log. Adam users are unaffected by the scale itself (a constant factor cancels in the update); on plain SGD a non-1.0grad_scaleacts as a learning-rate rescale. If you clip, either raise the threshold to suit this loss or setgrad_scalebelow1.0. See gradient magnitude in the design notes.
pip install torch-l1-snrpip install git+https://github.com/crlandsc/torch-l1-snr.gitOr, you can clone the repository and install it in editable mode for development:
git clone https://github.com/crlandsc/torch-l1-snr.git
cd torch-l1-snr
pip install -e .All loss functions in this package (L1SNRLoss, L1SNRDBLoss, STFTL1SNRDBLoss, and MultiL1SNRDBLoss) accept standard audio tensors of shape (batch, samples), (batch, channels, samples), or (batch, num_sources, channels, samples). For the time-domain losses, any 3D/4D input is flattened across all non-batch dimensions (e.g., sources, channels, and samples) into a single vector per example before the loss is computed. For the spectrogram-domain loss, inputs are reshaped to (batch, streams, samples) by flattening all non-time dimensions into a “stream” dimension (e.g., streams = channels or streams = num_sources * channels), and a separate STFT is computed for each stream.
The loss functions can be imported directly from the torch_l1_snr package.
The simplest loss function - pure L1SNR without regularization.
import torch
from torch_l1_snr import L1SNRLoss
# Create dummy audio signals
estimates = torch.randn(4, 2, 44100, requires_grad=True) # Batch of 4, stereo, 44100 samples
actuals = torch.randn(4, 2, 44100)
# Basic L1SNR loss
loss_fn = L1SNRLoss(name="l1_snr_loss")
# Calculate loss
loss = loss_fn(estimates, actuals)
loss.backward()
print(f"L1SNRLoss: {loss.item()}")Adds adaptive level-matching regularization to prevent silence collapse.
import torch
from torch_l1_snr import L1SNRDBLoss
# Create dummy audio signals
estimates = torch.randn(4, 2, 44100, requires_grad=True) # Batch of 4, stereo, 44100 samples
actuals = torch.randn(4, 2, 44100)
# Initialize the loss function with regularization enabled
# l1_weight=0.1 leans heavily toward L1SNR; see the calibration note below
loss_fn = L1SNRDBLoss(
name="l1_snr_db_loss",
use_regularization=True, # Enable adaptive level-matching regularization
l1_weight=0.1 # interpolation coefficient, not a behaviour fraction
)
# Calculate loss
loss = loss_fn(estimates, actuals)
loss.backward()
print(f"L1SNRDBLoss: {loss.item()}")Computes L1SNR loss across multiple STFT resolutions.
import torch
from torch_l1_snr import STFTL1SNRDBLoss
# Create dummy audio signals
estimates = torch.randn(4, 2, 44100, requires_grad=True) # Batch of 4, stereo, 44100 samples
actuals = torch.randn(4, 2, 44100)
# Initialize the loss function without regularization or traditional L1
# Uses multiple STFT resolutions by default: [512, 1024, 2048] FFT sizes
loss_fn = STFTL1SNRDBLoss(
name="stft_l1_snr_db_loss",
l1_weight=0.0 # Pure L1SNR (no regularization, no L1)
)
# Calculate loss
loss = loss_fn(estimates, actuals)
loss.backward()
print(f"STFTL1SNRDBLoss: {loss.item()}")Combines time-domain and spectrogram-domain losses into a single weighted objective.
import torch
from torch_l1_snr import MultiL1SNRDBLoss
# Create dummy audio signals
estimates = torch.randn(4, 2, 44100, requires_grad=True) # Batch of 4, stereo, 44100 samples
actuals = torch.randn(4, 2, 44100)
# Initialize the multi-domain loss function
loss_fn = MultiL1SNRDBLoss(
name="multi_l1_snr_db_loss",
weight=1.0, # Overall weight for this loss
spec_weight=0.6, # coefficients: 0.4 * time_loss + 0.6 * spec_loss
# (see Limitations: 0.5 is the paper-faithful default)
l1_weight=0.1, # applies to both domains; see the calibration note
use_time_regularization=True, # Enable regularization in time domain
use_spec_regularization=False # Disable regularization in spec domain
)
# Calculate loss
loss = loss_fn(estimates, actuals)
loss.backward()
print(f"Multi-domain Loss: {loss.item()}")Also exported is dbrms, the level measurement the regularizers are built on. It returns the RMS level in decibels for each element of a batch, flattening all non-batch dimensions:
import torch
from torch_l1_snr import dbrms
audio = torch.randn(4, 2, 44100) * 0.1 # (batch, channels, samples)
levels = dbrms(audio) # (4,) tensor of dBRMS values
print(levels)dbrms(x, eps=1e-8) computes 20 * log10(sqrt(mean(x**2) + eps)). The eps sits inside the square root, on a power quantity, and puts the floor for a digitally silent input at exactly -80 dB -- deliberately well below the lmin=-60 threshold the adaptive regularizer uses, so a silent target is correctly recognized as silent. (Before v0.2.0 a second epsilon was also added outside the root, on an amplitude; it could never prevent a log of zero and shifted the silence floor to -79.99913 dB.)
The goal of these loss functions is to provide a perceptually-informed and robust alternative to common audio losses like L1, L2 (MSE), and SI-SDR for training audio source separation models.
- Robustness: The L1 norm is less sensitive to large outliers than the L2 norm, making it more suitable for audio signals which can have sharp transients.
- Perceptual Relevance: The loss is scaled to decibels (dB), which more closely aligns with human perception of loudness.
- Adaptive Regularization: Prevents the model from collapsing to silent outputs by penalizing mismatches in the overall loudness (dBRMS) between the estimate and the target.
This package is motivated by, and largely follows, the objectives and regularizers described in the cited papers ([1–3]). Several novel algorithmic extensions have been included with the goal of increasing flexibility for improved performance depending on the specific task.
A key feature of L1SNRDBLoss is the adaptive regularization term, as described in [2]. This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (lambda) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
The STFTL1SNRDBLoss module applies the L1SNRDB loss across multiple time-frequency (spectrogram) resolutions. While not mentioned in the cited papers, by analyzing the signal with multiple different STFT window sizes and hop lengths, the loss function can capture a wider range of artifacts - from short, transient errors to longer, tonal discrepancies. This provides a more comprehensive error signal to the model during training. Using multiple resolutions for an STFT loss is common among many recent source separation works, such as the Band-Split RoPE Transformer.
A characteristic of these SNR-style losses that I experienced in many training experiments is that they encourage the model to make definitive, "all-or-nothing" separation decisions. This can be highly effective for well-defined sources (e.g. drums vs vocals), as it pushes the model to be confident in its estimations. However, this can also lead to "confident errors," where the model completely removes a signal component it should have kept. This poses a tradeoff for sources that may share greater similarities (e.g. speech vs singing vocals).
While the Level-Matching Regularization prevents a total collapse to silence, it does not by itself solve this issue of overly confident, hard-boundary separation. To provide a tunable solution, this implementation introduces a novel l1_weight hyperparameter. This allows you to create a hybrid loss, blending the decisive L1SNR objective with a standard L1 loss to soften its "all-or-nothing"-style behavior and allow for more nuanced separation.
Anecdotal, not measured. While this can potentially reduce the "cleanliness" of separations and slightly harm metrics like SDR, I found that re-introducing some standard L1 loss allows for slightly more "smearing" of sound between sources to mask large errors and be more perceptually acceptable for sources with many similarities. I have no hard numbers to report on this yet, just my experience.
So I recommend starting with no standard L1 mixed in (l1_weight=0.0), and then slowly increasing from there based on your needs.
l1_weight=0.0(Default): Pure L1SNR (+ regularization).l1_weight=1.0: Pure standard L1 loss.0.0 < l1_weight < 1.0: A weighted combination of the two.
The implementation is efficient at the endpoints: if l1_weight is 0.0 or 1.0, the unused component is not computed.
l1_weight is an interpolation coefficient, not a behaviour fraction: l1_weight=0.1 does not mean "10% L1 behaviour". How far the knob moves each update toward L1 depends on your target level relative to ref_level (default 0.05, the measured median for MUSDB-style stems), and it differs between the time and spectrogram domains. See docs/design_notes.md for the measured tables, the ref_level / spec_ref_level calibration, and the per-domain difference.
All loss functions work on CPU, CUDA, and MPS (Apple Silicon).
MPS note: PyTorch's MPS backend produces numerically incorrect gradients from torch.stft backward above an input length of 65,536 samples (2^16). The forward transform is correct to float32 precision, so the failure is silent. The error is not a simple function of size: a handful of specific lengths are exact while neighbouring ones are wrong by anywhere from 30% to 99%, so it cannot be predicted or avoided by choosing a particular window length. Batch sizes above 1 fail even at the lengths that are exact at batch 1. This affects STFTL1SNRDBLoss and MultiL1SNRDBLoss (which use STFT internally). As of v0.1.3, these losses automatically route STFT computation through CPU when on MPS (mps_cpu_fallback=True by default), producing correct gradients with negligible performance impact. Time-domain losses (L1SNRLoss, L1SNRDBLoss) are unaffected. CUDA and CPU users are completely unaffected by this change.
Typical audio training uses windows well above that threshold (6 seconds at 44.1 kHz is 264,600 samples), so leave the fallback enabled on Apple silicon unless you have verified on your own PyTorch version that the backward pass is correct.
To disable the workaround (e.g., if a future PyTorch release fixes the MPS bug):
loss_fn = STFTL1SNRDBLoss(name="stft_loss", mps_cpu_fallback=False)- The L1SNR loss is not scale-invariant. Unlike SI-SNR, it requires the model's output to be correctly scaled relative to the target.
- While the dB scaling and regularization are psychoacoustically motivated, the loss does not model more complex perceptual phenomena like auditory masking.
- The usable dynamic range collapses for quiet targets. Because
epssits in both the numerator and denominator, D1's floor at perfect reconstruction is10*log10(eps / (mean|y| + eps))rather than negative infinity. With the defaulteps=1e-3that floor is roughly -30 dB atmean|y|=1, -20 dB at 0.1, -10 dB at 0.01, and only -3 dB at 1e-3. A target near -58 dBFS RMS therefore has under 3 dB of total loss range to optimize within. This is inherited from the reference implementation rather than introduced here, but it means very quiet stems carry correspondingly little gradient signal, and it is worth checking your target levels before attributing poor performance on quiet sources to the model. The papers note the same constraint, that the loss is numerically stable forepsnot much smaller than the signal norm. spec_weightis a loss-value weight, not a gradient share. InMultiL1SNRDBLossthe combination is(1 - spec_weight) * time_loss + spec_weight * spec_loss, so the coefficients are exactly as documented. Butspec_lossinternally sums a real and an imaginary D1 term whiletime_lossis a single term, so equal coefficients do not mean the two domains contribute equally to the gradient. The default0.5is chosen for a specific reason: it is the value at which the time, real and imaginary terms all receive weight 0.5, reproducing the equal 1:1:1 weighting of the objective in [3]. Prefer to leave it there unless you have a reason to shift domain emphasis.- Finer numerical behaviours are documented separately. The spectrogram loss is not strictly monotone in reconstruction quality, the squared reductions overflow in float32 far below the dtype's range, and the level-matching regularizer exerts no gradient at exactly digital silence. See docs/design_notes.md for measured detail on each.
Contributions are welcome! Please open an issue or submit a pull request if you have any bug fixes, improvements, or new features to suggest.
This project is licensed under the MIT License - see the LICENSE file for details.
The loss functions implemented here are largely based on the work of the authors of the referenced papers. Thank you for your research!
The core D1 objective follows the authors' own reference implementations, not only the papers. In particular the mean-normalized form used here (rather than the summed L1 norm the papers write) matches their code, and this implementation is numerically equivalent to it. Those repositories are:
kwatcharasupat/bandit- Apache-2.0. Reference implementation for [1].kwatcharasupat/query-bandit- MIT, Copyright (c) 2024 Karn Watcharasupat. Reference implementation for [3].
No official implementation of [2] has been released, so the level-matching regularizer here follows the published equations alone.
[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, vol. 5, pp. 73-81, 2024. doi: 10.1109/OJSP.2023.3339428. arXiv:2309.02539
[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," arXiv:2501.16171. Preprint; not peer-reviewed at time of writing.
[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. arXiv:2406.18747