🤖 machine learning

A Mechanistic Diagnostic of Rank Collapse in Post-Norm Decoder Transformers

이 논문은 Post-Norm 디코더 전용 트랜스포머(decoder-only Transformer)에서의 랭크 붕괴(rank collapse)에 대한 2단계 기계론적 분석을 제공하며, 인과적 어텐션(causal attention)이 초기화 단계에서 토큰 유사성을 증폭시키는 한편 RMSNorm에 의한 그래디언트 수축(gradient contraction)이 학습 중 효과적인 복구를 방해함으로써, 결과적으로 빈도 기반 예측과 그래디언트 소실을 특징으로 하는 붕괴된 상태로 이어진다는 것을 입증한다.

원저자: Xingjian Wang, Qingyu Han, Xiaodong Luo, Yin Zhang

게시일 2026-08-11

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Xingjian Wang, Qingyu Han, Xiaodong Luo, Yin Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

컴퓨터가 수십억 개의 문장을 읽으며 이야기의 다음 단어를 추측하는 법을 배우는 세상을 상상해 보세요. 이를 위해 컴퓨터는 '트랜스포머(Transformer)'라고 불리는 특별한 뇌 구조와 같은 것을 사용합니다. 트랜스포머를 원시 데이터가 바닥으로 들어와 각 층에서 처리된 후, 꼭대기에서 스마트한 예측값으로 나오는 다층 공장이라고 생각해 보세요. 이 공장들이 가장 잘 작동하려면 매우 깊어야 합니다. 때로는 수십 층 높이까지 말이죠. 하지만 까다로운 문제가 하나 있습니다. 원래의 설계도를 사용하여 공장을 너무 높게 지으면, 아래층의 노동자들이 지시를 받지 못하게 됩니다. 꼭데기에서 오는 신호가 너무 약해져서 아래층들은 그냥 포기해 버리고, 전체 기계는 자신이 본 평균적인 단어를 계속 복사하기만 하는 지루한 루프에 빠지게 됩니다. 이는 마치 메시지가 끝에 도달할 때까지 너무 흐릿해져서 모두가 그저 똑같은 문구만 반복하게 되는 '전화기 게임(telephone game)'과 같습니다.

여러분이 읽게 될 논문은 왜 특정 유형의 공장 설계인 '포스트-놈(Post-Norm)'에서 이런 일이 발생하는지를 깊이 파고듭니다. 이 설계에서는 노동자들이 자신의 일을 마치고 메인 컨베이어 벨트에 결과물을 추가한 후에야 '검진(정규화, normalization)'을 받습니다. 저자들은 이 설정에 두 가지 치명적인 결함이 있다는 것을 발견했습니다. 첫째, 공장이 문을 열자마자, 과거를 바라보는 방식(어텐션, attention) 때문에 의도치 않게 모든 사람의 작업이 마치 복제 인간들처럼 똑같아 보입니다. 둘째, 일단 모두가 똑같아지면, 메시지를 뒤로 전달하는 공장의 내부 규칙(그래디언트, gradients)이 지시 사항을 완전히 사라질 때까지 축소시켜 버립니다. 그 결과, 기계는 새로운 것을 배울 수 없게 되고, 그저 단조로운 하나의 음만을 웅웅거리며 내뱉게 됩니다. 저자들은 단순히 추측한 것이 아니라, 어떻게 복제 인간들이 형성되는지 증명하기 위해 수학적 모델을 구축했고, 48층짜리 공장을 실제로 가동하여 실시간으로 일어나는 현상을 관찰했습니다.

거대한 복제 공장: 왜 깊은 AI는 갇히게 되는가

48층 규모의 거대한 마천루를 상상해 보세요. 각 층에는 이야기를 이해하려는 노동자 팀이 있습니다. 이것이 바로 AI 모델의 한 종류인 '포스트-노름 트랜스포머(Post-Norm Transformer)'입니다. 목표는 맨 위층이 문장의 다음 단어를 예측하는 것입니다. 이를 위해 정보는 바닥에서 위로 흐르고, 지시 사항(그래디언트)은 노동자들에게 더 잘하는 법을 가르치기 위해 아래로 흘러 내려옵니다.

