Single and Few-step Diffusion for Generative Speech Enhancement

TL;DR

This paper introduces a two-stage training method for diffusion-based speech enhancement that significantly reduces inference time from 60 to 5 function evaluations while maintaining high performance, outperforming traditional diffusion models in low evaluation scenarios.

Contribution

The paper proposes a novel two-stage training approach that enables single or few-step diffusion for speech enhancement, improving efficiency and robustness over existing methods.

Findings

Achieves comparable performance with only 5 function evaluations.

Maintains steady performance when reducing function evaluations.

Outperforms baseline diffusion models in low evaluation settings.

Abstract

Diffusion models have shown promising results in speech enhancement, using a task-adapted diffusion process for the conditional generation of clean speech given a noisy mixture. However, at test time, the neural network used for score estimation is called multiple times to solve the iterative reverse process. This results in a slow inference process and causes discretization errors that accumulate over the sampling trajectory. In this paper, we address these limitations through a two-stage training approach. In the first stage, we train the diffusion model the usual way using the generative denoising score matching loss. In the second stage, we compute the enhanced signal by solving the reverse process and compare the resulting estimate to the clean speech target using a predictive loss. We show that using this second training stage enables achieving the same performance as the baseline model using only 5 function evaluations instead of 60 function evaluations. While the performance of usual generative diffusion algorithms drops dramatically when lowering the number of function evaluations (NFEs) to obtain single-step diffusion, we show that our proposed method keeps a steady performance and therefore largely outperforms the diffusion baseline in this setting and also generalizes better than its predictive counterpart.