Transport and Merge: Cross-Architecture Merging for Large Language Models

1. 문제 정의 (Problem Statement)

• 배경: 대규모 언어 모델 (LLM) 은 방대한 데이터와 모델 용량 스케일링을 통해 뛰어난 성능을 보이지만, 실제 배포 환경 (저자원 언어, 특수 도메인 등) 에서는 제한된 파라미터 예산과 데이터 가용성으로 인해 작은 모델이 주로 사용됩니다.
• 현황: 고자원 (High-resource) 모델에서 저자원 (Low-resource) 모델로 지식을 전이하는 것은 중요한 과제입니다. 기존에는 지식 증류 (Distillation) 나 경사 기반 적응 (Gradient-based adaptation, 예: SFT) 이 주류였으나, 이는 많은 계산 비용이나 대규모 학습 데이터가 필요합니다.
• 제약 사항: 최근 주목받는 모델 머지 (Model Merging) 기법은 학습 데이터 없이 여러 전문가 모델의 가중치를 직접 결합하여 성능을 향상시킬 수 있지만, 대부분의 기존 방법은 동일한 아키텍처 (Architecture) 를 가진 모델 간에만 적용 가능합니다.
• 핵심 문제: 서로 다른 아키텍처 (예: LLaMA-3-8B 와 LLaMA-3-1B, 혹은 Qwen 과 LLaMA 등) 를 가진 고자원 모델과 저자원 모델 간에 직접적인 가중치 머징을 통해 지식을 전이할 수 있는가? 라는 질문이 제기됩니다.

2. 제안 방법론 (Methodology)

저자들은 최적 수송 (Optimal Transport, OT) 이론을 기반으로 한 크로스-아키텍처 머징 프레임워크를 제안합니다. 이 방법은 활성화 공간 (Activation Space) 에서의 정렬을 통해 서로 다른 아키텍처 간의 뉴런 대응 관계를 추론하고, 이를 가중치 공간 (Weight Space) 의 직접적인 결합으로 변환합니다.

주요 단계:

1. 활성화 기반 대응 관계 추론 (Activation-based Correspondence Inference):

   • 소스 모델 (고자원) 과 타겟 모델 (저자원) 에 동일한 작은 입력 데이터셋 ($D$) 을 통과시켜 각 레이어의 중간 활성화 (Intermediate Activations) 를 추출합니다.
   • 특징 수준 (Feature-level): 각 레이어 내에서 소스 활성화 채널과 타겟 활성화 채널 간의 상관관계를 비용 행렬 (Cost Matrix) 로 정의하고, 엔트로피 정규화된 최적 수송 (Entropic OT) 을 풀어 특징 전송 행렬 ($Q_{\ell m}$) 을 구합니다. 이는 타겟의 특정 뉴런이 소스의 어떤 뉴런들의 조합으로 표현될 수 있는지를 나타냅니다.
   • 레이어 수준 (Layer-level): 레이어 간 유사성을 집계하여 전역적인 레이어 매핑 행렬 ($P$) 을 OT 를 통해 추론합니다.

2. 가중치 공간 융합 (Weight-space Fusion):

   • 활성화 공간에서 추론된 대응 관계를 가중치 공간으로 승격 (Lift) 시킵니다.
   • 선택적 뉴런 교체 (Selective Top-k Neuron Replacement): 모든 뉴런을 무작정 합치는 대신, 타겟 데이터셋에서 가장 활발하게 반응하는 상위 $k$개의 뉴런만 선별합니다. 이는 이질적인 아키텍처 간의 불필요한 간섭을 줄이고 노이즈를 필터링하기 위함입니다.
   • 잔차 기반 융합 공식:
     $$W^{\ell, fused}_A = W^{\ell}_A + \alpha \cdot M_\ell \odot \left( \sum_m P_{eff}[\ell, m] Q^{out}_{\ell m} W^m_B (Q^{in}_{\ell m})^\top - W^{\ell}_A \right)$$
     • 여기서 $W_B$는 소스 가중치, $W_A$는 타겟 가중치, $Q$는 OT 를 통해 구한 전송 행렬, $M_\ell$은 상위 $k$개 뉴런을 선택하는 마스크입니다.
     • 이 공식은 소스 모델의 가중치를 타겟 모델의 좌표계로 변환하여 타겟 가중치에 주입하는 원리입니다.

