[논문 리뷰] UNIWAV: TOWARDS UNIFIED PRE-TRAINING FOR SPEECH REPRESENTATION LEARNING AND GENERATION

목차

1. Introduction

2. Method

3. Experiments

4. Related works

5. Conclusion

1. Introduction

최근 몇 년간 음성 처리 분야에서 사전 학습과 표현 학습의 중요성이 크게 증가했습니다. 그러나 기존의 사전 학습 기법은 주로 판별(discriminative) 또는 생성(generative) 작업 중 하나에 특화되어 있어, 각 작업에 맞는 별도의 기초 모델이 필요합니다. 본 논문에서는 이러한 두 가지 작업을 통합할 수 있는 최초의 사전 학습 프레임워크인 UniWav를 제안합니다. UniWav는 음성 인식, 텍스트-음성 합성, 음성 토큰화 등의 다양한 작업에서 뛰어난 성능을 발휘하며, 기존의 여러 기초 모델을 대체할 수 있는 단일의 범용 모델을 구축하는 데 기여합니다.

2. Method

UniWav는 인코더-디코더 프레임워크로 설계되어 있으며, 음성의 표면 특징(예: 인코더의 Vector 또는 멜 스펙트로그램)을 입력으로 받아 표현을 학습하고 음성을 생성합니다.

2.1 인코더

인코더는 Transformer 아키텍처를 기반으로 하였으며,24개의 레이어로 구성되어 있습니다.

입력으로는 음성의 표면 특징이 사용되며, Convolutional-Positional-Embedding을 통해 위치 정보를 인코딩합니다.

학습에서는 자기 지도 학습과 온라인 클러스터링, 그리고 EMA 기반 Teacher Model을 사용하여 더 안정적이고, 일반화된 학습을 가능하게 해줍니다.

자기 지도 학습 : 자기 지도 학습은 입력 데이터 자체에서 학습 목표를 도출하여 모델을 학습하는 방식입니다. 이는 마스킹(masking)이나 미래 예측(future prediction)과 같은 전처리 과정을 통해 이루어집니다.

온라인 클러스터링 : 온라인 클러스터링은 데이터가 스트리밍되면서 실시간으로 클러스터링을 수행하는 기술입니다. 이는 데이터의 수집이 어려운 환경에서 유용합니다. 이 논문에서는 DinoSR에서 온라인 클러스터링과 Self-distillation 기술을 가져와서 적용했다고 하며, 이로 인해서 인코더와 디코더를 처음부터 End-to-End로 함께 학습시키는 기술이 완성되었다고 합니다.

EMA 기반 교사 모델 : EMA(Exponential Moving Average) 기반 교사 모델은 학생 모델의 파라미터를 지속적으로 업데이트하여 교사 모델을 개선하는 방식입니다. 이는 반복적인 학습 과정을 통해 모델의 안정성을 높입니다. Moving Average는 이러한 시계열 혹은 자연어 같은 sequence 데이터에 대해서 이동한 길이의 평균을 말하며, 얼마만큼 데이터를 빠르게 넘길 것인지 이동 속도를 결정한다고 이해하셔도 될 거 같습니다.

즉, 이 모델에서는 자기 주도 학습을 위해 masking과 pseudo-labeling이 되어있는 Audio 데이터로, Teacher Model을 통해 Moving Average를 조정하며, 온라인 클러스터링을 통해 계속해서 데이터를 수집하며 바로바로 학습에 적용했다는 것이죠.

2.2 디코더

디코더는 Flow Matching 기반의 아키텍처를 사용하며, 12개의 레이어로 구성되어 있습니다.

디코더에서는 3가지 주요 기술이 사용되었습니다.

먼저, 디코더는 ODE(Ordinary Differential Equation)를 이용한 시간 정보를 사용하여, input sequence에 시간 정보를 추가합니다. 지는 sinusoidal embedding을 통해 구현됩니다.
그리고 이렇게 추가된 시간 정보를 위해서, ALiBi Self-Attention Bias를 적용합니다. input sequence의 시간 정보를 더 잘 반영합니다.
마지막으로, 디코더는 U-Net과 유사한 구조로 설계되어 skip-connection을 통해 레이어 간의 정보를 효과적으로 보존해줍니다.

ODE

ODE는 한국인에게는 '미분방정식', 혹은 $\frac{dx}{dt} = f'(x)$ 이라는 표현이 더 와닿을 거 같습니다. 이 모델에서는 시간에 따른 변화 과정을 모델링하는 방정식으로, 다음과 같이 표현될 수 있습니다.

Sinusoidal embedding

입력 시퀀스의 위치 정보를 인코딩하는 데 사용됩니다.

ALibi self-attention

ALiBi는 Attention with Linear Biases로, 입력 시퀀스의 시간 정보를 더 잘 반영하기 위해 사용됩니다. 이는 쿼리-키 어텐션 스코어에 선형적으로 시간 정보를 추가하여, 최근 입력에 대한 편향을 줄입니다.

2.3 Loss

전체 Loss function은 인코더와 디코더의 손실을 결합하여 계산하며, 인코더에서는 CTC Loss, 디코더에서는 정확히 명시되지는 않았지만 Flow Matching에서 주로 사용되는 Cross-Entropy Loss를 사용했을 것으로 추측됩니다.

CTC(Connectionist Temporal Classification) Loss

Cross-Entropy Loss

3. Experiments

UniWav는 LibriLight 데이터셋을 사용하여 사전학습되었습니다. 모델 평가 결과, 모델은 음성 인식, TTS, 음성 토큰화 3가지의 작업에서 뛰어난 성능을 보였습니다. 특히, 음성 토큰화 작업에서 UniWav는 Discrete Bottleneck 구조 덕분에 낮은 비트레이트에서도 높은 품질의 음성을 생성해낼 수 있었습니다.

Discrete Bottleneck : 연속적인 표현을 이산적인 표현으로 변환하여, 모델의 일반화 성능을 향상시키는 기술입니다. 연속적인 데이터를 이산화해서, 모델에게 더 정확하게 정보를 전달해주는 것이죠.

4. Related works

관련된 연구는 대부분 판별 작업(인코딩)에 초점이 맞추어져 있었으며, 생성(디코딩)은 별도의 모델로 처리되었습니다. UniWav의 진가는 이 두가지 작업을 단일 모델로 통합하려는 최초의 시도라는 점에 있습니다.
그리고 그것이 가능한 이유는 자기 지도 학습을 통한 인코딩, 그리고 Flow Matching을 통한 디코딩이라고 강조해주는 부분이네요.

5. Conclusion

UniWav는 Speech 도메인에서 전반적인 판별 및 생성 작업을 모두 해낼 수 있는 최초의 Pre-trained Framework로서, 기존의 여러 모델을 대체할 수 있는 단일의 범욜 모델을 추구합니다. 이는 즉, 각각의 Task에 대해서 동일한 모델을 사용함으로서 학습의 효율성과 성능을 높이고 비용을 절감할 수 있다는 뜻입니다. 그러나 아직은 개별 모델에 비해서 성능의 격차가 존재하며, 향후 연구에서 이를 해소하기 위한 방향을 도모할 예정이라고 합니다.