기술 요약: 곡률 근사를 위한 가중치 공간 대칭성 활용

문제 정의

손실 함수의 곡률(헤시안 행렬)을 정확하게 추정하는 것은 2차 최적화, 베이지안 딥러닝, 지속 학습(continual learning), 모델 압축을 포함한 많은 머신러닝 하위 분야에서 매우 중요합니다. 그러나 현대적인 딥 네트워크를 위해 전체 헤시안을 구축, 저장 및 역행렬을 구하는 것은 계산 및 메모리 측면에서 불가능에 가깝습니다. (E)KFAC, Shampoo, Soap와 같은 기존의 근사 기법들은 특정 블록 대각(block-diagonal) 또는 크로네커 인수 분해(Kronecker-factored) 구조를 가정함으로써 확장성을 달성했지만, 이 분야에는 왜 이러한 특정 구조들이 작동하는지, 혹은 모델 아키텍처에 기반하여 어떻게 체계적으로 유도할 수 있는지에 대한 원칙적인 이론적 프레임워크가 부족합니다. 또한, 기존 연구들은 신경망의 잘 알려진 가중치 공간 대칭성(예: 은닉 유닛의 치환 불변성)에 의해 부과되는 곡률의 구조적 제약을 대체로 간과해 왔습니다.

방법론

저자들은 가중치 공간 대칭성에 의해 생성되는 "궤도(orbit)"에 대해 분석적으로 평균을 내는 새로운 헤시안 근사 프레임워크를 제안합니다. 핵심 통찰은 많은 손실 함수가 특정 파라미터 변환(군 작용, group actions)에 대해 불변이라는 점이며, 이는 손실의 그래디언트(gradient) 역시 동일한 변환에 대해 공변성(equivariance)을 가짐을 의미합니다.

1. 궤도 평균화 및 그래디언트 공변성

현재의 파라미터 집합 $w$와 그래디언트 $g$가 주어졌을 때, 저자들은 대칭 군 $G$로부터 적용 가능한 대칭 변환들에 의해 도달할 수 있는 모든 $(w^{(i)}, g^{(i)})$ 쌍의 집합을 고려합니다. 손실 함수가 불변하기 때문에, 궤도 내 임의의 지점에서의 그래디언트는 원래의 그래디언트와 군 작용에 의해 연관됩니다. 이 방법은 궤도 전체에 대해 이러한 그래디언트들을 평균함으로써, 여러 번의 그래디언트 계산 없이도 곡률 정보를 집계합니다.

2. 구조화된 헤시안 근사

이 방법은 궤도 평균화된 파라미터 차이와 그래디언트 차이를 기반으로 섹턴트 조건(secant condition)을 해결함으로써 구조화된 헤시안 근사 $\hat{H}$를 유도합니다.

• 교환자 대수(Commutant Algebra): 결과적으로 도출된 헤시안 근사는 대칭 군의 교환자 대수에 속합니다. 이는 근사가 항등 행렬(identity)과 올-원즈 행렬(all-ones matrix)의 크로네커 곱으로 형성된 희소한 이진 텐서인 소수의 기저 요소들의 선형 중첩으로 표현될 수 있음을 의미합니다.
• 인수 분해: 이 방법은 전체 행렬을 저장하는 대신, 이러한 기저 요소들의 가중치인 스칼라 또는 행렬 "요소(factors)"만을 저장합니다. 이를 통해 계산 가능한 추정, 저장 및 역행렬 계산이 가능해집니다.
• 단순화: 저자들은 파라미터 차이의 궤도 평균이 항등 행렬에 비례한다고 가정하는 실용적인 단순화를 제안합니다. 이는 $\hat{H}_g = S_g^{1/2}$ 형태의 곡률 추정치로 이어지며, 여기서 $S_g$는 그래디언트 외적의 궤도 평균입니다.

3. "Symo" 옵티마이저

이 프레임워크는 Symo(Symmetry-based optimizer)라는 새로운 옵티마이저로 구현됩니다. 업데이트 규칙은 $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$입니다. 저자들은 네트워크의 대칭성에 대한 심볼릭 명세를 입력받아 필요한 요소들을 계산하는 효율적인 PyTorch 루틴을 생성하는 적시 컴파일러(just-in-time compiler)를 제공합니다.

주요 기여

1. 새로운 헤시안 근사법: 대칭 군 궤도에 대한 평균화를 통해 단 한 번의 그래디언트 계산으로부터 구조화된 헤시안 추정치를 도출하는 방법을 제시합니다. 이 근사치는 매우 구조화되어 있으며 교환자 대수에 존재함이 증명되었습니다.
2. Shampoo와 Muon의 이론적 통합: 이 프레 framework 내에서 특정 대칭 군을 선택함으로써 Shampoo와 Muon 옵티마이저의 업데이트 규칙을 정확히 복구할 수 있음을 증명했습니다.
   • 대칭 군을 특정 부호 치환 부분군(예: 입력/출력 레이어에는 항등 대칭을 적용하고 은닉 레이어에는 부호 치환을 적용)으로 제한함으로써, 유도된 곡률 행렬은 크로네커 인수를 가진 블록 대각 행렬이 되며, 이는 수학적으로 Shampoo/Muon과 동일합니다.
   • 이는 이 옵티마이저들이 효과적인 이유에 대한 원칙적인 설명을 제공합니다: 이들은 암묵적으로 가중치 공간 대칭성을 활용하며, 근사 정확도와 계산 비용 사이의 "스윗 스폿(sweet spot)"을 공략합니다.
