[논문 리뷰] NormFormer: Improved Transformer Pretraining with Extra Normalization

Sam Shleifer, Jason Weston외 1명|arXiv (Cornell University)|2021. 10. 18.

Topic Modeling참고 문헌 34인용 수 28

한 줄 요약

NormFormer는 Pre-LN 모델의 각 트랜스포머 레이어에 세 가지 경량 정규화 기반 연산을 추가하여 그래디언트 불일치를 줄이고 사전학습을 가속화하며, 인과 및 마스킹 언어 모델 전반에서 perplexity와 다운스트림 작업을 개선합니다.

ABSTRACT

During pretraining, the Pre-LayerNorm transformer suffers from a gradient magnitude mismatch: gradients at early layers are much larger than at later layers. These issues can be alleviated by our proposed NormFormer architecture, which adds three normalization operations to each layer: a Layer Norm after self attention, head-wise scaling of self-attention outputs, and a Layer Norm after the first fully connected layer. The extra operations incur negligible compute cost (+0.4% parameter increase), but improve pretraining perplexity and downstream task performance for both causal and masked language models ranging from 125 Million to 2.7 Billion parameters. For example, adding NormFormer on top of our strongest 1.3B parameter baseline can reach equal perplexity 24% faster, or converge 0.27 perplexity better in the same compute budget. This model reaches GPT3-Large (1.3B) zero shot performance 60% faster. For masked language modeling, NormFormer improves fine-tuned GLUE performance by 1.9% on average. Code to train NormFormer models is available in fairseq https://github.com/pytorch/fairseq/tree/main/examples/normformer .

연구 동기 및 목표

• Identify gradient magnitude mismatches in Pre-LN transformers during pretraining.
• Propose lightweight normalization-based additions to stabilize and accelerate training.
• Evaluate NormFormer on causal and masked language models across multiple scales.
• Demonstrate improvements in pretraining perplexity and downstream task performance.
• Provide ablations and analyses to understand the contribution of each addition.

제안 방법

• Introduce three additions per layer: head-wise scaling of MHA outputs (HeadScale), a LayerNorm after the attention module, and a LayerNorm after the first FFN layer.
• Apply an additional LayerNorm inside the MHA path and a second LN after FFN, with small learnable parameters γ per head and per residual path.
• Optionally include residual scaling (ResScale) on the FFN path, analyzed for its impact at different scales.
• Train causal and masked language models across sizes 125M, 355M, 1.3B, and 2.7B, comparing NormFormer to compute-matched baselines under equal compute budgets.
• Experiment with zero-shot evaluations on GPT-3-like tasks and GLUE benchmarks to assess generalization.

실험 결과

연구 질문

• RQ1 NormFormer의 추가 정규화 연산이 Pre-LN 트랜스포머를 안정화하고 레이어 간 그래디언트 간극을 줄이는가?
• RQ2 NormFormer의 이득이 125M에서 2.7B 매개변수까지 모델 규모에 걸쳐 지속되는가?
• RQ3 추가된 연산이 사전학습 perplexity와 다운스트림 작업 성능(GLUE) 측면에서 튜닝된 Pre-LN 베이스라인과 비교해 어떤 영향을 미치는가?
• RQ4 다양한 모델 규모에서 NormFormer의 잔차 스케일링(ResScale)의 효과는 어떠한가?
• RQ5 추가 구성 요소 중 어느 하나를 제거해도 이득이 유지되지 않는지에 대한 민감도 분석이 Robust한가?

주요 결과

• NormFormer는 125M–2.7B 규모의 인과 및 마스킹 언어 모델 모두에서 사전학습 perplexity와 다운스트림 작업 성능을 개선한다.
• 1.3B 모델의 경우 NormFormer는 베이스라인 perplexity에 비해 더 빠르게 수렴하며 계산 동등 조건에서 동일 perplexity에 도달하는 속도가 24% 더 빠를 수 있으며, 같은 계산 예산에서 0.27 perplexity 더 우수하게 수렴할 수 있다.
• 제로샷 평가에서 NormFormer가 테스트된 모든 작업에서 GPT-3를 능가하는 것으로 나타났다.
• GLUE 미세조정 결과에서 NormFormer MLM이 Pre-LN 베이스라인보다 작업 전반에서 우수한 성능을 보이며 평균 이득이 나타난다.
• 아블레이션 연구에서 어떤 추가 연산도 제거하면 성능이 저하되며, HeadScale과 Post-attn LN이 특히 큰 영향을 미친다.
• 학습된 스케일링 매개변수(γ)는 초기 레이어 FG 그래디언트를 줄이고 초기 FFN 입력을 다운스케일하며, HeadScale은 특정 헤드를 강조해 안정성과 성능을 도울 수 있다.