Skip to main content
공개 노트

GEGLU 출력 가중치만 0으로 두면 왜 빨라질까

Views 9약 3분생성형 AI 활용
Summary

GEGLU 출력 가중치만 0으로 두면 왜 빨라질까

트랜스포머 학습에서 초기화는 작은 구현 디테일처럼 보이지만, 실제로는 수렴 속도와 초반 안정성에 영향을 줄 수 있다. 최근 공유된 한 실험에서는 QKNorm, GEGLU, RMSNorm을 쓰는 바닐라 트랜스포머에서 geglu_out만 0으로 초기화했을 때, 같은 손실에 도달하는 속도가 약 10% 빨라졌다고 밝혔다.1

핵심은 “제로 초기화가 좋다”가 아니다. 더 정확히는 GEGLU 블록 전체가 아니라 출력 투영 가중치만 0으로 둔다는 점이다.

왜 “출력만 0”인가

일반적으로 모든 가중치를 0으로 초기화하는 것은 좋지 않다고 알려져 있다. 뉴런들이 같은 값을 갖고 같은 방향으로 업데이트되면 대칭성이 깨지지 않아 서로 다른 특징을 배우기 어렵기 때문이다.2

하지만 geglu_out 제로 초기화는 그 경우와 다르다. GEGLU 내부의 게이트와 중간 표현까지 모두 0으로 죽이는 것이 아니라, 마지막 출력 투영만 0에서 시작하게 한다. 즉, 블록 내부는 학습 가능한 상태로 두되, 학습 초기에 그 블록이 잔차 경로에 주는 영향을 거의 없게 만드는 방식이다.

이 관점에서 보면 이 초기화의 목적은 명확하다. 모델 초반에는 FFN 블록이 표현을 과하게 흔들지 않도록 하고, 학습이 진행되면서 필요한 만큼만 출력 기여도를 키우게 하는 것이다.

GEGLU는 단순한 활성화 함수 이상의 구현 이슈다

GEGLU는 GLU 계열 활성화로, 변환된 값과 게이트 값을 곱해 표현을 조절한다. SwiGLU와 함께 트랜스포머 FFN에서 자주 쓰이는 패턴이다.

이 패턴은 하드웨어 최적화에서도 중요하게 다뤄진다. NVIDIA cuDNN의 Grouped GEMM + GLU fusion API는 SwiGLU와 GeGLU를 모두 지원하며, GeGLU는 act_func="geglu"로 선택할 수 있다.3 이 API는 NVIDIA Blackwell GPU의 SM100 이상에서 동작하도록 설계됐고, MoE 워크로드를 대상으로 하며, CUTLASS/CUTE로 구현됐다.3

다만 이 API는 실험적 API로, 향후 변경될 수 있다고 명시되어 있다.3 따라서 GEGLU 출력 가중치 초기화를 평가할 때는 알고리즘 차이만 볼 것이 아니라, 실제 사용하는 커널과 실행 경로까지 함께 고려해야 한다.

실험할 때 확인해야 할 것

geglu_out 제로 초기화는 바로 모든 모델에 적용할 정답이라기보다, 검증해볼 만한 초기화 가설에 가깝다. 최소한 다음 항목은 함께 봐야 한다.

  • 같은 손실에 도달하기까지 필요한 토큰 수

  • 학습 초반 loss spike 여부

  • 최종 손실과 다운스트림 성능

  • 기존 초기화 대비 안정성 변화

특히 MoE나 저정밀 학습에서는 해석에 더 주의해야 한다. cuDNN의 Grouped GEMM + GLU fusion은 dense mode에서 모든 expert 가중치를 하나의 연속된 (N, K, L) 텐서로 패킹하고, discrete mode에서는 expert별 가중치 포인터를 사용해 가중치 스태킹이 필요 없다.3 초기화 효과를 비교하더라도 커널 경로나 레이아웃이 함께 바뀌면, 실제로 무엇이 성능 차이를 만들었는지 분리하기 어려워진다.

결론

GEGLU 출력 가중치 제로 초기화는 “제로 초기화는 항상 나쁘다”는 단순한 규칙의 예외가 될 수 있다. 중요한 차이는 전체 네트워크나 GEGLU 블록 전체를 0으로 만드는 것이 아니라, 블록의 초기 출력 영향력만 제한한다는 점이다.

따라서 이 방법의 실용적 가치는 초반 학습을 더 안정적으로 만들고, 같은 손실에 더 빨리 도달할 가능성을 실험해볼 수 있다는 데 있다. 다만 모델 구조, 정밀도, MoE 여부, 커널 구현이 결과에 영향을 줄 수 있으므로, 적용 전에는 동일 조건에서 수렴 속도와 최종 성능을 함께 비교하는 것이 안전하다.

1Zero Init On Geglu Out Weights Speeds Transformer Training By 10%

2Weight Initialization Techniques for Deep Neural Networks

3Grouped GEMM + GLU (SM100)#

이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.

틸노트는 사용자가 AI를 활용해 노트를 쉽게 작성하거나 편집할 수 있는 서비스입니다.

AI 스퀘어 단톡방에 함께해요

AI 활용과 틸노트 사용 팁을 나누는 오픈채팅방입니다. 입장 비밀번호는 tilnote입니다.

참여하기