기술적 요약: Sharpness-Aware Pretraining 은 Catastrophic Forgetting 을 완화합니다

문제 제기

현재의 사전 학습 최적화 전략들은 일반적으로 기본 모델의 품질을 최대화하도록 조정되며, 이는 사전 학습 손실 또는 벤치마크 성능으로 측정됩니다. 이러한 접근은 기본 모델이 강할수록 미세 조정 (post-training) 이나 양자화 같은 후속 수정을 거친 후에도 더 강력한 모델이 된다는 가정 하에 이루어집니다. 그러나 최근 연구는 이 가정이 실패할 수 있음을 시사합니다. 즉, 특정 지점 이상으로 사전 학습을 연장하면 사전 학습 손실은 개선되지만, post-training 이후의 하위 작업 성능은 저하될 수 있습니다.

핵심적으로 지적된 문제는 기본 모델 평가가 파라미터 업데이트 하에서 모델 능력의 안정성을 무시한다는 점입니다. 이러한 업데이트에 매우 민감한 모델들은 "Catastrophic Forgetting"을 겪어, 초기 강도와 상관없이 사전 학습된 능력을 상실합니다. 이 민감성은 손실 지형 (loss landscape) 의 기하학적 구조와 연결되어 있습니다. 구체적으로, 'sharp'한 최소점에 수렴하는 모델들은 하위 작업 동안 파라미터가 교란될 때 성능 저하에 더 취약합니다.

방법론

저자들은 모델을 flatter minima 쪽으로 편향시켜 국소 파라미터 교란에 대한 민감도를 낮추도록 설계된 세 가지 사전 학습 최적화 접근법을 조사합니다:

1. Sharpness-Aware Minimization (SAM): 손실 곡률을 명시적으로 패널티로 부과하는 최적화 방법입니다. SAM 은 반지름 $\rho$의 이웃 내에서 최대 손실을 최소화하는 파라미터 $\theta$를 찾아, 작은 교란 하에서도 낮은 손실을 유지하는 최소점을 효과적으로 발견합니다.
2. 증가된 피크 학습률: 'Edge of Stability' 현상에 기반한 암시적 접근법으로, 더 큰 학습률이 헤시안 (Hessian) 의 최대 고유값을 암시적으로 제한하여 손실의 sharpness 를 줄인다는 가설에 근거합니다.
3. 단축된 학습률 어닐링 (Annealing): 학습률 스케줄을 수정하여 훈련 말기의 감쇠 기간을 줄이는 방법입니다. 이는 훈련 말기 역학이 해의 최종 기하학에 상당한 영향을 미친다는 발견에 동기를 부여합니다.

본 연구는 20M 에서 150M 파라미터에 이르는 다양한 모델 크기에 걸쳐 이러한 개입들을 평가하며, 1B 파라미터 모델 (OLMo-2-1B) 에서 규모를 확장하여 검증합니다. 평가는 두 가지 주요 트레이드오프에 초점을 맞춥니다:

• 학습 - 망각 트레이드오프: 새로운 능력 획득 (미세 조정 손실) 과 사전 학습 능력 유지 (미세 조정 후 사전 학습 손실) 사이의 균형.
• 압축 - 망각 트레이드오프: 양자화 (예: 4 비트) 를 통해 모델을 압축할 때 발생하는 능력 저하.

저자들은 또한 이러한 방법들이 미세 조정 방향을 따라 방향성 sharpness 를 실제로 감소시키는지, 그리고 미세 조정 교란에 대해 손실의 2 차 테일러 근사가 유효한지 확인하기 위해 헤시안을 분석합니다.

주요 기여 및 결과

1. 개선된 학습 - 망각 트레이드오프

80 회 이상의 사전 학습 실행과 3,500 회 이상의 미세 조정 실험을 통해, 저자들은 sharpness-aware 개입들이 일관되게 학습 - 망각 트레이드오프를 개선한다는 사실을 발견했습니다:

