Weight Decay Improves Language Model Plasticity

기술 요약: 가중치 감쇠(Weight Decay)가 언어 모델의 가소성을 향상시킨다

문제 제기
현재 대규모 언어 모델(LLM)을 훈련하는 관행은 일반적으로 사전 훈련(pretraining)과 사후 훈련(post-training, 미세 조정 및 정렬)을 분리하여 다룹니다. 사전 훈련 단계에서의 하이퍼파라미터 최적화와 스케일링 법칙은 주로 베이스 모델의 검증 손실(교차 엔트로피)을 기준으로 평가되며, 이는 낮은 사전 훈련 손실이 더 유능한 다운스트림 모델로 이어진다는 가정에 기반합니다. 그러나 이러한 접근 방식은 모델 가소성(model plasticity), 즉 재초기화 없이 추가 훈련을 통해 새로운 데이터와 작업에 성공적으로 적응할 수 있는 능력을 간과합니다. 본 논문은 사전 훈련 손실을 최소화하기 위해 사전 훈련 하이퍼파라미터를 최적화하는 것이 의도치 않게 이후의 미세 조정 단계에서 모델의 적응력을 저해하는지 조사합니다. 특히, 저자들은 표준적인 규제 하이퍼파라미터인 **가중치 감쇠(weight decay)**에 초점을 맞추어, 사전 훈련 손실을 최소화하기 위한 최적의 값과 다운스트림 적응력을 최대화하기 위한 최적의 값이 서로 다른지 확인하고자 합니다.

방법론
저자들은 두 가지 모델 제품군(Llama-2 및 OLMo-2)과 다양한 모델 크기(최대 4B 파라미터)에 걸쳐 체계적인 경험적 연구를 수행합니다. 실험 설계는 다른 요인들을 통제하면서 사전 훈련 시 가중치 감쇠 하이퍼파라미터($\lambda$)를 변화시킵니다.

1. 훈련 레짐(Training Regimes): 모델은 두 가지 별개의 레짐으로 훈련됩니다:
   • 컴퓨트 최적(Compute-optimal): 20 TPP (tokens-per-parameter).
   • 과적합 훈련(Overtrained): 140 TPP.
2. 미세 조정 평가: 사전 훈련된 모델은 다음과 같은 다양한 다운스트림 작업 세트에 대해 미세 조정됩니다:
   • 6개의 사고 사슬(Chain-of-Thought, CoT) 추론 데이터셋 (예: MetaMathQA, MMLUProCoT).
   • 5개의 언어 이해 및 상식 추론 벤치마크 (예: HellaSwag, WinoGrande).
   • 1개의 안전 정렬(safety alignment) 작업.
3. 지표: 다운스트림 성능은 Greedy (Pass@1), Pass@16, Majority Vote (Maj@16), Reward Model scoring (RM@16), 그리고 Correct Ratio를 포함한 다수의 지표를 사용하여 평가됩니다.
4. 메커니즘 분석: 가중치 감쇠가 왜 가소성에 영향을 미치는지 이해하기 위해 저자들은 다음을 분석합니다:
   • 선형 프로빙(Linear Probing): 내부 표현의 선형 분리 가능성(linear separability) 측정 (감성 및 주제 분류).
   • 어텐션 랭크(Attention Rank): 쿼리-키($W_{QK}$) 및 밸류-프로젝션($W_{VP}$) 행렬의 의사 랭크(pseudo-rank) 분석.
   • 과적합(Overfitting): 훈련 손실과 검증 손실 사이의 격차(Train-Val Gap) 계산.

주요 기여

1. 가소성의 동력으로서의 가중치 감쇠: 본 논문은 가중치 감쇠가 모델 가소성을 형성하는 데 결정적인 요소임을 입증합니다. 표준적인 기본값과는 반대로, 사전 훈련 시 더 큰 가중치 감쇠 값은 미세 조정 후 다운스트림 성능을 일관되게 향상시킵니다.
2. 목적 함수의 분리: 본 연구는 사전 훈련 검증 손실을 최소화하는 하이퍼파라미터 값이 반드시 다운스트림 성능을 최대화하는 값은 아니라는 증거를 제공합니다. 여러 레짐에서, (더 큰 가중치 감쇠로 인해) 더 높은 사전 훈련 손실을 가진 모델이 미세 조정 후에는 더 낮은 손실을 가진 모델보다 더 우수한 성능을 보였습니다.
3. 메커니즘적 통찰: 저자들은 가중치 감쇠가 가소성을 강화하는 세 가지 구체적인 메커니즘을 식별했습니다:
   • 이는 선형 분리 가능한 표현을 장려하여, 내부 특징들이 다운스트림 작업에 더 용이하게 접근할 수 있도록 합니다.
   • 이는 어텐션 행렬을 규제하며, 특히 쿼리-키 곱($W_{QK}$)의 랭크를 감소시켜 더 압축되고 견고한 구조를 만듭니다.
   • 이는 사전 훈련 데이터에 대한 과적합을 줄여, 효과적으로 특정 훈련 노이즈를 "망각"하고 새로운 정보를 학습할 수 있는 모델의 용량을 보존합니다.

