기술 요약: RMSNorm 트랜스포머를 위한 부호-치환 좌표 전송 (Signed-Permutation Coordinate Transport)

문제 정의

현대 대규모 언어 모델(LLM) 워크플로우는 스티어링 벡터(steering vectors), 희소 오토인코더(SAE) 딕셔너리, top-k 뉴런 집합, LoRA 어댑터와 같이 좌표로 인덱싱된 산출물을 서로 다른 모델 체크포인트 간에 이동하는 작업에 점점 더 많이 의존하고 있습니다. 그러나 "이 뉴런은 무엇인가?"라는 질문은 고정된 게이지(좌표계) 없이는 정의될 수 없는 문제입니다. 표준 정렬 도구(예: Git Re-Basin, 가중치 매칭)는 순열만을 고려한 대칭군($S_d$)을 가정하지만, 이는 RMSNorm(Root Mean Square Layer Normalization)을 사용하는 모델에 대해서는 잘못된 가정입니다.

평균 중심화(mean-centering)를 포함하여 순열(전역적 부호 반전 포함)로 대칭성을 제한하는 LayerNorm과 달리, RMSNorm은 평균 중심화가 결 없습니다. 결과적으로 RMSNorm의 잔차 연결(residual stream)은 순열뿐만 아니라 모든 좌표에 대한 독립적인 부호 반전을 허용합니다. 따라서 RMSNorm의 고유한 이산 게이지는 $S_d$가 아니라 초팔면체 군(hyperoctahedral group) $B_d = S_d \ltimes \{\pm 1\}^d$ (부호 치환)입니다. 이 부호 성분을 무시하는 것($S_d$ 전용 정렬)은 좌표로 인덱싱된 산출물이 정렬되지 않아 성능이 크게 저하되거나 기능적 붕괴가 발생하는 구조적 실패 모드로 이어집니다.

방법론

1. 고유 게이지의 이론적 특성화

본 논문은 일반적인 per-channel gain $\gamma$를 가진 트랜스포머에 대해 다음과 같이 규명합니다:

• LayerNorm: $O(d)$의 최대 좌표 보존 부분군은 $S_d \times C_2$ (순열 및 전역적 부호 반전)입니다.
• RMSNorm: 최대 좌표 보존 부분군은 전체 부호 치환 군인 $B_d$입니다.
• 시사점: 연속적인 직교 회전(SliceGPT 또는 QuaRot에서 사용됨)은 함수 보존에는 유효하지만, 희소 좌표 정체성(sparse coordinate identity)을 파괴합니다. 오직 $B_d$만이 희소 객체(예: 특정 뉴런 또는 SAE 특징)의 정체성을 보존합니다.

2. 부호 마진화된 게이지 복구 (Sign-Marginalized Gauge Recovery)

쌍을 이룬 활성화 함수로부터 올바른 $B_d$ 게이지를 복구하기 위해, 본 논문은 표준 헝가리안 매칭 알고리즘의 수정을 제안합니다:

• 비용 함수: $\sum C_{i, \pi(i)}$인 부호 상관관계(signed correlation)를 최대화하는 대신, 절대 상관관계 $\sum |C_{i, \pi(i)}|$를 최대화합니다.
• 부호 복구: 최적의 순열 $\pi^*$가 절대 비용 행렬을 통해 결정되면, 부호는 $s_i = \text{sign}(C_{i, \pi^*(i)})$로 복구됩니다.
• 이론적 보장: 논문은 상관관계가 없는 좌표 분포 사례에서, 표준 부호 상관관계 매칭은 올바른 음의 상관 관계 항목을 피하기 때문에 약 50%(양의 부호 비율)의 구조적 정확도 천장에 머문다고 증명합니다. 부호 마진화 매칭은 이러한 장애물을 제거하여 100% 정확도로 수렴합니다.

3. 궤적 구성을 통한 평행 이동 (Parallel Transport via Trajectory Composition)

두 종단점만을 매칭하는 대신, 본 논문은 훈련 궤적을 따른 평행 이동을 도입합니다:

• 연속적인 체크포인트 사이에서 부호 마진화 매칭을 사용하여 로컬 $B_d$ 게이지를 추정합니다.
• 이러한 로컬 게이지들을 결합하여 먼 거리의 체크포인트 간(또는 베이스를 공유하는 실행 간) 글로벌 전송 맵을 형성합니다.
• 이 방식은 직접적인 종단점 매칭의 "홀로노미(holonomy)" 문제를 피하고, 연속적인 체크포인트가 가지는 높은 할당 마진을 활용합니다.

4. 게이지 공변 전송 규칙 (Gauge-Covariant Transport Rules)

본 논문은 $B_d$ 하에서의 다양한 좌표 인덱스 산출물에 대한 명시적 변환 규칙을 도출합니다 (표 1):