3. 복잡도 제어: 프레임워크는 헤시안 근사의 복잡도가 선택된 대칭 군의 크기에 의해 직접적으로 결정됨을 보여줍니다. 더 큰 군은 더 낮은 차원의 교환자 대수(더 적은 요소, 더 저렴한 계산)를 생성하지만 더 큰 궤도를 가집니다(잠재적으로 낮은 국소 정확도). 더 작은 군은 더 많은 항을 가지지만 더 높은 국소 충실도(local fidelity)를 가집니다.
4. 실증적 검증:
   • 헤시안 품질: 소규모 MLP에 대한 실험에서, 제안된 근사치(특히 단순화된 $\hat{H}_g$)가 음의 그래디언트 방향을 따라 실제 헤시안과 높은 코사인 유사도를 달성하며, 궤도 평균화된 파라미터 추정치보다 우수한 성능을 보이고 Shampoo의 성능과 일치함을 보여주었습니다.
   • 옵티마이저 등가성: MNIST 오토인코더 및 nanoGPT(소규모 언어 모델)에서 Symo 옵티마이저는 Muon과 거의 동일한 훈련 곡선을 생성하여, 특정 군 제약 조건 하에서의 이론적 등가성을 입증했습니다. 두 방식 모두 Adam보다 현저히 뛰어난 성능을 보였습니다.

결과

• 정확도 vs 비용 트레이드오프: 논문은 다양한 군 선택에 따른 실험적 탐색을 수행합니다. 모든 레이어(입력/출력 포함)의 치환을 허용하는 "더 큰 군"은 여전히 Adam보다 우수하지만, 내부 레이어에 치환을 제한하고 블록 대각 근사를 사용하는 "더 작은 군"보다는 효과가 떨어지며, 이는 Shampoo의 성능과 일치합니다. 이는 아키텍처의 특정 불변성에 맞춰 군을 제한하는 것이 최적의 성능을 내는 데 중요하다는 것을 시사합니다.
• 계수 결핍(Rank Deficiency): 저자들은 전체 궤도 평균 행렬 $S_w$가 계수 결핍 상태가 될 수 있어 역행렬 계산을 불안정하게 만든다는 점을 지적합니다. 제안된 단순화($S_w \approx I$)는 이 문제를 피하며, 전체 양의 준정부호(positive semi-definite) 솔루션보다 경험적으로 더 나은 곡률 추정치를 제공합니다.
• 확장성: 이 방법은 Transformer 아키텍처(nanoGPT)로 성공적으로 확장되어, 대칭성 인식 접근 방식이 단순 MLP를 넘어 적용 가능함을 입증했습니다.

의의 및 주장

본 논문은 기존의 2차 옵티마이저들을 바라보는 통합적인 이론적 렌즈를 제공한다고 주장합니다. Shampoo와 Muon을 대칭 기반 헤시안 근사의 특정 사례로 프레임화함으로써, 이 연구는 이러한 방법들을 휴리스틱한 엔지니어링에서 원칙적인 수학적 유도로 격상시켰습니다.

저자들은 본 프레임워크가 다음과 같은 기능을 제공한다고 명시합니다:

• 적절한 대칭 군을 선택함으로써 특정 계산 예산에 맞춘 곡률 근사를 구축할 수 있는 체계적인 레시피를 제공합니다.
• 모델의 대칭성으로부터 구조화된 곡률로 가는 원칙적인 경로를 제공하며, 이는 다른 모델 클래스를 위한 새로운 옵티마이저 설계에 영감을 줄 수 있습니다.
• 최적화 외에도 불확정성 추정, 지속 학습, 모델 압축 등의 분야에 적용될 가능성을 제시하지만, 이는 입증된 결과라기보다 잠재적인 미래 방향으로 제시되었습니다.

논문은 훈련된 가중치 주변에서의 테일러 전개 정확도가 완전히 경계 지어지지 않았다는 점과, 프리컨디셔너(preconditioner)와 실제 헤시안 사이의 고유벡터 정렬에 대한 공식적인 보장이 여전히 미결 과제로 남아 있다는 점을 인정하며 한계점에 대해 겸허한 태도를 유지합니다. 또한, 프레임워크가 요소(factors)의 유도를 자동화하지만, 임의의 새로운 아키텍처를 위해 최적의 군을 선택하는 것은 여전히 사용자의 몫으로 남아 있습니다.