같은 연산 예산에서 중간 층을 두 번 도는 MoE 스케일링 법칙 SMELT
루프 트랜스포머를 MoE에 적용한 SMELT 논문이 연산과 파라미터와 KV 캐시를 모두 맞춘 비교에서 학습 FLOPs를 6.8~18.0% 절약했다고 보고했습니다.
왜 지금 이 주제인가
루프 트랜스포머(looped Transformer)는 같은 층 묶음을 여러 번 통과시켜 실질 깊이를 늘리는 구조입니다. 아이디어 자체는 새롭지 않지만, 대부분의 비교가 모델 크기를 고정한 채 이뤄져서 구조의 이점인지 그냥 연산을 더 쓴 결과인지 구분되지 않는다는 지적이 따라다녔습니다. 9월 1일 arXiv에 올라온 SMELT 논문은 이 문제를 정면으로 다룹니다. 연산량과 파라미터와 KV 캐시를 모두 맞춘 상태에서도 루프가 이득인지를 스케일링 법칙으로 확인한 작업입니다.
핵심 내용
세 가지 예산을 동시에 맞춘 비교
SMELT는 Sparse MoE Transformer, middle layers Loop Twice의 약자입니다. 이름 그대로 Mixture-of-Experts 트랜스포머에서 가운데 절반의 층을 두 번 통과시킵니다. 12층 모델이라면 중간 6층이 두 번 실행되어 실질 깊이가 18층이 됩니다. 루프 구간의 잔차 갱신은 절반으로 축소합니다. 세 번, 네 번 도는 변형보다 두 번이 나았다는 것이 여러 절제 실험 끝에 정리된 결론입니다.
핵심은 구조가 아니라 비교 조건입니다. 저자들은 토큰당 FLOPs, 비임베딩 총 파라미터 수, KV 캐시 크기 세 가지를 기준 모델과 함께 맞췄습니다. 층을 더 실행하는 만큼 늘어나는 연산은 은닉 차원을 좁혀 상쇄하고, 그렇게 줄어든 파라미터는 층당 전문가 수를 늘려 되돌리며, KV 캐시는 헤드 크기와 GQA 비율로 조정합니다. 논문은 세 예산 모두에서 남는 오차가 4% 아래라고 밝힙니다.
스케일링 법칙과 하류 성능
이렇게 맞춘 뒤 네 가지 규모로 확장했고, 가장 큰 설정은 비임베딩 파라미터 54B입니다. 각 구조에 대해 Chinchilla 방식의 스케일링 법칙을 따로 적합했더니 SMELT 쪽 손실이 연산 증가에 따라 더 빠르게 떨어졌습니다. 계산 최적 경계 기준으로 학습 FLOPs를 6.8~18.0% 절약하는 수준이고, 프런티어 지수는 SMELT가 0.250, 기준이 0.237로 보고됐습니다. 반면 계산 최적 지점에서 파라미터와 토큰을 나누는 비율은 두 구조가 거의 같았습니다. 기존 학습 예산 배분 관행을 바꾸지 않고도 이 구조를 끼워 넣을 수 있다는 뜻이라, 실무에서는 이 부분이 절약 폭 자체만큼 중요할 수 있습니다.
이 이득은 검증 손실만으로 예측되는 것보다 하류 벤치마크에서 더 크게 나타났습니다. DCLM Completion에서는 대응하는 96쌍을 모두 이겼고, DCLM Core에서는 96쌍 중 83쌍, MMLU에서는 30쌍 중 29쌍을 이겼습니다. 도메인별로는 코드에서 격차가 가장 커서 연산 효율 이득이 20.4%로 집계됐습니다. 표본 길이가 길어지고 문맥에 넣는 예시가 많아질수록 이득이 커지는 경향, 그리고 규모가 커질수록 하류 성능 이득이 함께 커지는 경향도 보고됐습니다.
두 번째 통과가 하는 일
메커니즘 분석에서는 두 번째 통과가 어텐션 싱크를 줄이고 주의를 내용어 쪽으로 옮긴다고 설명합니다. 두 번의 통과가 비슷한 토큰을 보기는 하지만 값 벡터 쪽이 질의나 키보다 크게 갈라지고, Dyck 언어처럼 구조가 뚜렷한 과제에서는 두 번째 통과에서 싱크가 거의 사라졌다고 합니다.
의미와 한계
파라미터를 늘리는 대신 깊이를 재사용하는 방향이 예산을 맞춘 조건에서도 통한다는 점이 이 논문의 실질입니다. 총 파라미터와 KV 캐시까지 고정했기 때문에, 서빙 메모리를 늘리지 않고 학습 효율만 가져갈 수 있다는 해석이 가능합니다. 다만 실제 서빙에서는 같은 층을 두 번 실행하므로 지연 시간과 처리량 특성은 달라질 수 있는데, 토큰당 FLOPs를 맞췄다는 사실만으로 이 부분이 자동 해결되지는 않습니다.
한계도 분명합니다. 스케일링 법칙 적합은 10^19에서 10^21 FLOPs 구간에서 이뤄졌고 그 바깥은 외삽입니다. 실제 프런티어 모델의 학습 예산은 이보다 훨씬 큽니다. 학습은 시드 하나로 진행됐고, 데이터는 내부 코퍼스라 그대로 재현하기 어렵습니다. 문맥 길이도 4096 토큰 범위에서 분석했습니다. 어텐션 싱크 설명 역시 함께 관찰된 현상에 가까워서, 성능 향상의 원인이라고 단정하기는 이릅니다. 이 이점이 더 큰 규모에서도 유지될지는 아직 추측의 영역입니다.
한 줄 정리
연산과 파라미터와 KV 캐시를 모두 맞춘 조건에서도 중간 층을 두 번 도는 MoE가 학습 FLOPs를 6.8~18.0% 아꼈습니다.