SICRN: Advancing Speech Enhancement through State Space Model and Inplace Convolution Techniques

The recently proposed deep neural state-space model(SSM) [13] advances speech tasks by combining the properties of both CNNs and RNNs. The SSM is defined in continuous time using the following equations:

$$h^{\prime}(t)=Ah(t)+Bx(t)$$

(1)

$$y(t)=Ch(t)+Dx(t)$$

(2)

To be applied on a discrete input sequence (u0, u1, . . .) instead of continuous function u(t), (1) must be discretized by a step size $\operatorname{\Delta}$ that represents the resolution of the input. The discrete SSM is

$$x_{k}=\overline{\boldsymbol{A}}x_{k-1}+\overline{\boldsymbol{B}}u_{k}\quad y_{k}=\overline{\boldsymbol{C}}x_{k}$$

(3)

$$\overline{\boldsymbol{A}}=(\boldsymbol{I}-\Delta/2\cdot\boldsymbol{A})^{-1}(\boldsymbol{I}+\Delta/2\cdot\boldsymbol{A})$$

(4)

where $\overline{\boldsymbol{A}}$,$\overline{\boldsymbol{B}}$,$\overline{\boldsymbol{C}}$ are the discretized state matrices.According to the conclusion in [13], it can be seen that:

$$y_{k}=\overline{\boldsymbol{CA}}^{k}\overline{\boldsymbol{B}}u_{0}+\overline{\boldsymbol{CA}}^{k-1}\overline{\boldsymbol{B}}u_{1}+\cdots+\overline{\boldsymbol{CA}\boldsymbol{B}}u_{k-1}+\overline{\boldsymbol{C}\boldsymbol{B}}u_{k}$$

(5)

$$y=\overline{\boldsymbol{K}}*u$$

(6)

$$\overline{\boldsymbol{K}}=\left(\overline{\boldsymbol{C}\boldsymbol{B}},\overline{\boldsymbol{CAB}},\ldots,\overline{\boldsymbol{CA}}^{L-1}\overline{\boldsymbol{B}}\right)$$

(7)

In other words, (5) is a single (non-circular) convolution and can be computed very efficiently with FFTs, provided that $\overline{\boldsymbol{K}}$ is known. For the specific details of SSM, you can refer to [13, 14] to understand.

The S4 layer was developed for 1-D inputs, which limits its applicability. In [13, 14], the input dimension is 2-D, the shape is (H, T), and the S4 layer is designed as H independent parallel calculations. Since there are no correlations in the H dimension, this is limited in speech signal processing. So S4ND compensates for the correlation in the frequency dimension. In [12], the conventional S4 layer was extended to multidimensional signals by turning the standard SSM (1-D ODEs) into multidimensional partial differential equations (PDEs) governed by an independent SSM in each dimension.

Let $u=u(t^{(1)},t^{(2)})$ and $y=y(t^{(1)},t^{(2)})$ be the input and output which are signals $\mathbb{R}^{2}\rightarrow\mathbb{C}$, and $x=(x^{(1)}(t^{(1)},t^{(2)})$, $x^{(2)}(t^{(1)},t^{(2)}))\in\mathbb{C}^{N^{(1)}\times N^{(2)}}$ be the SSM state of dimension $N^{(1)}\times N^{(2)}$, where $x^{(\tau)}:\mathbb{R}^{2}\rightarrow\mathbb{C}^{N^{(\tau)}}$. The 2D SSM is the map u $\mapsto$ y defined by the linear PDE with initial condition x(0, 0) = 0 :

	$\displaystyle\frac{\partial}{\partial t^{(1)}}x(t^{(1)},t^{(2)})$	$\displaystyle=(\boldsymbol{A}^{(1)}x^{(1)}(t^{(1)},t^{(2)}),x^{(2)}(t^{(1)},t^{(2)}))+$		(8)
		$\displaystyle\phantom{=}\ \ \boldsymbol{B}^{(1)}u(t^{(1)},t^{(2)})$
	$\displaystyle\frac{\partial}{\partial t^{(2)}}x(t^{(1)},t^{(2)})$	$\displaystyle=(x^{(1)}(t^{(1)},t^{(2)}),\boldsymbol{A}^{(2)}x^{(2)}(t^{(1)},t^{(2)}))+$
		$\displaystyle\phantom{=}\ \ \boldsymbol{B}^{(2)}u(t^{(1)},t^{(2)})$
	$\displaystyle y(t^{(1)},t^{(2)})$	$\displaystyle=\langle\boldsymbol{C},x(t^{(1)},t^{(2)})\rangle$

