Tuning Language Models by Mixture-of-Depths Ensemble

기술 요약: Mixture-of-Depths Ensemble (MoDE)를 통한 언어 모델 튜닝

문제 정의
전통적인 트랜스포머 기반 대규모 언어 모델(LLM)은 미세 조정(fine-tuning) 과정에서 예측과 손실 계산을 위해 오직 최종 레이어의 표현(representation)에만 전적으로 의존합니다. "로짓 렌즈(logit lens)"와 같은 해석 도구들은 후기 레이어의 은닉 상태(hidden states)가 이미 충분히 형성된 작업 관련 예측을 포함하고 있음을 보여주었으나, 이러한 중간 표현들은 일반적으로 명시적인 학습 신호라기보다는 암묵적인 부산물로 취급됩니다. 결과적으로, 후기 레이어의 잠재적인 예측 능력은 제대로 활용되지 못하며, 표준 미세 조정은 오직 최종 레이어만을 최적화하므로 중간 표현들은 간접적으로만 최적화되는 한계를 가집니다. 또한, 중간 레이어를 활용하는 기존 방식들(예: 효율성을 위한 조기 종료(early exit) 또는 표현 스티어링(representation steering))은 레이어별 로짓을 최종 예측의 직접적이고 학습 가능한 구성 요소로 다루지 않습니다.

방법론: Mixture-of-Depths Ensemble (MoDE)
저자들은 **Mixture-of-Depths Ensemble (MoDE)**을 제안합니다. 이는 LLM의 마지막 $k$개 레이어 출력을 단일 모델 내의 "전문가(experts)" 앙상블로 취급하는 튜닝 프레임워크입니다. 서로 다른 모델으로 토큰을 라우팅하는 Mixture-of-Experts(MoE)와 달리, MoDE는 동일한 모델 내의 서로 다른 깊이에 있는 로짓들의 가중치 결합으로 최종 예측을 라우팅합니다.

이 프레임ка는 세 가지 핵심 구성 요소로 이루어집니다:

1. 후기 레이어 로짓 추출 (Late-Layer Logit Extraction): $n$개의 레이어를 가진 모델에 대해, 마지막 $k$개 레이어($K = \{L_{n-k+1}, \dots, L_n\}$)를 선택합니다. 각 레이어의 은닉 상태에 어휘 헤드(vocabulary head) $\phi(\cdot)$를 적용하여 로짓 $\ell_i(x)$를 생성합니다.
2. 학습된 라우팅 네트워크 (Learned Routing Network): 마지막 $k$개 레이어 직전의 은닉 상태($x = H_{n-k}$)를 기반으로 가중치를 계산하는 경량 라우터 $G(x)$(단일 선형 레이어와 소프트맥스로 구현)를 사용합니다 추정된 최종 로짓은 다음과 같이 가중 합산으로 계산됩니다:
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   라우터는 추론 시 조기 종료를 가능하게 하기 위해 희소(Sparse) 형태인 Top-K 방식으로 만들 수 있습니다.
3. 정규화 및 증류를 통한 적응 (Adaptation via Normalization and Distillation): 사전 학습된 헤드를 중간 레이어에 직접 적용하면 초기 손실(loss)이 높게 나타날 수 있습니다. 이를 해결하기 위해 무거운 파라미터를 추가하는 대신(per-layer affine matrix와 달리), 저자들은 다음을 도입합니다:
   • 레이어별 정규화 ($N_k$): 표현을 헤드의 입력 공간에 적응시키기 위해 각 후기 레이어에 학습 가능한 정규화 모듈(사전 학습된 아키텍처와 일치하는 LayerNorm 등)을 추가합니다.
   • 증류 손실 ($L_{distill}$): 최종 레이어가 "교사(teacher)" 역할을 합니다. KL-발산(KL-divergence) 손실을 적용하여 초기 후기 레이어들의 출력 분포를 최종 레이어의 분포로 끌어당김으로써, 학습을 안정화하고 앙상블을 작업 목표에 정렬시킵니다.

전체 손실은 $L_{MoDE} = L_{task} + \lambda L_{distill}$입니다. MoDE는 기존의 매개변수 효율적 미세 조정(PEFT) 방식(예: LoRA, DoRA) 위에 레이어 형태로 얹거나 독립적으로 사용할 수 있는 플러그인 메커니즘으로 설계되었습니다.

주요 기여

