기술 요약: 옵티마이저 유도 모드 연결성: AdamW 에서 Muon 으로

문제 제기

모드 연결성 (mode connectivity) 은 독립적으로 훈련된 신경망 모델들이 낮은 손실 경로를 통해 연결될 수 있다는 현상으로, 광범위하게 연구되어 왔습니다. 그러나 기존 문헌은 대부분 해 집단을 단일한 전체로 취급하며, 이러한 해에 도달하는 데 사용된 특정 최적화 역학과는 무관한 연결성을 가정하는 경향이 있습니다. 본 논문은 서로 다른 옵티마이저가 고유한 암시적 편향 (implicit biases)(정규화) 을 부과하여 해를 손실 지형의 서로 다른 영역으로 제한한다고 주장합니다. 다루는 핵심 질문은 다음과 같습니다: 특정 옵티마이저의 암시적 편향과 일치하는 해로 제한될 때 모드 연결성은 어떻게 행동하는가? 구체적으로, 저자들은 동일한 옵티마이저가 찾은 해들이 연결된 집단을 형성하는지 (옵티마이저 내부 연결성), 그리고 서로 다른 옵티마이저가 찾은 해들이 손실 장벽 없이 연결될 수 있는지 (옵티마이저 간 연결성) 를 조사합니다.

방법론

이론적 프레임워크

저자들은 옵티마이저 정규화 해 집합의 개념을 공식화합니다. 적응형 옵티마이저 (예: AdamW 및 Lion-K 계열) 에 대한 최근의 이론적 특성화에 기반하여, 이러한 옵티마이저의 극한점을 제약 최적화 문제의 해로 정의합니다:

• AdamW: $\ell_\infty$-노름 제약 ($\|\theta\|_{\max} \leq 1/\lambda$) 을 만족하는 해로 수렴합니다.
• Muon (및 핵노름을 사용하는 Lion-K): 연산자 노름 제약 ($\|\theta\|_{\text{op}} \leq 1/\lambda$) 을 만족하는 해로 수렴합니다.

논문은 두 가지 정의를 도입합니다:

1. 옵티마이저 내부 연결성: 특정 옵티마이저 $A$에 대한 해 집합 $O_A(\lambda)$가 연결되어 있다.
2. 옵티마이저 간 연결성: 해 집합의 합집합 $O_{A_1}(\lambda_1) \cup O_{A_2}(\lambda_2)$가 연결되어 있다.

이론적 분석은 보간 영역 (제로 훈련 손실 달성 가능) 에서의 2 층 ReLU 네트워크에 대해 수행됩니다. 저자들은 다양한 네트워크 폭 ($m$) 하에서 이러한 정규화 집합의 기하학을 분석합니다.

실험 설정

실증적 검증은 5 개의 데이터셋 (enwik8, WikiText-103, Stories, BookCorpus, One Billion Word) 에 대한 GPT-2 사전 훈련 작업에서 수행되었으며, 두 가지 모델 규모 (~30M 및 ~50M 파라미터) 를 사용했습니다.

• 옵티마이저: AdamW 및 Muon.
• 경로 구성: 저자들은 수정된 폴리체인 보간 (polychain interpolation) 알고리즘과 대칭 정렬 (어텐션 헤드와 MLP 뉴런의 치환 및 직교 정렬) 을 결합하여 사용합니다. 이 방법은 모델 간의 손실 장벽을 최소화하기 위해 학습 가능한 굽힘점을 가진 조각별 선형 경로를 구성합니다.
• 지표: 연구는 보간 경로 along validation loss 장벽과 특이값 스펙트럼 (특히 안정적 순위) 의 진화를 추적합니다.

주요 기여

1. 옵티마이저 내부 연결성에 대한 이론적 증명

논문은 충분히 넓은 2 층 ReLU 네트워크의 경우, 특정 옵티마이저에 의해 유도된 해 집합이 연결되어 있음을 증명합니다:

• 정리 3.1: 네트워크 폭 $m \geq 4P$일 때 (여기서 $P$는 서로 다른 초평면 배열 패턴의 수), Muon(및 모멘텀이 있는 정규화 GD) 에 대해 옵티마이저 내부 연결성이 성립합니다.
• 정리 3.2: $m \geq m^*$일 때 (여기서 $m^*$는 데이터와 정규화 매개변수에 의해 결정되는 임계 폭), AdamW(및 Signum) 에 대해 옵티마이저 내부 연결성이 성립합니다.
• 메커니즘: 증명 과정은 임의의 해를 정규화 집합 내의 "희소"또는 "균등화"된 표준형으로 축소하고, 이러한 형태가 암시적 노름 제약을 보존하는 연속 경로를 통해 연결될 수 있음을 보여줍니다.

