Music Source Separation with Deep Equilibrium Models

A weight-tied network is realized by employing the same transformation $f_{\theta}$, parameterized by $\theta$, in each layer of a typical skip-connection-based architecture as

$$\mathbf{z}^{[i+1]}_{1:T}=f_{\theta}\left(\mathbf{z}^{[i]}_{1:T};\mathbf{x}_{1:T}\right),$$

(1)

where $\mathbf{x}_{1:T}\in\mathbb{R}^{T\times p}$ is the input sequence and $\mathbf{z}^{[i]}_{1:T}\in\mathbb{R}^{T\times q}$ is the hidden sequence in the $i$-th layer [17]. Such transformation can also be interpreted as a sequence of identical transformations, where each iteration is identified by the index $i$. This sequence is typically terminated after a fixed number of iterations $L$ (i.e., $i=0,1,\dots,L-1$). The number of iterations is equivalent to the number of function evaluations (NFE), thus $L$ is also called NFE. TrellisNet [17], which is implemented by using TCN as $f_{\theta}$, outperformed other typical DNN-based approaches on some natural language processing (NLP) tasks without a noticeable increase of the number of parameters.

In [13], it was empirically shown that (1) tends to converge to an equilibrium point $\mathbf{z}^{\star}_{1:T}$ as the layer index $i$ increases. The equilibrium point $\mathbf{z}^{\star}_{1:T}$ satisfies the following condition:

$$\mathbf{z}^{\star}_{1:T}=f_{\theta}\left(\mathbf{z}^{\star}_{1:T};\mathbf{x}_{1:T}\right),$$

(2)

and the output of DEQ is the equilibrium point itself. DEQ showed further improvement over weight-tied networks on some NLP tasks [13].

The forward pass of DEQ is performed by finding the equilibrium point $\mathbf{z}^{\star}_{1:T}$. Some root-finding solvers for nonlinear equations such as Newton’s method or quasi-Newton methods can be utilized to find the equilibrium point as follows:

$$\mathbf{z}^{[i+1]}_{1:T}=\mathbf{z}^{[i]}_{1:T}-\alpha Bg_{\theta}\left(\mathbf{z}^{[i]}_{1:T};\mathbf{x}_{1:T}\right)\quad(i=0,1,\dots,L_{\text{stop}}-1),$$

(3)

where $g_{\theta}(\mathbf{z}_{1:T}^{[i]};\mathbf{x}_{1:T})=f_{\theta}(\mathbf{z}^{[i]}_{1:T};\mathbf{x}_{1:T})-\mathbf{z}^{[i]}_{1:T}$, $B$ is the inverse Jacobian of $g_{\theta}$ (or its low-rank approximation) evaluated at $\mathbf{z}_{1:T}^{[i]}$, which is obtained by a root-finding solver, $\alpha$ is the step size, and $L_{\text{stop}}$ is the NFE when the iteration is stopped. The iteration stops if the norm of $g_{\theta}$ becomes smaller than a tolerance $\epsilon$ or the maximum number of function evaluations $L_{\text{max}}$ is reached. Broyden’s method [22], which is a quasi-Newton method, is often used for DEQ as it computes the inverse Jacobian efficiently.

In DEQ, the backpropagation is performed on the basis of implicit differentiation [13], which doesn’t need to store the intermediate values of the solver. Let $l$ be the loss function for optimizing the parameter $\theta$. The derivative of $l$ with respect to $\theta$ can be calculated based on implicit differentiation as follows:

$$\frac{\partial l}{\partial\theta}=-\frac{\partial l}{\partial\mathbf{z}^{\star}_{1:T}}\left(J^{-1}_{g_{\theta}}|_{\mathbf{z}^{\star}_{1:T}}\right)\frac{\partial f_{\theta}\left(\mathbf{z}^{\star}_{1:T};\mathbf{x}_{1:T}\right)}{\partial\theta},$$

(4)

where $J^{-1}_{g_{\theta}}$ is the inverse Jacobian of $g_{\theta}$ evaluated at the equilibrium point $\mathbf{z}_{1:T}^{\star}$, which appears in the implicit differentiation theorem. The operation $-\frac{\partial l}{\partial\mathbf{z}^{\star}_{1:T}}\left(J^{-1}_{g_{\theta}}|_{\mathbf{z}^{\star}_{1:T}}\right)$ can be performed by solving the equation in terms of $\mathbf{y}$:

$$\left(J^{\top}_{g_{\theta}}|_{\mathbf{z}^{\star}_{1:T}}\right)\mathbf{y}^{\top}+\left(\frac{\partial l}{\partial\mathbf{z}^{\star}_{1:T}}\right)^{\top}=\mathbf{0}$$

(5)

with solvers such as Broyden’s method.

Jacobian-Free backpropagation (JFB) [15] was proposed as a more efficient technique for gradient calculation in DEQ. This method approximates the inverse Jacobian in (4) with an identity matrix, avoiding the Jacobian-based computation from (5) and reducing computations. In [15] it was shown that the gradient obtained by JFB is valid (i.e., a descent direction for the loss function) under certain conditions. Furthermore, since this method does not impose a condition on how close the solver’s output is to the equilibrium point, it is expected to be robust to estimation errors by the solver.

