기술 요약: 언어 모델 확장 시 웜스타팅이 효과적인 시점은 언제인가?

문제 제기

최첨단 딥러닝 모델을 학습시키려면 막대한 컴퓨팅 자원이 필요합니다. 모델이 수백만 개에서 수십억 개의 파라미터로 확장됨에 따라, 더 작은 모델을 학습하는 데 투자된 컴퓨팅 자원을 더 큰 모델의 학습을 가속화하는 데 재사용할 수 있는지에 대한 중요한 질문이 제기됩니다. 웜스타팅 (Warmstarting, WS) 또는 **모델 성장 (Model Growth)**이라고 불리는 이 기술은 학습된 작은 체크포인트에서 더 큰 모델을 초기화하는 것을 포함합니다. 상당한 연구적 관심을 받고 있음에도 불구하고, WS 는 생산 파이프라인에서 실제 채택이 제한적입니다.

저자들은 이러한 채택 부재의 원인을 세 가지 주요 장벽으로 꼽습니다:

1. 구현 복잡성: 기존 방법들은 종종 기능 보존을 위한 특정 복제 규칙 (예: 아키텍처별 클로닝 규칙) 과 같은 아키텍처별 적응을 요구하며, 이는 현대 구성 요소에 일반화되지 않습니다.
2. 불명확한 성장 한계: 실무자들은 체크포인트가 웜스타팅이 제로부터 학습하는 것보다 역효과를 내기 전에 얼마나 성장할 수 있는지에 대한 지침이 부족합니다.
3. 하이퍼파라미터의 불투명성: WS 방법, 하이퍼파라미터, 그리고 확장 예산 (모델 크기 대 데이터 토큰) 간의 상호작용이 충분히 탐구되지 않아, 효율성 향상을 상쇄하는 튜닝 오버헤드를 초래합니다.

또한, 기존 문헌은 종종 초기화 시 기본 모델의 정확한 출력을 보존하는 것 (Function Preservation, FP) 을 강조하는데, 이는 불필요하게 연산자 설계를 제한할 수 있습니다.

방법론

1. 이론적 프레임워크: Shrink-Zero-Perturb (SZP)

저자들은 **Shrink-Zero-Perturb (SZP)**라고 불리는 일반화되고 아키텍처에 구애받지 않는 성장 연산을 제안합니다. 그들은 지속적 학습 (continual learning) 의 "Shrink-and-Perturb (SnP)"공식을 모델 성장으로 확장합니다. 목표 가중치 $\theta_{target}$은 다음과 같이 유도됩니다:
$$\theta_{target} = \lambda_{shrink} \cdot G(\theta_{base}) + N(0, \sigma_{perturb}^2)$$
여기서:

• $G$는 기본 가중치를 더 큰 목표 형태로 매핑하는 성장 함수 (예: 제로 패딩 또는 클로닝) 입니다.
• $\lambda_{shrink}$는 기본 표현의 지배력을 줄이기 위해 상속된 가중치를 축소합니다.
• $\sigma_{perturb}$는 대칭성을 깨고 새로 추가된 뉴런을 활성화하기 위해 가우시안 노이즈를 추가합니다.

저자들은 초기화 시 성장된 모델의 출력이 기본 모델과 정확히 일치하는 엄격한 **기능 보존 (FP)**이 효과적인 웜스타팅에 필수적이지 않다고 주장합니다. 대신 SZP 는 신호 보존과 새로운 용량을 활용하는 유연성 사이의 균형을 잡는 원칙적인 기준선을 제공합니다.

2. 실험 설정

이 연구는 두 가지 주요 실험 설정을 사용합니다:

• 합성 회귀 (MLPs): 멱법칙 푸리에 스펙트럼을 가진 합성 회귀 작업에서 너비 확장된 다층 퍼셉트론 (MLPs) 을 사용한 통제된 테스트베드입니다. 이는 튜닝 전략과 성장 한계를 분석하기 위해 밀집된 하이퍼파라미터 스윕을 가능하게 합니다.
• 언어 모델링 (LLMs): SLIMPAJAMA 코퍼스로 학습된 디코더 전용 GPT-2 스타일 모델 (및 ablation 에서 LLaMA 스타일) 입니다. 실험은 14M 에서 1.2B 파라미터까지의 모델 크기를 다룹니다.

주요 실험 통제:

• 하이퍼파라미터 전이: 저자들은 **$\mu$P (Maximal Update Parameterization)**를 사용하여 기본 모델에서 목표 모델로 하이퍼파라미터 (학습률, 배치 크기) 를 전이하여 재튜닝 비용이라는 혼란 변수를 제거합니다.
• 학습 예산: 모델은 고정된 토큰당 파라미터 예산 (10, 20, 30 토큰/파라미터) 으로 학습되며, 20 토큰/파라미터는 Chinchilla 컴퓨팅 최적 프론티어를 나타냅니다.
• 비교: SZP 는 Scratch(제로부터 학습) 와 표준 Net2Net(기능 보존 기준선) 과 비교됩니다.

주요 기여

