Scaling Linear Mode Connectivity and Merging to Billion Parameter Pretrained Transformers

기술 요약: 수십억 파라미터 규모의 사전 학습된 트랜스포머를 위한 선형 모드 연결성(LMC) 및 병합의 확장

문제 정의
선형 모드 연결성(Linear Mode Connectivity, LMC)은 독립적으로 학습된 신경망들이 가중치 공간 내에서 낮은 손실을 갖는 선형 경로로 연결될 수 있으며, 이를 통해 앙상블이나 재학습 없이도 효율적인 모델 병합이 가능하다는 가설을 제시한다. 그러나 기존 방식들은 이 개념을 수십억 파라미터 규모의 사전 학습된 트랜스포머로 확장하는 데 어려움을 겪고 있다. 주요 장애물은 독립적으로 학습된 모델들이 함수 보존 대칭성(예: 뉴런 치환, 가중치 회전) 및 초기화 또는 학습 역학의 차이로 인해 가중치 공간에서 서로 멀리 떨어진 영역을 점유한다는 점이다. 두 모델 간의 단순한 선형 보간(interpolation)은 높은 손실 장벽(loss barriers)에 직면하게 된다. 가중치 매칭(Weight Matching)이나 학습된 매칭(Learned Matching)과 같은 기존의 정렬 기술은 작은 네트워크에서는 유망한 결과를 보여주었으나, 대개 일방향 최적화(하나의 참조 모델에 다른 모델을 맞추는 방식)에 의존하며 트랜스포머 구조 고유의 복잡한 대칭 그룹을 완전히 활용하지 못해 대규모 모델에서의 효과가 제한적이다.

방법론
저자들은 두 엔드포인트 모델의 대칭 변환을 공동으로 최적화함으로써, 대규모 사전 학습된 트랜스포머에서 낮은 손실 경로를 찾아내는 확장 가능한 프레임워크인 **Dual Learned Matching (LMC-DM)**을 제안한다.

1. 함수 보존 대칭성(Functionality-Preserving Symmetries): 본 논문은 단순한 뉴런 치환을 넘어 트랜스포머에 특화된 광범위한 대칭군을 체계적으로 공식화한다. 여기에는 다음이 포함된다:

   • 정규화 흡수(Normalization Absorption): 아핀 파라미터를 가진 LayerNorm/RMSNorm을 인접한 선형 레이어에 흡수시켜 파라미터가 없는 연산자로 변환한다.
   • 글로벌 잔차 공간 대칭(Global Residual-Space Symmetry): RMSNorm에 의해 유지되는 잔차 스트림에 대한 전역 직교 회전 및 스케일링을 적용한다.
   • 어텐션 헤드 치환(Attention Head Permutations): 어텐션 헤드의 순서를 재배열한다.
   • 회로 대칭성(Circuit Symmetries): 어텐션 헤드의 쿼리-키(QK) 및 밸류-아웃풋(OV) 회로 내의 내부 가역 변환.
   • FFN 대칭성: 은닉 채널 치환 및 GLU 스타일의 리스케일링 대칭.

2. 파라미터화(Parameterization): 이러한 대칭성을 경사 하강법을 통해 학습 가능하게 만들기 위해, 저자들은 연속적 변환을 다음과 같이 파라미터화한다:

   • 직교 행렬을 위한 케이리 변환(Cayley Transform) (효율적이고 안정적임).
   • 가역 행렬을 위한 극 분해(Polar Decomposition) (직교 성분을 위한 케이리 변환과 대칭 양의 정치 행렬 성분을 위한 촐레스키 분해의 결 조합).
   • 이산 대칭(헤드/FFN 치환)은 가중치 매칭을 통해 초기화된 후 학습 단계 동안 고정된다.