Fig. 1 is an example of a general network architecture for DNN-based audio source separation that receives a multi-channel (typically stereo in MSS) spectrogram obtained by a short-time Fourier transform (STFT) and outputs the spectrogram for the target audio source. UMX is a typical example of such architecture, which shows performance close to state-of-the-art [4].

In the pre-processing block of UMX, the input spectrogram is first cropped to a maximum frequency (typically 16 kHz). Then the cropped spectrogram is fed into frame-wise linear transformations whose parameters are set to normalize the spectrogram to have zero mean and unitary standard deviation. The encoder block is composed of a fully-connected (FC) layer with batch normalization (BN). The sequence modeling block is a key part of source separation as it deals with contextual information necessary to separate the mixed sources. Fig. 2 shows in detail the sequence model of UMX. The output of the encoder block is first normalized with a hyperbolic tangent function such that its range is between $-1$ and $1$, then fed into 3 consecutive BLSTM layers. The normalized output of the encoder block and the output of the BLSTM block are concatenated and sequentially fed into another FC layer with BN and rectified linear unit (ReLU) activation. The decoder block is composed of a FC layer with BN. The post-processing block is a linear transformation whose parameters are initialized to 1. The output of the post-processing block is used as a multiplicative mask on the input of the pre-processing module. Such architecture is trained to separate a specific instrument. Finally, the separated signals for each target instrument are refined by applying a multi-channel Wiener filter (MWF) [23, 2] to the output of each network.

We replace the sequence model of UMX with a DEQ-based sequence model while keeping the other blocks as they are. Fig. 3 describes the overview of the sequence model of the proposed method. In the sequence model, the input sequence normalized by the hyperbolic tangent function is fed into the DEQ block followed by a ReLU as shown in Fig. 3(a). Fig. 3(b) shows our proposed architecture for the function $f_{\theta}$ of the DEQ block, which is inspired by the original UMX. Specifically, we adopt a concatenation block to merge the output of the BLSTM and the input sequence. We define the size of hidden sequence $q$ as equivalent to that of the input sequence $p$ since the original UMX has the same size of output with the input. Then the $i$-th concatenated sequence $\mathbf{u}^{[i]}_{1:T}\in\mathbb{R}^{T\times 2p}$ is fed into a FC layer such that its dimension is the same as the original input of the DEQ block,

$$\mathbf{v}^{[i]}_{1:T}=\text{FC}(\mathbf{u}^{[i]}_{1:T}),$$

(6)

where $\mathbf{v}^{[i]}_{1:T}\in\mathbb{R}^{T\times p}$ is the $i$-th output sequence of the FC block. Then the output sequence $\mathbf{v}^{[i]}_{1:T}$ is normalized with a group normalization (GN) block whose number of groups is one and fed into a hyperbolic tangent function followed by a BLSTM block with one layer as follows,

$$f_{\theta}\left(\mathbf{z}^{[i]}_{1:T};\mathbf{x}_{1:T}\right)=\text{BLSTM}(\tanh(\text{GN}(\mathbf{v}^{[i]}_{1:T}))).$$

(7)

The choice of GN is a deviation from the original UMX, which uses BN instead. This modification is inspired by [14] as BN often causes a poor conditioning of the Jacobian matrix. We use Broyden’s method to compute Eq. (3) in the forward pass and obtain the hidden sequence $\mathbf{z}^{[i+1]}_{1:T}$. Note that the tolerance $\epsilon$ and the maximum number of function evaluations $L_{\text{max}}$ during inference need to be the same with those of training.

Aside from the architecture shown in Fig. 3(b), we preliminary evaluated two different architectures, where we substitute concatenation with weight-averaging or with masking. We found that the architecture with the concatenation block achieved the best performance, which agrees with the tendency highlighted by [24]. Therefore, we adopt the architecture shown in Fig. 3(b).

To provide further insights into our architecture, an analogy with guided source separation (GSS) [25, 26, 27] may be helpful. In the first iteration of the solver (i.e., $i=0$), $\mathbf{z}^{[0]}_{1:T}$ (which is an all-zero sequence) is fed into the function $f_{\theta}$ and concatenated with the input $\mathbf{x}_{1:T}$. This process can be interpreted as if the function $f_{\theta}$ attempts to extract the desired sequence without any guide information. In the first few iteration (e.g., $i=0,1,2$), although the calculated $\mathbf{z}^{[i+1]}_{1:T}$ is still insufficient for the separation performance, it is used as the updated guide information for the next iteration. The guide information gradually becomes more accurate with each iteration and finally a sufficiently informative $\mathbf{z}^{[L_{\text{stop}}]}_{1:T}$ can be obtained.

