수억 원의 클라우드 컴퓨팅 비용을 들여 거대 언어 모델(LLM)을 사전 학습시키던 중, 갑자기 Loss 값이 'NaN'으로 변하며 학습이 중단되는 상황을 겪어보셨나요? 혹은 특정 체크포인트 이후로 모델의 성능이 더 이상 개선되지 않고 오히려 진동하는 현상 때문에 밤잠을 설친 적이 있다면, 문제는 여러분의 데이터셋이나 학습률(Learning Rate)이 아니라 활성화 함수(Activation Function)의 구조적 특성에 있을 가능성이 높습니다. 현재 Llama 3나 Mistral 같은 주요 모델들이 채택하고 있는 SwiGLU는 강력한 성능을 보여주지만, 학습의 안정성 측면에서는 치명적인 약점을 내포하고 있습니다.
SwiGLU의 화려한 성능과 그 이면에 숨겨진 리스크
현재 LLM 아키텍처의 표준처럼 자리 잡은 SwiGLU(Swish-Gated Linear Unit)는 입력값이 클 때 2차 함수($x^2$)에 가까운 거동을 보입니다. 이러한 특성은 모델이 복잡한 비선형 관계를 학습하는 데 있어 매우 유리하게 작용하며, 모델의 표현력(Expressive Capacity)을 극대화하는 핵심 요소로 평가받습니다. 실제로 Llama 2와 Llama 3 모델은 SwiGLU를 통해 이전 세대 모델들보다 더 적은 파라미터로도 높은 벤치마크 점수를 기록한 바 있습니다 (출처: Meta AI Llama 3 Technical Report).
하지만 이 강력한 비선형성은 양날의 검입니다. 입력값이 특정 임계치를 넘어서는 순간, 출력값이 기하급수적으로 커지며 그래디언트 폭주(Gradient Explosion)를 유발하기 때문입니다. 특히 수조 개의 토큰을 학습하는 사전 학습(Pre-training) 단계에서는 이러한 수치적 불안정성이 누적되어 결국 전체 학습 과정을 망가뜨리는 원인이 됩니다. 개발자 입장에서는 성능을 위해 SwiGLU를 쓰자니 학습의 안정성이 걱정되고, ReLU 같은 단순한 함수를 쓰자니 모델의 지능이 떨어지는 딜레마에 빠지게 됩니다.
PowLU가 제시하는 새로운 대안: 안정성과 효율의 균형
최근 연구에서 제안된 PowLU(Power Linear Unit)는 바로 이 SwiGLU의 고질적인 불안정성을 해결하기 위해 등장했습니다. PowLU의 핵심 아이디어는 입력값이 커질 때의 증가율을 2차 함수보다 완만하게 제어하는 데 있습니다. 이를 통해 모델이 충분한 비선형성을 확보하면서도, 수치적으로 폭주하지 않는 안전장치를 마련한 것입니다.
PowLU는 수치적 안정성을 확보함으로써 학습률을 더 공격적으로 설정할 수 있게 해줍니다. 이는 결과적으로 모델의 수렴 속도를 높여 전체 학습 시간을 단축하는 효과를 가져옵니다. 반면 SwiGLU는 안정성을 위해 학습률을 보수적으로 잡아야 하거나, 학습 도중 발생하는 스파이크를 잡기 위해 별도의 정규화 기법을 추가해야 하는 번거로움이 있습니다.
| 비교 항목 | SwiGLU | PowLU |
|---|---|---|
| 주요 특성 | 큰 입력에서 $x^2$에 근사 | 제어된 파워 커브(Power Curve) |
| 학습 안정성 | 낮음 (Loss Spike 위험 존재) | 높음 (수치적 안정성 확보) |
| 표현력 | 매우 높음 | 높음 |
| 권장 학습률 | 보수적 설정 필요 | 상대적으로 공격적 설정 가능 |
조직의 상황과 목표에 따른 최적의 선택지
활성화 함수를 선택할 때는 단순히 최신 논문의 결과만을 따르기보다, 현재 가용한 자원과 목표하는 모델의 규모를 냉정하게 따져봐야 합니다.
첫째, 수천 장의 GPU를 동원하여 70B 이상의 거대 모델을 처음부터 사전 학습시켜야 하는 기업이라면 PowLU가 훨씬 경제적인 선택입니다. 거대 모델 학습에서 가장 큰 비용은 GPU 시간이 아니라, 학습 실패로 인해 버려지는 시간과 인건비입니다. PowLU를 사용하면 학습 도중 발생할 수 있는 'NaN' 오류 가능성을 현저히 낮출 수 있어, 전체 프로젝트의 리스크 관리에 유리합니다.
둘째, 기존의 Llama 계열 모델을 미세 조정(Fine-tuning)하거나 소규모 데이터셋으로 특정 도메인 모델을 만드는 경우에는 SwiGLU를 유지하는 것이 낫습니다. 이미 SwiGLU로 사전 학습된 가중치(Weights)들이 해당 함수의 특성에 최적화되어 있기 때문입니다. 활성화 함수를 중간에 바꾸는 것은 모델의 뇌 구조를 바꾸는 것과 같아 성능 저하를 초래할 수 있습니다.
셋째, 연구 목적이나 실험적인 소형 모델(7B 이하)을 개발 중이라면 두 함수를 모두 테스트해 볼 가치가 있습니다. 하지만 하이퍼파라미터 튜닝에 투입할 시간이 부족하다면, 기본적으로 더 넓은 안정 범위를 제공하는 PowLU를 기본값으로 설정하는 것이 정신 건강에 이롭습니다.
최종 판단: 안정성이 성능을 압도하는 순간
필자의 판단으로는, 향후 LLM 아키텍처의 흐름은 극단적인 성능 추구에서 '예측 가능한 학습 안정성'으로 옮겨갈 것입니다. SwiGLU가 보여준 놀라운 성능은 분명 매력적이지만, 대규모 클러스터에서 학습을 관리하는 엔지니어 입장에서 불확실성은 가장 큰 적입니다.
결론적으로, 새로운 아키텍처를 설계하거나 대규모 프리트레이닝을 계획하고 있다면 PowLU를 도입할 것을 강력히 권장합니다. SwiGLU가 제공하는 미세한 표현력의 우위보다, 중단 없는 학습이 가져다주는 시간적, 금전적 이득이 훨씬 크기 때문입니다. 기술적 화려함보다는 견고한 안정성이 결국 더 똑똑한 모델을 완성하는 지름길입니다. 지금 당장 여러분의 모델 설정 파일에서 활성화 함수 옵션을 다시 점검해 보시기 바랍니다.
참고: arXiv CS.LG (Machine Learning)