Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

기술 요약: AdamW를 넘어선 독점적 셀프 어텐션 (Exclusive Self-Attention)

문제 정의
트랜스포머의 표준 셀프 어텐션 메커니즘은 $i$ 위치에 있는 토큰의 출력이 자기 자신의 밸류 벡터($v_i$)를 포함한 밸류 벡터들의 가중치 혼합을 포함하도록 허용한다. 이는 표현력이 뛰어나지만, 잔차 경로(residual pathways)가 이미 로컬 토큰 정보를 보존하고 있음에도 불구하고 어텐션 출력이 $v_i$와 강하게 정렬된 상태를 유지하게 만드는 "표현적 지름길(representational shortcut)"을 생성할 수 있다. 이러한 현상은 "어텐션 유사성 편향(attention similarity bias)"이라 불리며, Zhai [2]는 토큰 자신의 밸류 벡터에 평행한 성분을 명시적으로 제거하는 독점적 셀프 어텐션(Exclusive Self-Attention, XSA)을 통해 이를 해결하였다.

그러나 원래의 XSA 연구는 해당 방법을 AdamW 옵티마이저를 사용하는 NanoGPT 훈련 스택 내에서만 배타적으로 평가하였다. 원 저자들은 Muon과 같은 대안적인 옵티마이저와의 호환성을 미해결 과제로 명시적으로 언급하였다. 이는 현대의 소규모 언어 모델 연구(예: nanochat 스택)가 행렬 파라미터에는 Muon을, 임베딩에는 AdamW를 사용하는 하이브리드 최적화 전략과 가치 임베딩(value embeddings) 및 회전 위치 인코딩(rotary position encoding)과 같은 구조적 특징을 점점 더 많이 활용하고 있다는 점에서 매우 중요한 공백이다. XSA의 기하학적 제약이 이러한 현대적인 Muon 기반 훈련 체계로 전이되었을 때도 안정적이거나 유익하게 유지될지는 불분명하다.

방법론
본 연구는 nanochat 훈련 스택을 사용하여 통제된 재현 및 전이 실험을 수행한다. 방법론은 다음과 같다:

1. 구현: 매개변수화된 형태의 XSA가 nanochat에 통합된다. 연산은 $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$로 정의되며, 여기서 $y_i$는 표준 어텐션 출력이고, $v_i$는 토큰의 밸류 벡터이며, $\alpha \in [0, 1]$은 투영 강도를 조절한다. $\alpha=1$일 때, 자기 정렬 성분이 완전히 제거된다.
2. 정확성 검증: 구현은 $\alpha=0$일 때 XSA 경로가 표준 어텐션과 동일함을 보장하도록 수치적으로 검증된다. 짧은 훈련 과정을 통해 베이스라인과 $\alpha=0$ 경로가 일치하는 초기 손실(loss) 거동을 보이는지 확인한다.
3. 실험 설정: 동일한 데이터 순서, 옵티마이저 계열, 학습률 스케줄, 컨텍스트 길이를 갖춘 매칭된 훈련 조건 하에서 두 가지 모델 깊이가 평가된다:
   • d12: 12개 레이어, 은닉 차원 768, 12개 헤드.
   • d24: 24개 레이어, 은닉 차원 1024, 16개 헤드.
   • 옵티마이저: 행렬 파라미터는 Muon으로 최적화되며, 나머지 파라미터는 AdamW를 사용한다.
   • 데이터: ClimbMix-400B.
4. 지표: 본 연구는 훈련 손실(training loss), 검증 bits-per-byte (bpb), 다운스트림 집계 점수 (Base CORE 및 ChatCORE), 그리고 어텐션 진단(출력과 자기 밸류 사이의 코사인 정렬, 대각 성분 거동)을 추적한다.

주요 기여