3. 잔차 고정 적응 (Residual-Frozen Adaptation):

   • 머징 후, 전이된 잔차 (Residual) 파라미터는 고정 (Freeze) 하고, 원래 타겟 모델의 베이스 파라미터만 미세 조정 (Fine-tuning) 합니다.
   • 이는 전이된 지식이 이후 최적화 과정에서 덮어쓰이지 않도록 보장하면서도, 타겟 모델이 하위 작업 분포에 적응할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

1. 이질적 아키텍처 간 직접 머징: 기존 머징 기법의 한계였던 "동일 아키텍처" 제약을 극복하고, 최적 수송을 활용하여 구조가 다른 모델 간 직접적인 가중치 전이를 가능하게 했습니다.
2. 적극적 지식 전이 메커니즘: 활성화 공간의 정렬을 통해 뉴런 수준의 대응 관계를 추론하고, 이를 가중치 공간의 선형 변환으로 해석하여 수학적 엄밀성을 갖췄습니다.
3. 데이터 효율성: 대규모 학습 데이터나 반복적인 경사 하강법이 필요 없으며, 소량의 입력 데이터셋 (Calibration Set) 만으로 지식을 전이할 수 있어 계산 비용과 데이터 요구 사항을 크게 절감합니다.
4. 실용성: 저자원 언어 및 특수 도메인 (금융, 의료) 에서 기존 베이스 모델 대비 일관된 성능 향상을 입증했습니다.

4. 실험 결과 (Results)

저자들은 다양한 저자원 언어 (말레이, 인도네시아어, 태국어, 광둥어) 와 전문 도메인 (금융, 의료) 에서 실험을 수행했습니다.

• 저자원 언어 전이 (Low-resource Language Transfer):
  • 말레이어 (MalayMMLU): 베이스 모델 대비 'Fused w/ Adaptation' 설정에서 모든 카테고리 (인문학, 언어, STEM 등) 에서 3.3~6.5%p의 정확도 향상을 보였습니다.
  • 인도네시아어, 광둥어, 태국어: 다양한 벤치마크 (Belebele, CMMLU, MGSM 등) 에서 머징만으로도 성능이 향상되었으며, 적응 (Adaptation) 을 거치면 추가적인 개선을 이루었습니다. 특히 태국어 MGSM 태스크에서는 머징만으로도 0.50 에서 0.64 로 크게 상승했습니다.
• 일반 도메인에서 전문 도메인으로의 전이:
  • 금융 및 의료: 일반 목적 LLM 과 전문 도메인 LLM 을 머징했을 때, 적응 단계를 거친 모델이 모든 벤치마크에서 최상의 성능을 기록했습니다.
• 비교 분석:
  • SFT 및 증류 (Distillation) 대비: 정규화된 점수 기준으로 SFT 나 증류보다 일관되게 우수한 성능을 보였습니다.
  • 소스 모델 크기 영향: 소스 모델의 크기가 커질수록 (1B → 8B → 32B) 전이 성능이 선형적으로 향상됨을 확인했습니다.
  • 일반 능력 보존: 머징 후에도 모델의 일반적 추론 능력 (Commonsense, Reasoning) 이 저하되지 않음을 확인했습니다.

5. 의의 및 결론 (Significance)

• 실용적 대안 제시: 저자원 환경에서 고자원 모델의 지식을 활용하기 위해, 고비용의 재학습이나 증류 없이 경량화된 머징 기법을 제공할 수 있습니다.
• 이론적 통찰: 서로 다른 아키텍처를 가진 모델들 사이에도 활성화 공간에서 정렬된 내부 표현 (Aligned Internal Representations) 이 존재한다는 것을 최적 수송을 통해 증명하고, 이를 가중치 결합으로 연결하는 이론적 틀을 마련했습니다.
• 확장성: 이 방법은 다양한 모델 패밀리 (LLaMA, Qwen 등) 와 아키텍처 조합에 적용 가능하며, 저자원 언어와 특수 도메인 모델 개발의 장벽을 낮추는 핵심 기술로 기대됩니다.

요약하자면, 이 논문은 최적 수송 (Optimal Transport) 을 활용하여 서로 다른 구조의 LLM 간에 뉴런 단위의 대응 관계를 찾아내고, 이를 통해 지식을 직접 주입하는 새로운 머징 프레임워크를 제안함으로써, 저자원 환경에서의 모델 성능 향상을 위한 효율적이고 강력한 솔루션을 제시했습니다.