[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=1}^K{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,\overline{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,\overline{P}=I_k-P$
    - $\mathbb{1}$ : $K$ size의 one vector
  • 이때 $P$는 평균값들의 subspace로 projection 되고 $\overline{P}$는 orthogonal complement를 의미
    - (Eq.1)에 대한 compact alternative는 $Ps=\overline{s}$로 표현될 수 있음

- SDE for Score-based Generative Modelling

• SDE를 통한 Score-based Generative Modelling (SGM)은 복잡한 data 분포를 모델링하는 방법임
  • Target 분포의 sample에서 시작해 Gaussian 분포로 수렴하는 diffusion process는 다음의 SDE로 표현됨:
    (Eq.5) $d{x}_t=f(x_t,t)dt+g(t)d{w}_t$
    - $x_t:{\mathbb{R}}^N\to {\mathbb{R}}^N$ : time $t\in \mathbb{R}$의 vector-valued function
    - $d{x}_t$ : $t$에 대한 도함수
    - $dt$ : infinitesimal time step
    - $f:{\mathbb{R}}^N\to {\mathbb{R}}^N$, $g:\mathbb{R}\to \mathbb{R}$ : $x_t$의 dirft, diffusion coefficient
    - $d{w}_t$ : standard Wiener process
  • 본 논문에서는 Diffusion process의 time $t$가 audio signal의 time과 관계없다고 가정함
  • (Eq.5)에서 주목할 점은, mild condition에서 corresponding reverse-time SDE가 존재한다는 것:
    (Eq.6) $d{x}_t=-[f(x_t,t)-g(t{)}^2{\mathrm{\nabla }}_{x_t}log{p}_t(x_t)]dt+g(t)d\overline{w}$
    - 이때 $t=T$에서 $t=0$으로 이동, $d\overline{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 {\mathrm{\nabla }}_{x}log{p}_t(x)$가 되도록 neural network $q_{\theta }(x,t)$를 학습시키는 것을 목표로 함
    - 충분히 근사되었다면 (Eq.6)을 수치적으로 풀어 ${\mathrm{\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) $d{x}_t=-\gamma \overline{P}{x}_{t}dt+g(t)dw,x_0=s$
    - $\overline{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{2log\left(\frac{{\sigma }_{max}}{{\sigma }_{min}}\right)}$
    -> (Eq.7)은 $t$가 커짐에 따라 marginal ${\mu }_t$의 평균이 $t=0$에서 분리된 signal vector에서 mixture vector $\overline{s}$로 변하는 속성을 가짐
  • [Theorem 1.] 따라서, $x_t$의 marginal 분포는
    (Eq.9) ${\mu }_t=(1-e^{-\gamma t})\overline{s}+e^{-\gamma t}s$, (Eq.10) ${\sigma }_t={\lambda }_1(t)P+{\lambda }_2(t)\overline{P}$
    를 평균, 공분산 matrix로 가지는 Gaussian 분포이고,
    이때 ${\xi }_1=0,{\xi }_2=\gamma ,\rho =\frac{{\sigma }_{max}}{{\sigma }_{min}}$이라고 하면,
    (Eq.11) ${\lambda }_k(t)=\frac{{\sigma }_{min}^2({\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}\overline{P}$
  • 이때 시간에 따른 $x_t$ 분포의 parameter 변화를 확인해 보면,
    - $\gamma$값을 조정하여 ${\mu }_t$와 $\overline{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\mathrm{\angle }x}$, $c^{-1}(x)=\beta |x{|}^{1/\alpha }{e}^{j\mathrm{\angle }x}$ 
    - 실수부, 허수부는 concatenate 됨

- Inference

• 추론 시 score-matching network $q_{\theta }(x,t,y)\approx {\mathrm{\nabla }}_{x}log{p}_t(x)$를 사용하여 (Eq.5)을 solve 함으로써 source separation을 수행
  • 이때 initial value는 아래 분포로부터 sampling 됨:
    (Eq.13) ${\overline{x}}_T\sim \mathcal{N}(\overline{s},{\mathrm{\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) ${\mathrm{\nabla }}_{x_t}logp(x_t)=-{\mathrm{\Sigma }}_t^{-1}(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|{|}_{{\mathrm{\Sigma }}_t}^2$
    (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_{ϵ},T)$
    - $x_0$ : dataset에서 무작위로 선택됨
    - ${\mathrm{\Sigma }}_t$ : weighting scheme으로 유도된 norm
• BUT, 위의 절차로만 학습을 수행할 경우 추론 성능이 크게 개선되지 않음
  
  1. (Eq.13)의 $\mathbb{E}[{\overline{x}}_T]=\overline{s}$와 ${\mu }_T$ 사이의 불일치가 존재하기 때문
     - $t=T$에서 true score function은 (Eq.14)처럼 단순화되지 않으므로,
     - 대신 ${\mathrm{\nabla }}_{x_T}logp({\overline{x}}_T)=-{\mathrm{\Sigma }}_T^{-1}(\overline{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}mi{n}_{\pi \in \mathcal{P}}||L_T{q}_{\theta }({\overline{x}}_T,T,y)+z+L_T^{-1}(\overline{s}-{\mu }_T(\pi ))|{|}^2$
     - ${\overline{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 실험 결과