• SAM 성능: SAM 으로 사전 학습된 모델들은 표준 AdamW 에 비해 우월한 파레토 프론티어 (Pareto frontier) 를 달성합니다. 예를 들어, StarCoder 데이터셋에서 SAM 은 동일한 미세 조정 손실 수준에서 망각을 80% 감소시킵니다.
• 확장 효과: 사전 학습 토큰 예산이 증가함에 따라 AdamW 대비 SAM 의 이점이 커집니다. AdamW 체크포인트는 훈련이 진행될수록 더욱 민감해져 (catastrophic overtraining) 성능이 저하되는 반면, SAM 은 이를 완화하여 높은 토큰 영역에서 가장 강력한 이점을 보여줍니다.
• 모델 크기: 토큰 - 당 - 파라미터 비율을 통제한 상태에서 모델 규모 (20M 에서 150M 파라미터) 가 커짐에 따라 이점은 지속되며 오히려 커집니다.

2. 미세 조정을 넘어선 견고성

sharpness-aware 사전 학습이 부여하는 견고성은 다른 하위 작업 교란으로도 확장됩니다:

• 양자화: SAM 으로 훈련된 모델들은 4 비트 양자화 하에서 성능 저하가 현저히 적습니다. 높은 토큰 예산 영역에서 SAM 은 기준선에 비해 양자화로 인한 손실 증가를 2~3 배 줄여줍니다.
• 가우시안 교란: SAM 으로 훈련된 모델들은 가중에 추가된 등방성 가우시안 잡음에 덜 민감하여, 더 견고한 국소 기하학을 나타냅니다.

3. 암시적 방법과 실용적 확장성

• 학습률 역학: 더 높은 피크 학습률과 더 짧은 어닐링 기간 모두, 기본 모델의 사전 학습 손실이 약간 더 나빠지더라도 학습 - 망각 트레이드오프를 개선합니다.
• 중간 훈련 개입: SAM 의 계산 비용 (단계당 계산량을 약 2 배 증가시킴) 을 해결하기 위해, 저자들은 학습률 어닐링 단계 동안에만 SAM 을 적용할 것을 제안합니다.
  • OLMo-2-1B 에 대한 실험에서, 훈련 마지막 10% (중간 훈련) 동안만 SAM 을 적용한 경우, 표준 기준선에 비해 MetaMath post-training 후 망각이 31% 감소했고, 4 비트 양자화 후 40% 감소했습니다. 이 접근법은 추가 계산 비용을 거의 들이지 않으면서 대부분의 이점을 포착합니다.

4. 헤시안 분석

본 논문은 헤시안을 분석하여 이론적 직관을 검증합니다:

• 2 차 근사: 손실의 2 차 테일러 전개는 특히 SAM 으로 훈련된 모델들의 경우 망각에 대한 합리적인 상한선 역할을 합니다.
• 방향성 Sharpness: SAM 과 더 큰 피크 학습률 모두 미세 조정 방향을 따라 방향성 sharpness ($\Delta_{FT}^\top H \Delta_{FT} / \|\Delta_{FT}\|^2$) 를 명시적으로 감소시킵니다. 망각의 감소는 미세 조정 단계 크기의 감소가 아니라 손실 지형의 평탄화 (flattening) 에 의해 주도됩니다.

중요성 및 주장

본 논문은 post-training 업데이트에 대한 파라미터 민감도가 사전 학습 역학과 하위 작업 성능 사이의 중요하지만 간과된 연결고리라고 주장합니다. 저자들은 기본 모델 품질만을 최적화하는 표준 관행은 불충분하다고 역설합니다.

주요 주장들은 다음과 같습니다:

• 사전 학습 목표의 재정의: 사전 학습 레시피는 기본 모델 품질과 하위 작업 파라미터 이동에 대한 낮은 민감도 모두를 고려하여 선택되어야 합니다.
• 파이프라인을 위한 최적화: 언어 모델 사전 학습은 기본 체크포인트를 고립적으로 최적화하는 것이 아니라, 미세 조정, 정렬, 압축을 포함한 전체 모델 개발 파이프라인을 위해 최적화되어야 합니다.
• 실용적 영향: 실무자들은 전체 사전 학습 과정 내내 적용하는 대신, 훈련 말기 (어닐링) 에 sharpness-aware 기법 (예: SAM) 을 선택적으로 배포함으로써 상당한 견고성 이점을 얻을 수 있으며, 이는 접근법의 확장성을 가능하게 합니다.

저자들은 적응성과 낮은 민감도를 1 순위 평가 기준으로 간주하면, 기본 모델 자체가 전통적인 사전 학습 손실 지표로는 약간 약해 보일지라도 사전 학습된 모델들이 향후 수정을 위한 더 강력한 출발점이 될 수 있다고 결론지었습니다.