• Muon-스택 재현: XSA를 Muon 기반 훈련 스택에 최초로 통합하여, 원래 연구에서 남겨진 옵티마이저 호환성에 관한 열린 질문을 직접적으로 다루었다.
• 매개변수화된 구현: 표준 어텐션, 완전한 XSA, 그리고 강도 파라미터 $\alpha$를 통한 부분 투영을 모두 허용하며, 레이어 선택 제어를 포함하는 유연한 구현을 제공한다.
• 정확성 및 도구화: $\alpha=0$ 극한값에 대한 검증과, 초기 훈련 비교를 왜곡했던 로깅 아티팩트(step-index 및 smoothing 관례)의 수정을 완료하였다.
• 두 가지 깊이 평가: 포괄적인 지표 세트를 사용하여 d12 및 d24 깊이에서 비교 분석을 수행하였다.
• 메커니즘 및 시스템 분석: 어텐션 자기 정렬에 대한 진단과 런타임 오버헤드에 대한 평가를 수행하였다.

결과
실증적 결과는 XSA가 Muon 기반 스택 하에서 안정적이지만, 모델 깊이와 특정 지표에 따라 엇갈린 다운스트림 이점을 보인다는 것을 나타낸다:

• 최적화 안정성: XSA와 베이스라인 모델의 훈련 손실 궤적은 두 깊이 모두에서 밀접하게 일치한다. 발산, NaN, 또는 불안정성의 증거는 없으며, 이는 XSA가 Muon의 업데이트 기하학을 방해하지 않음을 확인시켜 준다.
• 검증 성능 (Bits-per-Byte): XSA는 홀드아웃 검증 bpb를 일관되게 개선한다:
  • d12: 0.8484에서 0.8457로 감소 (상대적 개선율 0.32%).
  • d24: 0.7193에서 0.7171로 감소 (상대적 개선율 0.31%).
• 다운스트림 지표 (CORE):
  • d24: XSA는 Base CORE (0.2593에서 0.2606, +0.50%)와 ChatCORE (0.3638에서 0.3682, +1.21%)를 모두 개선한다.
  • d12: bpb는 개선되었으나, Base CORE는 0.1602에서 0.1508로 감소(-5.87%)하였다.
• 메커니즘 진단: 어텐션 진단 결과, XSA 하에서 자기 정렬(self-alignment)과 대각 거동(diagonal behavior)이 감소하였으며, 이는 문맥적 혼합(contextual mixing)을 촉진하려는 의도된 메커니즘과 일치한다.
• 런타임 오버헤드: 오버헤드는 미미한 것으로 특징지어진다. 초기 추정치는 약 5%의 오버헤드를 시사했으나, 반복적인 처리량(throughput) 측정 결과 차이가 짧은 실행의 시스템 수준 변동성 근처에 머물러 있어, 과도한 페널티는 없는 것으로 나타났다.

의의 및 주장
본 논문은 Exclusive Self-Attention이 Muon 기반 훈련 스택에 호환 가능한 구조적 귀납 편향(architectural inductive bias)이지만, 그 효용성이 보편적이지는 않다고 결론짓는다. 주요 의의는 XSA의 표현적 제약이 훈련을 불안정하게 만들지 않으면서 현대적인 최적화 체계로 깔끔하게 전이된다는 점을 입증한 데 있다.

그러나 저자는 다운스트림 이득이 규모 및 지표에 따라 달라진다는 점을 강조한다. 두 깊이 모두에서 검증 bpb가 개선된 것은 표현적 이점을 시사하지만, d12에서 Base CORE가 저하된 것은 자기 정렬을 제거하는 것의 이점이 모델의 깊이, 로컬 정보를 보존하는 잔차/밸류 경로의 능력, 또는 특정 평가 작업에 따라 달라질 수 있음을 나타낸다. 본 연구는 XSA가 표준 어텐션을 대체하는 보편적으로 우월한 교체 수단이라고 주장하는 것이 아니라, 특정 맥락(본 연구에서는 d24)에서 이점을 제공하는 안정적이고 저방해적인 제약임을 보여준다.

이 작업은 단순한 바이트 단위의 재현이라기보다 견고성 체크로서 기능하며, 현대 언어 모델 연구에서 구조적 수정, 옵티마이저 기하학, 그리고 평가 지표 사이의 복잡한 상호작용을 부각시킨다.