[Paper 리뷰] Diffusion-Based Generative Speech Source Separation

Diffusion-Based Generative Speech Source Separation

---

• Source separation을 위해 Stochastic Differential Equation을 활용할 수 있음
• DiffSep
  • 분리된 source에서 시작해 mixture를 중심으로 하는 Gaussian 분포로 수렴하는 continuous time diffusion-mixing proces를 활용
  • Diffusion-mixing process의 score function에 대한 marginal probability를 근사하는 neural network를 훈련
  • Neural network를 활용하여 mixture에서 source를 점진적으로 분리하는 reverse-time SDE를 solve
• 논문 (ICASSP 2023) : Paper Link

---

1. Introduction

• Source Separation은 mixture에서 관심 있는 signal을 분리하는 것을 의미함
  • 초기에는 Non-Negative Matrix Factorization (NMF)를 주로 활용했음
  • 이후 Deep Neural Network (DNN)의 등장으로 분리 성능이 크게 향상됨
    - DNN은 source cluster를 예측하는 것에서 mask value를 직접 예측하는 것으로 변화

• 위와 같은 data-driven 방식들은 fundamental source permutation ambiguity를 해결해야 함
  • Source에 대한 inherent preffered order가 없기 때문
  • 일반적으로 permutation ambiguity 해결을 위해 permutation invariant training (PIT)이 주로 활용됨
    - Source의 모든 permutation에 대해 objective function을 계산하고 최솟값을 선택해 gradient를 계산하는 방식
    - Time-domain에서는 Conv-TasNet, Dual-path transformer 등
    - Time-Frequency domain에서는 WaveSplit, TF-GridNet 등
  • 대부분 Scale-Invariant Signal-to-Distortion Ratio (SI-SDR)을 통해 학습됨
    -> NMF를 제외하고는 generative model을 separation 작업에 활용하는 경우가 거의 없음
• Generative modelling은 복잡한 data 분포를 쉽게 근사할 수 있음
  • 특히 Score-based Generative Modelling (SGM)이 좋은 성능을 보임
    - Target 분포를 Gaussian noise로 점진적으로 변환하는 forward process를 정의하고, log-probability에 대한 gradient인 score function을 통해 reverse process를 수행하여 target 분포를 sampling
    
  • 이때, score function은 두 가지 방식을 통해 DNN으로 근사할 수 있음
    1. Graphical model
    2. Stochastic Differential Equation (SDE)
  • 이 중 SDE를 활용하는 diffusion-based separataion은 아직 제시되지 않음

-> 그래서 diffusion-based source separation model인 DiffSep를 제안

 

• DiffSep
  • 분리된 signal로부터 mixture의 평균으로 수렴하는 SDE를 설계
  • Reverse-time SDE를 풀어 mixture로부터 개별적인 source를 복원
  • 추론 단계에서 source assignment의 ambiguity를 해소하는 score-matching network에 대한 학습방식 제시

< Overall of DiffSep >

• SDE solving을 통해 mixture에서 source를 점진적으로 분리하는 diffusion-based model
• 자연스럽고 높은 non-intrusive score를 달성할 수 있는 source separation 모델
• Noise를 extra source로 고려함으로써 speech enhancement로의 확장 가능성 제시

Diffusion-mixing process

2. Background

- Notation and Signal Model

• Vector $x$에 대한 norm : $||x|| = (x^{T} x) ^{1/2}$, $I_{N}$ : $N \times N$ identity matrix
• Audio signal은 $N$개 sample을 가지는 $\mathbb{R}^{N}$의 real valued vector로써 표현됨
  
  • $K$개 source에 대한 mixture는:
    (Eq.1) $y = \sum^{K}_{k=1} s_{k}$
    - 모든 k에 대해, $s_{k} \in \mathbb{R}^{N}$
  • 여기서 model은 extra source를 추가함으로써 noise, reverberation, degradation을 설명할 수 있음
    - 모든 source signal을 concatenating 하여 $\mathbb {R}^{KN}$의 vector를 얻을 수 있음
    - 이때 분리된 source의 vector와 그 평균값은:
    (Eq.2) $s = [s_{1}^{T}, ..., s_{K}^{T}]^{T}, \bar{s} = K^{-1} [y^{T}, ..., y^{T}]^{T}$
