SMELT: 반복형 MoE Transformer의 계산량 대비 확장 법칙

왜 중요한가

Looped Transformer의 성능 향상이 반복 실행으로 늘어난 계산량 때문인지, 구조 자체의 효과인지 기존 비교만으로는 분리하기 어려웠습니다. SMELT는 토큰당 FLOPs, 전체 비임베딩 파라미터, KV cache를 모두 맞춘 조건에서도 손실과 downstream 성능이 앞서며, C = 10²¹ FLOPs에서 동일 손실에 필요한 학습 계산량을 14.7–18.0% 줄였습니다. 이는 반복형 구조를 실제 모델 규모 확장과 추론 비용 관점에서 평가할 기준을 제공합니다.

핵심 기여

세 가지 예산을 맞춘 SMELT 설계

중간 50%의 층을 두 번 실행하고, 추가 실행 비용을 숨은 차원 축소로 상쇄한 뒤 전문가 수를 늘려 전체 파라미터를 회복했습니다. Head 크기와 GQA 비율도 조정해 토큰당 FLOPs 오차를 4% 미만, 파라미터 오차를 1% 미만, KV cache 오차를 4% 미만으로 유지했습니다.

계산 최적 손실의 개선

100M·200M·600M·1.6B 활성 비임베딩 파라미터 규모와 85%·95%·97% 희소성에서 Baseline과 SMELT를 비교했습니다. SMELT의 계산량 대비 손실 감소 지수는 0.250으로 Baseline의 0.237보다 높았고, C = 10²¹ FLOPs에서 희소성별 계산량 절약률은 14.7–18.0%였습니다.

구조화된 입력과 In-context Learning에서의 이득

SMELT는 DCLM Completion 96/96쌍, DCLM Core 83/96쌍에서 Baseline을 앞섰고, Baseline이 우연 수준보다 10%포인트 이상 높은 MMLU 비교에서는 29/30쌍에서 우세했습니다. Code 영역의 계산량 절약률은 20.4%로 다섯 영역 중 가장 높았으며, 긴 샘플과 시연 예제가 있는 조건에서 격차가 커졌습니다.

두 번째 방문의 내부 작동 원리 확인

두 번째 방문은 첫 번째 방문과 일부 전문가와 Attention 위치를 공유하면서도 더 큰 Residual update를 기록했습니다. 같은 물리 층의 두 방문에서 Residual update cosine 평균은 0.56이었고, Attention의 Q·K 유사도는 0.89–0.93, V 유사도는 0.65–0.74로 나타나 검색 위치보다 읽어오는 내용이 더 크게 바뀌었습니다.

핵심 아이디어 이해하기

일반 Transformer는 층마다 서로 다른 가중치를 사용해 토큰 표현을 한 단계씩 갱신합니다. 층 수를 늘리면 토큰이 거치는 계산 단계는 많아지지만, 새 층의 파라미터와 토큰당 FLOPs도 함께 증가합니다. 반대로 Looping은 같은 층을 재사용해 고유 파라미터 수를 억제하지만, 12층 모델의 중간 6개 층을 두 번 실행하면 실제 실행 깊이는 12층에서 18층으로 늘어납니다.

이 추가 실행은 계산량을 늘리므로 단순히 같은 파라미터 수의 모델과 비교하면 반복 구조의 효과와 추가 FLOPs의 효과가 섞입니다. SMELT는 먼저 숨은 차원 H를 줄여 늘어난 실행 비용을 상쇄하고, 줄어든 전문가별 파라미터는 전문가 수를 늘려 보충합니다. GQA의 Head 크기와 비율을 조정해 KV cache도 Baseline에 가깝게 유지하므로, 같은 저장 용량·토큰당 비용·서비스 메모리에서 구조만 비교할 수 있습니다.

두 번째 방문은 첫 번째 계산을 그대로 복사하지 않습니다. 같은 토큰이 같은 물리 층을 다시 통과할 때 Router는 일부 전문가를 재사용하면서 나머지를 바꾸고, Attention의 Q와 K는 비슷하게 유지해 참조 위치를 보존합니다. 반면 V와 최종 Attention 출력은 더 크게 달라져, 같은 위치에서 새로운 내용을 읽어 Residual stream에 추가합니다.