하지만 여기 함정이 있습니다. 이 특정 마천루 설계에서는 아래층의 노동자들이 침묵의 취급을 받고 있습니다. 논문은 이 비극이 마치 2막으로 구성된 연극처럼 두 단계로 진행된다고 설명합니다.

제1막: 개업 날의 "복제" 효과

공장이 막 문을 열었다고 상상해 보세요. 노동자들은 신선하고, 매니저들(AI의 어텐션 메커니즘)은 어떻게 협력할지 고민하고 있습니다. 포스트-노름 건물에서 매니저들은 이상한 습관이 있습니다. 그들은 사람들의 아이디어를 평균화하는 경향이 있습니다.

저자들은 시작 단계에서 공장의 '어텐션(Attention)' 브랜치가 **접두사 평균화 연산자(prefix-averaging operator)**처럼 작동한다는 것을 발견했습니다. 이는 마치 학생 개개인의 독특한 아이디어에 귀를 기울이는 대신, 지금까지 말한 모든 것의 평균을 내어 모두에게 그것을 복사하라고 말하는 교사와 같습니다. 건물이 매우 높기 때문에(48개 층), 신호가 위로 올라감에 따라 이 평균화 과정이 반복해서 일어납니다. 신호가 꼭대기 층에 도달할 때쯤이면, 모든 노동자의 출력값은 거의 동일해 보입니다. 그들은 모두 복제 인간이 되어버린 것입니다.

논문은 이것이 공장이 학습을 시작하기도 전부터 일어난다는 것을 보여줍니다. 마치 거울의 탑을 세웠는데, 첫 번째 반사부터 이미 모든 것을 똑같이 보이게 만든 것과 같습니다. 저자들은 이를 측정했으며, '어텐션' 부분이 노동자들을 복제 인간으로 만드는 주범임을 발견했습니다. 공장의 다른 부분인 'FFN(피드포워드 네트워크)'은 이를 밀어내고 차별성을 유지하려 노력하지만, 이는 허리케인을 막으려는 부드러운 미풍과 같아서 복제 현상을 막기에는 역부족입니다.

제2막: 사라지는 속삭임

이제 공장이 가동 중이고, 노동자들이 이미 복제 인간이 되었다고 가정해 봅시다. 꼭대기의 매니저들은 깨닫습니다. "이봐, 우리 모두 똑같아졌잖아! 이걸 고쳐야 해!" 그들은 행동을 바꾸기 위해 아래층으로 메시지를 보냅니다. 이것이 '역방향 패스(backward pass)' 또는 그래디언트입니다.

하지만 바로 여기서 포스트-노름 설계가 그들을 배신합니다. 이 건물에서는 노동자가 작업을 마칠 때마다, 그들의 출력이 얼마나 큰지에 따라 작업의 규모를 조절하는 '검진(RMSNorm)'을 받습니다. 노동자들이 문제를 해결하려고 노력할수록, 그들의 출력값은 점점 더 커집니다. 검진 기계는 이 거대한 출력을 보고, 관리 가능한 수준으로 유지하기 위해 이를 축소합니다.

문제는 이 축소가 노동자가 메인 컨베이어 벨트에 자신의 작업을 추가한 '후'에 일어난다는 점입니다. 따라서 "수정" 메시지가 뒤로 전달되려고 할 때, 이 메시지는 축소 기계를 통과해야 합니다. 저자들은 노동자들의 출력이 커짐에 따라 기계가 "수정" 메시지를 너무 공격적으로 축소하여, 메시지가 아래층에 도달할 때쯤에는 완전히 사라져 버린다는 것을 발견했습니다. 이는 마치 48층 높이의 복도에서 교정 명령을 외치는 것과 같지만, 층마다 소리를 흡수하는 벽이 점점 더 두꺼워지는 것과 같습니다. 외침이 바닥에 닿을 때쯤에는 그저 아무도 들을 수 없는 속삭임이 되어버립니다.

아래층 노동자들은 지시를 들을 수 없기 때문에, 복제 인간이 되는 것을 멈출 수 없습니다. 공장은 오직 훈련 데이터에서 가장 흔한 단어만을 예측하는 '고유사성 체제(high-similarity regime)'에 갇히게 됩니다. 저자들은 이를 '빈도 분포(frequency distribution)'라고 부릅니다. 이는 AI가 "무슨 말을 해야 할지 모르겠으니, 내가 들어본 것 중 가장 인기 있는 단어를 말하겠다"라고 말하는 방식입니다.