Note that (8) differs from the usual notion of multidimensional SSM, which is simply a map from $u(t)\in\mathbb{C}^{n}\mapsto y(t)\in\mathbb{C}^{m}$ for higher-dimensional n, m $>$ 1 but still with 1 time axis. However, (8) is a map from $u\left(t_{1},t_{2}\right)\in\mathbb{C}^{1}\mapsto y\left(t_{1},t_{2}\right)\in\mathbb{C}^{1}$ for scalar input/outputs but over multiple time axes. When thinking of the input $u(t^{(1)},t^{(2)})$ as a function over a 2D grid, (8) can be thought of as a simple linear PDE that just runs a standard 1D SSM over each axis independently For details, please refer to [12]. S4ND can be regarded as a convolution kernel with infinite receptive fields in N dimensions. In the frequency domain signal, the enhanced performance of S4ND-U-Net [15] is significant. Furthermore, both the parameter count and computational load are remarkably low. Importantly, S4ND adheres to the principles of causal networks and ensures real-time performance.

SICRN is shown in figure 1(a). The comprehensive network architecture adheres to the overall-detailed framework, employing the complex spectrum as its input. Initially, the feature channels are modified through inplace convolution. Subsequently, the SIC block undertakes preliminary extraction and integration of both global and local features from the real and imaginary components. Subsequently, the features undergo temporal modeling via a 2-layer LSTM. In the final step, the SIC block is employed for individual extraction and reconstruction of features from the real and imaginary components. Additionally, it estimates a complex mask, which is then applied through multiplication with the original real and imaginary parts to yield the enhanced feature spectrum.

The SIC block can function as a convolution kernel, effectively replacing standard convolution kernels, and excels at extracting both local and global features. Importantly, the absence of downsampling operations within the entire module ensures the preservation of the original features.

The SIC block is shown in figure 1(b). The input channel is bifurcated into two segments. The initial 1/2 of the channel employs the inplace convolution kernel to capture local information. Following 1D convolution, they serve as the original features and local attention features, respectively. The latter 1/2 of the channel feeds into the S4ND layer for global information extraction, which is then harmonized with local attention features. The attention map is derived via the sigmoid activation function and subsequently multiplied with the original feature, thereby facilitating the extraction and fusion of local and global features. The specific calculation process is as follows:

