[Paper 리뷰] FlowDec: A Flow-Based Full-Band General Audio Codec with High Perceptual Quality

FlowDec: A Flow-Based Full-Band General Audio Codec with High Perceptual Quality

---

• Lower bitrate에서도 동작하는 general full-band audio codec이 필요함
• FlowDec
  • Non-adversarial codec training과 conditional flow matching에 기반한 stochastic postfilter를 활용
  • Fine-tuning이나 distillation 없이 required postfilter evaluation을 절감
• 논문 (ICLR 2025) : Paper Link

---

1. Introduction

• Audio codec은 audio waveform을 compact, quantized representation으로 compress 하고, 해당 representation을 기반으로 audio waveform을 faithfully reconstruct 하는 것을 목표로 함
  • BUT, 기존 codec은 ad-hoc design과 extensive manual effort가 필요하므로 12kbit/s 이하의 lower bitrate에서 high-fidelity audio coding을 위한 end-to-end optimization이 어려움
    - 한편으로 SoundStream, DAC, AudioDec, EnCodec과 같은 End-to-End (E2E) Neural Codec은 8kbit의 lower bitrate에서도 우수한 audio quality를 달성함
  • 특히 ScoreDec과 같이 score-based diffusion이나 flow-based generative model을 도입하여 reconstruction quality를 개선할 수 있음
    - BUT, ScoreDec은 high-bitrate인 24kbit/s에서만 동작 가능하고 DNN evaluation으로 인한 Real-Time-Factor (RTF) 저하 문제가 있음

-> 그래서 lower bitrate에서도 효과적으로 동작하는 flow-based neural codec인 FlowDec을 제안

 

• FlowDec
  • 기존 score-based method를 Conditional Flow Matching (CFM) method로 확장
  • Fine-tuning, distillation 없이 DNN evaluation을 줄여 lower-bitrate, general full-band audio coding을 지원

< Overall of FlowDec >

• CFM에 기반한 full-band neural auido codec
• 결과적으로 RTF를 크게 개선하고 lower-bitrate에서 high-fidelity perceptual quality를 달성

2. Method

• 논문은 code $c:=E(x^{\ast })$가 주어진 clean audio $x^{\ast }\in {\mathbb{R}}^L$의 estimate $\hat{x}\in {\mathbb{R}}^L$을 reconstruct 하는 stochastic inference problem을 고려함 
  • 이때 model은 distribution에서 sampling을 통해 clean audio estimate $\hat{x}$를 제공하는 것을 목표로 함:
    (Eq. 1) $\hat{x}\sim p_{\text{data}}(\hat{x}|c),c=E(x^{\ast })\in {\mathbb{Z}}^{\ell },\ell \ll L$
    - $p_{\text{data}}(\cdot |c)$ : code $c$가 주어졌을 때 clean audio의 conditional distribution 
  • $x^{\ast }\in {\mathbb{R}}^L$을 lower-dimensional discrete representation $c$로 mapping 하는 모든 encoder $E$는 many-to-one mapping이므로 multiple $x^{\ast }$은 same code $c$를 가짐
    1. 따라서 $D$가 one-to-one mapping인 경우 ideal property $D(E(x^{\ast }))=x^{\ast }$를 fullfilling 하는 것이 formally impossible 함 
    2. 대신 (Eq. 2)를 minimizing 하여 $D$를 optimal estimator로 구성할 수 있음:
       (Eq. 2) $\underset{D}{min}{\mathbb{E}}_{x^{\ast }}[\text{dist}(D(E(x^{\ast })),x^{\ast })]$
       - $\text{dist}$ : $L^2,L^1$ distance와 같은 pairwise distance
    3. BUT, 해당 방식으로 training 되는 경우 domain-specific loss가 있더라도 perceptually pleasing signal을 생성하지 못함
       
  • 이를 해결하기 위해 SoundStream, EnCodec, DAC 등은 adversarial training loss를 도입하여 decoded signal distribution을 natural signal에 close 하도록 유도함
    - BUT, adversarial training은 interpretability가 부족하고 $p(\hat{x}),p(x^{\ast })$ 간의 distance를 properly minimize 하지 못함
  • 따라서 논문은 ScoreDec과 같이 $D$를 one-to-many mapping으로 구성함
    1. 즉, FlowDec은 deterministic pre-trained initial decoder $D_0$와 stochastic postfilter $\mathrm{\Omega }$를 combining 한 stochastic decoder $D_s(c)=\mathrm{\Omega }(D_0(c))$의 형태로 구성됨 
    2. 여기서 $y:=D_0(c)$라고 하면 $\mathrm{\Omega }$는 learned distribution $p_{\mathrm{\Omega }}(\cdot |y)$로부터 conditional sample $\hat{x}\sim p_{\mathrm{\Omega }}(\hat{x}|y)$를 생성함
       - 해당 sample은 statistical divergence $\mathcal{D}$를 mimimize 하여 intractable distribution $p_{\text{data}}(\cdot |y)$를 approximate 함