실험: 붕괴를 관찰하다

이것이 단순한 이론이 아님을 증명하기 위해, 저자들은 48층짜리 포스트-노름 공장을 구축하고 C4라는 거대한 데이터셋으로 훈련시켰습니다. 그들은 훈련이 진행되는 동안 면밀히 관찰했습니다.

복제의 성장: 그들은 노동자들의 출력값이 얼마나 유사한지 측정했습니다. 수학적 예측대로, 시작 직후 유사도가 급증했으며, 이는 '접두사 평균화' 어텐션이 모두를 복제 인간으로 만들고 있음을 확인해주었습니다.
축소 현상: 그들은 '축소 계수(RMSNorm backward factor)'를 추적했습니다. 훈련이 진행되고 노동자들의 출력이 커짐에 따라, 이 계수가 1 미만으로 떨어지는 것을 확인했습니다. 이는 공장이 수정 신호를 능동적으로 축소하고 있었음을 의미합니다.
사라지는 그래디언트: 그들은 아래층을 조사했고, 그래디언트(학습 신호)가 몇 자릿수 차이로 급격히 떨어지는 것을 보았습니다. 아래층은 사실상 꼭대기의 지시로부터 눈이 먼 상태였습니다.
막다른 길: 마지막으로, 그들은 손실(loss, 예측이 얼마나 나쁜지)을 확인했습니다. 공장이 붕괴되자, 손실은 개선을 멈추고 가장 흔한 단어만을 추측하는 기계의 이론적 최소치인 '빈도 손실(frequency loss)'에 머물렀습니다.

Pre-Norm은 왜 다른가?

여러분은 "왜 모든 AI 공장이 이런 문제를 겪지 않나요?"라고 물을 수 있습니다. 논문은 '프리-노름(Pre-Norm)'이라 불리는 대안 설계가 검진을 노동자가 작업을 컨베이어 벨트에 추가하기 전에 수행함으로써 이 문제를 해결한다고 설명합니다. 프리-노름에서는 축소 기계가 새로운 작업에만 영향을 미칠 뿐, 전체 컨베이어 벨트에는 영향을 주지 않습니다. 이는 "수정" 메시지가 짓눌리지 않고 복도를 따라 내려갈 수 있게 하여, 아래층이 깨어 있는 상태로 학습을 지속할 수 있게 해줍니다.

요점

이 논문은 단순히 "포스트-노름은 나쁘다"라고 말하는 데 그치지 않습니다. 왜 실패하는지에 대한 명확한 두 단계의 이야기를 제공합니다. 첫째, 어텐션 메커니즘이 시작 단계에서 실수로 복제 인간 무리를 만듭니다. 둘째, 메시지를 뒤로 전달하는 구조적 규칙이 일단 형성된 복제 현상을 되돌리는 것을 불가능하게 만듭니다. 공장은 자신이 아는 가장 흔한 단어만을 반복하는 루프에 갇히게 되며, 변화하라는 지시가 바닥에 도달하지 못하기 때문에 어떤 훈련으로도 깨울 수 없습니다.

저자들은 이 사실에 대해 확신하고 있습니다. 그들은 이러한 동작을 뒷받침하는 수학적 증명을 가지고 있으며, 48개 층 모델로부터 얻은 실험 데이터 또한 그들의 예측과 완벽하게 일치합니다. 심지어 어텐션에서 '평균화' 부분을 제거하면 복제가 멈춘다는 것을 보여줌으로써, 어텐션 메커니즘이 이 이야기의 진정한 악당임을 입증했습니다. 그러므로 깊고 스마트한 AI를 만들고 싶다면, 당신의 공장 설계가 실수로 노동자들을 침묵하는 똑같은 모습의 군중으로 만들지 않도록 주의해야 합니다.

기술 요약: 왜 Post-Norm 트랜스포머는 붕괴하는가

1. 문제 정의

