기술 요약: CLHA – 가중치 공유 Mixture-of-Experts 트랜스포머 양자화를 위한 교차 레이어 헤시안 집계 (Cross-Layer Hessian Aggregation)

1. 문제 정의

Megrez2와 같은 최신 효율적 Mixture-of-Experts (MoE) 트랜스포머 아키텍처는 **교차 레이어 가중치 공유 (CLWS)**를 채택하고 있습니다. 이 설계에서는 인접한 트랜스포머 레이어(예: $\ell$ 및 $\ell+1$)가 동일한 라우티드 전문가 가중치 텐서를 재사용하여, 토큰당 활성화되는 파라미터 수는 유지하면서 파라미터 점유량을 절반으로 줄입니다.

그러나 GPTQ 방식의 접근법을 포함한 기존의 사후 학습 양자화 (Post-Training Quantization, PTQ) 방법들은 이러한 공유 가중치를 단일 레이어에 속한 것으로 취급합니다. 이들은 하나의 "주요(primary)" 레이어에서 추출한 활성화 통계만을 사용하여 캘리브레이션을 수행합니다. 이는 근본적인 불일치를 초래합니다:

• 잔차 연결(residual stream) 표현은 레이어 간에 진화하며, 이에 따라 레이어 $\ell_a$에서 해당 공유 전문가를 활성화하는 토큰 분포는 레이어 $\ell_b$에서의 분포와 다릅니다.
• 기존 방식은 보조 레이어의 활성화 통계를 무시하므로, 특히 저비트 영역(예: 2-bit)에서 재구성 오차(reconstruction error) 최소화가 최적화되지 못합니다.

2. 방법론

2.1 이론적 기초: 교차 레이어 헤시안 집계

본 논문은 레이어 $\ell_1, \dots, \ell_n$에 걸쳐 공유되는 가중치 $W$를 양자화하기 위한 목적 함수를 정식화합니다. 올바른 목적 함수는 해당 가중치를 사용하는 모든 레이어에 걸쳐 총 재구성 오차를 최소화하는 것입니다:
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
여기서 $H_{\ell_i}$는 레이어 $\ell_i$의 헤시안(활성화 공분산)이며, $\alpha_i$는 각 레이어의 라우팅 빈도(토큰 비율)를 나타냅니다.

**명제 1(Proposition 1)**은 이 결합된 오차를 최소화하는 것이 **결합된 헤시안(combined Hessian)**을 적용하는 것과 수학적으로 동일함을 입증합니다:
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
이는 단일 레이어의 헤시안($H_{\ell_a}$)만을 사용하여 오차를 최소화하는 레이어별 독립 캘리브레이션 (Per-Layer Independent Calibration, PLIC) 방식과 대조됩니다.

2.2 CLHA 알고리즘

저자들은 세 단계의 패스를 거치는 캘리브레이션 알고리즘인 **CLHA (Cross-Layer Hessian Aggregation)**를 제안합니다:

1. 수집 (Collection): 각 캘리브레이션 시퀀스에 대해, 모든 활성 전문가에 대한 레이어별 헤시안과 토큰 수를 누적합니다.
2. 집계 (Aggregation): 각 공유 전문가 쌍 $(\ell_a, \ell_b)$에 대해 결합된 헤시안을 계산합니다. 본 논문은 두 가지 가중치 전략을 평가합니다:
   • 라우팅 빈도 가중치 ($\alpha = n_a / (n_a + n_b)$): 각 레이어가 수신한 실제 토큰 수에 따라 헤시안에 가중치를 부여합니다.
   • 동일 가중치 ($\alpha = 0.5$): 캘리브레이션 데이터가 제한적이거나 보조 레이어의 헤시안 추정치가 노이즈가 심할 때 유용한 정규화 역할을 합니다.
3. 양자화 (Quantization): 집계된 $H_{comb}$를 사용하여 표준 GPTQ 열 순차 업데이트(column-sequential updates)를 적용합니다.

2.3 CLHA-MP (혼합 정밀도)

**항상 활성화되는 공유 전문가(always-active shared experts)**는 라우티드 전문가(fraction $K/E$만큼만 수신)와 달리 모든 $N$개의 토큰을 받게 된다는 점을 인식하여, 이들의 양자화 민감도가 $E/K$배 더 높다는 점을 논문은 지적합니다. CLHA-MP는 라우티드 전문가를 $b$ 비트로 유지하면서 공유 전문가에는 $b+1$ 비트를 할당하여, 최소한의 저장 공간 오버헤드(전문가 파라미터의 10~20%)로 이를 처리합니다.

2.4 구현상의 오류 수정: 훅 모호성 (Hook Ambiguity)