1. SZP 로의 일반화: 저자들은 간단하고 아키텍처에 구애받지 않는 성장 전략인 SZP 를 소개합니다. 그들은 초기화 시 정확한 기능 보존이 강력한 최종 성능에 필수적이지 않음을 입증합니다. 실제로 SZP 는 특히 언어 모델에서 Net2Net 과 같은 복잡한 FP 기반 연산자보다 종종 더 나은 성능을 발휘합니다.
2. 성장 인자에 대한 경험적 상한 ($g_{upper}$): 이 논문은 성장 인자 $g$(목표 파라미터 대 기본 파라미터 비율) 에 대한 경험적 상한을 설정합니다. $g_{upper}$를 넘어서면, 성장된 모델을 제로부터 학습하는 것이 웜스타팅하는 것만큼 빠르거나 더 빠르게 수렴합니다. 이 상한은 MLP 와 LLM 모두에서 일관됩니다.
3. 스케일링 법칙 예측: 관찰 데이터를 기반으로 스케일링 법칙을 적합시켜, 저자들은 WS 가 효과적인 시점에 대한 예측 지침을 제공합니다. 그들은 제로부터 학습 대비 WS 가 수렴 속도 향상을 가져오는 "효율성 창 (efficiency window)"을 특징짓는데, 이러한 이득은 특정 토큰당 파라미터 예산 하에서 가장 두드러집니다.

주요 결과

1. SZP vs Net2Net vs Scratch

• 수렴 속도: SZP 는 제로부터 학습하는 것보다 일관되게 빠르게 수렴하며, 다양한 모델 너비와 성장 인자에서 Net2Net 을 능가합니다.
• 기능 보존은 필수적이지 않음: Net2Net(FP 기반) 은 기본 모델의 손실을 빠르게 회복하지만, 더 큰 성장 시나리오에서는 종종 초기 부분 공간에서 벗어나는 데 어려움을 겪어 SZP 에 비해 낮은 유효 랭크와 더 느린 장기 수렴을 초래합니다.
• 강건성: SZP 는 성장 인자에 더 강건하며 가중치 클로닝이나 노이즈 주입을 위한 아키텍처별 규칙이 필요하지 않습니다.

2. 성장 인자 한계 ($g_{upper}$)

• 한계의 존재: 저자들은 WS 가 이점을 상실하는 명확한 $g_{upper}$를 확인합니다. 많은 설정에서 이 한계는 대략 $g = 2$입니다.
• 체감 효과: 성장 인자가 증가함에 따라 기본 모델은 목표 아키텍처의 점점 더 작은 부분을 차지하게 됩니다. 초기화는 점점 더 무작위 초기화와 유사해지며, 상속된 구조의 제약이 이점을 압도합니다.
• 예산 의존성: WS 의 효과는 학습 예산과 밀접하게 연관되어 있습니다. WS 의 이득은 $\le$ 20 토큰/파라미터인 예산에서 가장 큽니다. 예산이 증가함에 따라 (Chinchilla 최적 프론티어에 접근하거나 초과함에 따라) WS 의 이점은 감소하고 제로부터 학습하는 것이 더 효율적이 됩니다.

3. 스케일링 법칙 분석

• 저자들은 파라미터 ($N$) 와 토큰 ($D$) 의 함수로서 손실을 예측하기 위해 $L(N, D) = E + A N^{-\alpha} + B D^{-\beta}$ 멱법칙 모델을 적합합니다.
• 등손실 등고선: WS 와 Scratch 간의 손실 차이를 시각화하면 경계가 드러납니다. 고정된 모델 크기의 경우, 더 적은 토큰/파라미터로 학습하는 것은 WS 를 선호하는 반면, 더 많은 토큰/파라미터로 학습하는 것은 Scratch 를 선호합니다.
• 예측 지침: 적합된 법칙을 통해 실무자는 주어진 성장 인자와 예산에 대해 WS 가 더 이상 유익하지 않은 교차점을 예측할 수 있습니다.

중요성 및 실무 가이드라인

이 논문은 웜스타팅의 효능에 대한 체계적이고 원리 기반의 조사를 제공하여 anecdotal 증거를 넘어 경험적 한계와 가이드라인을 제시합니다. 저자들은 다음과 같은 실무적 교훈을 주장합니다:

• 모델 확장: 하나의 모델 실행만 가능한 경우, 2$\times$ 성장 인자는 웜스타팅을 위한 가장 신뢰할 수 있는 확장 전략입니다.
• 효율성 창: 학습 예산이 컴퓨팅 최적 스케일링 (약 20 토큰/파라미터) 이하일 때 제로부터 학습하는 것보다 웜스타팅을 우선시해야 합니다.
• 하이퍼파라미터 전이: $\mu$P는 특히 더 큰 규모의 튜닝 예산이 제한될 때 웜스타팅에서 하이퍼파라미터를 전이하기 위한 원칙적인 전략으로 작용합니다.
• 표준화된 보고: 저자들은 향후 문헌이 공정한 비교를 촉진하고 비효율적인 성장 전략의 채택을 방지하기 위해 경험적 $g_{upper}$(WS 가 Scratch 성능과 일치하는 성장 인자) 를 명시적으로 보고해야 한다고 제안합니다.

이 연구는 웜스타팅이 자원 절감을 제공하지만 그 유용성은 제한적임을 결론짓습니다. 제안된 SZP 방법과 성장 인자에 대한 식별된 한계는 실무자가 진정한 효율성 향상을 보장하기 위해 모델을 언제 그리고 얼마나 성장시켜야 하는지 결정할 수 있는 프레임워크를 제공합니다.