이 과정에서 두 번째 Residual update는 첫 번째보다 커지고 방향도 대체로 일치합니다. 특히 Attention 질량이 초기 토큰에 몰리는 Attention Sink가 줄어들고 내용과 관련된 토큰 쪽으로 질량이 이동합니다. 따라서 SMELT의 이득은 단순한 깊이 증가보다, 공유 중간층이 첫 번째 표현을 덮어쓰지 않고 관련 정보 추출을 한 번 더 정제하는 데서 발생하는 것으로 해석됩니다.

방법론

모든 모델은 decoder-only Transformer와 sparse MoE Feed-Forward 층으로 구성했으며, 각 토큰은 top-8 전문가로 라우팅했습니다. Attention에는 GQA를 사용하고, 100M·200M·600M·1.6B 규모와 85%·95%·97% compute-equivalent sparsity를 조합했습니다. Baseline은 표준 비반복 구조이고, Looped Transformer는 연속된 중간 층을 정한 횟수만큼 반복합니다. 두 구조는 동일한 데이터, Optimizer, 평가 절차를 사용했으며 반복 범위·횟수와 그에 따른 폭·전문가·Head 설정만 달랐습니다.

예산 정합은 반복으로 늘어난 유효 깊이를 먼저 계산하고, 숨은 차원 H를 줄여 토큰당 학습 FLOPs를 맞추는 방식으로 진행했습니다. H 축소로 줄어든 전체 파라미터는 층별 전문가 수를 늘려 회복했고, KV cache는 Attention Head 크기와 GQA 비율로 보정했습니다. 200M 규모 예시에서 Baseline은 H=1280, 전문가 192개, 토큰당 FLOPs 1.33 × 10⁹, 전체 파라미터 3.87 × 10⁹이었고, SMELT는 H=1056, 전문가 288개, 1.37 × 10⁹ FLOPs, 3.89 × 10⁹ 파라미터를 기록했습니다.

학습에는 AdamW와 WSD 스케줄을 사용했습니다. 안정 구간은 196,075 steps, 전역 배치는 시퀀스 256개였고 step당 약 1M tokens를 처리했습니다. 안정 구간에서 205B tokens를 사용한 뒤, 10B tokens씩 추가하는 여섯 개 cosine-decay 분기를 만들었으며 최대 학습량은 약 215B tokens였습니다. 학습 시퀀스는 4096-token context에 packing했고, 문서별 segment-level Attention mask를 적용했습니다.

설계 탐색은 200M에서 세 번의 Ablation으로 진행했습니다. 12개 층 중 반복 범위를 0–100%로 바꿔 중간 50%를 선택했고, 물리적 깊이와 유효 깊이를 비교해 SMELT가 더 큰 유효 깊이 대 폭 비율에서 최적임을 확인했습니다. 마지막으로 두 번째·세 번째·네 번째 방문을 비교해 두 번 실행하는 설정을 고정했습니다. 이후 네 규모와 네 희소성 조건의 4 × 4 grid에서 32개 안정 학습과 96개 Baseline·SMELT 비교쌍을 구성했습니다.

검증 손실은 39개 validation source의 token-weighted cross-entropy로 계산했고, DCLM Core 22개 task, DCLM Completion, MMLU 5-shot을 별도로 평가했습니다. Scaling law는 Baseline과 SMELT에 각각 여섯 계수를 맞춘 뒤, 총 계산량 C를 토큰당 FLOPs와 학습 토큰 수로 나누어 동일 손실에 필요한 참조 계산량을 역산했습니다. 내부 작동은 1M-token held-out 샘플에서 전문가 선택, Residual update, Attention의 Q·K·V·가중치·출력을 방문별로 비교했습니다.

주요 결과

SMELT의 validation loss 곡선은 네 규모와 세 희소성에서 대부분 학습 구간 동안 Baseline 아래에 놓였습니다. 별도 Scaling law의 계산량 지수는 SMELT 0.250, Baseline 0.237이었으며, C = 10²⁰ FLOPs에서 SMELT의 계산량 절약률은 희소성별 6.8–10.0%, C = 10²¹ FLOPs에서는 14.7–18.0%였습니다. C = 10²² FLOPs의 19.6–23.5% 수치는 학습 지원 범위를 넘어선 외삽값입니다.

Downstream 평가에서도 손실 차이 이상의 격차가 나타났습니다. DCLM Completion은 96개 matched pair 전부에서 SMELT가 낮았고, DCLM Core는 96개 중 83개, MMLU는 Baseline이 chance보다 10%포인트 이상 높은 30개 중 29개에서 SMELT가 높았습니다. Validation loss에서 예상한 benchmark 점수와 실제 점수의 잔차는 모든 규모와 세 benchmark에서 양수였으며, 규모가 커질수록 DCLM Completion과 DCLM Core의 초과 이득이 증가했습니다.

