SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers

논문 링크

SMELT: 계산량을 맞춘 MoE 루프 트랜스포머의 스케일링 법칙

한 줄 요약 (TL;DR)

루프 트랜스포머는 레이어 블록을 반복 실행해 깊이를 늘리지만, 기존 연구는 대부분 파라미터 수만 고정한 채 FLOPs를 늘려 “공짜 점수"를 얻고 있었다. SMELT는 per-token FLOPs·총 파라미터·KV 캐시 세 예산을 동시에 맞추고도 루핑이 순수한 아키텍처 이득임을 증명한다(근거: §1). 그 결과 compute-optimal 프론티어에서 6.8~18.0%의 학습 FLOPs를 절약하고, 이 이득은 검증 손실이 예측하는 것보다 다운스트림에서 더 크게 나타난다(근거: §4.3, §5.1).

핵심 아이디어

SMELT는 Sparse MoE Transformer의 middle layers Loop Twice를 뜻한다(근거: §1). 핵심 주장은 한 문장으로 요약된다.

“저자들은 중간 50% 레이어를 2회 반복하는 MoE 루프 구조를 사용함으로써, **기존 루프 연구가 해결하지 못했던 ‘계산량·파라미터·KV 캐시를 동시에 고정한 비교’**를 수행하여, 동일 예산에서도 더 낮은 손실과 더 빠른 손실 감소율을 달성할 수 있다고 가정한다.”

이 주장을 뒷받침하는 세 가지 설계 규칙이 에이블레이션에서 도출된다(근거: §3.3–§3.5).

1. 중간 절반만 루프 — 전체 스택이 아니라 중간 50% 레이어를 반복한다(근거: §3.3).
2. 더 큰 유효 깊이-대-폭 비율 — Baseline보다 더 많은 실행 깊이를 허용한다(근거: §3.4).
3. 2회 반복 — 3회·4회보다 2회가 최적이다(근거: §3.5).

Key Numbers (요약)

$$ L(F,S,D) = E + \frac{A(1-S)^b}{F^a} + \frac{K}{D^c} $$

배경: 그들이 해결한 문제

루프 트랜스포머의 매력과 함정

루프 트랜스포머는 새 레이어를 쌓는 대신 공유 블록을 반복해 유효 깊이(effective depth)를 늘린다(근거: §2). Universal Transformer에서 시작된 이 아이디어는 이후 Huginn, Ouro 등으로 확장되었고, “파라미터 수 대비” 추론 성능이 강하다는 결과가 쌓여 왔다(근거: §2, Tab. 1).

그런데 저자들은 여기에 결정적 결함을 지적한다. 파라미터 수의 일부로 계산되는 것이 곧 비용의 일부는 아니라는 것이다.

루프로 12층을 24층까지 실행하면 가중치는 절반만 저장하지만, per-token FLOPs는 24층 모델 수준을 쓰고 KV 캐시도 24층 분량이 필요하다(근거: §1).

기존 평가들은 다음 두 갈래로 나뉘며, 둘 다 문제를 안고 있다.

• 파라미터 고정 + 반복 깊이 증가(Huginn, Ouro, Prairie et al.): 아키텍처 이득과 통제되지 않은 추가 FLOPs가 뒤섞인다.
• FLOPs 고정(Schwethelm et al.): FLOPs를 고정하면 루프 모델의 고유 파라미터가 줄어들어, 손실의 원인이 루핑 자체인지 파라미터 감소인지 분리할 수 없다. 이들은 반복 1회가 r^0.46 고유 블록에 해당한다는 결론을 내렸다(근거: §2).

즉 연구 공백은 명확하다. “루핑이 파라미터 효율성을 넘어서, 계산량까지 통제한 상태에서도 이득을 주는가?“라는 질문은 미해결이었다. 이를 답하려면 세 예산 — ① per-token FLOPs, ② 총 파라미터, ③ KV 캐시 — 을 동시에 고정해야 하는데, 기존 연구 중 이 셋을 모두 닫은 것은 없었다(근거: Tab. 1).

왜 MoE가 해법인가

세 예산을 동시에 맞추는 것이 가능해진 이유가 바로 MoE다. MoE는 총 파라미터와 per-token FLOPs를 분리한다(근거: §2). 루프 모델은 두 번째 방문 비용을 hidden dimension 축소로 지불하고, 잃은 용량은 expert 수 증가로 되돌릴 수 있다. KV 캐시는 head 크기와 GQA 비율 조정으로 별도로 맞춘다(근거: §3.2).

