AG-REPA: Causal Layer Selection for Representation Alignment in Audio Flow Matching

이 논문은 오디오 플로우 매칭에서 의미적 풍부함과 기능적 기여도가 불일치하는 현상을 규명하고, 예측 속도장의 인과적 기여도를 정량화하는 '포워드-온리 게이트 제거 (FoG-A)' 기법을 도입하여 표현 정렬을 최적화하는 AG-REPA 를 제안함으로써 기존 방법론의 성능을 향상시킵니다.

핵심

이 논문은 **"소리를 만드는 인공지능 (AI) 을 더 똑똑하고 빠르게 가르치는 새로운 방법"**에 대해 설명합니다.

기존의 AI 는 소리를 만들 때 중간 과정을 가르치기 위해 "어떤 층 (Layer) 을 봐야 할지"를 대충 짐작해서 정했습니다. 하지만 이 논문은 **"그건 틀릴 수 있다"**고 말하며, "무엇을 알고 있는지 (저장)"와 "무엇을 실제로 사용하는지 (작동)"는 다릅니다라는 놀라운 사실을 발견했습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 풀어보겠습니다.

---

1. 문제 상황: "책상 위가 정리된 학생" vs "시험을 잘 보는 학생"

소리를 만드는 AI 는 거대한 건물을 짓는 것과 같습니다. 이 건물에는 20~30 층의 층 (Layer) 이 있습니다.

• 기존의 생각 (REPA): "중간 층 (예: 8 층) 에 가면 소리에 대한 지식이 가장 풍부할 거야. 그래서 8 층을 선생님 (기존의 잘 만든 AI) 과 비교해서 가르치자!"
• 현실: 8 층은 소리에 대한 지식을 가장 많이 저장하고 있는 도서관 같습니다. 책이 꽉 차 있지만, 정작 건물을 짓는 **작업 (소리를 내는 일)**은 1 층이나 다른 층에서 이루어집니다.
• 결론: 도서관 (8 층) 을 열심히 가르쳐도, 실제 공사 현장 (1 층) 이 느리면 건물은 빨리 지어지지 않습니다.

2. 핵심 발견: "알고 있음"과 "행동함"의 분리 (SCD)

논문은 이 현상을 **'저장 - 기여 분리 (Store-Contribute Dissociation)'**라고 부릅니다.

• 저장 (Knowing): 깊은 층 (20~24 층) 은 소리의 모든 정보를 가장 잘 기억하고 있습니다. (도서관)
• 기여 (Doing): 하지만 소리를 만들어내는 가장 중요한 힘은 얕은 층 (1~3 층) 에서 나옵니다. (현장 지휘자)

비유:
마치 지휘자가 있습니다.

• 깊은 층은 악보의 모든 내용을 외우고 있는 사람입니다. (지식 풍부)
• 얕은 층은 지휘봉을 휘두르며 오케스트라의 소리를 실제로 만들어내는 사람입니다. (실질적 기여)
• 기존 방법은 "악보를 가장 잘 외운 사람 (깊은 층)"을 칭찬하고 가르쳤는데, 정작 소리를 만드는 건 "지휘봉을 휘두르는 사람 (얕은 층)"이었습니다. 그래서 효율이 떨어졌습니다.

3. 새로운 해결책: AG-REPA (의사결정 가이드)

저자들은 이제 **"무엇을 가르쳐야 할지"**를 대충 짐작하는 게 아니라, AI 가 실제로 어떤 층을 가장 많이 사용하는지를 측정해서 가르칩니다.

• 새로운 도구 (FoG-A): "이 층을 잠시 끄면 소리가 어떻게 변할까?"를 실험해봅니다.
  • 1 층을 끄면 소리가 완전히 망가집니다. → 이 층이 진짜 핵심이다!
  • 20 층을 끄면 소리는 비슷하게 나옵니다. → 이 층은 그냥 정보를 쌓아두는 곳이다.
• 새로운 전략 (AG-REPA): 이제 AI 를 가르칠 때, 정보를 많이 가진 깊은 층이 아니라, 실제로 소리를 만드는 얕은 층 (핵심 층) 위주로 집중해서 가르칩니다.

