Dust: 역전파 없이 Transformer를 사전학습하는 방법

TL;DR Highlight

활성값에 토큰별 노이즈를 줬더니 역전파 없이 Transformer 사전학습이 됐다

Who Should Read

LLM 학습이 역전파에 묶여 있는 게 늘 궁금한, LLM API·에이전트를 쓰는 백엔드 개발자

Core Mechanics

토큰마다 노이즈를 넣어 gradient를 추정한다
- 각 linear layer 출력에 토큰별로 독립적인 가우시안 노이즈를 더하고 forward pass를 돌린다.
- 토큰별 loss 변화를 보상으로 삼아 노이즈를 평균 내면 출력 오차 추정치가 나온다.
- 이 추정치와 layer 입력의 외적이 weight gradient가 된다.
토큰 하나가 가상 population 한 명이다
- 기존 weight-space ES(진화 전략, 가중치에 노이즈를 주고 좋은 쪽을 고르는 방식)는 forward pass 한 번에 population 한 명만 평가한다.
- Dust는 토큰 수천 개를 한 번에 평가해서 forward pass 한 번에 세 자릿수 이상 큰 population을 본다.
EGGROLL보다 10^3~10^4배 효율적이다
- 1M 토큰부터는 최신 ES 방법인 EGGROLL의 transformer 구현 대비 그 정도 효율이라고 한다.
- 이 수치는 저자들의 외삽(extrapolation) 결과다.
큰 모델일수록 population 효율이 좋다
- 243M 파라미터 모델이 120배 작은 모델보다 대부분의 population 크기에서 성능이 좋다.
- zeroth-order 방법은 큰 네트워크에 안 통한다는 통념과 반대다.
population을 키우면 backprop에 근접하고 일부는 넘어선다

Community Reactions

회의적인 반응: 신경망 목적함수는 매끄러워서 gradient가 더 유용하다는 의견
- 파라미터가 많을수록 gradient가 더 유용해지고, Lipschitz 볼록 최적화에서는 gradient 기반과 derivative-free 사이의 복잡도 격차가 증명됐다는 의견이다.
- 대안으로 Muon 같은 신경망 구조 특화 gradient 옵티마이저를 꼽았다.
nonconvexity를 논문이 다루지 않는다는 지적
- 비볼록성을 제거하면 Dust의 장점이 사라질 것 같고, 개선된 1차 방법에 밀릴 수 있다는 의견이다.
- zeroth-order가 이미 기본 선택지인 경우, 예를 들어 gradient를 안 주는 시뮬레이터를 호출하는 네트워크에는 쓸모 있다고 했다.
역전파가 약한 환경에서 쓸모 있다는 경험
- ZOO를 backprop 대체재가 아니라 backprop이 근본적으로 약한 곳에 쓰는 기법으로 본다는 의견이다.
- CD-RGE로 backprop-through-time 없이 큰 RNN을 학습한 논문이 인상적이었고, Dust를 개인 프로젝트에 써 보겠다는 댓글이다.
병렬화는 쉬워도 연산은 더 든다는 반응
- backprop보다 연산 효율은 낮지만 병렬화가 쉬운 것 아니냐는 질문이 나왔다.
- 이미 backprop으로 학습한 체크포인트를 Dust로 fine-tune하면 이득이 있을지 궁금해하는 댓글도 있었다.

How to Apply

논문 TL;DR과 Method 2장만 읽기
- 오늘 30분 안에 Introduction과 2장 Method를 읽고 Dust가 backprop과 다른 점을 정리하라.
- 수식 (1)~(3)에서 보상 가중 노이즈와 입력의 외적으로 gradient를 만드는 흐름만 잡아라.
backprop이 안 되는 학습 구간 목록 만들기
- 미분 불가능한 시뮬레이터나 외부 프로그램이 학습 루프에 들어가는 구간을 찾아 적어라.
- 그런 구간은 zeroth-order 방법의 후보가 되므로 댓글이 말한 쓰임새와 맞춰 보라.
공개 코드로 작은 모델에서 backprop과 비교하기
- 논문 상단의 Code 링크에서 저장소를 받아 작은 설정으로 돌려 보라.
- backprop과 같은 토큰 예산에서 validation loss를 비교하고 population을 64, 1k, 16k로 바꿔 보라.

Try It

복사해서 바로 해보기

# Claude나 ChatGPT에 그대로 붙여넣기 (논문 방법의 간이 버전: 코드 실행 없이 개념 확인용)

너는 ML 교수다. 아래 설명한 Dust 방식을 내가 이해했는지 검증해 줘.

[Dust 요약]
- 각 linear layer 출력 y_t에 토큰별 독립 노이즈 σ·a_t를 더하고 forward pass를 돌린다.
- 토큰별 loss 변화 c_s를 보상으로 쓰고, 보상 가중 노이즈를 draw K개에 걸쳐 평균 내서 출력 오차를 추정한다.
- 이 추정 오차와 layer 입력 x_t의 외적을 토큰에 걸쳐 합하면 weight gradient가 된다.

[내 이해]
{{내가 이해한 내용을 3~5문장으로}}

위 요약과 비교해서 틀린 부분을 짚어 주고, backprop과 다른 점을 개발자 비유로 설명해 줘.

이어서 해 볼 것: 'population이 커지면 왜 gradient 추정의 노이즈 비율이 줄어드는지'를 초보자 수준으로 물어보기

Terminology

Backprop(역전파)loss를 출력에서 입력 방향으로 거슬러 올라가며 각 가중치의 책임량(gradient)을 정확히 계산하는 방법이다. 채점 결과를 보고 각 단계에 정확히 점수를 돌려주는 것과 같다.

Zeroth-order 방법gradient를 직접 계산하지 않고 값을 흔들어 보며 loss 변화만 관찰해 방향을 추정하는 방법이다. 눈을 가리고 여기저기 발을 디뎌 보며 내리막을 찾는 것과 비슷하다.

Node perturbation가중치 대신 뉴런의 출력(활성값)에 노이즈를 넣는 방식이다. 가중치는 수가 많아서, 출력 쪽을 흔드는 게 탐색 공간이 작다는 논리다.

ES (Evolution Strategies)파라미터에 무작위 변형을 여러 개 만들어 평가하고, 성적 좋은 쪽으로 평균을 옮기는 진화식 최적화다. EGGROLL은 여기에 low-rank 변형을 써서 싸게 만든 방법이다.

Virtual population가중치 복사본을 실제로 만들지 않고, 토큰마다 활성값에 다른 노이즈를 주어 population 효과를 내는 것이다. 한 번의 forward pass에서 수천 개 후보를 동시에 평가한다.