새로운 접근법: SMELT

예산 매칭 = 계산 배분 문제

저자들은 Baseline(비루프 MoE)과 Looped Transformer를 비교할 때 세 예산을 <1~4% 오차로 맞춘다(근거: Tab. 2).

  flowchart LR
    A[루프로 인한<br/>추가 레이어 실행] --> B[FLOPs 증가]
    B --> C[hidden dimension H 축소]
    C --> D[총 파라미터 감소]
    D --> E[expert 수 증가로 복구]
    C --> F[KV 캐시 변화]
    F --> G[head size / GQA 비율 조정]
    E --> H[3개 예산 모두 매칭]
    G --> H

루프 구간 내 각 서브레이어의 residual 업데이트는 1/r(루프 횟수)로 스케일링되는데, 이는 weight-tied 업데이트가 잔차 스트림을 반복 방문마다 부풀리는 것을 막기 위함이다(근거: §3.1).

“계산 등가 스파스티티"라는 새 지표

저자들은 sparsity를 FLOPs 비율로 재정의한다. 완전 활성 컨트롤(모든 expert 활성)의 per-token FLOPs를 F₀, 파라미터를 N₀라 하면(근거: §3.1):

$$ N_{\text{act}}^{eq} = \frac{F}{F_0} N_0, \qquad S = 1 - \frac{N_{\text{act}}^{eq}}{N} $$

이 지표 덕분에 dense와 sparse를 하나의 축 위에서 비교할 수 있고, 루프 모델을 “동일 계산 강도"의 Baseline과 정렬할 수 있다. 보고되는 매칭 쌍은 S≈85%, 95%, 97% 세 수준이다(근거: §3.1).

작동 원리: 구체적인 예시로 살펴보기

200M, S≈95% 셀을 toy example로 삼자(근거: §3.2).

Baseline은 물리 깊이 L=12, hidden dimension H=1280, 레이어당 192 experts, top-8 라우팅이다. per-token FLOPs는 1.33×10⁹, 총 파라미터는 3.87×10⁹이다.

SMELT는 12층 중 중간 6층을 2번 실행해 총 18층을 실행한다. 그런데 6번의 추가 레이어 실행은 FLOPs를 잡아먹는다. 그래서:

1. H를 1280 → 1056으로 축소해 per-token FLOPs를 되돌린다.
2. H 축소는 모든 expert의 FFN을 얇게 만들어 총 파라미터를 줄이므로, expert 수를 192 → 288로 늘려 복구한다.
3. 결과: per-token FLOPs는 1.37×10⁹ (+2.9%), 총 파라미터는 3.89×10⁹ (+0.4%), KV 캐시는 4% 이내.

즉 같은 가격에 더 깊이 실행되는 모델을 얻는다. 이것이 SMELT의 정수다.

세 에이블레이션이 레시피를 잠근다

200M 스케일에서 세 질문을 순서대로 답한다(근거: §3.3–§3.5).

① 어느 레이어를 루프할까? 루프 구간을 0층(비루프)부터 12층(전체)까지 스윕한 결과, 50% span에서 validation loss가 최소가 된다(근거: Tab. 3). 첫·마지막 레이어는 독립 파라미터가 유리한 특화 역할을 하므로, 중간 블록을 반복하는 편이 낫다는 기존 발견과도 일치한다(근거: §3.3).

② 얼마나 깊게? 깊이-대-폭 비율을 스윕하자, Baseline은 물리 깊이 12에서 최적이지만 Looped Transformer는 12/18(물리 12, 실행 18)에서 최적이다(근거: Tab. 4). 저자들은 공유 레이어가 여러 방문으로부터 gradient 기여를 받아, 후반부의 짧은 경로 신호가 더 깊은 직렬 계산을 최적화하기 쉽게 만든다고 해석한다(근거: §3.4).

③ 몇 번 루프? 2회(18층), 3회(24층), 4회(30층)를 비교하면 2회가 셋 모두에서 승리한다. 3·4회는 matched FLOPs 상한 때문에 모델이 얇아져 오히려 퇴보한다(근거: Tab. 5).

이 세 규칙이 SMELT를 정의한다. 이후 모든 실험은 이 레시피를 4개 규모(100M/200M/600M/1.6B) × 4개 sparsity(0%/85%/95%/97%)의 4×4 그리드로 확장한다(근거: §4.1).

성능 검증: 주요 결과

별도의 스케일링 법칙을 피팅한다

