Wav2Vec 2.0
🍋 Key Takeaways
- 1️⃣ Automatically learn discrete speech units for transformer using codebooks, Gumbel-Softmax
- 2️⃣ Apply entropy concept to diversity loss (to use codewords equally)
- 3️⃣ PT: SSL(Contrastive Learning) - FT: SL
Overview
- 1️⃣ multi-layer convolutional feature encoder $f : x \rightarrow Z$
- $X$: input raw audio
- $Z = z_1, z_2, …, z_T$: latent speech representations
- 2️⃣ transformer $g : Z \rightarrow C$
- $c_1, c_2, …, c_T$: capturing information from entire sequence(contextualized representations)
- 3️⃣ Quantization module: multiple codebooks과 codewords
- 2️⃣, 3️⃣은 순차적이 아닌, 병렬적으로 이루어진다.
2. Model
Encoder NW: Feature encoder
- Input: $X$: raw audio
- Output: $Z = z_1, z_2, …, z_T$: latent speech representations
- Action: model이 음성 데이터를 효율적으로 처리할 수 있게 변환한다.
Layer Normalization과 GELU를 포함한 Temporal CNN block들로 이루어져 있다. 먼저 raw audio $X$는 평균 0, 분산 1로 정규화된다. 변환 과정에서 데이터의 dimensionality가 감소한다. 실험에서는 7개의 CNN block, 512개의 channel을 사용했다. CNN의 kernel width와 stride는 더 깊어질수록 감소한다. sampling rate는 16kHz이다.
Context NW: Contextualized representations with Transformers
- Input: $Z = z_1, z_2, …, z_T$(latent speech representations)
- with Masking
- Output: $c_1, c_2, …, c_T$(capturing information from entire sequence(contextualized)
- Action: input의 맥락적 정보를 파악한다.
먼저 NW에 입력되기 전, latent speech representations에 대해 masking을 수행한다. Masking은 self-attention과 연관되며, 자세한 내용은 이 글의 Training 섹션에서 다룬다.
원래 transformer에서는 positional embedding을 absolute로 사용하지만, 이 논문에서는 convolution 연산으로 relative positional embedding의 역할을 하게 하였다. 또한 CNN의 ouput에 GELU를 적용하여 input에 더하고, layer normalization을 수행하였다.
이를 통해 continuous speech representations 위에 context representations를 뽑아내었다. 또한 transformer의 self-attention은 dependencies over the entire sequence of latent representations end-to-end를 잡아낼 수 있게 된다.
Quantization module
- Input: $z_t \in \R^d$(latent speech representations)
- without Masking
- Output: $q_t \in R^f$(final speech unit, quantized representation)
- Action: discrete speech unit을 자동으로 학습한다.
text와 다르게 speech는 continuous nature를 가지고 있어, 자연스럽게 word나 sub-words로 분할하기 어렵다. 따라서 transformer를 사용하는 것에 어려움이 있다. discretization은 가능하긴 하지만, 사람이 하나하나 labeling해야 하기 때문에 현실적으로 cost가 너무 크다. 이 논문에서는 discrete speech units 을 자동으로 학습하는 방법을 제안한다.
어떻게 discrete speech unit을 자동으로 학습할 수 있을까? 바로 Gumbel-Softmax 에서 sampling하면 된다.
- codebook: $G \times V$ ($e \in \R^{V \times d / G}$ 이 $G$ 개)
- embedding matrix
- codeword: $V$ 개의 codeword가 하나의 codebook을 이룬다.
- embedding vector
- codebook과 codeword는 모두 학습 가능한 parameter이다.
$z_t$ 는 $G$ 개의 codebooks 각각에서 하나의 codeword와 대응된다, 각 줄에서 codeword가 선택되면 이를 모두 concatenate하여 quantized representation $q_t \in \R^f$ 가 만들어진다. 논문에서는 $G=2, V=320$ 을 사용했다. 따라서 discrete speech units의 종류는 최대 $320 \times 320 = 102,400$ 개가 된다.
즉, latent representations는 quantization matrix과 곱해져 logit을 생성한다. 이 logit은 각 codebook에서 각 codewords에 대한 score로 볼 수 있다. 그 후 Gumbel-Softmax를 통해 logit을 확률로 변환하여 각 codebook에서 하나의 codeword를 sampling하도록 한다.
Gumbel-Softmax: \(p_{g, v} = \frac{\exp(l_{g, v}+n_v)/\tau}{\sum_{k=1}^{V}{\exp(l_{g, k}+n_k)/\tau}}\)
3. Training
3.1 Masking
BERT의 Masking과 유사하다. latent representation space에서 일정 비율의 timestamp를 masking한다. 추후에 loss function에서 contrastive loss를 통해 self-attention의 효과를 확인해볼 수 있는데, 틀린 것들 사이에서 masked된 부분에 들어갈 올바른 discrete unit을 식별할 수 있는지가 학습 목표이다.
1️⃣ 먼저 비복원 추출로 masking 시작 index를 일정 비율 $p$로 random sampling한다.
2️⃣ 그 후 sampling된 각 index에서 $M$ 개의 연속적인 time step을 masking한다. masking된 구간은 겹칠 수 있다.
이때 $p$ 와 $M$ 은 hyper-parameter이다. masking된 부분은 모두 동일하게 trained feature vector로 대체된다.
3.2 Objective
Pretrain stage에서, contrastive learning을 이용하여 audio의 representation을 학습한다. \(L = L_m + \alpha L_d\)
- $L_m$: Negative discrete units 중에서 masking된 부분에 들어갈 true quantized latent speech representation 를 식별해야만 한다.
- $L_d$: 이때 Diversity Loss를 통해 model이 각 codebook에서 각각의 codeword을 균등하게 사용하도록 한다.
- $\alpha$: hyperparameter
Contrastive Loss
General하게, contrastive learning의 목적은 positive pair의 embedding은 서로 가깝게, negative pair의 embedding은 서로 멀게 하는 것이다.
- $t$: masking이 수행된 time step
- $c_t$: masking이 수행된 time step에서 추출된 context representation으로, 이 시점에서 전체 sequence에 대한 맥락 정보가 반영되어 있다.
- $\tilde q ∈ Q_t$: $K$개의 distractor(negative)와 정답 역할을 하는 1개의 qt로 구성되어, 총 $K + 1$개의 candidate quantized representation을 형성한다.
- 이때 $K$개의 distractor는 동일 발화의 다른 masking된 time step에서 랜덤하게 추출된 값들이다.
- $κ$: temperature(상수값)
- $sim()$: cosine similarity
따라서 contrastive loss는 masking된 $t$ 시점에서의 context representation $c_t$가 주어졌을 때, 정답에 해당하는 $q_t$를 다른 오답 후보들 중에서 구별하는 역할을 한다.
Diversity Loss(kind of regularization)
\[L_d=\frac{1}{GV}∗(−H(\bar p_g))=\frac{1}{GV}\sum^G_{g=1}\sum^V_{v=1} \bar p_{g,v} \log(\bar p_{g,v})\]- 모델의 codebook 활용 문제
- 모델이 모든 codeword 조합의 확률을 고르게 고려하고 quantized representation을 만든다는 보장이 없다.
- codeword 선택에 대한 경우의 수가 많지만, 이를 제대로 활용하지 않으면 codebook을 사용하는 의미가 사라진다.
Entropy in information theory \(H(X)=−∑xP(x)log(P(x))H(X) = - \sum_x P(x) \log(P(x))H(X)=−x∑P(x)log(P(x))\)
- entropy: 확률 분포의 불확실성 정도를 나타내며, 균등한 확률 분포에서 가장 큰 값을 가진다.
- 저자는 Diversity loss에서 entropy를 극대화시켜 이를 해결했다.
- 이를 통해 model은 codebook 내의 모든 codeword를 고르게 활용할 수 있게 되었다.
Fine-tuning: ASR
기존의 모델 위에 무작위로 초기화된 linear projection layer를 추가하여, 음성 인식에 필요한 특정 vocabulary에 맞게 모델을 조정한다. 이 projection layer는 해당 task의 vocabulary를 나타내는 CCC개의 클래스에 데이터를 맞추는 역할을 한다.
Librispeech 데이터셋의 경우, 문자 타겟을 위해 29개의 token과 하나의 word boundary token을 사용한다. 모델은 CTC loss를 최소화하는 방식으로 최적화되며, 이는 음성 신호와 텍스트 간의 정렬 없이 음성 인식을 가능하게 한다.
또한, 학습 중에 SpecAugment라는 기법을 사용하여 성능을 향상시킨다. SpecAugment는 time-steps와 channels를 무작위로 masking하는 방식으로, 모델이 특정 정보에 과도하게 의존하는 것을 방지하고, overfitting을 지연시킨다. 이 기법은 특히 레이블이 적은 데이터셋에서 최종 error rates를 크게 개선하는 데 효과적이다. Libri-light와 같은 레이블이 적은 데이터셋에서는 이 기법이 특히 중요한 역할을 한다.