기술 요약: 모델 병합이 라우팅을 깨뜨릴 때: MoE를 위한 학습 없는 보정 (Training-Free Calibration)

1. 문제 정의

모델 병합은 공동 재학습의 막대한 비용 없이 여러 거대 언어 모델(LLM)의 능력을 통합하는 비용 효율적인 전략이 되었습니다. 기존 기술들(예: Task Arithmetic, TIES-Merging)은 선형 파라미터 산술을 통해 밀집(dense) 아키텍처를 성공적으로 병합하지만, 혼합 전문가(Mixture-of-Experts, MoE) 아키텍처에 적용될 때는 실패합니다.

저자들은 **라우팅 붕괴(routing breakdown)**라고 명명된 결정적인 실패 모드를 식별했습니다. MoE 모델에서 라우터는 비선형 소프트맥스(softmax) 활성화 함수와 그에 따른 이산적 Top-k 선택을 사용하여 토큰을 특정 전문가 집합으로 동적으로 배정합니다. 본 논문은 표준적인 선형 병합 방법이 다음과 같은 이유로 이 설정에 부적합하다고 주장합니다:

1. 비선형성 불일치: 라우터 가중치의 선형 보간($W_m = \sum \lambda_i W_i$)은 소프트맥스 함수의 볼록성(convexity)으로 인해 라우팅 확률의 선형 보간으로 이어지지 않습니다.
2. 민감도: 라우터 가중치의 작은 섭동(perturbation)은 소프트맥스와 Top-k 메커니즘에 의해 기하급수적으로 증폭됩니다.
3. 부하 분산 제약: 로드 밸런싱을 강제하는 사전 학습 목표는 전문가 할당을 게이팅 변화에 매우 민감하게 만듭니다.

경험적 증거: 두 개의 OLMoE 모델(수학 및 코드용 미세 조정 모델)을 병합한 예비 실험 결과, 병합 직후 전문가 선택의 50.11%가 소스 모델과 달라졌습니다. 이 불안정성은 네트워크 깊이를 통해 누적되어, 16개 레이어 전체에서 99.98%의 토큰이 라우팅 발산을 경험했습니다. 결정적으로, 교차 평가를 통해 병합된 모델에 원래의 라우터를 복구했을 때는 성능이 회복되었으나, 병합된 라우터를 원래 모델에 적용했을 때는 성능이 저하됨을 확인했습니다. 이는 일관된 라우팅 동작 없이는 전문가 파라미터를 보존하는 것이 무의미함을 입증합니다.

2. 방법론: 헤시안 인식 라우터 보정 (Hessian-Aware Router Calibration, HARC)

라우팅 붕괴를 해결하기 위해, 저자들은 2차 곡률 정보를 사용하여 병합된 라우터를 재정렬하는 학습 없는 프레임워크인 HARC를 제안합니다.

2.1. 이론적 정식화

목표는 모든 태스크에 대해 병합된 라우터의 출력 분포와 소스 라우터들의 분포 사이의 쿨백-라이블러(KL) 발산을 최소화하는 통합 라우터 가중치 행렬 $W_m$을 찾는 것입니다.
$$W^*_m = \arg \min_{W_m} \sum_{i=1}^D \mathbb{E}_{x \sim \mathcal{D}_i} [\text{KL}(r_i(x) \parallel r_m(x; W_m))]$$

직접적인 최적화는 소프트맥스의 비선형성으로 인해 다루기 어렵습니다. 저자들은 2차 테일러 전개(Lemma 4.1)를 사용하여 KL 발산의 **이차 근사(quadratic approximation)**를 도출합니다:
$$\text{KL}(r_i \parallel r_m) \approx \frac{1}{2}(z_m - z_i)^\top H_i (z_m - z_i)$$
여기서 $H_i = \text{diag}(r_i) - r_i r_i^\top$는 로그 분할 함수(log-partition function)의 헤시안(Hessian)입니다. 이 헤시안은 전문가 간의 경쟁적 역학(비대각 성분)과 활성 전문가의 민감도(대각 성분)를 포착합니다.

2.2. 폐쇄형 해법 (Closed-Form Solution)

선형 관계 $z = Wx$를 이차 근사에 대입하면 문제를 가중 최소 제곱 시스템으로 변환합니다. 최적의 가중치$w^*_m$(벡터화된 형태의$W_m$)은 다음 선형 시스템을 만족합니다:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
여기서 $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$는 태스크별 정밀도 행렬(precision matrix)입니다.

2.3. 효율적 계산 (Matrix-Free)

$\Sigma$를 명시적으로 구성하고 역행렬을 구하는 것은 계산적으로 매우 부담이 큽($O(K^2d^2)$ 공간 복잡도). HARC는 확장성을 보장하기 위해 두 가지 핵심 전략을 사용합니다:

1. 대각 정규화 (Diagonal Regularization): 희소한 Top-k 게이팅으로 인한 헤시안의 계수 결핍(rank-deficiency)을 처리하기 위해, 저자들은 전체 정밀도 행렬과 그 대각 근사 사이를 보간하는 정규화 항을 도입합니다. 이는 전문가 간의 독립성을 가정하지 않고도 솔루션을 안정화합니다.
2. 행렬 없는 공액 경사법 (Matrix-Free Conjugate Gradient, CG): $\Sigma$를 직접 생성하는 대신, 알고리즘은 크로네커 곱(Kronecker product) 항등식을 사용하여 온더플라이(on the fly)로 행렬-벡터 곱 $\Sigma w$를 계산하는 행렬 없는 연산자를 사용합니다. 이를 통해 공간 복잡도를 $O(Kd)$로 줄이고 공액 경사법(CG)을 효율적으로 사용할 수 있게 합니다.

