Neural Discrete Codec은 어떻게 학습되지?
Audio codec VQ VAE를 예로 들자. 일반적으로 codec은 auto-encoder 구조이다. 여기서 intput-output은?
- encoder
- input - audio mel 혹은 raw waveform (mono 24kHz라면 초당 24000개의 floating point 값을 가진 array)
- output - Language tokenizer의 output은 단일 정수의 array로, 각 정수 값은 vocab에서의 index이다.
encoder의 output은 모델 구조에 따라 다를 수 있는데 여기서는 quantized vector라고 하자.
- decoder
- input - encoder의 output
- output - encoder의 input
- inference 과정
- 우선 data sample을 downsampling + channel을 늘리도록 여러 layer를 거치고, [latent_seq_len,
dim] 형태의 latentz가 된다. - codebook의 각 임베딩과 L2 distance를 계산해서 가장 가까운 임베딩은 1, 나머지는 0인 K 길이 categorical distribution을 얻는다.
- 그리고 위에서 가장 가깝다고 계산된 k-index의 codebook 임베딩을 선택, 사용해서 decoder로 넘기거나, quantized latent로 사용한다.
raw waveform이 들어왔다.
- training 과정
- encoder
- decoder
- codebook, nn.Embedding(K, dim)
- codebook 임베딩을 업데이트 하기 위한 Codebook loss
- encoder output z가 코드북 벡터 근처로 매핑되도록 강제
- 어쨌든 encoder-decoder 구조의 학습의 목표는 decoder를 통과하여 얻은 출력이 입력을 reconstruction하도록 하는 것이다.
- 이 방법으로 하면 인코더는 업데이트 되는데, 코드북은 업데이트 되지 않는다.
- 그런데 현재 식에서는 VAE의 KLD 처럼 prior 분포에 대한 제약이 없다.
- decoder: reconstruction loss로만 업데이트
- encoder: reconstruction + commitment loss
- codebook(embedding): codebook loss로만 업데이트
- 두 loss에 가중치를 다르게 주고싶어서라면 이해됨.
- loss 자체는 같지만, 각 loss의 의미는 다르긴하다. 결국 합쳐서 하나로 하면 학습이 비효율적이게 됨.
학습되어야 하는건 크게 3가지 모듈이다.
codebook =
dim dimension을 가진 임베딩이 num_embeddings(= K)개 있다. 구현으로는 nn.Embedding(K, dim)초기화는 평균 0, 범위는 +- 1/
K로, uniform하게 초기화해서 초기 편향이 없게 한다.loss는 3가지로 구성된다. 각 loss는 가중합한다.
목적 : 코드북이 encoder의 출력을 따라가도록 학습한다.
F.mse_loss(z_q, z.detach())여기서 z_q는 선택된 임베딩이고, encoder output z는 detach를 해서 현재 loss를 통해 가까워지도록 할 때 encoder는 가만히 있고 코드북만 업데이트 되도록 한다.
encoder가 제멋대로 여러가지 넓은 space로 output을 보내지 않도록 한다.
F.mse_loss(z_q.detach(), z)여기서는 반대로 z_q를 detach
encoder를 태우는건 위 inference 과정과 같다. 거리를 기반으로 k번째 index를 선택한다는 연산 자체는 미분 불가능 하기 때문에 decoder의 output으로 계산한 loss의 gradient를 back-propagation으로 encoder까지 흘려줄 수 없다.
따라서 z_q를 선택한다는 그 연산 자체를, 로 바꿔 써서(straight-through estimator 방법) 미분 가능하게 그것도 gradient가 그대로 전달되게 한다.
그래서 인코더 output을 따라가도록 codebook loss를 추가한다.
즉, encoder의 output이 무지하게 커지거나 분산이 커져도 상관없어진다. encoder는 reconstruction loss 때문에 활발히 움직이는데, codebook이 이걸 따라가려다 latent space가 과하게 커질 수 있다. 그러다보면 학습이 불안정해진다.
따라서 이걸 막기위해 이번에는 반대로 encoder output이 codebook과 비슷해지도록(= 현재 codebook과 과하게 멀어지지 않도록) commitment loss를 추가한다.
즉:
→ 궁금 : 어짜피 서로서로 비슷해지게 할거면 commitment, codebook loss를 detach해서 두개 쓰는게 아니라 |e - z_q| 로 하나만 두는 것과 같지 않나?
하지만 detach로 분리하지 않으면 encoder가 업데이트 될 때 recon loss에 의해서만 이동해도 codebook이 따라오기 때문에 코드북이 deocder에게 최적화된다.
하지만 여기서 문제는?
- 기본적으로 다양한 auxiliary loss와 STE와 같은 방식이 필요하다는건 학습이 불안정하다는 뜻이다. 그걸 완화하기 위해 여러 모듈을 추가한 것.
- 학습이 불안정한 이유 중 하나는 codebook이 5000개라도 그중 1000개만 사용하게될 수 있기 때문이다.(= codebook collapse) 표현력이 제한되고 학습 데이터에만 피팅될 수 있다.
- 이런 이유로 코드북을 키워도 성능이 커지지 않는 경우가 많다.
→ 이런 이유로 FSQ 방식이 제안되었다.
Finite Scalar Quantization은 기존 방식으로 quantize하지 않는다.
- 우선 벡터를 아주 작은 d dimension으로 만든다.
- 각 차원의 값들이 가질 수 있는 값 범위를 정한다. [-K, K]
- 그리고 각 차원의 값을 bounding 처리한다. → 각 차원의 값들이 범위가 [-K ~ K]인 실수 값이 됨.
- -K ~ K 사이의 정수 값으로 반올림한다.
→ 그럼 자연스럽게 코드북 수는
FSQ는 벡터 전체를 하나의 벡터로 quantize가 아니라, 각 feature dimension을 미리 정해진 finite set 중 하나로 quantize한다. → 코드북의 코드 수가 많아져도 안정적이고,, 다시말해 여러 다양한 조합을 사용할 수 있게 된다.
FSQ 논문에서의 VQ와 성능 비교.
특히 codebook size가 큰 경우에 차이가 발생한다.
이렇게되면 encoder의 출력이 [-1, 1, -2, -3, 2, 0] 과 같은 형태기 때문에 language model에 사용할 수 있는 단일 index에 1을 가지는 형태의 분포는 아닌데, index화를 통해 변형해서 사용할 수 있다.
하지만 FSQ가 아닌 Residual Vector Quantization도 Sound Stream을 시작으로 많이 쓰인다.
Share article