두 아키텍처에 대해 각각 Chinchilla 스타일의 표면을 피팅한다(근거: §4.2). 이때 파라미터 N 대신 per-token FLOPs F를 쓰고, capacity 항을 sparsity S로 스케일한다. 피팅 계수는 다음과 같다(근거: Tab. 6).

핵심은 두 지수 모두 SMELT가 크다는 것이다. 이를 결합한 프론티어 지수는:

$$ \gamma = \frac{ac}{a+c}, \quad \gamma_{\text{base}} = 0.237,\ \gamma_{\text{SMELT}} = 0.250 $$

즉 SMELT의 손실은 계산량 한 단위당 5.5% 더 빠르게 떨어진다(근거: §4.3). E의 0.005 nats 차이는 피팅 RMSE보다 작아 피팅 노이즈로 해석된다(근거: §4.3).

CE Gain: 프론티어에서 실제 절약되는 계산량

동일 손실에 도달하기 위해 필요한 계산량의 역산으로 compute efficiency gain을 정의한다(근거: §4.2):

$$ \text{CE Gain} = 1 - \frac{C_{\text{tgt}}}{C_{\text{ref}}} $$

† 피팅 윈도우를 벗어난 외삽값(근거: Tab. 7).

γ가 크기 때문에 프론티어 격차는 계산 범위 내에서 복리로 벌어진다. 흥미롭게도 최적 배분(Tokens Per Parameter)은 두 아키텍처가 거의 같다(6% 이내, 근거: Tab. 8). 즉 SMELT의 절약은 예산을 다르게 배분해서가 아니라 같은 배분에서 더 낮은 손실에 도달하는 데서 온다(근거: §4.3).

다운스트림 이득은 검증 손실 예측을 초과한다

단순 승률은 96쌍 중 DCLM Completion 96승, DCLM Core 83승, MMLU 29/30승이다(근거: §5.1). 그러나 저자들은 여기서 멈추지 않고, Baseline 96개 endpoint로 loss→score 시그모이드를 피팅한 뒤 SMELT의 잔차(residual)를 계산한다(근거: §5.1). 결과적으로 잔차는 세 벤치마크 모두 양수이며 규모에 따라 단조 증가한다 — SMELT의 다운스트림 이득은 검증 손실 개선만으로 설명되는 양을 초과한다는 뜻이다(근거: Fig. 9a).

이 이득은 구조화된 데이터에 집중된다. 도메인별 CE Gain은(근거: Fig. 10a):

• Code 20.4% > Finance 16.8% > Math/STEM 16.6% > Knowledge 14.9% > Web 14.8%

또한 이득은 긴 샘플(512–4096 토큰 구간의 이득이 32–256 구간의 1.52배, 근거: Fig. 11a)과 많은 인-컨텍스트 예시(k=0에서 0.9pp, 데모 제공 시 1.9pp로 확대, 근거: Fig. 11b)에서 커진다. Dyck 언어 과제에서는 k=32에서 29.8% vs 26.4%로 격차가 계속 벌어진다(근거: Fig. 11c).

두 번째 방문 내부에서는 무슨 일이 일어나는가

메커니즘 분석은 일관된 시그니처를 보여준다(근거: §6).

• 라우팅: 두 방문이 expert의 핵심 부분집합을 재사용하고 나머지를 다양화한다. S≈97%에서도 overlap이 2~3개로 우연 수준을 크게 웃돈다(근거: Fig. 12).
• 잔차 쓰기: 두 번째 방문의 업데이트 norm이 1.2~3.5배 크고, 방향도 정렬되어 있다(cross-visit 코사인 평균 0.56 vs 비매칭 0.16, 근거: Fig. 13–15). 두 번째 방문은 첫 방문을 덮어쓰는 것이 아니라 증폭한다.
• 어텐션: Q·K는 0.890.93으로 유지되지만 V는 0.650.74로 분기한다(근거: Fig. 16). 즉 검색 좌표는 유지하고 읽는 내용만 바꾼다.
• 어텐션 싱크: 두 번째 방문에서 싱크가 크게 줄고 질량이 콘텐츠 관련 토큰으로 이동한다. Dyck 사례에서 BOS 질량이 0.60→0.02로 감소하고 데모 정답 질량이 0.24→0.85로 증가한다(근거: Fig. 18). 이는 일반 데이터에서도 재현되며, 깊이가 깊어질수록 싱크가 강해지는 일반적 깊이 경향을 역전시킨다(근거: Fig. 19).