데이터 유형별 C = 10²¹ FLOPs, 95% 희소성 조건의 계산량 절약률은 Code 20.4%, Finance 16.8%, Math/STEM 16.6%, Knowledge 14.9%, Web 14.8%였습니다. 샘플 길이별로는 512–4096 tokens 구간의 평균 정규화 이득이 32–256 tokens 구간의 1.52배였습니다. 16개 few-shot task의 평균 격차는 예제가 없을 때 0.9 percentage points에서 예제 1개 이후 1.9 percentage points로 커졌고, Dyck Languages에서는 32-shot 기준 SMELT 29.8%, Baseline 26.4%였습니다.

내부 측정에서 두 번째 방문의 각 구성요소 norm 비율은 16개 규모·희소성 cell 모두 1보다 컸고 범위는 1.2–3.5배였습니다. 같은 물리 층의 top-8 전문가 중 겹치는 수는 97% 희소성에서 2–3개였지만 무작위 라우팅 기대치보다 높았습니다. 두 방문의 top-8 Attention token 집합은 56–66% 겹쳤고, 다른 Head를 짝지은 통제 조건은 28–34%였습니다. Q·K cosine은 0.89–0.93, V는 0.65–0.74였으며, 이는 두 번째 방문이 검색 위치보다 읽은 표현을 더 크게 갱신한다는 결과입니다.

기술 상세

SMELT는 물리적 깊이 12층인 200M 규모 예시에서 중간 6층을 두 번 실행해 유효 깊이 18층을 만듭니다. 반복 횟수를 r이라고 할 때 loop 내부 각 Sublayer의 Residual update를 1/r로 조정해 같은 가중치가 여러 번 더해지며 Residual stream을 과도하게 키우는 현상을 억제했습니다. 추가 실행으로 생긴 FLOPs는 H=1280에서 H=1056으로 줄여 상쇄했고, 전문가 수는 192개에서 288개로 늘려 전체 파라미터를 3.87 × 10⁹에서 3.89 × 10⁹으로 회복했습니다. 최종 토큰당 FLOPs는 1.33 × 10⁹ 대비 1.37 × 10⁹으로 +2.9%였고 KV cache 차이는 4% 이내였습니다.

Scaling law는 각 구조에서 검증 손실을 계산량 F, compute-equivalent sparsity S, 학습 토큰 D의 함수로 맞췄습니다. F는 평균 문서 길이에서 측정한 토큰당 학습 FLOPs이고, S가 높을수록 같은 토큰당 비용으로 더 많은 비활성 전문가를 저장합니다. 전체 학습 계산량 C가 F와 D의 곱으로 고정되면 F를 늘려 모델 측 손실을 낮출수록 D가 줄어 데이터 측 손실이 커지므로 두 항의 균형점이 계산 최적 배분이 됩니다. SMELT의 capacity exponent는 0.3892로 Baseline의 0.3703보다 높고, data exponent는 0.7011로 Baseline의 0.6594보다 높아 두 축에서 감소 속도가 더 컸습니다.

두 번째 방문의 Residual update는 첫 번째보다 절대값과 Residual 대비 비율 모두 대체로 컸습니다. 1.6B 모델에서 같은 물리 층의 두 방문 update cosine 평균은 0.56이었고, 일치하지 않는 방문 간 비교 평균은 0.16이었습니다. 이는 두 번째 쓰기가 첫 번째 신호를 지우기보다 비슷한 방향으로 증폭한다는 뜻입니다. Attention에서는 Q와 K가 0.89–0.93으로 유지되어 참조 위치가 보존됐지만 V는 0.65–0.74로 더 크게 변했고, PV와 최종 O는 각각 0.56–0.69와 0.49–0.64로 더 낮아졌습니다.

MoE Router는 두 방문에서 일부 top-8 전문가를 재사용하고 나머지를 다르게 선택했습니다. 97% 희소성에서는 후보 전문가 풀이 크게 늘어났는데도 같은 토큰의 겹치는 전문가가 2–3개로 무작위 기대치보다 높았습니다. Attention의 top-8 token 집합은 56–66% 겹쳤고 다른 Head 통제 조건은 28–34%에 머물렀습니다. 따라서 두 번째 방문은 새로운 위치를 폭넓게 찾기보다 비슷한 위치의 Value 표현을 다시 계산해 내용 전달을 바꾸는 경로에 가깝습니다.

