Fast Speech Foundation Model Distillation Using Interleaved Stacking

기술 요약: 인터리브드 스태킹(Interleaved Stacking)을 이용한 빠른 음성 기초 모델 증류

문제 정의
지식 증류(Knowledge Distillation, KD)는 HuBERT와 같은 대규모 음성 기초 모델(Speech Foundation Models, SFMs)로부터 효율적인 학생 모델을 배포하여 추론 지연 시간과 계산 비용을 줄이기 위해 사용하는 표준적인 기술이다. 기존의 접근 방식들은 메모리 제약 조건 하에서 성능을 유지하기 위해 주로 "깊고 좁은(deep-and-narrow)" 학생 아키텍처를 채택하지만, 이러한 모델들은 제한된 병렬화로 인해 "얕고 넓은(shallow-and-wide)" 아키텍처에 비해 학습 속도가 느리다는 단점이 있다. 모델 스태킹(모델 깊이를 점진적으로 증가시키는 방식)을 통한 단계별 학습은 다른 도메인에서 학습 가속화를 위해 탐구되어 왔으나, SFM 증류에 대한 적용은 아직 연구되지 않았다. 또한, 기존의 스태킹 방법들(예: 점진적 스태킹(Gradual Stacking), MIDAS)은 훈련 단계 전반에 걸쳐 레이어 위치의 일관성을 깨뜨리기 때문에 SFM에서 성능 저하를 일으킨다. SFM은 레이어별로 구별되는 고유한 지식을 인코딩하므로, 스태킹 과정 중 레이어를 재배치하는 것은 이러한 지식의 효과적인 전달을 방해한다.

방법론
저자들은 SFM 증류를 위해 특별히 설계된 새로운 학습 가속 프레임워크인 **인터리브드 스태킹(Interleaved Stacking)**을 제안한다. 이 방법은 다음과 같은 메커니즘을 통해 작동한다:

1. 인터리빙을 통한 단계별 학습: 학습 과정은 $B$개의 단계로 나뉜다. 초기에는 $K$개의 레이어($K = N/B$, 여기서 $N$은 목표 깊이)를 가진 얕은 학생 모델을 학습시킨다. 각 단계가 끝날 때, 현재 모델의 상단이나 중간에 복사된 레이어를 추가하는 대신, 현재 $bK$레이어 모델에서 매$b$번째 레이어를 선택하여 복사한 뒤, 원래 레이어 바로 뒤에 삽입한다. 이는 다음 단계를 위한$(b+1)K$ 레이어 모델을 생성한다.
2. 레이어 위치 일관성: 레이어를 마지막에 복사하는 점진적 스태킹(Gradual Stacking)이나 중간 블록을 복사하는 MIDAS와 달리, 인터리브드 스태킹은 훈련 단계 전반에 걸쳐 레이어가 상대적인 비례 위치를 유지하도록 보장한다(초기 레이어는 계속 초기에, 후기 레이어는 계속 후기에 위치함).
3. 중간 단계 수준 KD와의 통합: 이 방법은 중간 단계 수준의 지식 증류 손실을 자연스럽게 지원한다. 레이어 위치가 일관되게 유지되기 때문에, 학생의 $k$번째 레이어는 모든 단계에서 특정 교사 레이어(예: 매 $M/K$번째 교사 레이어)와 일관되게 정렬될 수 있다. 이를 통해 레이어 인덱스가 이동하는 다른 스태킹 방법에서 관찰되는 학습 불안정성 없이 레이어 간 평균 제곱 오차(MSE) 손실을 사용할 수 있다.
4. 손실 함수: 전체 손실은 출력 레벨 KD 손실(교사와 투영된 학생 출력 간의 MSE)과 중간 단계 수준의 KD 손실을 결합한다. 중간 손실은 학생의 현재 깊이와 관련된 고정된 교사 레이어 인덱스를 대상으로 하여 안정적인 감독을 보장한다.

