아키텍처는 모델 성장, 재귀 깊이, 경계 연산자로 스케일링 지수를 바꾼다.

무엇인가

스케일링 법칙은 연산량이 늘어날 때 손실이 어떻게 줄어드는지를 예측한다. Hoffmann 등이 쓴 L(N,T) = E + (N/N0)^-α + (T/T0)^-β 형태에서 상수는 손실 곡선을 위아래로 옮기고, 지수는 곡선의 모양 자체를 바꾼다. 통념은 아키텍처 변경이 상수만 건드린다는 것이었고, 실제로 Muon 옵티마이저 같은 사례도 상수 개선으로 분류된다. 이 논문은 그 통념에 정면으로 반박한다. 사전학습 중 아키텍처 개입만으로 지수를 바꿀 수 있고, 지수 개선은 연산이 커질수록 이득이 복리로 커지기 때문에 상수 개선과는 질적으로 다른 결과라는 주장이다.

어떻게 동작하나

저자들이 쓰는 뼈대는 Geiping 등의 prelude-core-coda 정식화다. 입력을 임베딩하는 prelude P, K번 반복 적용되는 core R, 출력을 만드는 coda C로 나뉘고, e = P(s), h_k = R_θk(φ(h_{k-1}, e)), y = C(ρ(h_K, e))로 쓴다. 여기서 φ와 ρ가 경계 연산자이고, 논문은 BO(h,e) = Norm(h) + α·e를 쓴다. 두 부분으로 이루어지는데, Norm은 pre-norm 트랜스포머에서 깊이가 쌓일수록 잔차 스트림이 커져 각 블록의 기여가 희석되는 현상을 막아 매 패스가 온전한 가중치로 쓰이게 하고, α·e 항은 prelude 출력을 다시 주입해 모든 패스가 입력에 조건화되게 한다. 이 위에 세 축이 얹힌다. 가중치를 공유하면(tied) 하나의 core를 K번 적용하는 루프 트랜스포머가 되고, 공유하지 않으면(untied) 같은 FLOPs와 같은 계산 그래프를 가지면서 저장 블록 수가 늘어나는 깊은 트랜스포머가 된다. 세 번째 축은 모델 성장으로, 학습 도중 루프 수 K를 늘린다. tied에서는 기존 core를 더 여러 번 적용할 뿐 파라미터가 늘지 않고, untied에서는 학습된 core를 복사해 각각 따로 학습시킨다. 이렇게 해서 Vanilla, Operator-1, Loop-2, Untied-2, Loop-Grow, Untied-Grow 여섯 변형과 연산자 없는 대조군 두 개를 같은 조건에서 비교한다.

무엇과 다른가

실험은 FineWeb과 GPT-2 토크나이저로 진행하고, 폭-깊이 비율을 128로 고정한다. 핵심은 모든 아키텍처에 대해 별도의 compute-optimal 레시피를 적합했다는 점이다. Vanilla의 레시피를 Operator-1에 그대로 옮기면 d8 깊이에서 손실이 7.8×10^-3 나빠지고 지수 개선이 사라진다. 그래서 d8에서의 기본 하이퍼파라미터, 저장 파라미터당 토큰 수, 성장 시점, 학습률 스케일링 규칙의 네 단계를 아키텍처마다 따로 맞췄다. 저장 파라미터당 토큰 수는 Vanilla 5, 경계 연산자 사용 시 6, 성장 사용 시 7~8로 나왔고, 성장 시점은 학습의 1/2에서 4/5 사이가 최적이었다. core 패스 수는 고정 변형이 2, 성장 변형이 2에서 4로 늘어난다. 나머지 하이퍼파라미터를 다시 스윕해도 compute multiplier는 3% 이내로만 움직였다고 밝힌다.

어떻게 쓰나

결과는 compute multiplier, 즉 같은 손실에 도달하는 데 표준 트랜스포머가 몇 배의 연산을 필요로 하는지로 제시된다. 10^20 FLOPs에서 untied 가중치로 성장한 모델은 1.55배, tied 루프로 성장한 모델은 1.36배, 경계 연산자만 추가한 Operator-1은 1.25배의 연산 효율 이득을 보였다. 이 배수는 스케일이 커질수록 커지는데, 이것이 지수 개선의 특징이다. 7.4B Untied-Grow 모델은 1.23×10^21 FLOPs에서 CORE 0.3865를 기록해 GPT-3 13B의 참조값 0.3852와 비슷한 수준에 도달했고, GPT-3 13B가 쓴 2.31×10^22 FLOPs와 비교하면 약 20배 적은 연산이다. 다만 저자들은 두 모델이 다른 데이터로 학습됐고 GPT-3 참조값이 별도 평가 파이프라인의 추정치라서 이 20배는 통제된 비교가 아니라 지표로서의 의미라고 명시한다. FineWeb-Edu로 적합한 법칙에 따르면 Vanilla 대비 배수는 10^20 FLOPs에서 1.6배, 1.23×10^21에서 1.8배, 10^25 FLOPs에서는 2.7배로 늘어난다.