깊은 디코더 전용(decoder-only) 트랜스포머는 기존의 초기화 방식 하에서 Post-Norm 아키텍처가 웜업 스케줄과 학습률에 매우 민감하기 때문에, 흔히 이를 Pre-Norm 변형으로 대체한다. 선행 연구들은 랭크 붕괴(rank collapse, 토큰 표현이 거의 동일해져 유효 차원이 감소하는 현상) 및 그래디언트 소실(gradient vanishing)과 같은 증상들을 식별해 왔으나, 인과적 디코더 전용 트랜스포머에서 이러한 현상들을 연결하는 인과적 메커니즘은 여전히 명확히 밝혀지지 않았다. 특히, 인과적 어텐션(causal attention)이 어떻게 고유사성 표현을 생성하는지, 그리고 왜 일단 발생한 유사성을 학습 역학이 복구하지 못하는지에 대해서는 불분명하다. 이러한 간극은 실질적인 해결책(Pre-Norm으로의 전환)과 Post-Norm 표현이 왜 붕괴 영역으로 진입하며 왜 최적화가 이를 탈출하기 어려운지에 대한 이론적 이해 사이의 단절을 야기한다.

2. 방법론

저자들은 토큰 유사성($tsim$)을 스칼라 상태 변수로 활용하여 Post-Norm 랭크 붕괴를 분석하는 2단계 분석을 제안한다. 이 분석은 초기화 시점의 순전파(forward pass) 역학과 훈련 중의 역전파(backward pass) 역학으로 나뉜다.

이론적 프레임워크

순전파 (초기화): 저자들은 무작위 어텐션 행렬을 접두사 평균 연산자(prefix-averaging operator, $C_n$ )로 근사하고, 토큰 그람 행렬(Gram matrix)을 등상관 폐쇄(equal-correlation closure) 모델로 모델링한다. 또한 어텐션 브랜치와 SwiGLU 피드포워드 네트워크(FFN) 모두에 대한 기대 1단계 토큰 유사성 변화에 대한 폐쇄형 수식을 도출한다.
역전파 (훈련): 분석은 증가하는 Pre-normalization 잔차 노름(residual norms)과 RMSNorm 자코비안(Jacobian) 사이의 상호작용에 집중한다. 저자들은 서브레이어 그래디언트 기여 계수( $\alpha$ )를 정의하고, 잔차 노름이 성장할 때 RMSNorm 역전파 계수가 그래디언트를 어떻게 수축시키는지 분석한다.
붕괴 상태 특성화: 논문은 완전히 붕괴된 네트워크의 특성을 수학적으로 규명하며, 특히 최적 확률 분포와 파라미터 그래디언트의 거동을 분석한다.

실험 설정

모델: $d=512$ , $d_{ff}=1536$ , 4개의 어텐션 헤드, SwiGLU FFN, RMSNorm을 갖춘 48계층 디코더 전용 트랜스포머(Llama-2 스타일).
데이터셋: 시퀀스 길이 2048인 C4 데이터셋.
프로토콜: 초기화 분산과 학습률(측정 범위: $6 \times 10^{-4}$ ~ $1.8 \times 10^{-3}$ )을 변화시키는 통제된 실험. 본 연구는 매칭된 설정 하에서 Post-Norm과 Pre-Norm 아키텍처를 비교한다.
개입: 저자들은 초기화 시 어텐션 행렬에서 접두사 평균 성분을 실험적으로 제거하여, 이것이 유사성 증가에 미치는 영향을 격리한다.

3. 주요 기여

A. 붕괴에 대한 2단계 설명

본 논문은 Post-Norm 랭크 붕괴가 두 가지 뚜렷한 메커니즘에 의해 발생함을 입증한다:

어텐션 증폭 (순방향): 초기화 시, 인과적 어텐션은 레이어를 가로질러 토큰 유사성을 증가시키는 대략적인 접두사 평균 연산자로 작용한다. SwiGLU FFN은 미미한 댐핑(damping) 효과만을 제공한다. Post-Norm에서는 "어텐션의 양"이 레이어 전반에 걸쳐 일정하게 유지되므로, 유사성이 급격히 누적된다. 반면, Pre-Norm은 깊이에 따라 어텐션의 양이 감소하여 유사성 증가를 늦춘다.
그래디언트 복구 실패 (역방향): 네트워크가 고유사성 영역에 진입하면, Pre-normalization 잔차 노름이 성장한다. 이 성장은 RMSNorm 역전파 계수를 수축적(1 미만)으로 만든다. 만약 어텐션과 FFN 서브레이어의 그래디언트 기여가 이러한 수축을 충분히 보상하지 못한다면, 초기 레이어에 도달하는 그래디언트는 깊이에 따라 기하급수적으로 감소한다. 이는 최적화 알고리즘이 초기 레이어의 유사성을 복구하는 것을 방해한다.