비유:
기존에는 "책상 위에 책이 가장 많이 쌓인 학생 (8 층)"을 칭찬하며 가르쳤다면,
이제 **"시험 문제를 가장 잘 풀고 있는 학생 (1 층)"**을 찾아내어 그 학생에게 집중적으로 코칭을 해주는 것입니다.

4. 결과: 얼마나 좋아졌나요?

이 방법을 적용하자 놀라운 변화가 일어났습니다.

• 소리의 질: 더 선명하고 자연스러워졌습니다. (FAD 점수 16~18% 향상)
• 학습 속도: 같은 품질을 내는데 훨씬 더 빨리 학습이 끝났습니다.
• 적용 범위: 사람 목소리 (TTS) 이든, 배경음악이나 효과음 (일반 오디오) 이든 모두 잘 작동했습니다.

5. 한 줄 요약

"AI 가 무엇을 '알고 있는지'보다, AI 가 무엇을 '실제로 하고 있는지'를 파악해서 가르쳐야 더 똑똑해진다."

이 논문은 AI 의 내부 구조를 단순히 '정보 저장소'로 보지 않고, **'어떤 부분이 실제로 일을 하는지'**를 찾아내는 과학적인 방법을 제시함으로써, 앞으로 더 빠르고 훌륭한 소리 생성 AI 를 만드는 길을 열었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

• 배경: Flow Matching (FM) 모델은 오디오 생성 (음성 합성, 일반 오디오 생성) 에서 강력한 성능을 보이지만, 학습 비용이 매우 높습니다. 이를 가속화하기 위해 사전 학습된 Teacher 모델의 특징과 중간 은닉 상태를 정렬하는 REPresentation Alignment (REPA) 기법이 도입되었습니다.
• 현황의 한계: 기존 REPA 기반 방법들은 주로 휴리스틱 (heuristic) 한 접근법 (예: 고정된 중간 레이어, 예를 들어 레이어 8 등) 에 의존하여 정렬 레이어를 선택합니다.
• 핵심 문제: 오디오 생성에서 "가장 풍부한 의미 (semantic) 정보를 저장하는 레이어"가 반드시 "생성 과정 (velocity field) 을 주도하는 데 가장 크게 기여하는 레이어"와 일치하지 않을 수 있습니다. 기존 연구는 이 저장 (Storage) 과 기여 (Contribution) 의 불일치를 간과하여, 생성 역학에 실제로 영향을 미치지 않는 레이어에 학습 신호를 보내는 비효율적인 학습을 초래합니다.

2. 핵심 발견: Store-Contribute Dissociation (SCD)

저자들은 토큰 조건부 (token-conditioned) 오디오 FM 모델에서 다음과 같은 현상을 발견하고 이를 Store-Contribute Dissociation (SCD, 저장 - 기여 분리) 이라고 명명했습니다.

• 저장 (Storage): 깊은 레이어 (Deep layers) 는 Teacher 모델과 높은 유사도를 보이며 풍부한 의미 및 음향 정보를 저장합니다 (Representation Storage).
• 기여 (Contribution): 그러나 생성을 구동하는 속도장 (velocity field) 에 대한 인과적 기여 (causal contribution) 는 주로 얕은 레이어 (Shallow layers, 특히 레이어 1) 에서 발생합니다.
• 메커니즘: 잔차 네트워크 (Residual Network) 구조와 ODE 관점에서 볼 때, 초기 레이어의 변화는 전체 네트워크를 관통하는 야코비안 (Jacobian) 곱을 통해 증폭되어 최종 출력에 결정적인 영향을 미칩니다 (나비 효과). 반면, 깊은 레이어는 정보를 축적하지만 생성 동역학에는 상대적으로 수동적인 역할을 합니다.

3. 방법론 (Methodology)

이러한 통찰을 바탕으로 제안된 AG-REPA (Attribution-Guided REPA) 는 휴리스틱 선택 대신 인과성 기반 (causality-driven) 의 레이어 선택 전략을 사용합니다.

3.1. 진단 도구 (Diagnostic Toolkit)

학습 효율성을 높이기 위해 네트워크의 '무엇을 아는가 (Representation)'와 '무엇을 사용하는가 (Function)'를 구분하는 세 가지 도구를 개발했습니다.