Overall of FlowDec

- Flow Matching

• Flow Matching은 tractable dsitribution $q_0(x_0)$에서 intractable data distribution $q_1(x_1)=p_{\text{data}}$로 sample을 transport 하는 model을 학습하는 것을 목표로 함
  • 이때 Flow Matching은 sample $x_0\sim q_0$에서 시작하여 다음의 Ordinary Differential Equation (ODE)를 solve 함:
    (Eq. 4) $\frac{d}{dt}{\varphi }_t(x)=u_t({\varphi }_t(x)),{\varphi }_0(x)=x_0$
  • 여기서 ${\varphi }_t:[0,1]\times {\mathbb{R}}^N\to {\mathbb{R}}^N$를 flow라 하고, time-dependent vector field $u_t:[0,1]\times {\mathbb{R}}^N\to {\mathbb{R}}^N$를 통해 $p_{t=0}=q_0,p_{t=1}=q_1$인 probability density path $p_t:{\mathbb{R}}^N\to {\mathbb{R}}_{>0}$을 생성함
  • 그러면 $v_{\theta }$는 다음의 CFM training loss를 통해 학습됨:
    (Eq. 5) ${\mathcal{L}}_{\text{CFM}}:={\mathbb{E}}_{x,t,p_t(x|x_1)}[||v_{\theta }(x,t)-u_t(x|x_1)|{|}_2^2]$
    - $x_1\sim q_1$
    - 특히 conditional (Eq. 5)는 intractable unconditional flow matching objective와 동일한 gradient를 가지고 correct unconditional probability path $p_t(x)$와 flow field $u_t(x)$를 marginalize 함

- Joint Flow Matching for Signal Enhancement

• Original flow matching formulation에서 $x_0,x_1$은 zero-mean Gaussian $q_0=\mathcal{N}(0,{\sigma }^{2}I)$에서 independently sampling 됨
  • $q_0$가 standard Gaussian인 경우 conditional path $p_t(x|x_1)$은 $q_0$에서 $q_1$으로의 Optimal Transport (OT)를 fulfill 하지만 modeled marginal probabilitiy path $p_t(x)$는 OT를 fulfill 하지 않음
    1. 결과적으로 learned marginal flow field $v_{\theta }$에서 high-variance training과 lower-straightness가 발생하므로 inefficient inference와 suboptimal sample quality가 나타남
    2. 이를 해결하기 위해 각 training batch $\{(x_{b,0},x_{b,1}){\}}_{b=1}^B$에서 pairing을 reorder 하고 각 batch에서 OT algorithm을 통해 optimal coupling을 결정하는 per-batch approximation을 도입할 수 있음
       - 즉, $(x_0,x_1)\sim q(x_0,x_1)$을 independently sample 하지 않고 jointly sampling 함
       - 특히 $(x_0,x_1)$를 jointly sampling 하면 OT solver나 extra computation이 필요하지 않음
  • 이때 initial estimate $y=D_0(c)=D_0(E(x^{\ast }))$에 access 할 수 있으므로, 논문은 다음의 probability path를 choice 함:
    (Eq. 6) $p_t(x_t|x_1,y)=\mathcal{N}(x_t;{\mu }_t,{\sigma }_t):=\mathcal{N}(x_t;y+t(x_1-y),(1-t{)}^2{\mathrm{\Sigma }}_y)$
    - ${\mathrm{\Sigma }}_y=\text{diag}({\sigma }_y^2)$ : diagonal covariance matrix
  • 해당 probability path는 $y,x_1$ 간의 linear interpolation에 해당하고, noise는 ${\sigma }_y$에서 $0$으로 linearly decrease 하므로 $x_0,x_1$ 간에 $y$를 통한 coupling이 나타남 
    1. 즉, $q_0(x_0|x_1,y)=\mathcal{N}(x;y,{\mathrm{\Sigma }}_y)$인 경우 $x_0$의 mean은 $0$에서 $y$로 shift 됨
    2. 그러면 marginalized $q_0(x_0)$는 아래 그림과 같이 variance가 ${\sigma }_y^2$이고 training data에 centered 된 Gaussian mixture로 볼 수 있음 
    3. 이때 ${\sigma }_y$가 well-chosen 되어 Gaussian이 negligible overlap 되면 per-batch coupling을 통해 optimal 하다고 가정할 수 있으므로, mini-batch OT가 필요하지 않음 
       - 따라서 ${\sigma }_y$의 choice는 output quality에 큰 영향을 미침

