Dust, 활성 공간 Perturbation으로 backpropagation과 경쟁하는 zeroth-order 학습법 공개

핵심 혁신: 활성 공간 Perturbation

Dust는 backpropagation을 활성 공간 perturbation으로 대체하는 zeroth-order 최적화 알고리즘이다. 가중치를 변경하는 기존 Evolution Strategies(ES)와 달리, Dust는 각 토큰에 대해 선형 레이어 출력에 독립적인 Gaussian noise를 추가한다. 이를 통해 각 토큰이 독립적인 개체로 작용하는 virtual population을 형성하며, 단일 forward pass로 대규모 population을 병렬 평가할 수 있다. 이 방식은 population 크기와 forward pass 횟수를 분리하여, EGGROLL 같은 weight-space ES 방법보다 같은 계산 예산 내에서 최소 세 자릿수 더 큰 population을 평가할 수 있게 한다.

Backpropagation 및 EGGROLL과의 성능 비교

FineWeb 데이터셋에서 GPT 스타일 transformer를 사용한 실험에서 Dust는 zeroth-order 방법이 backpropagation과 경쟁력 있음을 입증했다.

• 효율성: 100만 토큰부터 Dust는 최신 weight-space ES 방법인 EGGROLL의 transformer 구현체보다 10^3에서 10^4배 더 효율적인 것으로 추정된다.
• 손실 지표: 100k 및 1M 토큰 예산에서 Dust는 작은 population을 사용한 backpropagation보다 낮은 test loss를 달성했다. 더 큰 예산(10M 및 20M 토큰)에서는 population 크기가 증가함에 따라 backpropagation과의 격차가 크게 좁혀졌으며, 이 격차는 계속 줄어들 것으로 예측된다.
• ES 비교: Dust는 주어진 계산 예산 내에서 효율성과 손실 감소 측면에서 weight-space ES 방법을 크게 능가한다.

확장성 및 Gradient 정렬

zeroth-order 방법이 분산 문제로 인해 대규모 네트워크에서 실패한다는 통념과 달리, Dust는 더 큰 모델이 더 높은 population 효율성을 보임을 입증했다.

• 모델 크기: 243M 파라미터 모델은 대부분의 population 크기에서 2M 파라미터 모델과 유사하거나 더 나은 성능을 보여, zeroth-order 방법이 작은 네트워크에 한정된다는 생각을 반박했다.
• Gradient 정렬: population 크기가 증가함에 따라 Dust의 추정 gradient와 실제 backpropagation gradient 사이의 cosine similarity가 cos(K) = c_max / sqrt(1 + c/K)라는 fitted law에 따라 개선된다. 이 정렬은 1B 토큰까지 다양한 레이어 유형과 토큰 수에서 유지된다.

한계 및 향후 방향

Dust는 backpropagation을 근사하고 손실 측면에서 경쟁력이 있지만, 계산 효율성 제약으로 인해 backpropagation의 실용적인 대체재는 아직 아니다. 저자들은 Dust가 현재 규모에서 backpropagation과 경쟁하려면 계산 효율성이 몇 자릿수 개선되어야 한다고 지적했다. 그러나 Dust는 end-to-end differentiability 요구 사항을 제거하여 recurrent 및 looped computations 등 표준 backpropagation으로 학습하기 어려운 아키텍처 탐색에 새로운 가능성을 열었다.