• 학습 가능한 후기 레이어 신호: 논문은 후기 레이어 로짓이 증류(distillation)를 통해 개별적으로 최적화될 때, 최종 레이어와 유사한 손실에 수렴하며 상보적인 문제들을 해결할 수 있음을 입증했습니다.
• MoDE 프레임워크: 학습된 라우터를 통해 후기 레이어 로짓을 앙상블하는 단순하고 오버헤드가 적은 튜닝 메커니즘입니다. 이는 모델 크기의 약 +0.04%에 불정하는 극히 적은 학습 파라미터를 추가하면서도, 훨씬 더 큰 학습 모듈(예: 전체 레이어 LoRA)을 대체할 수 있는 성능을 보여줍니다.
• 효율성 및 유연성: MoDE는 희소 Top-K 라우팅을 지원하여, 최소한의 정확도 손실로 최대 1.6배의 생성 속도 향상과 조기 종료를 가능하게 합니다. 또한 기존 PEFT 방법들과 결합하여 성능을 높일 수 있습니다.

실험 결과
실험은 LLaMA2-7B 및 LLaMA3-8B를 대상으로 산술 추론, 상식 추론, 지시 이행(instruction following) 작업에서 수행되었습니다.

• 산술 추론 (Arithmetic Reasoning): MoDE는 기본 방법들(LoRA, DoRA)의 성능을 일관되게 개선했습니다. 경량 변형(마지막 $k$개 레이어를 제외한 모든 레이어에 LoRA 적용 + MoDE)은 훨씬 적은 파라미터를 사용하면서도 전체 레이어 LoRA와 0.3~0.5 포인트 이내의 성능 차이를 보였습니다.
• 상식 추론 (Commonsense Reasoning): MoDE는 가중치 기반 방식(LoRA, DoRA)에 대해 부가적인 이득을 제공했습니다. 표현 기반 방식(ReFT)이 단독으로도 강력한 성능을 보였으나, MoDE는 가중치 기반 접근법과 결합했을 때 신뢰할 수 있는 개선을 제공했습니다.
• 지시 이행 (Instruction Following): MoDE는 MT-Bench에서 작지만 일관된 이득을 보였으나, 추론 작업에 비해 개선 폭이 완만했습니다. 이는 지시 이행이 앙상블만으로는 완전히 포착하기 어려운 후기 레이어의 포맷팅(formatting)에 크게 의존함을 시사합니다.
• 절제 연구 및 분석 (Ablation and Analysis):
  • 라우터는 자연스럽게 마지막 한두 개의 레이어에 가중치를 집중시키지만, 증류와 정규화 구성 요소 덕분에 초기 후기 레이어들도 효과적으로 기여할 수 있습니다.
  • 앙상블 크기는 $k \approx 3\text{--}4$일 때 최적입니다. 더 큰 $k$는 쉬운 데이터셋에는 도움이 되지만 복잡한 추론 작업에는 해가 됩니다.
  • 희소 Top-K 라우팅은 대부분의 정확도를 유지하면서 상당한 속도 향상을 가능하게 합니다.

의의 및 주장
저자들은 MoDE를 후기 레이어 로짓을 직접 유용하게 만들 수 있음을 보여주는 메커니즘 연구로 규정합니다. 논문의 주장은 다음과 같습니다:

1. 후기 레이어 표현은 단순한 적응(정규화 및 증류)을 통해 잠금을 해제할 수 있는 상당한 예측 잠재력을 가지고 있습니다.
2. 이러한 중간 신호는 훨씬 더 큰 학습 모듈(예: 전체 레이어 LoRA)을 유사한 성능으로 대체할 수 있는 매개변수 효율적인 대안이 될 수 있습니다.
3. MoDE는 추론 시점의 프로빙(probing)이나 가중치 편집(weight-editing) 방식과는 구별되는, 중간 레이어 표현을 실행 가능한 학습 신호로 전환하는 보완적인 방향을 제공합니다.

본 논문은 고정된 $k$의 사용(동적 선택 불가), 지시 이행에서의 상대적으로 작은 이득, 디코더 전용 LLM에 국한된 평가 등 한계점을 언급하며 그 범위를 겸허하게 설정하고 있습니다. 향-후 연구로는 동적 $k$ 선택, 더 넓은 모델 제품군, 그리고 더 큰 규모의 모델에 대한 적용이 제안되었습니다.