본 논문은 Transformer의 다양한 구성 요소(embedding layer, normalization layer, self-attention mechanism, point-wise feedforward network) 간의 상호작용과 차이점을 분석하여, 각 블록별로 초기 학습 단계부터 지속되는 명확한 Sharpness Disparity를 발견했습니다. 이를 바탕으로, 각 블록의 sharpness에 맞춰 학습률을 조정하는 Blockwise Learning Rate (LR) 전략을 제안합니다. AdamW에 Blockwise LR을 통합하여, 기존 AdamW 대비 최종 손실 감소 및 약 2배의 학습 속도 향상을 달성했습니다. GPT-2와 LLaMA (0.12B ~ 2B 파라미터) 모델과 OpenWebText, MiniPile, C4 데이터셋을 사용하여 다양한 실험을 통해 성능 향상을 검증했습니다. 또한, 메모리 효율적인 Adam 변형체인 Adam-mini에 Blockwise LR을 적용하여 학습 속도 2배 향상과 메모리 사용량 2배 절감을 동시에 달성했습니다. 결과적으로, sharpness disparity를 활용하여 LLM 학습을 개선할 가능성을 보여줍니다.
