[Paper 리뷰] SpeechFlow: Generative Pre-Training for Speech with Flow Matching

SpeechFlow: Generative Pre-Training for Speech with Flow Matching

---

• Single pre-trained generative model을 다양한 downstream task에 활용할 수 있음
• SpeechFlow
  • Flow Matching과 masked condition을 사용하여 untranscribed speech로 pre-training을 수행
  • Pre-trained generative model을 task-specific data로 fine-tuning 하여 다양한 task에 적용
• 논문 (ICLR 2024) : Paper Link

---

1. Introduction

• Discriminative model은 speech recognition, enhancement, separation, text-to-speech (TTS) 등의 다양한 speech task에서 사용되고 있음
  • BUT, Wav2Vec 2.0, HuBERT와 같은 pre-trained foundation model은 data distribution $p(\text{speech})$ modeling 보다 discriminative task에 대한 useful representation을 학습하는데 중점을 둠
  • 한편으로 discriminative model과 달리 generative model을 활용하면 data distribution에 대한 sampling이 가능함
    - BUT, 기존의 generative speech model은 task-specific conditioning이나 distribution mapping에만 사용됨

-> 그래서 generative method를 활용한 speech foundation model인 SpeechFlow를 제안

 

• SpeechFlow
  • Masked audio modeling과 Flow Matching을 combine 하여 generative model을 구성
    - 이때 masked audio에서 speech conditioning의 underlying distribution을 estimate 하기 위해 unlabled speech로 training 됨
  • Unlabled speech data로 train 된 generative model을 기반으로, labeled data를 통해 task-specific fine-tune 함으로써 다양한 downstream speech task에 적용

< Overall of SpeechFlow >

• Flow Matching과 Masked audio modeling을 활용한 generative speech foundation model
• 결과적으로 다양한 downstream task에서 우수한 성능을 달성

2. Method

- Background: Flow Matching for Generative Modeling

