[Paper 리뷰] ZipCodec: Simple and Pretrained-Model-Free Speech Tokenizer via Flow-Matching

ZipCodec: Simple and Pretrained-Model-Free Speech Tokenizer via Flow-Matching

---

• Speech tokenizer는 massive pre-trained model이나 complex multi-stage training pipeline에 의존함
• ZipCodec
  • Zipformer-based flow-matching decoder에 acoustic reconstruction을 offloading
  • Encoder Consistency Regularization loss를 도입하여 encoder의 semantic capacity를 향상
• 논문 (INTERSPEECH 2026) : Paper Link

---

1. Introduction

• Speech tokenizer는 continuous audio waveform을 compact discrete unit으로 convert 함
  • BUT, 대부분의 neural speech codec은 acoustic-level reconstruction을 목표로 하므로 semantic richness 측면에서 한계가 있음
  • 이때 semantic distillation/guidance나 HuBERT, Whisper와 같은 pre-trained model을 활용하면 semantic modeling을 향상할 수 있음
    - BUT, 해당 방식은 extra pre-trained model과 complex multi-stage pipeline에 의존적임

-> 그래서 single-stage, pretrained-model-free speech tokenizer인 ZipCodec을 제안

 

• ZipCodec
  • Flow-Matching (FM) 기반의 ZipFormer backbone을 채택하여 lightweight pipeline을 구성
  • Encoder Consistency Regularization (ECR) loss를 활용해 robust semantic extraction을 보장

< Overall of ZipCodec >

• FM backbone과 ECR loss를 활용하여 semantic capacity를 향상한 lightweight speech tokenizer
• 결과적으로 기존보다 우수한 성능을 달성

2. Method

- Overview