본 논문은 중요한 엔지니어링 함정을 식별했습니다. 공유된 ExpertFFN 모듈에 부착된 표준 PyTorch 포워드 훅(forward hooks)은 해당 모듈을 호출하는 어떠한 레이어에 대해서도 실행됩니다. 이로 인해 서로 다른 레이어의 활성화 값이 단일 헤시안 누적기에 혼합되어, 레이어별 귀속(attribution)이 왜곡됩니다.

• 해결책: 저자들은 MoELayer.forward() 시그니처를 통해 전달되는 **레이어 인지 콜백 메커니즘(layer-aware callback mechanism)**을 제안합니다. 이는 레이어 인덱스를 콜백에 명시적으로 제공함으로써, 하위 가중치 모듈이 공유되더라도 $H_a$와 $H_b$를 별도로 분리하여 누적할 수 있도록 보장합니다.

3. 주요 기여

1. 이론적 명제: 공유 가중치에 대한 최적의 헤시안은 레이어별 헤시안의 라우팅 빈도 가중 합계라는 것을 정식화했습니다.
2. CLHA 알고리즘: GPTQ 프레임워크 내에서 결합된 헤시안 캘리브레이션을 구현하는 실용적인 방법을 제시했습니다.
3. CLHA-MP: 매우 민감한 공유 전문가에 추가 비트를 할당하는 혼합 정밀도 확장 방식을 제안했습니다.
4. 구현 수정: 공유 모듈 아키텍처에서 캘리브레이션을 조용히 망가뜨리는 "훅 모호성" 문제를 식별하고 해결했습니다.
5. 실증적 검증: CLHA의 이득이 교차 레이어 라우팅 발산(cross-layer routing divergence)과 상관관계가 있음을 입증했습니다.

4. 실험 결과

실험은 WikiText-2로 처음부터 학습된 34.5M 파라미터 규모의 CLWS-MoE 모델을 대상으로 수행되었습니다.

• 2-bit 성능: 미캘리브레이션(miscalibration)이 중첩되는 임계 영역입니다.
  • CLHA-rt (라우팅 가중치 적용)는 PLIC 대비 퍼플렉시티(PPL)를 4.3% 감소시켰습니다.
  • CLHA-eq (동일 가중치 적용)는 PPL을 5.4% 감소시켰습니다.
  • CLHA-MP는 PLIC 대비 **18.6%**의 개선을 달왔습니다.
  • 4-bit 및 3-bit에서는 차이가 오차 범위 내(<1%)였습니다.
• 캘리브레이션 vs. 테스트 오차: 흥미롭게도, PLIC는 해당 분포에 직접 최적화되므로 훈련 세트에서의 캘리브레이션 헤시안 오차(CHE)는 더 낮았습니다. 그러나 CLHA가 테스트 세트에서 더 잘 일반화되었는데, 이는 결합된 헤시안이 단일 레이어 관점보다 공유 가중치의 실제 분포를 더 잘 포착함을 나타냅니다.
• 라우팅 날카로움(Routing Sharpness) 실험: 라우터 로짓을 인위적으로 날카롭게 만들어(교차 레이어 발산을 증가시켜) 관찰한 결과:
  • 발산이 증가함에 따라 PLIC 성능은 5% 저하되었습니다.
  • CLHA는 안정적으로 유지되었습니다(<1% 저하).
  • CLHA와 PLIC 사이의 성능 격차는 자연스러운 라우팅 시 **+4.3%**에서 높은 발산 시 **+8.0%**로 커졌으며, 이는 CLHA의 이득이 교차 레이어 발산에 따라 커진다는 이론적 예측을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 교차 레이어 가중치 공유가 적용된 MoE 모델에 대한 최초의 PTQ 분석을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 목적 함수의 교정: 공유 가중치를 단일 레이어 개체로 취하는 것이 수학적으로 최적이 아니며, 헤시안 집계가 필수적인 교정임을 입증했습니다.
• 이득의 확장성: 모델이 커지고 라우팅 패턴이 전문화됨에 따라(교차 레이어 발산 증가), CLHA의 이점이 커질 것이며, 이는 미래의 대규모 CLWS-MoE(예: Megrez2, DeepSeek, Qwen)에 필수적임을 보여줍니다.
• 실무적 필요성: 기존 PTQ 도구를 CLWS 아키텍처에 나이브하게 적용하는 것은 근본적으로 잘못되었음을 지적하고, 이 분야의 향후 캘리브레이션 방법을 위한 필수적인 해결책을 제공했습니다.

저자들은 모델 규모가 작고(34.5M) 학습 단계가 제한적이라는 점을 언급하며 절대적인 성능에 대해서는 겸손한 태도를 유지하지만, 표준 캘리브레이션 대비 상대적인 개선 효과가 견고하고 이론적으로 근거가 있음을 강조합니다.