기술 요약: 서로 다른 레이어, 서로 다른 매니폴드

문제 정의

가중치 공간(weight-space)의 기하학적 구조는 신경망 최적화의 핵심적인 요소이지만, 현재의 매니폴드 제약 기반 훈련 방식은 일반적으로 트랜스포머의 모든 가중치 행렬에 대해 단일하고 균일한 기하학적 제약(예: 직교성 또는 Stiefel 제약)을 적용합니다. 본 연구는 트랜스포머의 서로 다른 모듈—구체적으로 어텐션(attention) 레이어와 MLP(Multi-Layer Perceptron) 블록—이 서로 다른 가중치 공간 기하학을 선호하는지 조사합니다. 저자들은 이러한 모듈들의 계산적 역할과 비선형 메커니즘이 크게 다르다고 가정하며, 따라서 균일한 매니폴드 할당은 최적이 아니거나 심지어 불안정성을 초래할 수 있다고 주장합니다.

방법론

본 연구는 은닉층 가중치에 행렬 정규화 업데이트를 적용하면서 이를 특정 행렬 매니폴드에 구속하는 옵티마이저인 Manifold Muon을 활용합니다. 저자들은 두 가지 주요 제약 계열을 비교합니다:

1. Stiefel 제약 (Stiefel Constraint): 열 직교성($W^\top W = I$)을 강제하여 특이값(singular values)을 엄격하게 제한하고 강력한 스펙트럼 제어를 제공합니다.
2. DGram 제약 (DGram Constraint): 그램 행렬(Gram matrix)의 비대각 성분이 0이어야 한다는($\text{Off}(W^\top W) = 0$) 더 약한 그램 공간 제약입니다. 이는 열 직교성은 유지하지만, 열의 노름(norm) 및 이에 따른 특이값이 변할 수 있도록 허용하여 추가적인 스케일 자유도를 제공합니다.

저자들은 OpenWebText로 사전 학습된 GPT-2 small 모델(124M 파라미터)에 대해 다섯 가지의 서로 다른 할당 전략을 평가합니다:

• UNCONSTRAINED: 매니폴드 제약 없음.
• ALL-STIEFEL: 어텐션과 MLP 가중치 모두에 Stiefel 제약 적용.
• ALL-DGRAM: 어텐션과 MLP 가중치 모두에 DGram 제약 적용.
• HETERO: 어텐션에는 Stiefel 제약, MLP에는 DGram 제약 적용.
• HETERO-INV: 어텐션에는 DGram 제약, MLP에는 Stiefel 제약 적용.

모든 설정은 매니폴드 할당의 효과만을 분리하기 위해 동일한 모델 아키텍처, 데이터, 훈련 스케줄 및 옵티마이저 하이퍼파라미터를 공유합니다.

주요 결과

실험 결과, 서로 다른 모듈이 매니폴드 제약에 반응하는 방식에서 명확한 비대칭성이 드러났습니다:

• 최적의 구성: HETERO 할당(어텐션에 Stiefel, MLP에 DGram)은 가장 낮은 안정적인 검증 손실(3.3544)을 달 기록하며, 비제약 베이스라인과 균일한 ALL-STIEFEL 구성을 모두 능가했습니다.
• DGram-Attention에서의 불안정성: 공유된 하이퍼파라미터 설정 하에서 어텐션 레이어에 DGram 제약을 할당한 구성(ALL-DGRAM 및 HETERO-INV)은 불안정해졌습니다. 이 실행들은 검증 손실의 폭발적인 증가를 보였으며 회복에 실패했습니다.
• DGram-MLP에서의 안정성: 반대로, (Stiefel 어텐션과 결합되었을 때) MLP 레이어에 DGram 제약을 할당하는 것은 안정적으로 유지되었으며 균일한 Stiefel 구성보다 성능을 개선했습니다.

메커니즘 분석

본 논문은 DGram 제약이 적용된 어텐션에서 관찰된 불안정성에 대해 메커니즘적 설명을 제공합니다:

1. 특이값 성장 (Singular Value Growth): Stiefel 제약과 달리, DGram은 가중치 행렬의 특이값이 임의로 커지는 것을 허용합니다. 어텐션 레이어에서는 쿼리($W_Q$)와 키($W_K$) 투영이 곱해져 어텐션 로짓($Z = X W_Q W_K^\top X^\top / \sqrt{d}$)을 형성합니다.
2. 로짓 증폭 (Logit Amplification): $W_Q$와 $W_K$의 스펙트럼 노름(spectral norms)에서 발생하는 통제되지 않은 성장은 어텐션 로짓의 크기를 직접적으로 증폭시킵니다.
3. 소프트맥스 포화 (Softmax Saturation): 커진 로짓 크기는 소프트맥스 함수를 포화 영역으로 밀어넣어, 출력이 원-핫 벡터(one-hot vector)에 가까워지게 만듭니다. 이는 소프트맥스의 자코비안(Jacobian)을 지수적으로 소멸시켜 그래디언트 퇴화(gradient degeneracy)와 훈련 불안정성을 초축합니다.
4. MLP의 내성 (MLP Tolerance): MLP 레이어는 활성화 함수(예: GELU)가 요소별(element-wise)로 작동하며 정규화된 확률 심플렉스를 통해 토큰 점수를 결합하지 않기 때문에 DGram 제약을 견뎌낼 수 있습니다. 결과적으로 MLP 가중치의 스케일 증폭이 모듈 전체를 포화된 라우팅 체제로 몰아넣지 않습니다.

기여점

저자들은 세 가지 주요 기여를 식별했습니다:

1. 경험적 평가: GPT-2 사전 학습 과정에서 Manifold Muon의 모듈별 매니폴드 할당을 평가하여, 균일한 제약만으로는 불충분함을 입증했습니다.
2. 비대칭성 발견: 어텐션과 MLP 레이어가 서로 다른 가중치 공간 기하학을 선호하며, 어텐션은 강력한 스펙트럼 제어(Stiefel)를 요구하고 MLP는 약한 스케일 보존 자유도(DGram)로부터 이득을 얻는다는 것을 보여주었습니다.
3. 불안정성 메커니즘: 특이값 성장과 그에 따른 소프트맥스 포화가 DGram 제약 어텐션의 불안정성을 일으키는 타당한 메커니즘임을 확인했습니다.

의의 및 한계

본 논문은 트랜스포머를 위한 대칭 인지 및 기하학 인지 최적화는 균일한 방식이 아닌 모듈별로 이루어져야 한다고 주장합니다. 이는 가중치 공간의 기하학이 현대 네트워크의 훈련 안정성과 역학을 개선하기 위한 실질적인 설계 축임을 시사합니다.

저자들은 다음과 같이 신중한 입장을 유지합니다:

• DGram 어텐션이 본질적으로 모든 조건에서 불안정한 것이 아니라, 테스트된 공유 하이퍼파라미터 설정 하에서 추가적인 스케일 제어 없이는 불안정한 경로를 노출한다는 점을 명시했습니다.
• 실험은 GPT-2 small과 특정 시드 세트에 국한되어 있으며, 최적의 헤테로 구성과 균일 구성 간의 검증 손실 차이는 완만합니다.
• 메커니즘 분석은 관찰된 불안정성에 대한 설명을 제공하는 것이지, (예: 학습률 재조정이나 명시적 스케일 규제를 통해) DGram 어텐션을 안정화할 수 없다는 것을 증명하는 완전한 증명이 아닙니다.

향후 연구로는 다양한 모듈, 아키텍처 및 규모에 걸쳐 매니폴드 제약을 자동으로 할당하는 방법을 개발하는 것이 제안되었습니다.