1. BiT-C (Bi-Stream Teacher Cosine Alignment): 음성 (Whisper Teacher) 과 일반 오디오 (BEATs Teacher) 두 가지 Teacher 모델을 사용하여 레이어별 표현의 의미/음향 정렬 정도를 측정합니다.
2. LASP (Layer-wise Analysis via Shared Projection): 공유된 프로젝션 헤드를 사용하여 각 레이어가 Teacher 특징과 얼마나 유사한지 (정보 저장량) 를 정량화합니다.
3. FoG-A (Forward-only Gate Ablation): 핵심 도구. 특정 레이어의 게이트를 닫아 (ablation) 예측된 속도장 (velocity field) 에 발생하는 변화를 측정합니다. 이는 해당 레이어가 생성에 인과적으로 얼마나 필수적인지를 정량화합니다.

3.2. AG-REPA 전략

• 레이어 선택: 고정된 레이어 대신, FoG-A 점수가 높은 상위 K 개의 레이어 (인과적으로 중요한 레이어) 를 동적으로 선택합니다.
• 가중치 할당: 선택된 각 레이어에 FoG-A 점수에 비례하는 가중치 ($\lambda_k$) 를 부여하여 정렬 손실 (Alignment Loss) 을 적용합니다.
• 목적 함수: Flow Matching 손실과 함께, 인과적으로 중요한 레이어에만 집중된 정렬 손실을 추가하여 학습 효율을 극대화합니다.

4. 주요 결과 (Results)

LibriSpeech (음성) 와 AudioSet (일반 오디오) 을 활용한 통합 학습 실험에서 AG-REPA 는 기존 REPA 베이스라인을 압도적으로 능가했습니다.

• 성능 향상:
  • FAD (Fréchet Audio Distance): 음성 생성에서 18%, 일반 오디오에서 16% 개선.
  • WER (Word Error Rate): 3.45 로 감소 (기존 5.82 대비).
  • MOS (Mean Opinion Score): 4.12 로 향상.
• 비교 실험:
  • Deep REPA (정보 저장 풍부한 레이어 정렬): 성능 향상 미미.
  • Shallow REPA (초기 레이어 정렬): 상당한 개선.
  • AG-REPA: Shallow REPA 보다 더 정교하게 상위 인과 레이어를 선택하고 가중치를 부여하여 최고의 성능 달성.
• 일반화 능력: Voicebox, CosyVoice, F5-TTS 등 다양한 아키텍처의 Flow Matching 모델에 적용 시에도 일관된 성능 향상을 보였습니다.
• 수렴 속도: 인과적 레이어에 정렬을 적용할 경우, 무작위 레이어나 정보 저장량이 많은 레이어에 정렬하는 것보다 3.3 배 빠른 수렴을 보였습니다.

5. 의의 및 기여 (Significance & Contributions)

1. 이론적 통찰: 오디오 생성 모델에서 "정보 저장 (Representation)"과 "기능적 기여 (Function)"가 분리될 수 있음을 최초로 체계적으로 증명하고, 이를 SCD로 정의했습니다.
2. 방법론적 혁신: 휴리스틱에 의존하던 레이어 선택 방식을, 인과적 기여도 (FoG-A) 를 기반으로 한 데이터 주도적 (data-driven) 인 동적 선택 방식으로 전환했습니다.
3. 실용적 가치: AG-REPA 는 생성 모델의 학습 효율성을 높이고, 더 낮은 WER 과 높은 음질 (FAD, MOS) 을 달성하여, 해석 가능성 (Interpretability) 이 실제 학습 전략으로 직접 연결될 수 있음을 입증했습니다.
4. 미래 지향성: "무엇을 아는가 (Knowing)"보다 "무엇을 사용하는가 (Doing)" 에 초점을 맞춘 정렬 전략이 생성 AI 의 투명성과 제어 가능성을 높이는 새로운 패러다임을 제시합니다.

결론적으로, 이 논문은 오디오 생성 모델의 학습 효율성을 극대화하기 위해, 단순히 정보를 많이 담고 있는 레이어가 아닌, 실제 생성 역학을 주도하는 인과적 레이어를 대상으로 정렬해야 함을 증명하고 이를 실현 가능한 알고리즘 (AG-REPA) 으로 제안했습니다.