Using BLSTM in the function $f_{\theta}$ is a key point. If a convolutional neural network (CNN)-based architecture such as TrellisNet [17] was assigned to $f_{\theta}$, it would require many NFE to obtain a sufficiently large receptive field. In fact, TrellisNet requires hundreds of iterations to converge, and dozens of iterations are performed in practice [17, 13]. In contrast, we can expect that using BLSTM in the function $f_{\theta}$ enables us to utilize all information of the input sequence $\mathbf{x}_{1:T}$ even if the NFE is small.

Although a small NFE contributes to reducing the computational cost, it could be a problem when using implicit differentiation for the backward pass because it is not guaranteed that the obtained hidden sequence satisfies the equilibrium condition sufficiently. Therefore, since JFB is expected to be more robust to the estimation error of the solver, we use JFB for an efficient gradient calculation instead of using implicit differentiation-based backpropagation.

We use the same training setup (i.e., loss function, learning rate scheduler, etc…) as in the original UMX. Following [13], the parameters of the DEQ-UMX are pre-trained as weight-tied network using the typical backpropagation with a pre-defined number of epochs for the stability of the training. NFE of the weight-tied network and the number of epochs for pre-training are determined on the basis of the validation loss.

We use the MUSDB18 dataset, which consists of 150 professionally recorded songs [28]. For each song, the clean waveform of vocals, drums, bass, and other together with their mixture are available. Each audio track is stereo with a sampling frequency of 44.1 kHz. Following the official split, we used 86 songs for the training, 14 songs for the validation, and 50 songs for the test set. All networks in this experiment were trained on 6 seconds long segments using the Adam [29] optimizer and a weight decay of $10^{-5}$. The learning rate was initially set to 0.001 and decreased by 70% if the average of the loss function on the validation set did not improve in 80 consecutive epochs. The training was interrupted when the average of the loss function on the validation set did not improve in 300 consecutive epochs. For inference, the full-length signals were processed with the trained networks. MWF was also applied to the full-length signals. Signal-to-distortion ratio (SDR) [30] was used for the evaluation. We compute the SDR by taking the median over all frames of a song and then calculating the median over all songs, which is the standard way of SDR computation on MUSDB18.

We first investigated the feasibility of DEQ-based source separation using the weight-tied network with the architecture described in Fig. 3, which we refer to as WT-UMX hereafter. Fig. 4 shows the relationship between the NFE of WT-UMX and the average SDR. The average SDR reaches the best SDR score when the NFE, $L$, is set to four, while it gets worse as the NFE increases. This tendency is similar to previous work in other tasks [17, 13, 16]. Although weight-tied networks and DEQ usually require hundreds of NFE to converge, the models that are evaluated are trained with dozens of NFE [13]. The relationship between NFE and performance, which clearly shows a maximum performance around a specific NFE value, is compatible with another approach [16], which uses a closed-form representation. This shows that the best NFE depends on the task and the network architecture. In our case, the model tends to reach the best performance with a relatively small NFE. This implies that using BLSTM as $f_{\theta}$ helps the network to utilize all the information of the input sequence early in the iterations, as hypothesized in sec. 3.1. We can conclude that it is possible applying weight-tying to UMX and that the required NFE is relatively small comparing to the experiment conducted in other works.

The evaluation results are shown in Table 1. The number of parameters, the number of multiply-and-accumulate operations (MACs), the SDR values for each instrument, and the average SDR values over all the instruments are compared among several variants of UMX, WT-UMX (from the preliminary study), and our proposed DEQ-UMX. The MACs do not include the computations for the STFT and MWF. “UMX large” is a variant of UMX with more layers in the BLSTM block than the original UMX. “UMX small” is a variant with smaller hidden size (410) than that of UMX (512).

First, we can observe that even if the number of layers of UMX is simply increased, there is almost no improvement in terms of the average SDR. Also, if we reduce the hidden size of UMX, then the number of parameters will decrease while the average SDR is penalized by 0.09 dB. WT-UMX with $L=4$ reduces the number of parameters by 30% and improves the average SDR, with a particular improvement of the SDR for drums. The proposed DEQ-UMX achieved the best average SDR. We can also see that JFB contributes to improving the SDR values for all instruments with respect to the model trained with implicit differentiation. In particular, the proposed DEQ-UMX improved the average SDR by 0.29 dB with respect to “UMX small”, which has the comparable number of parameters. We also compared the SDR distributions of both methods and observed that the distributions have a similar shape, i.e., there are no extreme outliers for DEQ-UMX but its distribution is shifted to a better SDR score. These results suggest that the equilibrium point obtained by our proposed method allows for a better separation performance than the ones obtained by the weight-tied network with typical backpropagation and the DEQ-UMX trained by backpropagation with implicit differentiation. In our proposed method, $L$ used in pre-training was 4 for all instruments, and the number of epochs for pre-training was 400 for vocals, 0 for drums, and 300 for bass and other. After pre-training, $L_{\text{max}}$ was set to 6 for all instruments. During inference, the NFE always reached $L_{\text{max}}$ in the test set. Although the MACs of DEQ-UMX are higher than UMX, the number of parameters of DEQ-UMX is lower than for UMX by around 30%, which enables the implementation of DEQ-UMX in devices with limited memory.