• Time-invariant mixing operation을 $(A \otimes I_{N})$으로 정의할 수 있음
  
  • $A$ : $K \times K$ matrix, $\otimes$ : Kronecker product
  • Vector $v \in \mathbb{R}^{KN}$에 위의 matrix를 곱하는 것은, interval $N$에서 $v$로부터 가져온 element와 length $K$의 모든 sub-vector를 $A$로 곱하는 것과 동일함:
    (Eq.3) $((A \otimes I_{N}) v)_{kN+n} = \sum_{\ell=1}^{K} A_{k\ell}v_{\ell N+n}$
    - $k = 1, ..., K$, $n = 1, ..., N$
• 이후 식을 간단히 하기 위해 regular matrix product $Av \triangleq (A \otimes I_{N})v$ 라고 하고,
  • Projection matrix는 아래와 같이 정의됨:
    (Eq.4) $P = K^{-1}\mathbb{1} \mathbb{1}^{T}, \bar{P} = I_{k} - P$
    - $\mathbb{1}$ : $K$ size의 one vector
  • 이때 $P$는 평균값들의 subspace로 projection 되고 $\bar{P}$는 orthogonal complement를 의미
    - (Eq.1)에 대한 compact alternative는 $Ps = \bar{s}$로 표현될 수 있음

- SDE for Score-based Generative Modelling

• SDE를 통한 Score-based Generative Modelling (SGM)은 복잡한 data 분포를 모델링하는 방법임
  • Target 분포의 sample에서 시작해 Gaussian 분포로 수렴하는 diffusion process는 다음의 SDE로 표현됨:
    (Eq.5) $dx_{t} = f(x_{t}, t)dt + g(t) dw_{t}$
    - $x_{t} : \mathbb{R}^{N} \rightarrow \mathbb{R}^{N}$ : time $t \in \mathbb{R}$의 vector-valued function
    - $dx_{t}$ : $t$에 대한 도함수
    - $dt$ : infinitesimal time step
    - $f : \mathbb{R}^{N} \rightarrow \mathbb{R}^{N}$, $g: \mathbb{R} \rightarrow \mathbb{R}$ : $x_{t}$의 dirft, diffusion coefficient
    - $dw_{t}$ : standard Wiener process
  • 본 논문에서는 Diffusion process의 time $t$가 audio signal의 time과 관계없다고 가정함
  • (Eq.5)에서 주목할 점은, mild condition에서 corresponding reverse-time SDE가 존재한다는 것:
    (Eq.6) $dx_{t} = - [ f(x_{t}, t) - g(t)^{2} \nabla_{x_{t}} log \, p_{t} (x_{t}) ] dt + g(t)d\bar{w}$
    - 이때 $t =T$에서 $t=0$으로 이동, $d\bar{w}$ : reverse Brownian process, $p_{t} (x)$ : $x_{t}$의 marginal 분포
• Forward process 동안 $x_{0}$가 알려지면, 일반적으로 $p_{t}(x)$에 대한 closed-form을 가짐
  • Affine $f(x,t)$는 Gaussian이고 해당 parameter가 평균, 공분산 matrix이므로 tractable 함
    - Reverse process에서 $p_{t}(x)$는 unknown이므로 (Eq.6)은 directly solve가 불가능함
  • 따라서, SGM은 $q_{\theta} (x,t) \approx \nabla_{x} log \, p_{t} (x)$가 되도록 neural network $q_{\theta} (x,t)$를 학습시키는 것을 목표로 함
    - 충분히 근사되었다면 (Eq.6)을 수치적으로 풀어 $\nabla_{x} log \, p_{t}(x)$를 $q_{\theta}(x,t)$로 대체함으로써 $x_{0}$ 분포에서 sample을 생성할 수 있음

3. Diffusion-Based Source Separation