$$\begin{split}X_{0\sim\frac{c}{2}}^{L}=C_{1d}(\operatorname{IC}(X_{0\sim\frac{c}{2}})\end{split}$$

(9)

$$\begin{split}X_{\frac{c}{2}\sim c}^{R}=\operatorname{S}(X_{\frac{c}{2}\sim c})\end{split}$$

(10)

$$\begin{split}ATmap=\sigma(C_{1d}(\operatorname{IC}(X_{0\sim\frac{c}{2}}))+X_{\frac{c}{2}\sim c}^{R})\end{split}$$

(11)

$$\begin{split}\mathrm{X}=X_{0\sim\frac{c}{2}}^{L}\cdot ATmap\end{split}$$

(12)

Where, $X_{0\sim\frac{c}{2}}^{L}$ denotes the local features, while $X_{\frac{c}{2}\sim c}^{R}$ represents the global features. $X_{0\sim\frac{c}{2}}$ corresponds to the features of the first half of the channels, and $X_{\frac{c}{2}\sim c}$ pertains to the features of the latter half of the channels. $\operatorname{IC}(\cdot)$ denotes the inplace convolution operation, and $\operatorname{S}(\cdot)$ signifies the convolution kernel used in the S4ND layer. The symbol $\sigma$ represents the sigmoid activation function. ’$ATmap$’ denotes the attention map and $C_{1d}(\cdot)$ denotes the conv1d.

Figure 1(c) shows the S4ND block. Initially, the S4ND is employed to extract global features. This extraction process is followed by passage through an ELU activation function, and output via a linear layer. Subsequently, a residual connection is implemented to address potential problem related to gradient vanishing or exploding. Finally, a batch normalization layer is applied to generate the final output. The rationale behind choosing S4ND for global feature modeling is as follows:

1. The global modeling capacity of S4ND surpasses that of LSTM, while maintaining a smaller parameter count and computational load.

2. Given that S4 processes elements independently within the frequency dimension, it isn’t ideally suited for processing frequency domain information. This limitation is addressed by the utilization of S4ND.

Furthermore, as demonstrated in [12], experimental comparisons between S4ND and 2D convolutions reveal that S4ND outperforms the latter. Consequently, we opt for S4ND as the method to extract global feature information.

We apply a scale-invariant signal-to-noise ratio (SI-SNR) [16] loss, which is a time domain loss function as follows:

$$\mathbf{s}_{\text{target }}=\frac{\langle\hat{\mathbf{s}},\mathbf{s}\rangle\mathbf{s}}{\|\mathbf{s}\|^{2}}$$

(13)

$$\mathbf{e}_{\text{noise }}=\hat{\mathbf{s}}-\mathbf{s}_{\text{target }}$$

(14)

$$\mathcal{L}_{\text{si-snr }}=10\log_{10}\frac{\left\|\mathbf{s}_{\text{target }}\right\|^{2}}{\left\|\mathbf{e}_{\text{noise }}\right\|^{2}}$$

(15)

where $\hat{\mathbf{s}}\in\mathbb{R}^{1\times T}$ and $\mathbf{s}\in\mathbb{R}^{1\times T}$ refer to the estimated and clean sources, respectively, and $\|\mathbf{s}\|^{2}=\langle\mathbf{s},\mathbf{s}\rangle$ denotes the signal power.

We evaluated the SICRN on the DNS Challenge (INTERSPEECH 2020) dataset [17]. The clean speech set includes over 500 hours of clips from 2150 speakers. The noise dataset includes over 180 hours of clips from 150 classes. To make full use of the dataset, we simulate the speech-noise mixture with dynamic mixing during model training. In detail, before the start of each training epoch, 75% of the clean speeches are mixed with randomly selected room impulse responses (RIR) from (1) the Multichannel Impulse Response Database [18] with three reverberation times (T60) 0.16s, 0.36s, and 0.61 s. (2) the Reverb Challenge dataset [19] with three reverberation times 0.3 s, 0.6 s and 0.7 s. After that, the speech-noise mixtures are dynamically generated by mixing the clean speech (75% of them are reverberant) and noise with a random SNR in between -5 and 20 dB. The DNS Challenge provides a publicly available test dataset, including two categories of synthetic clips, i.e., without and with reverberations. Each category has 150 noisy clips with SNR levels distributed in between 0 dB to 20 dB. We use this test dataset for evaluation.

For STFT ,we adopt 510/160 for win-length/hop-length,the analysis is Hanning window. We use 510-point discrete Fourier transform (DFT) to extract 256-dimensional complex spectra for 16 kHz sampling rate. The model is optimized by Adam. The initial learning rate is 0.0002 and halved when the validation loss of four consecutive epochs no longer decreased. When training the model, the specific settings are as follows: there are 2 layers of SIC blocks, with channel sizes of 16 and 32, respectively, for each layer. Additionally, there are 2 LSTM layers in the middle. Within the SIC block, there are 3 inplace convolution layers, and 4 layers of S4ND blocks. For the integrity of the experiment, the 4-layer S4ND blocks in SIC block are replaced with 4-layer inplace convolution, which is named IICRN.

In order to verify the effectiveness of the proposed method, we select FullSubNet as the baseline models. At the same time, compare according to the same evaluation standard of FullSubNet. In addition, we also compared with the topranked methods in the DNS challenge (INTERSPEECH 2020), including NSNet [20], DTLN [21], Conv-TasNet [22], DCCRN [23] and PoCoNet [24].