3. 이중 학습 절차(Dual Learning Procedure): 하나의 모델($\Theta_B$)을 고정된 참조 모델($\Theta_A$)에 맞추는 기존의 "단방향" 접근 방식과 달리, LMC-DM은 양방향 최적화를 도입한다. 두 모델 모두 공유된 선형 보간 경로를 향해 각자의 변환 파라미터($\pi_A$와 $\pi_B$)를 학습한다:
   $$\Theta_{interp} = \lambda \pi_A(\Theta_A) + (1 - \lambda) \pi_B(\Theta_B)$$
   대칭 파라미터는 보간 경로를 따라 손실을 최소화하도록 직접 최적화되며, 이를 통해 한 모델이 다른 모델의 좌표계에 강제로 맞춰지는 것이 아니라 두 모델 모두가 공유된 저손실 베이슨(low-loss basin)을 향해 이동할 수 있게 한다.

핵심 기여

• 포괄적인 대칭성 공식화: 정규화 흡수, 잔차 공간 회전, 특정 어텐션/FFN 회로 대칭을 포함하여 트랜스포머를 위한 광범위한 함수 보존 변환군을 정의하였다.
• 이중 학습 매칭(Dual Learned Matching): 두 엔드포인트가 모두 변환되는 양방향 최적화 프레임워크를 도입하여, 단방향 방식보다 더 잘 조건화된 최적화 지형과 더 낮은 보간 장벽을 구현하였다.
• 확장성 입증: 수십억 개의 파라미터를 가진 모델까지 LMC 기반 병합을 성공적으로 적용하였으며, 이는 근접 장벽 없는 연결성(near-barrier-free connectivity) 측면에서 이전에 탐구되지 않았던 규모이다.

실험 결과
저자들은 14M에서 6.9B 파라미터에 이르는 Vision Transformer(ViT-S, ViT-B, ViT-L) 및 언어 모델(Pythia 및 OLMo)에 대해 LMC-DM을 평가하였다.

• 손실 장벽(Loss Barriers): WikiText 데이터셋에서 LMC-DM은 160M 파라미터 모델에 대해 0.18의 손실 장벽을 달로, 가중치 매칭(0.27) 및 단방향 학습된 매칭(0.27)을 크게 상회하는 성능을 보였다. 1B 파라미터 모델에서도 장벽은 0.23으로 낮게 유지되었다.
• 비전 도메인: ViT-L은 전체 보간 경로에 걸쳐 69% 이상의 ImageNet top-1 정확도를 유지하였는데, 이는 급격히 하락하는 단순 보간 방식 및 단방향 방식에 비해 상당한 개선이다.
• 확장 동작(Scaling Behavior): 최대 6.9B 파라미터까지의 Pythia 모델은 작은 손실 장벽을 보이는 반면, OLMo-7B는 더 높은 장벽을 보였다. 이는 LMC가 대규모 모델에서 규모뿐만 아니라 학습 데이터 및 최적화 역학 등 다른 요인의 영향을 받는다는 것을 시사한다.
• 절제 연구(Ablation Studies): QK 또는 OV 대칭을 제거하면 특히 ViT에서 손실 장벽이 크게 증가한다. 이중 학습 방식은 단방향 변형보다 일관되적으로 우수한 성능을 보였으며, 직교 대칭을 위한 케이리 기반 파라미터화가 제약 없는 방식이나 행렬 지수(matrix-exponential) 방식보다 더 효과적임을 입증하였다.

의의 및 주장
본 논문은 중간 규모 파라미터(160M)를 가진 언어 모델에 대해 최초로 장벽이 거의 없는 선형 연결성을 입증하였으며, 적절히 정렬될 경우 수십억 파라미터 규모의 LLM 또한 작은 손실 장벽만을 가지고 연결될 수 있음을 보여준다고 주장한다.

저자들은 이러한 결과가 대규모 사전 학습된 트랜스포머가 가중치 공간에서 고립된 점이 아니라, 적절히 정렬되고 병합될 수 있는 "대칭 관계의 표현(symmetry-related representations)"임을 시사한다고 논한다. 이중 최적화를 통해 파라미터 대칭성을 해결함으로써, 이 방법은 선형 모드 연결성이 다루기 불가능하다고 여겨졌던 규모에서도 신뢰할 수 있는 모델 병합 및 구성을 가능하게 한다. 이 연구는 이론적인 손실 지형 기하학과 실질적인 모델 재사용 사이의 간극을 메우며, 대규모 사전 학습 모델의 확장 가능한 구성(composition)을 위한 토대를 제공한다.