$q_0(x_0)$ vs. $q_0(x_0❘x_1)$

• 결과적으로 conditional $u_t$는,
  • Flow Matching을 따라 다음과 같이 derive 됨:
    (Eq. 7) $u_t(x|x_1,y)=\frac{x_1-x_t}{1-t}$
  • $x_t$는 $x_0$를 통해 나타낼 수 있으므로:
    (Eq. 8) $x_t=t{x}_1+(1-t)x_0,x_0\sim \mathcal{N}(x_0;y,{\mathrm{\Sigma }}_y)$
    (Eq. 9) $=t{x}_1+(1-t)y+(1-t){\sigma }_tϵ,ϵ\sim \mathcal{N}(0,I)$
    (Eq. 10) $x_0=y+{\sigma }_yϵ,ϵ\sim \mathcal{N}(0,I)$
  • (Eq. 6)을 따라 $x_1=x^{\ast }$을 대입하면 simple joint flow matching loss를 얻을 수 있음:
    (Eq. 11) ${\mathcal{L}}_{\text{JFM}}:={\mathbb{E}}_{t\sim \mathcal{U}(0,1),(x^{\ast },y)\sim \mathfrak{D},ϵ\sim \mathcal{N}(0,I),x_t\sim p_t(x_t|x_0)}\left[{\left||v_{\theta }(x_t,t,y)-(\underset{=x_1}{\underbrace{x^{\ast }}}-\underset{=x_0}{\underbrace{(y+{\sigma }_yϵ)}})|\right|}_2^2\right]$
    - $\mathfrak{D}$ : training dataset
  • 해당 loss는 $x_0,x_1$에 대해 reparameterize 함으로써 (Eq. 7)의 $t\approx 1$ 주변의 numerical stability를 제거함
    1. 특히 ${\sigma }_y>0$을 choice 하면 flow field가 contractive mapping이 되도록 force 할 수 있음
       - 이를 통해 inference를 위한 ODE는 numerically stable 하고 locally converge 됨
    2. $p_t$에 대한 choice는 trajectory가 $x^{\ast }$에 exactly reach 하도록 하여 SGMSE를 개선함
       - 기존 SGMSE는 correct $q_0$를 modeling 하지 않으므로 fail 할 수 있음
  • 추가적으로 논문은 multiple hyperparameter를 가지는 Stochastic Differential Equation (SDE) 대신, 하나의 hyperparameter ${\sigma }_y$ 만을 사용하는 data-based heuristic을 도입함:
    (Eq. 12) ${\sigma }_y=\frac{1}{3}\sqrt{Q(|{\mathbf{X}}^{\ast }-\mathbf{Y}{|}^2,0.997)}$
    - $Q$ : quantile operation
  • 한편으로 independent CFM formulation을 사용하여 constant ${\sigma }_t=\sigma$를 얻고 sampled noise에 대해 target flow field $u_t$가 independent 하도록 할 수 있음
    1. BUT, 해당 방식은 ${\sigma }_1=\sigma >0$ 이므로 아래 그림과 같이 non-contractive flow field와 estimate에 residual noise가 나타남
    2. 반면 FlowDec은 ${\sigma }_1=0$이므로 postfiltering task에서 더 나은 quality를 달성할 수 있음

$t=0.7$에서 Flow Field