• Source separation을 위해 SGM을 도입하여 시간이 지남에 따라 source 전반에 걸쳐 diffusion과 mixing이 발생하는 forward SDE를 설계
  • Process의 각 step은 infinitesimal noise를 더하고 infinitesimal mixing을 수행하는 것으로 볼 수 있음
  • 이를 SDE로 공식화하면:
    (Eq.7) $dx_{t} = - \gamma \bar{P} x_{t} dt + g(t) dw, \, x_{0}=s$
    - $\bar{P}$ : (Eq.4), $s$ : (Eq.2)
    - Diffusion coefficient로써 Variance Exploding SDE를 사용:
    (Eq.8) $g(t) = \sigma_{min} \left ( \frac{\sigma_{max}} {\sigma_{min}} \right )^{t} \sqrt{2 log \left ( \frac {\sigma_{max}} {\sigma_{min}} \right ) }$
    -> (Eq.7)은 $t$가 커짐에 따라 marginal $\mu_{t}$의 평균이 $t=0$에서 분리된 signal vector에서 mixture vector $\bar{s}$로 변하는 속성을 가짐
  • [Theorem 1.] 따라서, $x_{t}$의 marginal 분포는
    (Eq.9) $\mu_{t} = (1-e^{-\gamma t} )\bar{s} + e^{-\gamma t }s$, (Eq.10) $\sigma_{t} = \lambda_{1} (t) P + \lambda_{2} (t) \bar{P}$
    를 평균, 공분산 matrix로 가지는 Gaussian 분포이고,
    이때 $\xi_{1} = 0, \xi_{2} = \gamma, \rho = \frac{\sigma_{max}}{\sigma_{min}}$이라고 하면,
    (Eq.11) $\lambda_{k}(t) = \frac{ \sigma^{2}_{min} (\rho^{2t} - e^{2\xi_{k}t} log \, \rho) } {\xi_{k} + log \, \rho}$
    이다.
  • 결과적으로 Sample $x_{t}$에 대한 explicit expression은:
    (Eq.12) $x_{t} = \mu_{t} + L_{t}z, \, z \sim \mathcal{N} (0, I_{KN})$
    - $L_{t} = \lambda_{1}(t)^{1/2}P + \lambda_{2} (t) ^{1/2} \bar{P}$
  • 이때 시간에 따른 $x_{t}$ 분포의 parameter 변화를 확인해 보면,
    - $\gamma$값을 조정하여 $\mu_{t}$와 $\bar{s}$의 차이를 작게 만들 수 있음
    - Mixing process로 인해 source에 더해지는 noise의 correlation이 시간에 따라 증가함 

시간에 따른 SDE (Eq.7)의 parameter 변화

• Speech enhancement를 위해 일반적으로 STFT와 같은 non-linear transform을 활용하여 diffusion process를 적용함
  • BUT, (Eq.8)은 source의 linear mixing을 modelling 하므로 non-linear transform을 적용할 수 없음
  • 대신 time-domain에서 diffusion process를 수행하고, non-linear STFT domain에서 network를 동작
    - Noise Conditioned Score Matching Network (NCSN)을 활용
    1. STFT와 iSTFT layer가 network 앞뒤에 위치
    2. STFT 이후에 non-linear transform $c(x)$을 추가하고, iSTFT 이전에 그 역 연산을 수행함:
    $c(x) = \beta^{-1} |x|^{\alpha} e^{j \angle x}$, $c^{-1}(x) = \beta |x| ^{1/\alpha} e^{j \angle x}$ 
    - 실수부, 허수부는 concatenate 됨

- Inference

• 추론 시 score-matching network $q_{\theta} (x,t,y) \approx \nabla_{x} log \, p_{t}(x)$를 사용하여 (Eq.5)을 solve 함으로써 source separation을 수행
  • 이때 initial value는 아래 분포로부터 sampling 됨:
    (Eq.13) $\bar{x}_{T} \sim \mathcal{N} (\bar{s}, \Sigma_{T} I_{KN})$
    - 이후 (Eq.5)를 solve 하기 위해 predictor-corrector approach를 채택
  • Prediction step은 reverse diffuse sampling을 통해 수행되고, correction은 annealed Langevin sampling을 통해 수행됨