주요 결과

• 최적 가중치 감쇠의 이동: 컴퓨트 최적(20 TPP) 훈련의 경우, 다운스트림 성능을 위한 최적의 가중치 감쇠는 표준 기본값인 0.1보다 현저히 큽니다. 예를 들어, 이 레짐에서 Llama-2 및 OLMo-2 모델의 최적값은 0.5–1.0에서 발견되었습니다. 과적합 훈련(140 TPP) 레짐에서도 다운스트림 성능을 위한 최적값(0.3)은 사전 훈련 손실을 최소화하는 값(0.1)보다 높게 유지되었습니다.
• 역설적인 트레이드오프: 본 논문은 "더 나쁜" 베이스 모델(높은 사전 훈련 손실)이 미세 조정 후 "더 나은" 최종 모델이 될 수 있다는 트레이드오프를 강조합니다. 예를 들어, $\lambda=0.3$ 또는 $1.0$으로 사전 훈련된 OLMo-2-1B-140x 모델은 $\lambda=0.1$ 베이스라인보다 높은 사전 훈련 손실을 가졌음에도 불구하고 눈에 띄게 더 나은 다운스트림 정확도를 달-성했습니다.
• 강건성: 이러한 결과는 모델 제품군, 크기, 그리고 사전 훈련 및 미세 조정 중 다른 하이퍼파라미터(학습률, 배치 크기)가 변하는 경우에도 일관되게 나타납니다.
• 메커니즘적 상관관계:
  • 높은 가중치 감쇠는 더 구조화된 표현을 나타내는 높은 선형 프로빙 정확도와 상관관계가 있습니다.
  • 높은 가중치 감쇠는 어텐션 행렬(특히 $W_{QK}$)의 낮은 의사 랭크와 상관관계가 있으며, 이는 저차원(low-rank)의 견고한 가설로의 이동을 시사합니다.
  • 높은 가치 감쇠는 더 작은 Train-Val Gap과 상관관계가 있으며, 이는 과적합의 감소를 나타냅니다.

의의 및 주장
본 논문은 하이퍼파라미터 최적화를 오직 사전 훈련 검증 손실에만 의존하여 수행하는 관행을 재고해야 한다고 주장합니다. 저자들은 모델 가소성이 독특하고 중요한 속성이며, 순수하게 사전 훈련 손실만을 최적화할 때 종종 희생된다고 주장합니다.

본 연구의 의의는 다음과 같습니다:

1. 기본값의 재평가: 표준적인 하이퍼파라미터 선택(예: $\lambda=0.1$)이 LLM 개발의 전체 생애 주기, 특히 다운스트림 적응력이 궁극적인 목표인 경우에 서브옵티멀(suboptimal)할 수 있음을 시사합니다.
2. 교차 엔트로피의 한계: 교차 엔트로피 손실이 다운스트림 능력의 불완전한 대리 지표임을 강조합니다. 유사한 손실을 가진 모델이라도 가소성 측면에서는 매우 다를 수 있기 때문입니다.
3. 하이퍼파라미터의 다면적 역할: 가중치 감쇠와 같은 단일 하이퍼파라미터가 일반화뿐만 아니라 미래의 학습을 용이하게 하는 구조적 특성(선형성, 랭크)에 영향을 미치는 복잡하고 다면적인 역할을 수행함을 강조합니다.

저자들은 자신들의 결과가 테스트된 규모 전반에서 견고하지만, 현재 실험 범위 너머의 극도로 거대하거나 과하게 훈련된 시나리오에서는 최적값이 변할 수 있음을 언급하며 겸허한 태도를 유지합니다. 또한, 안정성-가소성(stability-plasticity) 트레이드오프와 이러한 발견이 다른 파운데이션 모델 모달리티에 적용될 수 있는지에 대한 추가 연구를 촉구합니다.