[Paper 리뷰] TS3-Codec: Transformer-based Simple Streaming Single Codec

TS3-Codec: Transformer-based Simple Streaming Single Codec

---

• 대부분의 neural audio codec은 convolution을 기반으로 함
• TS3-Codec
  • Transformer와 linear layer로만 구성된 simple streaming single codec
  • Convolution layer를 fully eliminate 하여 simplicity와 expressiveness를 향상
• 논문 (INTERSPEECH 2025) : Paper Link

---

1. Introduction

• Neural Audio Codec (NAC)는 audio signal을 discretized code로 compress 하는 것을 목표로 함
  • 특히 NAC는 AudioLM, VALL-E 등의 Speech Language Model (SLM)에서 주로 활용됨
    - BUT, 대부분 convolutional layer에 의존하므로 Transformer의 benefit을 활용하기 어려움
  • 한편으로 SLM에서 NAC를 효과적으로 활용하기 위해서는 다음을 고려해야 함
    1. Streaming : seamless real-time interaction을 위해 low latency로 동작해야 함
    2. Single codebook : Residual Vector Quantization (RVQ)는 SLM에서 additional complexity가 발생하므로 single codebook-based model이 필요함
    3. Low computation : model의 computational demand를 reduce 해야 함
    4. Low token rate : long sequence는 SLM training을 unstable하게 만듦

-> 그래서 위 고려사항을 반영한 Transformer 기반의 simple single-codebook codec인 TS3-Codec을 제안

 

• TS3-Codec
  • Convolution 없이 Transformer layer stack 만으로 구성
  • Hyperparameter tuning과 large computation을 줄여 expressiveness를 향상

< Overall of TS3-Codec >

• Transformer layer를 사용한 simple single-codebook codec
• 결과적으로 기존보다 우수한 성능을 달성

2. Method

- Model Architecture

• TS3-Codec은 encoder, quantizer, decoder로 구성되고, Generative Adversarial Network (GAN) framework를 사용하여 training 됨 
  • Encoder는 2개의 linear layer와 left context에 sliding window attention을 적용하는 Transformer layer로 구성됨 
    1. Signal length $T$에 대해 input audio signal이 $\mathbb{R}^{T}$의 shape를 가진다고 하면, 해당 audio signal은 먼저 $\mathbb{R}^{F\times N}$ shape의 2-dimensional tensor로 reshape 됨
       - $F$ : window size, $N=\frac{T}{F}$ : windowed frame 수
    2. Activation이 없는 2개의 linear layer는 reshaped windowed frame을 $\mathbb{R}^{D\times N}$ shape로 convert 함
    3. 이후 $D$ embedding size를 가진 Transformer layer stack을 통해 $\mathbb{R}^{D\times N}$ shape의 encoder embedding을 output 함
       - 구조적으로 encoder는 $8$ Transformer layer, $16$ attention head, $4096$ feed-forward dimension을 가짐
  • Transformer layer는 self-attention을 위해 left-context에 대해 $16$ 또는 $32$ size를 가진 sliding window를 활용함 
    - 이를 통해 Transformer의 computational complexity가 $N$에 대해 quadratically blow up 하는 것을 방지하고, model의 generalizability를 크게 향상할 수 있음
  • Encoder module 이후에는 factorized VQ layer를 적용하여 encoder output을 low-dimensional space로 discretize 함 
    - 이때 codebook dimension은 $8$ 또는 $16$이고 codebook size는 $65536$ 또는 $131072$
  • Decoder는 encoder의 symmetric structure로써 2개의 linear layer를 가진 Transformer layer stack으로 구성됨
    1. Transformer output은 $\mathbb{R}^{D\times N}$ shape를 가지고, 2개의 linear layer를 통해 $\mathbb{R}^{F\times N}$ shape tensor로 convert 됨
    2. 최종적으로 해당 tensor는 $\mathbb{R}^{T}$ shape의 single-dimensional tensor로 reshape 됨

Overview

- Design Principle

• 논문은 다음의 이점으로 인해 Transformer를 backbone으로 사용함
  • Convolution은 동일한 parameter size의 model에 대해 Transformer 보다 더 많은 computation을 요구함
  • Convolution은 inherent bias를 가짐
    1. Convolution은 서로 다른 time stamp의 all intermediate feature map에 fixed weighted-sum을 적용함
    2. Transformer는 self-attention을 사용하여 각 weight를 dynamically determine 함
       - 추가적으로 positional embedding을 사용하여 서로 다른 time stamp의 서로 다른 feature map에 distinct embedding을 적용할 수 있음
  • Transformer는 model design의 simplicity를 제공함
    - Convolution의 up-/down-sampling을 avoid 할 수 있음

- Training Objective

• TS3-Codec은 GAN framework를 사용하여 training 됨
  • 이때 논문은 다음과 같은 multiple loss를 사용함:
    1. Reconstruction loss : multiple scale의 spectral domain의 $L_{1}$ distance로 compute 되는 Multi-Scale Mel-spectrogram Reconstruction loss
    2. Least-Square GAN loss : HiFi-GAN의 Multi-Period Discriminator (MPD), EnCodec의 Multi-Scale STFT Discriminator (MS-STFT)로 얻어지는 GAN loss
    3. Feature loss : discriminator feature에 대한 $L_{1}$ feature matching loss
    4. VQ loss : quantization before/after feature 간 $L_{1}$ loss
       - 이 과정에서 stop-gradient operation이 적용되고, $0.25$ weight의 commitment loss가 추가됨
       - Reconstruction, GAN, feature loss에 대한 weight는 각각 $15.0, 1.0, 1.0$으로 설정됨
  • 추가적으로 논문은 codebook size에 비례하여 VQ loss를 설정함
    - Codebook size가 $8192$이면 $4.0$, codebook size가 $65536$이면 $32.0$

3. Experiments

- Settings

• Dataset : LibriLight
• Comparisons : EnCodec, DAC, SpeechTokenizer, WavTokenizer, Mimi, BigCodec

Comparisons

- Results

• 전체적으로 TS3-Codec의 성능이 가장 우수함

Model 성능 비교

• 600bps 이하에서도 우수한 성능을 보임

600bps 이하에서의 성능

• 비슷한 budget을 가지는 BigCodec과 비교하여 TS3-Codec은 더 뛰어난 성능을 보임

BigCodec과의 비교