우리의 관점: 강점, 한계, 그리고 이 연구가 중요한 이유

강점

가장 큰 기여는 비교 프로토콜 자체의 정직성이다. 기존 루프 연구의 “파라미터 효율” 내러티브가 실제로는 통제되지 않은 추가 FLOPs에 기댔을 수 있음을 명확히 짚고, MoE의 파라미터-FLOPs 분리 특성을 이용해 세 예산을 동시에 닫았다(근거: §2, Tab. 1). 여기에 아키텍처별 별도 스케일링 표면을 피팅해, 점 관측을 외삽 가능하고 귀속 가능한 주장으로 승격시킨 것도 의미 있다(근거: §4.2). 블록 수준 루핑에서 세 예산을 동시에 닫은 최초의 멀티스케일 연구라는 점은 독창성 요건을 충족한다(근거: §2).

메커니즘 분석이 단순 성능 보고를 넘어 “왜"에 대한 설명적 단서를 제공한다는 점도 강점이다. “루핑 = refinement 단계"라는 프레이밍은 단순한 capacity 추가가 아니라는 주장에 설득력을 준다(근거: §7).

한계

저자 스스로 인정한 한계는 세 갈래다(근거: §7).

1. 설계 에이블레이션이 200M 스케일에 고정되었다. 최적 span·루프 횟수가 더 큰 규모에서 달라질 수 있다.
2. 가장 단순한 루핑만 연구했다. 완전 가중치 공유의 연속 블록 반복만 다루었고, per-visit LoRA, 적응형 재귀 깊이, learned halting 등은 다루지 않았다.
3. FLOPs·파라미터·KV 캐시를 맞췄을 뿐 wall-clock 비용은 아니다. 직렬 재실행과 sparse 라우팅은 하드웨어 효율 격차를 낳을 수 있다.

분석적으로 볼 때 잠재적 한계도 있다. CE Gain의 계산량 절약은 validation loss라는 단일 지표에 기반하며, 가장 큰 값(18.0%)조차 10²¹ FLOPs에서의 부트스트랩 구간이 [8,28]%로 넓다(근거: Tab. 7). 10²² FLOPs 값들은 피팅 윈도우 밖의 외삽이라 신뢰 구간이 [0,51]%까지 벌어진다. 또한 잔차 분석의 sigmoid 매핑은 Baseline endpoint에만 피팅된 것이므로, 손실 외 요인(예: 능력별 차등 이득)을 구조적으로 분리하지는 못한다(근거: §5.1). 훈련 데이터와 내부 모델 패밀리가 비공개라 재현성이 제한된다는 점도 실무 관점의 아쉬움이다(근거: Appx. A).

왜 중요한가

이 논문은 “루핑은 정말 파라미터 효율 이상을 주는가?“라는 오래된 질문에 그렇다는 첫 엄밀한 증거를 제시한다. MoE가 이미 주류가 된 현재, 깊이 재사용은 width·depth·expert count에 이은 네 번째 스케일링 축으로 자리매김할 잠재력이 있다(근거: §7). 특히 코드·구조화 데이터·ICL에서 이득이 크다는 점은, 정확히 현대 LLM의 실사용 부하가 높은 영역과 겹친다(근거: §5.2–§5.3).

다음 단계는?: 앞으로의 길

저자가 제안하는 방향은 (1) 더 큰 규모에서의 설계 재검증, (2) richer 루프 변형(per-visit LoRA, adaptive depth, learned halting, block-selective sharing, cross-token state reuse)의 budget-matching 검증, (3) wall-clock 격차를 메울 시스템 최적화, (4) 메커니즘의 설명에서 인과적 규명으로의 이행이다(근거: §7).

합리적인 다음 단계로는 다음을 덧붙일 수 있다. 첫째, 설계 규칙(중간 50%·2회)의 규모 의존성을 실제 대규모에서 확인하는 것은 로우-행잉 프루트다. 둘째, CE Gain이 wall-clock과 실제 토큰 처리량에서도 유지되는지 — 특히 직렬 재실행의 지연 비용이 희석되지 않는지 — 를 검증해야 실배포 주장이 완성된다. 셋째, 어텐션 싱크 감소가 관찰된 ICL 이득의 인과적 원인인지를 개입 실험(예: 싱크 억제 어블레이션)으로 분리하면, “루핑 = refinement"라는 가설이 검증 가능한 이론으로 격상될 것이다(근거: §6, §7).