2. 옵티마이저 간 연결성의 특성화

저자들은 서로 다른 옵티마이저에 의해 유도된 영역 간의 연결성이 네트워크 폭과 정규화 강도에 크게 의존함을 확립합니다:

• 큰 폭 영역: 두 정규화 집합의 합집합은 집합들이 겹칠 때 (즉, 제로 손실 집합이 각각의 노름 공의 교차점을 통과할 때) 연결됩니다. 정규화 강도가 증가함에 따라 연결된 구성 요소의 수가 1 개에서 2 개로 전환될 수 있습니다.
• 유한 폭 구성: 논문은 제로 손실 해 집합이 여러 개의 불연속 구성 요소로 분해되는 특정 데이터셋과 폭 ($m=2$) 을 구성합니다. 중요하게도, 적절한 정규화 매개변수를 선택함으로써 AdamW 와 Muon 을 서로 다른 구성 요소로 강제할 수 있음을 증명합니다. 이 시나리오에서 AdamW 해와 Muon 해 사이에는 제로 손실 경로가 존재하지 않으며, 어떤 연결 경로도 증명 가능한 손실 장벽 (제곱 손실의 경우 적어도 $1/2$) 을 수반해야 합니다.

3. 실증적 스펙트럼 전이

GPT-2 실험은 뚜렷한 스펙트럼 행동을 보여줍니다:

• 동일 옵티마이저 경로: 동일한 옵티마이저로 훈련된 두 모델 간 보간 (예: AdamW-AdamW) 은 특이값 스펙트럼을 보존합니다 (예: AdamW 는 아웃라이어 유지, Muon 은 등방성 유지) 그리고 낮은 손실을 유지합니다.
• 교차 옵티마이저 경로: 서로 다른 옵티마이저 간 보간 (예: AdamW-Muon) 은 부드러운 스펙트럼 위상 전이를 초래합니다. 스펙트럼은 경로 along AdamW 서명에서 Muon 서명으로 진화합니다. 손실은 상대적으로 낮게 유지되지만, 중간 모델들은 양쪽 끝점과 구별되는 스펙트럼 프로파일을 보입니다.

4. 분포 외 (OOD) 일반화

논문은 교차 옵티마이저 보간 경로가 OOD 일반화를 향상시킬 수 있음을 보여줍니다. AdamW-Muon 경로 along 샘플링된 모델들은 OOD 데이터셋에서 음의 손실 장벽을 나타내며 (끝점의 선형 보간보다 성능이 우수함), 이는 서로 다른 옵티마이저 편향 사이의 공간을 탐색함으로써 구조적으로 다양한 모델을 생성하여 우수한 일반화 특성을 얻음을 시사합니다.

결과 및 의의

연구 결과 요약

• 옵티마이저 의존적 구조: 손실 지형은 낮은 손실 해의 단일 연결된 덩어리가 아니라, 암시적 정규화에 의해 정의된 옵티마이저별 "분지"또는 영역을 포함합니다.
• 조건부 연결성: 넓은 네트워크에서는 동일 옵티마이저 해가 연결되어 있지만, 유한 폭 설정에서는 서로 다른 옵티마이저가 손실 장벽으로 분리된 불연속 구성 요소로 수렴할 수 있습니다.
• 스펙트럼 지문: 특이값 스펙트럼은 옵티마이저 유도 영역에 대한 강력한 대리 지표 역할을 합니다. 동일 옵티마이저 경로는 이 지문을 보존하는 반면, 교차 옵티마이저 경로는 전이를 통과합니다.

주장된 의의

본 논문은 고전적인 모드 연결성 문헌에서는 보이지 않는 옵티마이저 의존적 구조를 드러낸다고 주장합니다. 암시적 편향을 고려하여 모드 연결성의 정의를 정제함으로써, 저자들은 다음을 보여줍니다:

1. 옵티마이저 선택은 도달 가능한 해 공간의 위상을 근본적으로 형성합니다.
2. 교차 옵티마이저 보간은 해 공간의 다양한 영역에 접근하는 모델 병합의 실현 가능한 메커니즘이며, 양자화 및 일반화와 같은 하류 특성을 개선할 수 있습니다.
3. 연결성에 대한 이론적 보장은 네트워크 아키텍처뿐만 아니라 옵티마이저가 유도하는 특정 정규화자를 고려해야 합니다.

이 연구는 이론적 암시적 편향 분석과 실용적 모델 융합 간의 간극을 연결하며, 모델 간의 "최적"경로가 서로 다른 옵티마이저 유도 정규화의 경계를 탐색하는 데 의존할 수 있음을 시사합니다.