주요 기여

• SFM 증류를 위한 스태킹의 첫 번째 조사: 본 논문은 SFM 증류를 위해 스태킹 기반 학습 가속의 효과를 최초로 탐구하였다.
• 인터리브드 스태킹 알고리즘: 저자들은 레이어 위치 일관성을 보존하는 스태킹 전략을 도입하여, 기존 방법들에서 발생하는 레이어 불일치로 인한 성능 저하 문제를 해결하였다.
• 중간 감독과의 호환성: 제안된 방법은 SFM 성능에 결정적인 중간 단계 수준의 KD 손실을 원활하게 통합할 수 있게 하며, 기존의 스태킹 베이스라인들이 학습 발산 없이 이를 수행하는 데 어려움을 겪는 것과 대조된다.
• 실험적 검증: 제안된 방법은 네 가지 다운스트림 작업인 음소 인식(PR), 자동 음성 인식(ASR), 슬롯 필링(SF), 화자 식별(SID)에 대해 SUPERB 벤치마크에서 검증되었다.

실험 결과
저자들은 HuBERT base 교사 모델로부터 증류된 12개 레이어 Transformer 학생 모델을 사용하여 방법을 평가하였다.

• 베이스라인 대비 성능: 인터리브드 스태킹은 모든 작업과 스케줄링 전략(Equal 및 Prop-1)에서 기존 스태킹 베이스라인(Gradual Stasting 및 MIDAS)보다 성능이 현저히 뛰어났다. 예를 들어, Equal 스케줄링을 적용한 PR 작업에서 인터리브드 스태킹은 9.08의 PER(Perplexity Error Rate)을 달이트었으며, 이는 Gradual Stacking의 11.50 및 MIDAS의 10.75와 비교된다.
• 전체 학습 대비 성능: 놀랍게도, Prop-1 스케줄링 전략 하에서 제안된 방법은 스태킹 없이 학습된 전체 깊이 모델(Full L2L)과 대등하거나 더 나은 성능을 PR, SF, SID 작업에서 달성하였으며, 동시에 약 1.16배의 학습 시간 단축을 보였다.
• 중간 손실의 영향: 중간 단계 수준의 KD 손실을 포함하는 것은 성능을 크게 향상시켰다. 이 방법은 중간 손실 없이도 베이스라인보다 우위를 유지했지만, 중간 손실을 포함했을 때 그 이득이 가장 두드러졌다. 반면, Gradual Stacking에 중간 손실을 적용하면 학습 불안정과 성능 저하가 발생했다.
• 레이어 유사도 분석: 교차 레이어 유사도 분석을 통해 SFM의 인접 레이어들이 높은 상관관계를 보임을 확인하였다. 인터리브드 방식은 이러한 관계를 보존하여, 복사된 레이어와 원본 레이어가 높은 유사성을 유지하는 "블록 형태"의 유사도 구조를 형성하였으며, 이는 인접 삽입이라는 설계 선택을 입증한다.

의의 및 주장
본 논문은 인터리브드 스태킹이 SFM 증류에서 학습 효율성과 모델 성능 사이의 트레이드오프에 대한 실질적인 해결책을 제공한다고 주장한다. 레이어 위치를 일정하게 유지함으로써, 이 방법은 다른 가속 기술에서 흔히 소실되는 SFM 고유의 레이어별 지식을 보존할 수 있게 한다. 저자들은 이 접근 방식이 학습 비용을 줄임으로써 효율적인 학생 모델의 배포를 가속화할 뿐만 아니라, 단계별 학습과 관련된 성능 저하를 완화한다고 주장한다. 이 연구는 학습 효율성이 실제 SFM 배포에서 중요하지만 아직 충분히 탐구되지 않은 요소임을 강조하며, 학습 중 세심한 아키텍처 조작이 학습 속도가 빠르면서도 매우 효과적인 모델을 만들어낼 수 있음을 보여준다.