• Deep generative model은 neural network에 의해 parameterize되는 distribution $p(x)$를 사용하여 real world $d$-dimensional data $x\in\mathbb{R}^{d}$의 unknown distribution $q(x)$를 estimate 하는 것을 목표로 함
  • Sampling을 위해서는 normal distribution과 같은 simple prior distribution $p_{0}(x)$을 starting point로 사용하여 $p_{1}(x)\approx q(x)$가 되는 neural transport map $p_{1}=F_{\theta}(p_{0})$를 finding함
    1. 이때 Generative Adversarial Network, Variational AutoEncoder와 같은 approach는 $x_{0}\sim p_{0}(x), x_{1}\sim q(x)$에서 $x_{1}=f_{\theta}(x_{0})$를 directly modeling 함
       - 즉, network $f_{\theta}$를 사용하여 noise로부터 data를 predict 하는 방식
    2. Diffusion model의 경우 noise $x_{0}$에서 data $x_{1}$까지 step size $\Delta t$로 traverse 하는 iterative denoising model $x_{t+\Delta t}=f_{\theta,t,\Delta t}(x_{t})$를 사용해 generation quality를 향상함
  • 따라서 SpeechFlow는 Continuous Normalizing Flow (CNF)에 해당하는 Flow Matching을 활용하여 neural transport map $p_{1}=F_{\theta}(p_{0})$를 구성함 
  • 먼저 CNF는 time-dependent probability density function $p_{t}:[0,1]\times \mathbb{R}^{d}\rightarrow \mathbb{R}_{>0}$을 통해 simple prior $p_{0}$와 target distribution $p_{1}$ 간의 path를 정의함
    1. Path를 따르는 $x$의 flow $\phi_{t}:[0,1]\times \mathbb{R}^{d}\rightarrow \mathbb{R}^{d}$는 Ordinary Differential Equation (ODE)를 통해 정의됨:
       (Eq. 1) $\frac{d}{dt}\phi_{t}(x)=v_{t}(\phi_{t}(x));\,\,\, \phi_{0}(x)=x$
       - $v_{t}:[0,1]\times \mathbb{R}^{d}\rightarrow \mathbb{R}^{d}$ : time-dependent vector field
       - 여기서 time-dependent probability density function $p_{t}$는 variable change를 통해 $p_{0}\left(\phi_{t}^{-1}(x)\right)\det\left[\frac{\partial \phi_{t}^{-1}}{\partial x}(x)\right]$와 같이 derive 됨
    2. Flow Matching objective는 $p_{t}(x)$에 해당하는 target vector field $u_{t}(x)$가 주어졌을 때, $\theta$로 parameterize 된 neural network를 사용하여 vector field $v_{t}$를 predict 하여 얻어짐:
       (Eq. 2) $\mathcal{L}_{FM}(\theta)=\mathbb{E}_{t\sim\mathcal{U}[0,1],x\sim p_{t}(x)}\left|\left| v_{t}(x;\theta)-u_{t}(x)\right|\right|^{2}$
  • BUT, $\mathcal{L}_{FM}(\theta)$는 $p_{t},u_{t}$로 인해 intractable 함 
    1. Real data $x_{1}$에 대한 $p_{t},u_{t}$ conditioning은 generative model training을 위한 indetical gradient $\theta$를 제공함
       - 즉, Conditional Flow Matching obejctive $\mathcal{L}_{CFM}(\theta)$를 활용할 수 있음
    2. 따라서 Optimal Transport Conditional Path를 채택하여 mean $\mu_{t}(x)=tx_{1}$, standard deviation $\sigma_{t}(x)=1-(1-\sigma_{\min})t$이 time에 따라 linearly change 한다고 가정할 수 있음
    3. 그러면 sufficiently small $\sigma_{\min}=\text{1e-5}$에 대해 tractable $p_{t}(x|x_{1}) =\mathcal{N}(x|\mu_{t}(x_{1}),\sigma_{t}(x_{1})^{2}I), u_{t}(x|x_{1})=\frac{(x_{1}-(1-\sigma_{\min})x)}{(1-(1-\sigma_{\min})t)}$를 얻을 수 있음
       - 이때 $p_{1}(x|x_{1})$는 $x_{1}$을 center로 가짐
  • 결과적으로 reparameterization을 적용하여 얻어지는 Conditional Flow Matching obejctive는:
    (Eq. 3) $\mathcal{L}_{CFM}(\theta)=\mathbb{E}_{t,q(x_{1}),p_{0}(x_{0})}\left| \left| v_{t}\left(\psi_{t}(x_{0});\theta\right)-\left(x_{1}-(1-\sigma_{\min})x_{0}\right)\right| \right|^{2}$
    - $\psi_{t}(x_{0})=\sigma_{t}(x_{1})x_{0}+\mu_{t}(x_{1})$, $t$ : $[0,1]$에서 uniformly sampling 됨

SpeechFlow

- Generative Pre-Training of SpeechFlow with Unlabled Speech

• 논문은 VoiceBox를 따라 Flow Matching을 사용하여 unlabled speech로 generative model을 pre-training 함
  • 즉, $L$-frame, $d$-dimensional mel-spectrogram을 가지는 acoustic feature $x\in\mathbb{R}^{d\times L}$에 대한 $q(x)$ modeling을 고려함
    - 이때 simple prior $p_{0}$가 normal distribution이라고 가정함
  • 특히 generative model은 unsupervised/self-supervised이므로 Flow Matching model은 pure speech 만으로 training 될 수 있음 
• Masked Audio Condition
  • Wav2Vec 2.0, HuBERT와 같은 masked prediction을 활용한 self-supervised speech representation leraning을 따라, training 중에 $p_{\text{cond}}$의 chance를 가지는 partially masked target audio $x_{\text{mask}}$로 $v_{t}$를 additionally conditioning 함
    1. 이는 model이 fully masked $x_{\text{mask}}$를 receive 할 chance가 $1-p_{\text{cond}}$인 것으로 볼 수 있음
    2. Masked condition $x_{\text{mask}}$는 $l_{\text{mask}}$의 minimum masking span length로 masked frame의 $n_{\text{mask}}$를 randomly select하여 얻어짐
  • 위 과정을 통해 conditional generative model을 얻을 수 있지만, $x_{\text{mask}}$는 unlabled speech $x_{1}$에서 directly derive 되므로 SpeechFlow는 self-supervised로 동작할 수 있음
    - 특히 condition이 없는 vanilla flow matching model은 $p_{\text{cond}}<1$인 pre-training stage에서 계속 사용 가능함
  • 이러한 auxiliary condition은 time step $t$에 관계없이 $v_{t}$를 predict 하기 위한 more context를 model에 제공함
    - 이때 pre-training stage에 auxiliary condition을 도입하면 다양한 task에 대해 intuitively fine-tuning 할 수 있음
