같은 블록 반복으로 LLM 깊이를 늘리는 LOOM 기법
원제 Looping Beyond Twice: A Scalable Recipe for Looped Mixture-of-Experts
추천 16댓글 2
Key Point
같은 블록을 몇 번까지 반복할 수 있는지, 그리고 MoE 라우터의 동작 특성을 고려해 어떻게 효율적으로 설계할지가 실무 모델 개발에 직결된다.
핵심 요약
- Looped Transformers는 같은 Transformer 블록을 반복 적용해 파라미터를 늘리지 않고 모델의 유효 깊이를 증가시킨다.
- 기존 연구는 반복 횟수 증가에 따른 성능 이득이 빠르게 감소해 2번 반복으로 제한했으며, MoE 모델에서 특히 이 문제가 심했다.
- LOOM은 두 가지 병목을 진단했다: 반복할 때마다 hidden-state 분산이 증가해 깊은 모델이 불안정해지는 '깊이의 저주' 현상, 그리고 라우터가 매번 같은 expert를 선택해 계산 다양성이 떨어지는 'expert 선택 붕괴' 현상이다.
- LOOM은 각 반복에서 새로운 계산을 추가하면서도 상태를 안정화하도록 설계했다: residual 업데이트 크기를 조절해 분산 성장을 제한하고, 매 반복마다 입력 임베딩을 재주입하며, 루프별로 다른 expert를 선택하는 라우터와 이전 출력을 전달하는 Looping Residual을 사용한다.
- 100M~1.7B 규모 모델에서 9~12번 반복까지 안정적으로 확장된다.
- FLOP을 맞춰 비교할 때 700M 모델은 5번 반복에서 최적이며, 기준 대비 perplexity를 18.36에서 16.54로 낮추고 zero-shot 평균 정확도를 38.84%에서 39.53%로 향상시켰다.
- FLOP 제약 없이 1.7B 모델을 60B 토큰으로 학습했을 때 9번 반복에서 최고 성능을 보였으며, perplexity는 9.62에서 7.77로, zero-shot 평균 정확도는 42.4%에서 47.7%로 개선되었다.
- 코드는 공개되어 있다.