언어 모델 사전 학습을 위한 점진적 잔차 웜업 (ProRes)

대규모 언어 모델 학습 시 모든 레이어가 동시에 학습을 시작하면 깊은 레이어가 불안정한 초기 신호를 증폭시켜 학습이 불안정해지는 문제가 있다. 이 논문은 얕은 레이어부터 순차적으로 활성화하는 간단한 스케줄링만으로 학습 안정성을 높이고 최종 성능을 개선할 수 있음을 입증했다. 복잡한 초기화 기법 없이도 더 깊은 모델을 효과적으로 학습시킬 수 있는 실용적인 방법론이다.

Transformer는 수십 개의 레이어를 쌓아 올린 구조이며, 각 레이어는 Residual Connection(잔차 연결)을 통해 이전 레이어의 정보를 전달받고 자신의 출력을 더한다. 초기 학습 단계에서는 모든 레이어의 가중치가 무작위로 설정되어 있어 출력값이 매우 불안정하다. 특히 깊은 레이어는 아직 정돈되지 않은 얕은 레이어들의 출력을 입력으로 받아 학습을 진행하므로, 노이즈를 증폭시키고 전체 최적화 경로를 방해하는 경향이 있다.

ProRes는 "얕은 레이어가 먼저 배운다"는 직관을 실제 구현에 옮긴 것이다. 각 레이어의 잔차 경로에 alpha라는 가중치를 곱하는데, 학습 시작 시점에는 이 값을 0으로 설정하여 모델이 마치 입력이 그대로 출력되는 Identity Mapping처럼 동작하게 만든다. 이후 시간이 흐름에 따라 얕은 레이어부터 alpha를 1까지 서서히 높여 잔차 신호를 반영하기 시작하며, 깊은 레이어일수록 이 웜업 기간을 더 길게 가져가 상위 레이어가 안정될 때까지 기다리게 한다.

이 방식은 깊은 레이어가 초기 단계의 혼란스러운 업데이트에 휘둘리지 않게 보호한다. 결과적으로 모델 전체의 활성화 값 성장이 억제되어 학습이 안정화되며, 얕은 레이어에서 구축된 탄탄한 기초 위에 깊은 레이어의 표현력이 점진적으로 쌓이게 된다. 이는 단순한 수렴 속도 향상을 넘어, 모델이 더 나은 일반화 성능을 갖춘 지점에 도달하도록 돕는다.

ProRes의 핵심은 잔차 연결 공식 x_{l+1} = x_l + alpha(l, t) * F(Norm(x_l))에 있다. 여기서 alpha(l, t)는 레이어 인덱스 l과 학습 스텝 t에 따라 결정되는 스칼라 값이다. 학습 초기(t=0)에는 모든 레이어의 alpha가 0이 되어 모델 전체가 Identity 함수로 시작하며, 이는 초기 그래디언트 흐름을 원활하게 하고 활성화 값의 폭발적 증가를 막는 역할을 한다.

구체적인 스케줄링으로 Linear Schedule을 사용한다. alpha(l, t) = min(t / (T * l), 1) 공식을 적용하는데, [학습 스텝 t를 레이어 인덱스 l과 기본 웜업 길이 T의 곱으로 나누는 연산]을 수행한다. 이 결과값은 0에서 시작해 1에서 포화되며, 레이어 번호 l이 클수록(즉, 깊은 레이어일수록) 분모가 커져 1에 도달하는 데 더 많은 시간이 걸리게 된다.

이 기법은 Pre-LN뿐만 아니라 Post-LN, Sandwich-LN, DeepNorm 등 다양한 변형 구조에도 적용 가능하다. 각 구조의 잔차 결합 방식에 맞춰 alpha를 곱하는 위치를 조정하며, 모든 경우에서 레이어별 기여도를 시간적으로 제어한다는 동일한 원칙을 유지한다. 실험에서는 T=1000과 같은 고정된 하이퍼파라미터만으로도 튜닝 없이 우수한 성능을 확보했다.

C4 데이터셋을 이용한 사전 학습 실험에서 130M, 350M, 1.3B 규모의 모든 모델이 ProRes 적용 시 Perplexity가 유의미하게 감소했다. 특히 Post-LN 구조에서 가장 큰 폭의 개선이 관찰되었는데, 이는 Post-LN이 본래 깊은 레이어의 기여도를 과도하게 높게 잡는 경향을 ProRes가 효과적으로 보정한 결과이다.

제로샷(Zero-shot) 추론 벤치마크 평가 결과, 1.3B 모델 기준 평균 1.27%의 정확도 향상을 기록했다. 특히 LAMBADA 데이터셋에서는 정확도가 2.89% 상승하며 장기 의존성(Long-range dependency) 처리 능력이 강화되었음을 입증했다.

학습 안정성 측면에서 모델 깊이를 120 레이어까지 확장했을 때, 기존 Pre-LN은 Perplexity가 정체되거나 불안정해지는 반면 ProRes는 깊이가 깊어질수록 성능이 지속적으로 향상되는 우수한 Depth Scaling 특성을 보였다. 또한 학습 중 발생하는 손실값 및 그래디언트 노름의 스파이크 현상이 거의 발생하지 않아 매우 안정적인 학습 곡선을 유지했다.

ProRes는 Transformer의 잔차 경로를 제어하여 학습 초기 단계의 Identity Behavior를 보장한다. 이는 초기화 시점에 모델 업데이트의 크기를 제한해야 한다는 이론적 원칙을 학습 전 과정으로 확장한 것이다. 기존의 정적 초기화 기법과 달리, ProRes는 학습 진행에 따라 제약 조건을 동적으로 완화하며 모델의 표현 용량을 점진적으로 개방한다.

Pre-LN 구조에서 흔히 발생하는 지수적 활성화 값 증가(Exponential Activation Growth) 문제를 해결한다. 분석 결과, 일반적인 Pre-LN은 레이어가 깊어질수록 활성화 노름이 기하급수적으로 커지지만, ProRes는 이를 선형적인 증가로 억제하여 상위 레이어의 그래디언트 소실이나 노이즈 증폭을 방지한다.

레이어별 표현 진화(Representation Evolution) 분석을 통해, ProRes가 적용된 모델은 얕은 레이어부터 순차적으로 수렴하며 깊은 레이어로 갈수록 부드러운 코사인 유사도 변화를 보임을 확인했다. 이는 모든 레이어가 동시에 경쟁적으로 업데이트되는 기존 방식보다 더 효율적인 최적화 경로를 형성함을 시사한다.