• Objective
  • Masked feature $x_{\text{mask}}$에 대한 predicted time-dependent vector field $v_{t}$ conditioning을 통해, SpeechFlow의 generative pre-training obejective는 (Eq. 3)을 다음과 같이 modify 하여 얻어짐:
    (Eq. 4) $\mathbb{E}_{t,q(x_{1}),p(x_{0})}\left|\left| v_{t}\left( \psi_{t}(x_{0}),x_{\text{mask}};\theta\right)- \left( x_{1}-(1-\sigma_{\min})x_{0}\right)\right|\right|^{2}$
  • 구조적으로는 learnable parameter $\theta$를 가지는 Transformer encoder를 사용하여 vector field $v_{t}$를 predict 함
    1. 여기서 masked input $x_{\text{mask}}$는 frequency axis를 따라 $\phi_{t}(x_{0})$로 concatenate 된 다음, model dimension $d_{\theta}$와 match 되도록 project 됨
       - 그리고 time step $t$의 Sinusoidal Positional Encoding을 append 하여 $\mathbb{R}^{d_{\theta}\times (L+1)}$의 model input을 얻음
    2. 결과적으로 model output은 predicted vector field $v_{t}\in\mathbb{R}^{d\times L}$이 됨

(좌) Model Architecture (우) Task 별 Condition

- Supervised Fine-Tuning SpeechFlow on Different Taks

• Task-Specific Condition
  • Pre-trained SpeechFlow를 사용하면 $p_{1}(x)$에서 new data를 sampling 할 수 있지만, speech application을 위해서는 어느 정도의 control이 필요함
  • 따라서 논문은 audio $x_{1}$의 task-specific condition $y\in\mathbb{R}^{d_{y}\times L_{y}}$를 사용하여 controllable generation을 위한 fine-tuning을 수행함
    - 이때 논문은 $y$와 $x_{1}$이 align 되는 $L_{y}=L$에 focus 함 
• Objective
  • Pre-training stage를 따라 fine-tuning objective는 masked condition $x_{\text{mask}}$를 task-specific condition $y$의 pre-training으로 swapping 하여 derive 됨:
    (Eq. 5) $\mathbb{E}_{t,q(x_{1}),p(x_{0})}\left|\left| v_{t}\left( \psi_{t}(x_{0}),y;\theta\right)-\left( x_{1}-(1-\sigma_{\min})x_{0}\right)\right|\right|^{2}$
  • Fine-tuning에서는 pre-training stage의 $\theta$를 reuse 하는 것이 중요함
• Inference
  • Training 이후 speech generation은 먼저 simple prior $p_{0}(x)$에서 $x_{0}$를 sampling 하고,
  • ODE solver를 통해 $d\phi_{t}(x_{0})/dt=v_{t}(\phi_{t}(x_{0})), \, \phi_{0}(x_{0})=x_{0}$가 주어졌을 때 $\phi_{1}(x_{0})$를 solve 함
  • 이후 mel-spectrogram $x_{1}$의 time-domain으로부터 audible speech를 생성함

3. Experiments

- Settings

• Dataset
  - Speech Enhancement : VoiceBank-Demand, WSJ0-CHiMe3, DNS2020
  - Speech Separation : WHAM, LibriMix
  - Zero-Shot Text-to-Speech : LibriLight
• Comparisons :
  - Speech Enhancement : Conv-TasNet, MetricGAN+, SGMSE+
  - Speech Separation : SepFormer 
  - Zero-Shot Text-to-Speech : VALL-E, VoiceBox, YourTTS

Training Details

- Results

• Speech Enhancement
  • 전체적으로 SpeechFlow의 성능이 가장 뛰어남

Speech Enhancement 성능

• Speech Separation
  • Separation task에 대해서도 SpeechFlow가 가장 우수한 성능을 달성함

Speech Separation 성능

• Zero-Shot Text-to-Speech
  • Text-to-Speech에서도 SpeechFlow는 우수한 성능을 보임

Zero-Shot Text-to-Speech 성능

• Multi-Task Fine-Tuning
  • Multi-task SpeechFlow를 활용해도 single-task 수준의 성능을 달성할 수 있음

Single-Task vs. Multi-Task Model