2.4. 알고리즘

HARC는 보정되지 않은 병합 모델에 대한 후처리 단계로 작동합니다. 각 MoE 레이어에 대해:

1. 병합된 모델에 보정 데이터를 통과시켜 현재 레이어까지의 입력 특징(features)을 생성합니다.
2. 소스 모델들에 대한 라우팅 출력과 헤시안을 계산합니다.
3. 선형 시스템을 형성하기 위해 통계량을 누적합니다.
4. 행렬 없는 CG 솔버를 사용하여 보정된 라우터 가중치를 구합니다.
5. 병합된 모델의 라우터 가중치를 업데이트합니다.

3. 주요 기여

1. 라우팅 붕괴 식별: 본 논문은 전문가 파라미터의 간섭보다 라우팅 불안정성이 MoE 병합 시 성능 저하의 주요 원인임을 확립했습니다.
2. HARC 프레임워크: 2차 헤시안 정보를 활용하여 라우터 분포를 재정렬하는 새로운 학습 없는 보정 방법을 제안하며, 비선형 병합 문제에 대한 폐쇄형 해법을 제공합니다.
3. 확장 가능한 최적화: 헤시안을 명시적으로 구현하지 않고도 대규모 MoE 아키텍처에 대응할 수 있는 행렬 없는 공액 경사법 솔버를 유도했습니다.
4. 경험적 검증: 다양한 병합 베이스라인(Weight Averaging, TIES-Merging, WUDI 등)과 태스크(Math, Code)에 대해 HARC가 효과적으로 라우팅 붕괴를 완화함을 입증했습니다.

4. 실험 결과

실험은 수학적 추론(GSM8K, MATH500)과 코드 생성(HumanEval+, MBPP+)에 미세 조정된 변형 모델들을 병합한 OLMoE-1B-7B-0125 모델을 대상으로 수행되었습니다.

• 지속적인 개선: HARC는 테스트된 모든 병합 베이스라인에서 일관되게 성능을 향상시켰습니다. 예를 들어, 강력한 WUDI-Merging 베이스라인에 적용했을 때 전체 점수가 38.08에서 38.96으로 상승했습니다.
• 도메인 민감도 완화: 단순 병합은 코드 생성 성능을 심각하게 저하시켰습니다(예: Weight Averaging 적용 시 HumanEval+가 33.50에서 20.27로 하락). HARC는 이 손실을 부분적으로 회복하여, 전문가 선택을 교정함으로써 도메인 특화 능력을 보존할 수 있음을 보여주었습니다.
• 절제 연구 (Ablation Studies):
  • 헤시안 구조: 전체 헤시안을 사용하는 것이 대각 또는 항등 근사를 사용하는 것보다 우수한 성능을 보였으며, 이는 전문가 간의 경쟁적 비대각 역학을 모델링하는 것이 중요하다는 것을 확인시켜 줍니다.
  • 정규화: 오프-대각 성분을 억제하는 댐핑 인자 $\alpha$가 대각 스케일링 $\gamma$보다 더 안정적이었습니다.
  • 데이터 구성: 필터링된 올바른 소스 응답(SrcCrt)을 사용했을 때 가장 좋은 결과를 얻었습니다. 놀랍게도 프롬프트만 사용하는 것이 전체 응답을 사용하는 것과 유사한 성능을 보였는데, 이는 초기 라우팅 결정이 매우 중요함을 시사합니다.
• 데이터 효율성: HARC는 단 **40%**의 보정 샘플만으로도 최적에 가까운 성능에 도달하여, 2차 통계량의 빠른 수렴을 보여주었습니다.
• 레이어별 분석: 베이스라인 방법들의 KL 발산 시각화 결과, 라우팅 오류가 깊은 레이어로 갈수록 기하급수적으로 누적됨을 확인했습니다. HARC는 이러한 추세를 성공적으로 억제하여 중요한 깊은 레이어에서의 정렬을 유지했습니다.

5. 의의 및 주장

본 논문은 희소 아키텍처에서 라우팅 로직을 보존하는 것이 전문가 파라미터를 정렬하는 것만큼 필수적임을 주장합니다. 라우팅 붕괴라는 특정 실패 모드를 해결함으로써, HARC는 재학습 없이 MoE 모델을 통합할 수 있는 견고하고 효율적인 경로를 제공합니다.

저자들은 본 방법이 다음과 같음을 강조합니다:

• 전문가에 대한 그래디언트 역전파나 재학습을 요구하지 않습니다.
• 가벼운 후처리 단계로서 다양한 병합 전략(선형 산술, 최적화 기반)에 걸쳐 일반화됩니다.
• 여러 전문가 모델을 신뢰성 있게 재사용할 수 있게 함으로써, 유능한 모델을 배포하는 데 드는 계산 비용과 장벽을 낮춥니다.

결론적으로, 기존의 선형 병합 가설(Linear Mode Connectivity)은 밀집 레이어에는 유효하지만, MoE 라우터에는 취약합니다. HARC는 MoE 아키텍처 고유의 비선형 게이팅 역학을 명시적으로 고려함으로써 이 간극을 메웁니다.