• FlowDec은 $x^{\ast },y$를 invertible feature extractor $\mathrm{\Phi }$의 feature representation ${\mathbf{X}}^{\ast },\mathbf{Y}$로 replace 하여 feature domain의 flow를 학습함
  • 이때 $\mathrm{\Phi }$는 compression exponent $\alpha =0.3$의 amplitude-compressed complex STFT를 채택함
    - 추가적으로 input에서 channel-wise concatenate를 통해 $\mathbf{Y}$에 $v_{\theta }$의 conditioning을 제공함
  • Training 이후 flow model $v_{\theta }$는 (Eq. 4)의 ODE와 함께 conditional distribution $p_{\mathrm{\Omega }}({\mathbf{X}}^{\ast }|\mathbf{Y})$를 modeling 함
    1. 이때 clean feature estimate $\hat{\mathbf{X}}\sim p_{\mathrm{\Omega }}$를 생성하기 위해, initial state (latent) ${\mathbf{X}}_0\sim q_0({\mathbf{X}}_0|\mathbf{Y})$를 sample 함
    2. 이후 $t=0$에서 $t=1$ 까지 $v_{\theta }$를 사용하여 numerical ODE solver를 통해 flow (Eq. 4)를 solve 하여 ${\hat{\mathbf{X}}}_1$을 구함
       - 논문은 3-step Mid-point Solver ($\text{NFE}=6$)를 사용
    3. 최종적으로 feature extractor $\mathrm{\Phi }$의 inverse를 사용하여 waveform estimate $\hat{x}={\mathrm{\Phi }}^{-1}({\hat{\mathbf{X}}}_1)$을 생성함

- Non-Adversarial Codec Training

• Effective phase loss 없이 spectral loss 만으로 training 된 NAR audio generative model은 unsynchronized phase로 인한 buzzy noise가 발생함 
  • 이때 adversarial training을 도입하면 해당 문제를 해결하고 natural-sounding audio를 얻을 수 있음
    - BUT, unstable training, mode-collapse, handcrafted multi-discriminator design 등의 문제가 발생함
  • 따라서 논문은 adversarial training을 제거하는 대신 generative postfilter를 도입함
    1. 즉, adversarial loss 없이 deterministic neural codec을 initial decoder $D_0$로 training 하고 stochastic postfiler $\mathrm{\Omega }$를 통해 output audio와 clean audio의 distribution을 matching 함  
    2. 구조적으로는 DAC와 같은 neural codec을 $D_0$로 채택하고 adversarial loss term과 관련된 모든 component를 제거하여 사용함

- Underlying Codec: Improved Non-Adversarial DAC

• FlowDec의 stochastic postfilter는 ScoreDec과 같이 any underlying codec에 대해 training 되어 waveform estimate를 향상함
  • 이때 다른 sampling rate, bitrate에 대한 adaptibility가 우수한 DAC를 underlying codec의 basis로 채택함
    - 추가적으로 adversarial loss를 제거하고 아래 표와 같이 configuration을 modify 함
  • 한편으로 해당 non-adversarial loss를 training 할 때 $-30\text{dB}$의 Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) value가 발생하는 경우가 있음
    1. 이는 low-frequency ($\le 2\text{kHz}$)가 badly modeling 되기 때문으로, 논문은 Multiscale Constant-Q Transform (CQT) loss를 도입하여 해당 문제를 해결함
    2. 추가적으로 DAC의 multiscale Mel loss와 같이 amplitude와 log-amplitude의 difference를 모두 사용하고, SI-SDR과 phase error를 반영하는 $L^1$ waveform-domain loss를 추가함

Underlying Codec Variant

- Frequency-Dependent Noise Levels

• ${\sigma }_y$ choice는 output quality에 큰 영향을 미침
  • 이때 single scalar ${\sigma }_y$를 사용하면 added Gausian noise가 high-frequency를 dominate 할 때 over-smoothing이 나타날 수 있음
  • 따라서 논문은 각 STFT frequency band에 대해 (Eq. 12)의 heuristic quantile calculation을 independently performing 하여 frequency-dependent curve ${\sigma }_y(f)$를 calculate 함
    

3. Experiments

- Settings

• Dataset : 아래 표 참조
• Comparisons : DAC, EnCodec, ScoreDec

Dataset

- Results

• 전체적으로 FlowDec이 가장 우수한 성능을 보임

Objective Evaluation

• Perception-Distortion trade-off 측면에서도 FlowDec은 DAC 보다 더 robust 함

Perception-Distortion

• ScoreDec과의 비교에서도 FlowDec이 더 뛰어남

FlowDec vs. ScoreDec

• Spectrogram 측면에서도 FlowDec은 더 나은 reconstruction이 가능함

Spectrogram 비교

• Listening Test
  • Subjective evaluation을 위한 listening test parameter는 아래 표와 같이 설정됨

Listening Test Parameters

• Subjective evaluation 측면에서도 FlowDec이 가장 우수함

Subjective Evaluation

• Real-Time-Factor (RTF)
  • $\text{NFE}=6$의 default setting에서 FlowDec-75의 total RTF는 0.2285, FlowDec-25는 0.2235와 같음
  • 즉, FlowDec은 ScoreDec의 1.707 RTF와 비교하여 상당한 RTF 개선을 보임