한계점

Scaling law는 85%·95%·97% 희소성 grid만 사용했으며 S = 0% dense-reference control은 별도 regime으로 판단해 적합에서 제외했습니다. 지원되는 누적 학습 계산량 범위는 1.3 × 10¹⁹–2.2 × 10²¹ FLOPs이고, C = 10²² FLOPs의 절약률은 이 범위를 넘어선 외삽값입니다. 각 설정은 한 번만 학습했기 때문에 표준오차에는 학습 실행 간 변동이 포함되지 않습니다.

실무 활용

SMELT는 MoE 언어 모델에서 고유 층을 추가하지 않고 중간층 재사용으로 유효 깊이를 늘리려는 설계 선택지입니다. 기존 모델과 비교할 때는 반복 횟수만 늘리지 말고 폭, 전문가 수, GQA, KV cache를 함께 조정해야 하며, 논문 수치상 긴 문서·Code·few-shot 처리에서 우선 평가할 근거가 있습니다.

토큰당 추론 FLOPs와 전체 파라미터를 동시에 제한해야 하는 sparse MoE 언어 모델 설계
긴 문서와 In-context examples를 처리하는 모델에서 중간층 반복 구조의 효과 비교
Code·Math/STEM처럼 구조와 장거리 의존성이 있는 학습 데이터의 도메인별 성능 평가

코드 공개 여부: 비공개

키워드

SMELT(중간층을 두 번 반복하는 sparse MoE Transformer)Mixture-of-Experts(전문가 혼합 모델)Looped Transformer(반복형 Transformer)Scaling Law(모델 확장 법칙)Attention Sink(초기 토큰 집중 현상)KV Cache(Key·Value 캐시)In-context Learning(문맥 내 학습)

용어 해설

반복형 Transformer(Looped Transformer): — 기존 Transformer 층을 새 가중치로 쌓는 대신 동일한 층 블록을 여러 번 실행하는 구조입니다. 저장하는 고유 파라미터 수를 늘리지 않고 토큰이 거치는 순차 계산 깊이를 키우지만, 반복 횟수만큼 FLOPs와 KV cache 부담이 커지므로 공정한 비교에는 비용 조정이 필요합니다.
Mixture-of-Experts: — 각 토큰을 전체 Feed-Forward 전문가가 아니라 일부 전문가에만 라우팅하는 구조입니다. 전체 전문가를 저장해 파라미터 용량을 확보하면서 토큰당 활성화되는 전문가 수를 제한해 계산량을 분리할 수 있어, SMELT에서는 반복 실행으로 늘어난 FLOPs와 줄어든 폭을 함께 조정하는 데 활용됩니다.
KV 캐시(KV Cache): — Autoregressive 추론에서 이미 계산한 Attention의 Key와 Value를 저장해 이전 토큰의 투영을 다시 계산하지 않게 하는 메모리 구조입니다. 캐시 크기는 레이어 수와 Head 차원에 영향을 받으므로, 반복 실행 모델의 실제 서비스 가능 문맥 길이를 비교할 때 파라미터와 FLOPs와 함께 맞춰야 합니다.
Attention Sink: — 토큰의 의미와 무관하게 Autoregressive Transformer가 시퀀스의 초기 토큰에 Attention 질량을 과도하게 모으는 현상입니다. SMELT에서는 같은 중간 층의 두 번째 방문이 이 질량을 낮추고 내용과 관련된 토큰으로 재분배해, 긴 입력과 In-context Learning에서의 이득과 연결됩니다.
Chinchilla 스케일링 법칙(Chinchilla Scaling Law): — 모델의 계산량, 파라미터 또는 활성 계산량, 학습 토큰 수와 검증 손실 사이의 관계를 거듭제곱 형태로 근사하는 방법입니다. 총 학습 계산량을 고정하면 모델 계산과 데이터 양을 어떻게 나눌지 결정할 수 있으며, 이 논문은 Baseline과 SMELT에 별도 법칙을 맞춰 동일 손실에 필요한 계산량을 비교합니다.

SMELT: 반복형 MoE Transformer의 계산량 대비 확장 법칙

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

지식과 추론을 분리한 Mobius

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

지식과 추론을 분리한 Mobius