LatentLoop은 루프 MoE의 첫 방문을 압축해 연산을 재배분한다
LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation
무엇인가
루프 트랜스포머는 같은 층을 반복 적용해 저장 파라미터를 늘리지 않고 유효 깊이를 키우지만, 방문마다 연산은 그대로 든다. 그래서 파라미터 수만 고정한 비교는 루프 모델에 유리하게 기울고, 연산만 맞추면 저장 파라미터가 줄어 불리해진다. 이 논문은 SMELT의 루프 MoE 레시피(중간 절반 층을 두 번 방문, 엔트리·엑시트 블록은 한 번)를 출발점으로 삼아, 저장 파라미터·KV 캐시·스텝당 FLOPs·총 학습 FLOPs를 1% 이내로 맞춘 상태에서 루프 내부의 연산 배분 자체를 개선하는 문제를 다룬다.
어떻게 동작하나
출발점은 관찰이다. 두 방문에 각각 학습 가능한 잔차 계수 α1, α2를 붙여 모델 가중치와 함께 최적화하면, 1.3B에서 (0.27, 0.87), 3B에서 (0.36, 1.09)로 첫 방문의 계수만 크게 작아진다. 고정 계수 실험도 같은 비대칭을 뒷받침한다. 첫 방문 계수를 0.125~0.5로 줄이면 검증 NLL이 유지되거나 오히려 개선되지만, 두 번째 방문에 같은 축소를 적용하면 NLL이 올라간다. 테스트한 조합 중 (0.25, 1)이 두 크기 모두에서 최저 NLL을 기록해 (1, 1)보다 각각 약 0.008, 0.003 개선됐고, 방문마다 1/N로 균등 축소하는 대칭 설정 (0.5, 0.5)보다도 좋았다. 학습 파라미터나 연산 예산을 늘리지 않고 얻은 이득이다.
무엇과 다른가
이 비대칭이 제안 방법 LatentLoop의 근거가 된다. 첫 방문은 시퀀스를 압축 인자 k=2로 절반 길이의 잠재 시퀀스로 줄여 처리하고, 두 번째 방문은 원래 해상도로 전체 시퀀스를 처리한다. 절약된 연산은 예산 정합 제약 안에서 모델 폭으로 재배분한다. 인코더는 인접한 두 상태를 평균 내는 파라미터 없는 mean aggregation을 쓰고, 첫 방문이 만든 잠재 증분은 학습 가능한 writeback 계수 w_s로 채널·수신 위치별로 스케일을 조정해 원래 시퀀스에 되쓴다. 이때 인과 정렬을 지켜 증분은 위치 2j와 2j+1에만 더해지고, 시퀀스 첫 토큰은 잠재 업데이트를 받지 않는다. 추가 파라미터는 2d뿐이다. 파라미터 공유는 선택적으로 바꾼다. MoE 라우터와 전문가 가중치는 두 방문이 공유하되, 서로 다른 시퀀스 해상도를 다루는 어텐션 블록만 분리한다. 늘어난 어텐션 파라미터는 저장 전문가 수를 층당 하나 줄여(3B 설정에서 76→75) 상쇄하고, 활성 전문가 수는 그대로 둬서 전문가 FLOPs를 유지한다.
어떻게 쓰나
실험은 1.3B·3B·7B에 걸친 15개 예산 정합 설정에서 이뤄졌고, 최대 학습 예산은 10^21 FLOPs를 넘는다. 데이터는 FineWeb-Edu 고정 스트림, TinyLlama 토크나이저, 0.1B 토큰 검증셋을 쓴다. 베이스라인은 층마다 64개 전문가를 저장하고 각 층을 한 번만 실행하는 MoE Base와, SMELT식 루프 구조에 잔차 계수 1을 쓰는 MoE Loop다. 15개 설정 전부에서 LatentLoop이 두 베이스라인보다 낮은 검증 NLL을 기록했다. 각 아키텍처에 손실 곡면을 따로 피팅한 결과, NLL 1.988 지점에서 LatentLoop은 1.85×10^20 FLOPs로 도달하며, 이는 MoE Base 대비 19.0%, MoE Loop 대비 7.9%의 연산 절약에 해당한다.
전제와 한계
다운스트림 평가는 LAMBADA OpenAI·Standard, HellaSwag, PIQA, WinoGrande, ARC-Easy, ARC-Challenge, SciQ, MMLU(continuation 형식) 아홉 개 벤치마크를 0-shot과 5-shot으로 돌렸다. 크기별 최대 토큰 예산(1.3B는 40B, 3B는 70B, 7B는 140B)에서 LatentLoop이 여섯 개 크기·샷 조합 모두 평균 최고 점수를 냈다. MoE Loop 대비 평균 이득은 5-shot에서 0.73~1.18%p, 0-shot에서 0.17~0.67%p이고, MoE Base 대비로는 5-shot 0.67~1.82%p, 0-shot 0.80~1.39%p다.
절제 실험은 3B/40B 기준으로 이뤄졌다. 첫 방문 연산을 줄이는 네 가지 방식(잔차 스트림 절반 폭 투영, 전문가 중간 폭 절반, 활성 전문가 8→4, 시퀀스 압축) 중 시퀀스 압축이 가장 낮은 NLL을 냈다. 각 잠재 위치에서 전체 폭과 전문가 활성 수를 유지하는 점이 다른 전략 대비 이점으로 해석된다. 인코더·디코더 비교에서는 직접 writeback 2.0268, 채널 컨볼루션 인코더 2.0219, scaled decoder 2.0199, 둘 다 학습 2.0209, 계수 RMS 정규화 2.0228, 단일 스칼라만 학습 2.0210, SpiralFormer식 softmax 집계 2.0264였다. 파라미터 공유 절제에서는 일반 MoE Loop가 어텐션을 분리하고 전문가를 92→91로 줄이면 NLL이 2.0293에서 2.0319로 나빠지지만, LatentLoop은 76→75로 줄일 때 2.0303→2.0268(mean+direct), 2.0237→2.0199(mean+scaled decoder)로 개선됐다. 압축된 첫 방문에서만 어텐션 분리가 이득이라는 뜻이다.
분석 파트는 비대칭이 두 방문에 국한되지 않음을 보인다. 방문을 3회·4회로 늘려도 계수는 방문 순서대로 커지고, 학습 90% 시점의 연속 계수 비율은 3회에서 약 2.4배, 4회에서 약 2.0배였다. 학습된 writeback 계수는 두 수신 위치 모두 초깃값 1보다 훨씬 작은 채널 평균 약 0.199와 0.160으로 수렴했고, 소스 쌍의 상태를 포함하는 가까운 수신 위치 2j가 더 먼 2j+1보다 일관되게 큰 계수를 학습했다.
개발자 관점에서 이 논문이 주는 실무적 신호는 명확하다. 루프 MoE를 설계할 때 두 방문에 같은 연산을 배분하는 기본값을 그대로 쓰지 말고, 첫 방문의 잔차 계수나 writeback 계수를 로그로 찍어 실제 업데이트 강도를 확인하라는 것이다. 계수 하나를 0.25로 낮추는 것만으로 검증 손실이 개선된다는 결과는 재학습 비용 대비 시도해볼 만한 저비용 개입이다. 다만 검증 범위는 최대 7B 파라미터, 10^21 FLOPs 규모이고 데이터는 영어 FineWeb-Edu 중심이며 MMLU는 continuation 형식으로 평가됐다는 점을 감안해야 한다. 저자들도 밝히듯 이 비대칭이 왜 생기는지에 대한 기제는 아직 불완전하게 이해된 상태이고, 특히 3회·4회 방문에서 나타난 계수 패턴은 설명하지 못한 채 향후 과제로 남겨져 있다.