B. 붕괴된 네트워크의 특성화

저자들은 붕괴된 네트워크의 거동을 다음과 같이 규명한다:

손실 하한선 (Loss Floor): 랭크 붕괴 하에서 네트워크가 달성할 수 있는 최적 확률 분포는 빈도 분포(frequency distribution, 토큰 확률이 훈련 데이터 내 출현 빈도에 비례하는 분포)이다. 이는 상대적으로 높은 손실 하한선을 초래한다.
그래디언트 소실: 출력값이 빈도 분포와 일치하고 은닉 상태가 붕괴되었을 때($tsim=1$), 붕괴된 레이어의 파라미터 그래디언트는 정확히 소멸한다. 이를 초기 레이어의 그래디언트 기하적 감소와 결합하면, 네트워크는 최적화가 정체되는 근사 정지점(near-stationary point)에 도달하게 된다.

C. 실험적 검증

초기화: 1단계 유사성 증분 측정값은 정리 3.2(어텐션)와 정리 3.4(FFN)의 이론적 예측과 일치한다. 접두사 평균 성분을 제거하면 유사성 증가가 억제되며, 이는 해당 성분이 주요 동인임을 확인시켜 준다.
훈련 역학: 붕괴가 발생하는 실행 과정에서 저자들은 잔차 노름의 급격한 상승, 그에 따른 RMSNorm 수축 계수( $c < 1$ )의 하락, 그리고 초기 레이어에서의 그래디언트 기하적 감소를 관찰한다.
손실 거동: 붕괴된 훈련 실행은 예측된 빈도 손실 근처에서 안정화되며, 이는 이론적 하한선을 검증한다.

4. 결과

유사성 성장: Post-Norm 네트워크는 초기화 시 토큰 유사성이 빠르게 성장하여 조기에 포화되는 반면, Pre-Norm 네트워크는 더 완만한 성장을 보인다.
그래디언트 수축: 붕괴로의 전이 과정 동안 그래디ਆਂ트 기여 계수( $\alpha$ )는 약간 증가하지만, RMSNorm 수축 계수( $\sqrt{d}/\|y\|$ )가 크게 떨어져 대부분의 레이어에서 순 수축( $c < 1$ )이 발생한다.
그래디언트 소실: 전이 구간 동안 초기 레이어로 향하는 그래디언트는 몇 자릿수(orders of magnitude)만큼 급락하는 반면, 후기 레이어의 그래디언트는 상대적으로 안정적으로 유지된다.
손실 수렴: 붕괴된 훈련 손실은 훈련 데이터로부터 계산된 빈도 손실과 밀접하게 일치하며, 이는 네트워크가 한 번 붕괴되면 통계적 토큰 빈도 이상을 학습할 수 없음을 확인해 준다.

5. 의의 및 한계점

의의:
본 논문은 Post-Norm 실패에 대한 통합적인 이론적 설명을 제공하며, 순방향 유사성 증폭(네트워크를 붕괴로 몰아넣음)과 역방향 복구 불능(복구를 방해함)을 구분한다. 또한 Pre-Norm이 왜 더 견고한지를 명확히 한다. 즉, Pre-Norm은 초기화 시 유사성의 지속적인 증폭을 피하고, RMSNorm 자코비안이 전체 역방향 신호(스킵 연결 포함)를 수축시키는 것을 방지한다.

한 한계점:
저자들은 유사성이 왜 증가하고 그래디언트가 왜 소실되는지는 설명하지만, 특정 옵티마이저 스텝에서 마지막 레이어의 토큰 유사성과 훈련 손실이 급격히 상승하는 구체적인 트리거(trigger)는 설명하지 못한다고 명시했다. 이 전이를 시작하는 정확한 조건을 식별하는 것은 향후 과제로 남겨두었다. 또한, 이론적 분석은 등상관 가정 및 어텐션 행렬에 대한 접두사 평균 대용치(surrogate)와 같은 근사치에 의존하고 있다.