StoRM: A Diffusion-based Stochastic Regeneration Model for Speech Enhancement and Dereverberation의 핵심을 짚고 추가 아이디어 만들기
논문 StoRM: A Diffusion-based Stochastic Regeneration Model for Speech Enhancement and Dereverberation [1] (줄여서 StoRM)의 기본아이디어를 정리하고 추가적으로 StoRM 에서 제시된 공식들을 검증해보고 좀 더 일반화 해보겠다.
StoRM 의 아이디어
StoRM 은 SDE프레임워크 하에서 score based model을 추정 하는 방법을 사용한다. [2] \mathbf{y} 를 corrupted speech 라고 하고 \mathbf{x}(0) 를 clean speech 라고 할 때, 다음의 SDE를 통해서 forward process를 학습하고, 그것의 반대과정인 reverse process를 이용해서 \mathbf{y} 로부터 clean speech를 생성한다.
\begin{align}d \mathbf{x} = \gamma( \mathbf{y} - \mathbf{x} ) dt + g(t)^2 d \mathbf{w}, \quad 0\leq t \leq 1, \quad\end{align} \begin{align} g(t) = \sigma_{min} \left( \frac{ \sigma_{max}} {\sigma_{min}} \right)^\tau \sqrt { 2 \log \left( \frac{\sigma_{max}}{\sigma_{min}} \right) } \end{align}여기서 \mathbf{x}, \mathbf{w} \in \mathcal{R}^d 이고 \mathbf{w} 는 Brownian Motion 이다.
논문에서는 audio signal을 Time-Frequency domain에서 복소신호로 보고 처리하고 있지만 핵심아이디어를 설명하기 위해서는 실수 (real number)에서만 논의를 해도되니 앞으로 나오게 될 식들은 다 실수를 다룬다는 것을 양해해 주길 바란다.
StoRM 에서 score 추정 방식
score p( \mathbf{x}(t) | \mathbf{y} ) 를 추정하기 위해 p(\mathbf{x} | \mathbf{x}(0), \mathbf{y})을 계산 하였는데 아래와 같은 결과가 나온다.
\begin{align}p_{0, t} (\mathbf{x}(t) | \mathbf{x}(0), \mathbf{y} ) = \mathcal{N} ( \mathbf{x}(t) ; \mathbf{\mu}(\mathbf{x}(0), \mathbf{y}, t), \sigma(t)^2 \mathbf{I}) \end{align} \begin{align} \mu(\mathbf{x}(0), \mathbf{y}, t) = e^{-\gamma t} \mathbf{x}(0)+ (1 - e^{-\gamma t}) \mathbf{y} \end{align} \begin{align}\sigma(t)^2 = \sigma_{\min}^2 \log (\sigma_{\max}/\sigma_{\min}) \frac{ \left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^{2\tau} - e^{-2\gamma\tau} } { \gamma + \log (\sigma_{max} /\sigma_{min})} \end{align}StoRM 의 아이디어 검증
위의 식 (3), (4), (5)등에 대해 검증을 해보겠다. (3)처럼 가우시안 분포로 나온다는 것은 이 글(Linear SDE의 transition pdf)에서 이미 알아보았고 (4), (5)를 유도하기 위한 mean과 covariance에 대해 미분방정식을 세우는 방법은 이 글(Linear SDE의 transition pdf)에서 알아봤다. 여기서 중요한것은 (4), (5)를 유도하는 것이기 때문에 이 글(Linear SDE의 transition pdf)을 토대로 미분방정식을 유도해보겠다. 좀 더 일반적인 상황에 대해 증명하기 위하여 식 (1)에서 \gamma 대신 t에 대한 함수인 \gamma(t) 대해 논의해보겠다. 즉 아래의 SDE (6)에 대해 mean 과 covariance를 유도하겠다는 의미이다.
\begin{align}d \mathbf{x} = \gamma(t)( \mathbf{y} - \mathbf{x} ) dt + g(t)^2 d \mathbf{w}, \quad 0\leq t \leq 1, \quad\end{align}\mathbf{x}(0), \mathbf{y} 가 주어졌을 때 기호를 편하게 하기 위해 \mathbf{\mu}(t) := \mathbf{\mu}(\mathbf{x}(0), \mathbf{y}, t) 라 표기하자, 그리고 covariance를 \Sigma (t) = E[ (\mathbf{x}(t) - \mathbf{\mu}(t))(\mathbf{x}(t) - \mathbf{\mu}(t))^T ] 라 하자. 이 글(Linear SDE의 transition pdf)에 의해 아래의 미분 방정식이 유도된다.
\begin{align} \frac{d\mathbf{\mu}(t)}{dt} = -\gamma(t) \mathbf{\mu} + \gamma(t) \mathbf{y} \end{align} \begin{align} \frac{ d \mathbf{\Sigma}}{dt} = -2 \gamma(t) \Sigma(t) + g^2(t) \mathbf{I} \end{align} \begin{align} \text{ 초기조건: } \mathbf{\mu}(0) = \mathbf{x}(0), \mathbf{\Sigma}(0) = \mathbf{0} \end{align}식 (7) 에서 \mathbf{\mu} , \mathbf{y} 는 벡터인데 이것의 i번째 원소에 대응되는 미분 방정식은 아래와 같다.
\begin{align} \frac{d \mathbf{\mu}_i}{dt} = -\gamma(t) \mathbf{\mu}_i + \gamma(t) \mathbf{y}_i \end{align}식 (10)은 1계 선형미분방정식이므로 1계 선형미분방정식의 풀이법과 초기조건 \mathbf{\mu}(0) = \mathbf{x}(0) 이용하면 아래와 같이 해를 구할 수 있다.
\small{\begin{align} \mathbf{\mu}(t)_i = \exp \left(- \int_{0}^t \gamma (\tau) d\tau\right) \mathbf{x}(0)_i + \exp\left(-\int_0^t \gamma(\tau) d\tau\right) \left[ \int_0^t \exp\left(\int_0^\tau \gamma(s) ds\right) \gamma(\tau) d\tau \right] \mathbf{y}_i \end{align} }식 (11) 또한 1계 선형미분방정식이므로 1계 선형미분방정식의 풀이법과 초기조건 \mathbf{\Sigma}(0) = \mathbf{0} 을 이용하면 아래와 같이 \mathbf{\Sigma} 의 각원소값 \mathbf{\Sigma}_{ij}를 구할 수 있다.
\begin{align}\mathbf{\Sigma}(t)_{ij} = \sigma(t)^2 \delta[i-j]\end{align} \begin{align} \sigma(t)^2 =\exp{\left( -2 \int_0^t \gamma(\tau) d\tau\right)} \left[ \int_0^t \exp{ \left( 2 \int_0^\tau \gamma(s) ds \right) g^2(\tau) d\tau} \right] \end{align}
식 (11), (12)는 \mathbf{\mu}, \mathbf{\Sigma} 의 구성하는 원소에 대한 식인데, 이것을 벡터형태와 행렬형태로 기술하면 아래와 같이 나온다.
\begin{align}\mathbf{\mu}(t) &= \exp \left(- \int_{0}^t \gamma (\tau) d\tau\right) \mathbf{x}(0) + \exp\left(-\int_0^t \gamma(\tau) d\tau\right) \left[ \int_0^t \exp\left(\int_0^\tau \gamma(s) ds\right) \gamma(\tau) d\tau \right] \mathbf{y}\\ \mathbf{\Sigma}(t) &= \sigma^2(t) \mathbf{I}\end{align}식 (14), (15)에 \gamma(t) = \gamma 를 대입하면 식 (4), (5)가 유도된다.
StoRM 에서 아이디어 추가하기
첫번째 아이디어: 식 (1)에 \gamma 대신 \gamma(t) 를 이용해 논의를 전개했었다. \gamma(t) 를 어떻게 전개하냐에 따라 mean과 covariance가 달라지니까 \gamma(t) 를 잘 잡는 것이 중요해 보인다. \gamma(t) 를 조정하다보면 적분 (11), (13)을 구해야 되는데, 적분을 구성하는 요소인 \gamma(t) 를 세심하게 잡아야 겠고 마찬가지로 적분의 요소인 diffuse term g(t) 또한 건들여 볼만하다.
두번째 아이디어: 첫번째 아이디어에서 스칼라 함수 \gamma(t) 를 사용하는 방법에 대해 생각해보았는데, Matrix 함수인 \mathbf{\Gamma}(t) 를 이용해 SDE를 만들고 논의를 전개해도 괜찮을수도 있겠는데, 목적에 맞게 잡아야 겠지!
Reference
[1] Lemercier, J. M., Richter, J., Welker, S., & Gerkmann, T. (2022). StoRM: A diffusion-based stochastic regeneration model for speech enhancement and dereverberation. arXiv preprint arXiv:2212.11851.
[2]Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2020, October). Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations.