트랜스포머 안정성의 비밀: Pre-LN과 Post-LN, 레이어 정규화의 위치가 중요한 이유

트랜스포머 안정성의 비밀: Pre-LN과 Post-LN, 레이어 정규화의 위치가 중요한 이유

인공지능 모델의 세계는 마치 정교한 부품들로 조립된 고성능 엔진과 같습니다. 그중에서도 트랜스포머(Transformer) 아키텍처는 자연어 처리(NLP) 분야에서 압도적인 성능을 자랑하며 수많은 모델의 기반이 되어왔습니다. 하지만 이 강력한 엔진을 제대로 길들이는 데에는 한 가지 까다로운 조건이 따랐습니다. 바로 '학습률 예열(learning rate warm-up)'이라는 과정입니다. 모델 훈련 초기에 학습률을 매우 낮게 시작하여 점진적으로 높여가는 이 과정은 안정적인 학습을 위해 필수적이었지만, 동시에 훈련 시간을 늘리고 추가적인 하이퍼파라미터 튜닝을 요구하는 번거로움이 있었습니다. 그런데 만약 이 예열 과정을 생략하고도 엔진을 안정적으로 구동할 수 있다면 어떨까요? 2020년에 발표된 "On Layer Normalization in the Transformer Architecture" 논문은 이 문제의 해답이 아키텍처의 아주 사소해 보이는 디테일, 레이어 정규화(Layer Normalization)의 위치에 있음을 밝혀냈습니다. 이 글에서는 기존 트랜스포머(Post-LN)가 왜 예열 단계를 필요로 했는지, 그리고 레이어 정규화의 위치를 바꾸는 것만으로 어떻게 이 문제를 해결하고 더 빠르고 효율적인 훈련을 가능하게 했는지(Pre-LN) 그 원리를 자세히 살펴보겠습니다.

---

기존 트랜스포머의 딜레마: Post-LN과 불안정한 그래디언트

오리지널 트랜스포머 아키텍처는 이제 Post-LN (Post-Layer Normalization) 구조로 불립니다. 이름에서 알 수 있듯이, 어텐션이나 피드포워드 신경망 같은 서브레이어(sub-layer)를 거친 후 잔차 연결(residual connection)이 이루어지고, 그 이후(Post)에 레이어 정규화가 적용되는 방식입니다.

• 개념: Post-LN의 정보 흐름은 입력 -> 서브레이어 -> + (잔차 연결) -> 정규화 -> 출력 순서로 이루어집니다. 즉, 연산이 끝난 결과물을 마지막에 정돈하는 방식입니다.
• 비유: 통제 없이 질주하는 경주차
  Post-LN 구조의 훈련 초기 상태를 이제 막 출발선을 떠난 고성능 경주차에 비유할 수 있습니다. 훈련 시작 시점(initialization)에 모델의 파라미터들은 무작위(Random)로 설정됩니다. 이 상태에서 계산되는 그래디언트(gradient, 모델이 어느 방향으로 학습해야 할지 알려주는 신호)는 특히 출력층에 가까워질수록 그 크기가 매우 커지는 경향이 있습니다. 이는 마치 경험 없는 드라이버가 처음부터 가속 페달을 끝까지 밟는 것과 같습니다. 차는 엄청난 힘으로 튀어 나가지만, 방향을 제어할 수 없어 트랙을 이탈하거나 충돌할 위험이 매우 큽니다. 마찬가지로, 큰 그래디언트에 높은 학습률을 그대로 적용하면 파라미터가 비정상적으로 크게 업데이트되어 훈련 과정이 불안정해지고 결국 모델이 수렴하지 못하는 문제가 발생합니다.
• 문제 해결을 위한 '예열': 학습률 예열(warm-up)은 바로 이 문제를 해결하기 위한 실용적인 해결책이었습니다. 훈련 초반에는 아주 낮은 학습률을 사용해 경주차가 천천히 출발하도록 속도를 제어하고, 모델이 어느 정도 안정적인 방향을 찾으면 점차 학습률을 높여 속도를 내는 것입니다. 이 방식은 효과적이었지만, 훈련을 더디게 만들고 예열 기간과 목표 학습률 등 신경 써야 할 변수를 늘리는 단점이 있었습니다.

---

혁신적인 해법: Pre-LN과 안정적인 학습

이러한 Post-LN의 근본적인 불안정성을 해결하기 위해 등장한 것이 바로 Pre-LN (Pre-Layer Normalization) 구조입니다. 이 구조의 핵심은 레이어 정규화의 순서를 바꾸는 단순하지만 강력한 아이디어에 있습니다.

• 개념: Pre-LN에서는 서브레이어에 입력이 들어가기 이전(Pre)에 먼저 레이어 정규화를 적용합니다. 정보의 흐름은 입력 -> 정규화 -> 서브레이어 -> + (잔차 연결) -> 출력 순서가 됩니다. 연산을 시작하기 전에 입력 데이터부터 깔끔하게 정돈하고 시작하는 것입니다.
• 비유: 숙련된 코치의 사전 전략 브리핑
  Pre-LN 방식은 경기를 시작하기 전에 숙련된 코치가 선수들에게 명확한 전략을 브리핑하는 것과 같습니다. 선수(입력 데이터)들은 경기에 투입되기 전 코치의 지시(레이어 정규화)를 통해 각자의 역할을 명확히 인지하고 정돈된 상태가 됩니다. 이렇게 잘 정돈된 상태에서 경기를 시작하면, 선수들은 경기 초반부터 혼란 없이 자신의 역할을 안정적으로 수행할 수 있습니다. 즉, Pre-LN 구조에서는 훈련 시작부터 그래디언트가 폭발적으로 커지지 않고 "얌전하게(well-behaved)" 유지됩니다. 그래디언트가 안정적이기 때문에, 처음부터 높은 학습률을 적용해도 훈련이 불안정해질 위험이 없습니다. 따라서 학습률 예열이라는 복잡한 과정 자체가 필요 없게 되는 것입니다.
• 예시: "The cat sat on the mat"이라는 문장을 처리한다고 가정해 봅시다. Post-LN에서는 'cat'이라는 단어 벡터가 어텐션 메커니즘을 거친 후, 원래의 'cat' 벡터와 더해진 다음에야 그 값이 정규화됩니다. 만약 이 과정에서 값이 너무 커지면 다음 레이어로 전달될 때 불안정성을 야기할 수 있습니다. 반면, Pre-LN에서는 'cat'이라는 단어 벡터를 어텐션 메커니즘에 넣기 전에 먼저 정규화합니다. 깨끗하게 정제된 데이터를 가지고 연산을 시작하기 때문에, 결과값이 급격하게 변동할 위험이 현저히 줄어듭니다.

---

Pre-LN이 가져온 명확한 이점

Pre-LN 구조를 채택함으로써 트랜스포머 훈련은 다음과 같은 명확한 이점을 얻게 됩니다.

• 훈련 속도 향상: 예열 단계를 완전히 제거할 수 있어 전체 훈련 시간을 크게 단축시킬 수 있습니다.
• 하이퍼파라미터 튜닝 간소화: 더 이상 예열 스케줄을 신경 쓸 필요가 없어 모델 최적화 과정이 훨씬 단순해집니다.
• 안정적인 성능: 예열 과정 없이도 기존의 Post-LN 모델과 비슷한 수준의 성능을 달성할 수 있음이 다양한 실험을 통해 입증되었습니다.

---

마무리하며

트랜스포머 아키텍처에서 레이어 정규화의 위치를 바꾸는 것은 사소한 수정처럼 보일 수 있지만, 그 결과는 모델의 훈련 역학을 근본적으로 바꾸는 나비효과를 불러왔습니다. Post-LN의 불안정성을 해결하고 예열 단계를 제거한 Pre-LN의 등장은 어떻게 작은 설계 변화가 모델의 안정성과 효율성에 지대한 영향을 미칠 수 있는지를 명확히 보여주는 사례입니다. 이는 우리가 더 강력하고 효율적인 인공지능을 만들기 위해 모델의 거시적인 구조뿐만 아니라, 그 안을 구성하는 부품들이 어떻게 상호작용하는지에 대한 깊은 이해가 얼마나 중요한지를 다시 한번 일깨워 줍니다. Pre-LN은 트랜스포머를 더 빠르고, 더 쉽게, 그리고 더 안정적으로 만드는 중요한 핵심이 되었습니다.

---

https://arxiv.org/abs/2002.04745

On Layer Normalization in the Transformer Architecture