Source-linked AI summary
TSTNN: Two-stage Transformer based Neural Network for Speech Enhancement in the Time Domain
Kai Wang, Bengbeng He, Wei-Ping Zhu
TL;DR
Speech-enhancement systems often rely on time-frequency processing with transform overhead and problematic noisy-phase handling. TSTNN performs end-to-end monaural denoising in the time domain using a two-stage transformer between encoder and decoder, and experiments show strong performance across time- and frequency-domain comparisons with low complexity.
Problem
Time-frequency speech-enhancement methods incur Fourier-transform overhead and commonly ignore noisy phase information during denoising.
Method
TSTNN uses an encoder, two-stage transformer module, masking module, and decoder to extract local and global context from noisy speech and reconstruct enhanced speech.
Results
TSTNN outperforms most state-of-the-art methods across most evaluation metrics while using substantially fewer trainable parameters.
Takeaways & Limitations
The results support a time-domain transformer architecture that combines local and global contextual extraction with lower model complexity than competing models.
Abstract
from arXiv · showhide
In this paper, we propose a transformer-based architecture, called two-stage transformer neural network (TSTNN) for end-to-end speech denoising in the time domain. The proposed model is composed of an encoder, a two-stage transformer module (TSTM), a masking module and a decoder. The encoder maps input noisy speech into feature representation. The TSTM exploits four stacked two-stage transformer blocks to efficiently extract local and global information from the encoder output stage by stage. The masking module creates a mask which will be multiplied with the encoder output. Finally, the decoder uses the masked encoder feature to reconstruct the enhanced speech. Experimental results on the benchmark dataset show that the TSTNN outperforms most state-of-the-art models in time or frequency domain while having significantly lower model complexity.
1. INTRODUCTION
Speech enhancement methods commonly operate in the time-frequency domain, but Fourier-transform overhead and noisy-phase handling remain limitations. TSTNN addresses these issues with a time-domain transformer that models both local and global context and reports strong benchmark performance with relatively low complexity.
- Limitations of time-frequency methods: Time-frequency speech-enhancement methods typically use STFT magnitudes and noisy phase for time-domain reconstruction, adding transform overhead and often ignoring phase during denoising.The paper notes that phase information is important for speech quality.
- Limitations of existing time-domain methods: Time-domain CNN approaches require additional convolutional layers for long-range context, while RNN approaches cannot parallelize efficiently and may miss contextual information.The paper presents these constraints as motivation for transformer-based sequence modeling.
- Proposed approach: TSTNN is an end-to-end monaural time-domain speech-enhancement model that inserts a two-stage transformer module between encoder and decoder.The module is designed to learn local and global contextual information from long-range speech sequences.
- Experimental outcome: Benchmark experiments report that TSTNN outperforms state-of-the-art methods on most evaluation criteria while retaining relatively light model complexity.This is the paper’s stated overall experimental finding.
2. TWO-STAGE TRANSFORMER
The paper develops an improved transformer encoder and a two-stage transformer block for speech features. The block separates local processing from global dependency fusion to extract contextual information efficiently.
- Improved transformer: The improved transformer removes positional encoding because the authors consider it unsuitable for acoustic sequences.The model uses only the transformer encoder because noisy and enhanced speech sequences have the same length.
- Improved transformer: Multi-head attention maps input features into queries, keys, and values, computes attention for each head, concatenates the heads, and applies projection, normalization, and residual connection.The attention operations are described across the corresponding equations and prose.
3. PROPOSED MODEL
TSTNN combines segmentation, an encoder, a two-stage transformer module, masking, decoding, and overlap-add reconstruction for time-domain speech enhancement. Its transformer blocks process local and global contextual information, while joint time-domain and time-frequency losses supervise reconstruction.
- Architecture: TSTNN comprises segmentation, an encoder, a two-stage transformer module, a masking module, a decoder, and overlap-add reconstruction.The segmentation stage frames the raw mixture, and overlap-add recovers the enhanced waveform.
- Transformer encoder: The improved transformer encoder uses multi-head attention followed by a position-wise feed-forward network with residual connections and layer normalization.The model removes positional encoding and uses four parallel attention layers.
- Two-stage transformer: The two-stage transformer block applies a local transformer across frame-size dimension F, then a global transformer across frame dimension N.The input is a 3-D tensor [C, N, F], and the two stages extract local and global contextual information respectively.
- Encoder: The encoder uses two convolutional layers with an intermediate four-layer dilated dense block to transform framed mixtures into feature representations.The first convolution increases channels to 64, while the second halves frame size using stride (1, 2).
- Masking and decoder: The masking module estimates a denoising mask from TSTM features and obtains the final masked encoder feature through element-wise multiplication.The decoder then reconstructs enhanced speech using a dilated dense block, sub-pixel convolution, and overlap-add.
- Loss function: The training objective combines time-domain mean square error with a time-frequency-domain loss based on clean and enhanced spectrograms.The time-frequency loss is described as providing supervision for speech intelligibility and perceptual quality; α is set to 0.2.
4. EXPERIMENTS
Experiments evaluate TSTNN on a standard Voice Bank-derived benchmark using objective speech-quality, intelligibility, MOS, signal-to-noise, and parameter-size measures. TSTNN achieves strong performance against waveform- and T-F-domain methods while using substantially fewer parameters, and also outperforms a comparison architecture in ablation tests.
- Experimental setup: The benchmark contains 11,572 training utterances from 28 speakers and 824 test utterances from two speakers, with seen and unseen noise conditions across multiple SNRs.Training uses 10 noise types at 15, 10, 5, and 0 dB; testing uses five unseen noise types at 17.5, 12.5, 7.5, and 2.5 dB.
- Experimental setup: Evaluation uses PESQ, STOI, three MOS measures—CSIG, CBAK, and COVL—SSNR, and trainable parameter count.PESQ ranges from -0.5 to 4.5, STOI from 0 to 1, MOS measures from 1 to 5, and SSNR from -10 to 35.
- Comparison with time-domain methods: TSTNN outperforms most waveform-based methods in PESQ, achieves the best STOI score among existing time-domain models, and leads on three MOS evaluations.Its PESQ performance is comparable to DEMUCS with large configuration while using about 36 times fewer parameters.
- Comparison with T-F methods: TSTNN outperforms all T-F methods on most evaluation metrics, especially PESQ, with 0.92 million parameters versus 2.3 million for DCUNet-16.The reported parameter count is 2.5 times lower than DCUNet-16 and is described as extremely low among the listed T-F models.
- Influence of two-stage transformer block: In the ablation comparison, TSTNN obtains better scores on all evaluation metrics while using 2.6 times fewer parameters than the comparison model.The comparison model uses two transformer blocks and more encoder and decoder layers, whereas TSTNN uses its two-stage transformer-block design.
5. CONCLUSION
The paper proposes TSTNN for monaural speech enhancement in the time domain, using two-stage transformer processing to extract local and global contextual information from long-range speech sequences. Experiments report strong performance across most evaluation metrics with substantially fewer trainable parameters than current models.
- TSTNN efficiently extracts both local and global contextual information for long-range speech sequences in monaural time-domain speech enhancement.
- Experimental results show that TSTNN outperforms most state-of-the-art methods on most evaluation metrics.
- TSTNN uses much fewer trainable parameters than other current models.