- Permutation and Mismatch-aware Training Procedure

• Score network $q_{\theta} (x,t,y)$를 학습시키기 위해 modified score-matching procedure를 도입
  • [Theorem 1.]에 의해 score function은 closed form을 가지므로 계산식은:
    (Eq.14) $\nabla_{x_{t}} log \, p(x_{t}) = - \Sigma^{-1}_{t} (x_{t} - \mu_{t}) = - L_{t}^{-1}z$
  • 이때 Training loss는:
    (Eq.15) $\mathcal{L} = \mathbb{E}_{x_{0,z,t}} || q_{\theta}(x_{t}, t,y) + L_{t}^{-1}z || ^{2}_{\Sigma_{t}}$
    (Eq.16) $= \mathbb{E}_{x_{0,z,t}} || L_{t} q_{\theta} (x_{t}, t, y) + z ||^{2} $
    - $z \sim \mathcal{N} (0, I_{KN})$, $t \sim \mathcal{U} (t_{\epsilon}, T)$
    - $x_{0}$ : dataset에서 무작위로 선택됨
    - $\Sigma_{t}$ : weighting scheme으로 유도된 norm
• BUT, 위의 절차로만 학습을 수행할 경우 추론 성능이 크게 개선되지 않음
  
  1. (Eq.13)의 $\mathbb{E} [\bar{x}_{T}] = \bar{s}$와 $\mu_{T}$ 사이의 불일치가 존재하기 때문
     - $t = T$에서 true score function은 (Eq.14)처럼 단순화되지 않으므로,
     - 대신 $\nabla_{x_{T}} log \, p(\bar{x}_{T}) = - \Sigma_{T}^{-1} (\bar{s} + L_{T}z - \mu_{T})$를 사용
  2. Network는 source를 출력할 order를 결정해야 하기 때문
     - Model mismatch를 포함하는 PIT objective로 network를 학습하면 됨
     1. 확률 $1-p_{T}, p_{T} \in [0,1]$인 각 sample에 대해, (Eq.16)을 최소화하는 regular score-matching procedure를 적용
     2. 확률 $p_{T}$에서 $t=T$로 두고, 아래의 alternative loss를 최소화:
     $\mathcal{L} = \mathbb{E}_{x_{0}, z} min_{\pi \in \mathcal{P}} || L_{T} q_{\theta} (\bar{x}_{T}, T, y) + z + L^{-1}_{T} (\bar{s} - \mu_{T} (\pi) ) ||^{2}$
     - $\bar{x}_{T}$ : (Eq.13)에서 sampling 되는 값, $\mathcal{P}$ : source permutations의 집합, $\mu_{T}(\pi)$ : permutation $\pi$에 대해 계산된 $\mu_{T}$
     -> 이를 통해 score network는 noise와 model mismatch를 모두 제거하는 방법을 학습할 수 있고, 결과적으로 speech enhancement시 source order에 대한 ambiguity를 줄일 수 있음

4. Experiments

- Settings

• Datasets : WSJ0_2mix, Libri2Mix, VoiceBank-DEMAND
• Comparisons : Conv-TasNet, CDiffuse, SGMSE

- Results

• Separation
  • SI-SDR, PESQ, STOI 측면에서는 DiffSep의 성능이 Conv-TasNet보다 낮게 측정됨
    - 비슷한 소리를 내는 speaker에 대해 natural sounding permutation이 발생하기 때문
    
  • OVRL 측면에서는 DiffSep가 우수한 성능을 보임

Separation 실험 결과

• Clean source와 분리된 source의 spectrogram을 비교해 보면, 낮은 품질의 sample에서 source에 대한 block permutation이 발생함

Spectrogram 비교

• Enhancement
  
  • DiffSep는 Conv-TasNet에 비해서는 상대적으로 낮은 성능을 보이지만, 다른 diffusion-based 방법들과 비교했을 때 더 나은 성능을 보임
  • OVRL 측면에서 DiffSep는 높은 OVRL 점수를 나타내어 좋은 enhancement 품질을 가진다고 볼 수 있음

Enhancement 실험 결과