• ZipCodec은 quantization bottleneck을 포함한 flow-matching-based diffusion autoencoder architecture를 따름
  • 구조적으로는 speech encoder, Finite Scalar Quantization (FSQ) module, FM decoder, down/upsampling module로 구성됨
  • 논문은 encoder/FM decoder에 ZipFormer backbone을 적용하고, mel-spectrogram을 input/output domain으로 사용하여 Vocos vocoder를 통해 waveform reconstruction을 수행함
    1. Speech encoding 시 mel-spectrogram $\mathbf{x}_{1}\in\mathbb{R}^{T\times d_{mel}}$은 convolution module stack을 통해 $4\times$ downsample 됨:
       (Eq. 1) $\mathbf{h}=\text{CNNs}(\mathbf{x}_{1},\theta_{CNNs})$
       - $\mathbf{h}\in\mathbb{R}^{T'\times d_{enc}}$, $T'=T/4$
    2. 이후 encoder는 $\mathbf{h}$를 hidden embedding vector $\mathbf{e}\in\mathbb{R}^{T'\times d_{enc}}$로 preporcess 함:
       (Eq. 2) $\mathbf{e}=\text{Enc}(\mathbf{h},\theta_{enc})$
    3. FSQ module은 $\mathbf{e}$를 discrete speech token으로 quantize 함:
       (Eq. 3) $\mathbf{z}=\text{round}\left(\mathcal{T}\left(\text{tanh}\left( \mathbf{eW}^{\top}_{in}+\mathbf{b}_{in}\right)\right)\right)$
       - $\mathbf{z}=[\mathbf{z}_{1},...,\mathbf{z}_{T'}]$, $\mathbf{z}_{j}\in\prod_{i=1}^{d_{FSQ}}\mathcal{L}_{i}$, $\mathcal{L}_{i}=\{0,1,2,...,L_{i}-1\}$
       - $L_{i}$ : FSQ에서 $i$-th channel의 level 수, $\mathbf{W}_{in}\in\mathbb{R}^{d_{FSQ}\times d_{enc}},\mathbf{b}_{in}\in\mathbb{R}^{d_{FSQ}}$ : FSQ 이전의 linear down projection layer parameter, $\mathcal{T}$ : fixed linear mapping $\mathcal{T}:[-1,1]\rightarrow [0,L_{i}]$
    4. Continuous value는 다음과 같이 recover 됨:
       (Eq. 4) $\mathbf{e}'=\mathcal{T}^{-1}(\mathbf{z})\mathbf{W}_{out}^{\top}+\mathbf{b}_{out}$
  • Decoding 시 FM decoder는 $\mathbf{e}'$을 condition으로 하여 Conditional Flow Matching (CFM) process를 통해 mel-spectrogram을 recover 함
    - 이때 original length $T$에 match 하기 위해 $\mathbf{e}'$은 time axis에 대해 $4\times$ repeat 됨
  • CFM은 simple prior distribution을 complex data distribution으로 transport 하는 vector field를 학습함
    1. $\mathbf{x}_{0}\in\mathbb{R}^{T\times d_{mel}}$이고 $\text{vec}(\mathbf{x}_{0})\sim\mathcal{N}(0,I)$를 standard Gaussian noise라고 하자
    2. Optimal Transport (OT) formulation을 통해 $t\in[0,1]$에 대해 probability density path $\mathbf{x}_{t}=(1-t)\mathbf{x}_{0}+t\mathbf{x}_{1}$을 얻을 수 있음
       - 해당 path는 constant target vector field $\mathbf{u}_{t}(\mathbf{x}_{t}|\mathbf{x}_{1})=\mathbf{x}_{1}-\mathbf{x}_{0}$로 drive 됨
  • Decoder network $\text{Dec}(\cdot\cdot\cdot,\theta_{dec})$은 $\mathbf{e}'$에 condition 되어 해당 field를 approximate 하도록 training 되고, Mean Squared Error를 통해 optimize 됨:
    (Eq. 5) $\mathbf{v}_{t}=\text{Dec}(\mathbf{x}_{t},t,\mathbf{e}',\theta_{dec})$
    (Eq. 6) $\mathcal{L}_{CFM}(\Theta)=\mathbb{E}_{t,\mathbf{x}_{0},\mathbf{x}_{1}}\left[\left|\left| \mathbf{v}_{t}-\mathbf{u}_{t} \right|\right|^{2}\right]$
    - $\Theta=\{\theta_{enc}, \theta_{dec},\theta_{CNNs},...\}$ : ZipCodec model의 모든 parameter
  • Generated speech와 conditioning token 간의 alignment를 strengthen 하기 위해 논문은 Classfier-Free Guidance (CFG)를 도입함
    1. Training 시 condition은 probability $p_{drop}$으로 randomly zero-replace 되고 추론 시에는 predicted vector field를 extrapolate 함:
       (Eq. 7) $\tilde{\mathbf{v}}_{t}=\mathbf{v}_{t}+\gamma\left(\mathbf{v}_{t}-\mathbf{v}_{t}^{uncond}\right)$
       - $\gamma$ : guidance scale
    2. 여기서 $\mathbf{v}_{t}^{uncond}$는:
       (Eq. 8) $\mathbf{v}_{t}^{uncond}=\text{Dec}(\mathbf{x}_{t},t,0,\theta_{dec})$
  • Mel-spectrogram은 $\mathbf{x}_{0}$에서 start 하여 Ordinary Differential Equation (ODE) solver를 통해 reconstruct 됨:
    (Eq. 9) $\mathbf{x}_{t_{n+1}}=\mathbf{x}_{t_{n}}+(t_{n+1}-t_{n})\tilde{\mathbf{v}}_{t_{n}}$
    - $0=t_{0}\leq ...\leq t_{n}\leq t_{n+1}\leq ...\leq t_{N}=1$, $N$ : Number of Function Evaluation (NFE)

Overview

- Encoder Consistency Regularization

• 논문은 Encoder Consistency Regularization (ECR)을 사용하여 semantic capacity를 향상함
  • Clean speech signal $\mathbf{y}\in\mathbb{R}^{T_{wav}}$가 주어지면 Gaussian noise를 add 하여 $\mathbf{y}^{aug}=\mathbf{y}+\beta\epsilon$과 같이 augment 함
    - $\epsilon\sim\mathcal{N}(0,I)$
  • Clean speech $\mathbf{y}$, augmented speech $\mathbf{y}^{aug}$는 mel-spectrogram으로 convert 되어 encoder에 전달되고, 논문은 output $\mathbf{e},\mathbf{e}^{aug}$에 대해 consistency regularization loss를 적용함:
    (Eq. 10) $\mathcal{L}_{ECR}(\mathbf{e},\mathbf{e}^{aug})=\frac{1}{2}\left(\text{sim}\left( \text{sg}[\mathbf{e}],\mathbf{e}^{aug}\right)+\text{sim}\left( \mathbf{e},\text{sg}[\mathbf{e}^{aug}]\right)\right)$
    - $\text{sg}[\cdot]$ : stop gradient operation
  • $\text{sim}$은 cosine similarity loss로써:
    (Eq. 11) $\text{sim}(\mathbf{a},\mathbf{b})=1-\frac{\mathbf{a}^{\top}\mathbf{b}}{||\mathbf{a}||_{2}|| \mathbf{b}||_{2}}$
  • 결과적으로 total training loss는 다음과 같이 얻어짐:
    (Eq. 12) $\mathcal{L}=\frac{1}{2}(\mathcal{L}_{CFM}+\mathcal{L}_{CFM}^{aug})+\lambda\mathcal{L}_{ECR}$
    - $\mathcal{L}_{CFM}^{aug}$ : augmented branch의 CFM loss, $\lambda$ : weight

- Flow-Distillation for Fast Inference

• 논문은 ZipVoice를 따라 NFE를 reduce 하고 model이 CFG를 explicitly learn 하도록 ZipCodec-distill을 구성함
  • 이를 위해 FM decoder를 modify 하여 guidance scale $\gamma$를 additional embedding으로 incorporate 함
  • 이때 trained frozen ZipCodec은 teacher model로 사용되고, teacher의 2-step ODE inference result는 student model의 target이 됨:
    (Eq. 13) $\mathbf{x}_{t_{1}}=(1-t_{1})\mathbf{x}_{0}+t_{1}\mathbf{x}_{1},$
    $\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\, \mathbf{x}_{t_{2}}^{(t)}=\mathbf{x}_{t_{1}}+(t_{2}-t_{1}) \tilde{\mathbf{v}}\left( \mathbf{x}_{t_{1}},\mathbf{e}',t_{1},\gamma,\Theta^{(t)}\right),$
    $\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\,\mathbf{x}_{t_{3}}^{(t)}=\mathbf{x}_{t_{2}}^{(t)}+(t_{3}-t_{2})\tilde{\mathbf{v}}\left( \mathbf{x}_{t_{2}}^{(t)},\mathbf{e}',t_{2},\gamma,\Theta^{(t)}\right)$
    - $\tilde{\mathbf{v}}(\cdot\cdot\cdot)$ : CFG를 포함한 1-step ODE inference, $\mathbf{x}_{t_{2}}^{(t)},\mathbf{x}_{t_{3}}^{(t)}$ : parameter $\Theta^{(t)}$를 가지는 teacher model에 대한 inference result
    - $t_{1}\sim\mathcal{U}(0,1)$은 uniform distribution에서 sample 되고 $t_{2}\sim \mathcal{U}(t_{1},1),t_{3}\sim \mathcal{U}(t_{2},1)$은 $t_{n+1}-t_{n}\leq 0.3$을 만족하도록 sequentially sample 됨
    
  • 결과적으로 distillation loss function은:
    (Eq. 14) $\mathbf{v}_{t_{1}}^{(t)}=\frac{\mathbf{x}_{t_{3}}^{(t)}-\mathbf{x}_{t_{1}}}{t_{3}-t_{1}}$
    (Eq. 15) $\mathbf{v}_{t_{1}}^{(s)}=\text{Dec}^{(s)}\left(\mathbf{x}_{t_{1}},t_{1},\gamma,\mathbf{e}',\theta^{(s)}_{dec}\right)$
    (Eq. 16) $\mathcal{L}_{distill}(\Theta^{(s)})=\mathbb{E}_{\mathbf{x}_{0},\mathbf{x}_{1},\gamma, t_{1},t_{2},t_{3}}\left[\left|\left| \mathbf{v}_{t_{1}}^{(s)}-\mathbf{v}_{t_{1}}^{(t)}\right|\right|^{2}\right]$
    - $\Theta^{(s)}$는 student model parameter로써 teacher model에 의해 initialize 되고, student model의 encoder parameter $\theta_{enc}^{(s)}$는 training 시 freeze 됨 

3. Experiments

- Settings

• Dataset : Emilia, LibriTTS
• Comparisons : EnCodec, SpeechTokenizer, X-Codec, XY-Tokenizer, BiCodec

- Results

• 전체적으로 ZipCodec의 성능이 가장 우수함

Model 성능 비교

• ECR loss weight $\lambda=0.25$일 때 최적의 성능을 달성함

ECR Loss Weight $\lambda$ 별 성능

• Downstream TTS task에서도 가장 우수한 성능을 보임

Downstream TTS 성능