전제와 한계

데이터가 제약된 다중 에폭 구간에서는 그림이 달라진다. 100M개의 고유 FineWeb 토큰을 10 에폭 반복해 1B 토큰을 쓰는 설정에서, 최적 루프 수는 연산이 커질수록 늘어나 가장 작은 예산에서 약 1, 가장 큰 예산에서 약 7이 된다. 고정 루프 수 곡선은 결국 과적합으로 상승 전환하는데, 루프 수가 클수록 전환점이 늦고 최솟값도 낮다. Operator-1 대비 손실 격차는 0.00에서 0.11로 벌어진다. 더 중요한 것은 루프 수 확장이 모델 크기 확장보다 연산 효율적이라는 점이다. 모델 크기 확장에 맞춰 weight decay를 매번 재조정해도, 고정 모델 크기와 고정 weight decay에서 루프 수만 늘린 쪽이 최적 튜닝된 모델 크기 확장과 같은 손실을 2.2배 적은 연산으로 달성한다. 루프 수 확장 위에 weight decay를 추가 튜닝해도 손실 개선은 0.02 이하다. 최적 weight decay는 깊이에 따라 커지지만 루프 수에 대해서는 거의 평평해서, 한 루프 수에서 튜닝한 값이 더 큰 루프 수로 전이된다. untied 루프는 같은 깊이를 더 많은 파라미터로 얻기 때문에 Operator-1을 이기지 못한다.

저자들은 이 결과를 computational depth, 즉 주어진 연산 예산에서 예측 분포에 실질적으로 영향을 주는 블록의 수라는 관점으로 해석한다. 실행 깊이와 computational depth는 다를 수 있고, 낭비가 두 가지 있다고 본다. 하나는 블록이 실행되지만 예측을 바꾸지 않는 경우(curse of depth)이고, 다른 하나는 학습 후반에만 필요한 깊이를 처음부터 계속 지불하는 경우다. 경계 연산자는 전자를, 모델 성장은 후자를 줄인다는 가설이다. 이를 KL effective depth로 측정하는데, logit lens로 각 블록 뒤 잔차 스트림을 디코딩해 최종 출력과 KL 임계값 안에 들어오는 첫 블록을 기록한다. 10^20 FLOPs에서 Untied-2는 24, 연산자 없는 Deep Vanilla는 20으로 나와 같은 블록을 실행해도 경계 연산자가 더 많은 깊이를 쓴다는 것을 보여준다. Untied-Grow는 36으로 Untied-2의 24보다 훨씬 높다. 데이터 반복 구간에서는 루프 수를 늘리는 쪽이 K=1에서 깊이를 늘리는 쪽보다 KL effective depth를 빠르게 올려, 가장 큰 예산에서 18 대 14를 기록한다. 폭만 늘리고 깊이는 고정하면 Operator-1의 지수 개선이 상수 개선으로 바뀐다는 점도 깊이 스케일링이 지수 개선의 필요조건임을 시사한다.

개발자 입장에서 실용적인 시사점은 세 가지다. 첫째, 단일 에폭 compute-optimal 사전학습에서는 메모리가 허락한다면 untied 가중치를 쓰는 것이 좋다. tied는 파라미터 저장이 우선일 때 여전히 유용하고, 성장에서 오는 지수 개선은 유지하면서 연산 효율을 상수 배수만큼 내주는 선택이다. 둘째, 데이터가 반복되는 구간에서는 루프 수를 늘리는 것이 파라미터를 늘리는 것보다 낫고, weight decay를 매 스케일마다 재조정할 필요도 줄어든다. 셋째, 아키텍처별 하이퍼파라미터 튜닝이 필수다. 표준 트랜스포머의 레시피를 그대로 재사용하면 지수 개선이 가려진다. 저자들이 밝힌 한계도 분명하다. KL effective depth는 computational depth의 대리 지표일 뿐이고, 임계값 안에 들어온다고 해서 이후 블록이 예측을 바꾸지 않는다는 보장은 없으며 네트워크 중간의 미사용 블록은 탐지되지 않는다. GPT-3 13B와의 비교는 데이터와 평가 파이프라인이 달라 통제된 실험이 아니다. 또한 compute-optimal 레시피는 아키텍처마다 따로 적합해야 하며, 그 비용과 민감도는 부록에 의존한다.