• 벡터 (스티어링, 프로브, SAE 인코더): $v \mapsto vG$로 변환됩니다.
• 행렬 (SAE 디코더, LoRA WRITE): $W \mapsto W G^\top$로 변환됩니다.
• 옵티마이저 상태 (AdamW): 1차 모멘트 $m_W$는 가중치와 함께 변환되며, 2차 모멘트 $v_W$는 제곱 형태이므로 부호에 무관하며 순열을 통해서만 변환됩니다.

주요 결과

1. 복구 정확도

• 부호 마진화 vs. 부호 포함: 7B/8B RMSNorm 모델 실험에서, 단순 부호 상관관계 매칭은 순열의 약 49.5%, 부호의 약 50%만을 복구했습니다. 반면 부호 마진화 매칭은 **100%**의 복구율을 달림을 확인했습니다.
• 전송 vs. 종단점 매칭: Qwen2.5-1.5B 미세 조정 벤치마크(1500 스텝)에서, 베이스 구성을 통한 전송은 **91.1%**의 좌표를 복구한 반면, 직접적인 종단점 매칭은 **60.3%**만을 복구했습니다. 이 개선은 베이스만을 참조했기 때문이 아니라(베이스 경유 종단점 매칭은 11.4%였음), 많은 수의 고마진 로컬 게이지를 합성했기 때문입니다.

2. 산출물 전송 성능

$B_d$-정확한 전송은 여러 산출물 클래스에 걸쳐 $S_d$-전용 전송보다 현저히 우수한 성능을 보였습니다:

• LoRA 어댑터: Qwen2.5-1.5B (SST-2)에서, $S_d$ 전송은 어댑터 정확도를 비적응 베이스 수준으로 붕괴시킨 반면($0.523$ vs $0.525$),$B_d$전송은 이득의 **99.7$0.944$).
• SAE 재구성: TinyLlama의 경우, $S_d$ 전송은 재구성 NMSE 1.08(실패)을 초래했으나, $B_d$ 전송은 0.004(완벽에 가까움)를 달성했습니다.
• 스티어링 벡터: Qwen 감성 스티어링은 $B_d$ 하에서 **95.8%**의 효과를 유지했으나, $S_d$ 하에서는 **17.2%**에 그쳤습니다. 거절(refusal) 스티어링 벡터는 $S_d$ 하에서 부호가 반전되어 행동을 역전시켰습니다.

3. 상태 기반 훈련 및 병합

• 옵티마이저 상태: $S_d$-전용으로 전송된 AdamW 상태(순열은 맞지만 부호는 +1로 강제됨)로 훈련을 재개하면, 기능적으로 동일한 체크포인트에서 시작했음에도 불구하고 훈련 궤적이 참조 모델에서 벗어나게 됩니다. $B_d$ 전송은 수치적 허용 오차 범위 내에서 궤적을 보존합니다.
• 모델 병합: 좌표 보존 병합에서, $S_d$-전용 정렬은 파괴적인 부호 불일치를 유발하여 높은 "피크-어보브-코드(peak-above-chord)" 장벽을 생성했습니다. 부호 정렬은 이러한 장벽을 제거했습니다(게이지 스크램블 테스트에서 0.00으로 감소).

의의 및 주장

본 논문은 RMSNorm 기반 아키텍처(Llama 2, Llama 3, Qwen 등 널리 사용되는 모델 포함)에 대해 좌표 정체성은 부호 치환까지로만 잘 정의된다고 주장합니다.

• 재현성: 좌표 인데스 기반의 해석 가능성 주장(예: "뉴런 X는 거절 뉴론이다")은 특정 $B_d$ 게이지를 고정하고 적용하지 않으면 실행 간에 재현될 수 없습니다. 본 논문은 이러한 주장을 검증하기 위한 부정 대조군으로서 "게이지 스윕(gauge sweeps)"을 제안합니다.
• 도구 전송: $S_d$-전용 방식의 실패는 단순한 오류가 아니라 구조적 붕괴입니다. 스티어링 벡터, SAE 딕셔너리, 옵티마이저 상태를 성공적으로 이동시키기 위해서는 부호 전송이 필수적입니다.
• 한계: 본 논문은 $B_d$가 좌표 정체성을 보존하지만, 좌표 정체성이 주요 대상이 아닌 밀집 융합(dense fusion)이나 압축의 경우 연속적인 직교 방법이 여전히 유용할 수 있음을 언급합니다. 또한, FFN 은닉 유닛은 잔차 연결 게이지와 구별되는 독자적인 순열 대칭성($S_{d_{ffn}}$)을 가집니다.

요약하자면, 본 연구는 RMSNorm 모델을 올바르게 정렬하기 위한 이론적 토대와 실질적인 알고리즘을 제공하며, 부호 성분을 무시하는 게이지 설정이 좌표 기반 워크플로우를 근본적으로 망가뜨린다는 것을 입증합니다.