MoR(Mixture-of-Recursions)은 무엇인가?

라우팅 메커니즘

MoR이 재귀의 깊이(Recursion Depth)를 정하는 방법에는 두 가지가 있습니다.

1) 전문가 선택 라우팅(Expert-choice Routing)

이 방식은 일종의 다이나믹한 Early Exit이라고 할 수 있겠습니다. 어려운 토큰(Token)은 계속해서 레이어를 더 통과하고, 쉬운 토큰은 더 일찍 멈추고 재귀에서 빠져나오게 됩니다.

MoR은 각각의 재귀의 깊이(Recursion Depth)를 하나의 “전문가(Expert)”로 취급합니다.

각 단계마다 라우터는 토큰의 히든 상태(Hidden State)를 바탕으로 점수를 매기고, 상위 k개의 토큰만 다음 단계로 계속 진행시킵니다. 이전 단계에서 살아남은 토큰만 다음 단계에서 다시 평가되기 때문에, 자연스럽게 계층적으로 필터링(Hierarchical Filtering)이 일어나고, 깊이가 깊어질수록 후보 집합은 점점 더 줄어듭니다.

이 방식의 장점은 단계별로 연산 예산(Compute Budget)이 고정되어 있고 예측 가능하다는 점이지만, 학습 과정에서는 라우터가 미래의 토큰을 간접적으로 “보게 되는” 정보 누출(Information Leakage) 문제가 발생할 수 있습니다. 즉, 공정하게 학습을 했는가에 대한 이슈가 제기될 수 있고, 모델의 일반화 성능에 악영향을 줄 지도 모른다는 거죠. 그래서 이런 문제를 보완하기 위해서 보조 라우터(Auxiliary Router)나 보조 손실 함수(Auxiliary Loss) 같은 추가적인 해결책이 필요합니다.

2) 토큰 선택 라우팅(Token-choice Routing)

이 방식에서는 라우터가 각 토큰의 재귀의 깊이를 처음에 한 번만 결정합니다. 즉, 토큰마다 시작 시점에서 히든 상태를 기반으로 깊이(예: 1, 2, 3)를 할당받고, 그에 맞춰 정해진 횟수만큼 재귀 단계를 거친다는 거죠 - 더 많지도, 더 적지도 않게요.

전문가 선택 라우팅(Expert-choice Routing)과 비교했을 때, 토큰 선택 라우팅(Token-choice Routing)은 정보 누출(leakage) 문제가 없습니다. 하지만 반대로 깊이(Depth)별로 연산 부담을 균형 있게 유지하기가 어렵다는 단점이 있습니다.

그래서 이 방식을 제대로 활용하려면, 밸런싱 손실 함수(Balancing Loss Function)나, 깊이별 연산 분포를 고르게 해 주는 특화된 알고리즘(Specialized Algorithm)이 함께 필요합니다.

KV 캐싱 전략

MoR의 KV 캐싱(KV Caching)은 왜 특별한 요령이 필요할까요? 이유는 간단합니다. 어떤 토큰이 일찍 빠져나가면(Early Exit), 더 깊은 단계에서 사용해야 할 그 토큰의 KV(Key–Value) 쌍이 아예 존재하지 않게 됩니다. 이렇게 되면 나중에 들어오는 토큰들의 어텐션(Attention)이 깨질 수 있죠.

그래서 MoR은 이 문제를 해결하면서도 메모리 사용량을 줄이고 전체 속도까지 높여주는 두 가지 해법을 제안합니다.

MoR은 실제로 더 많은 단계가 필요한 토큰에 대해서만 메모리(Key–Value 쌍)를 저장하고 불러오는 방식을 사용합니다. 구체적으로는 두 가지 전략이 있습니다.

• 재귀 단위 KV 캐싱(Recursion-wise KV Caching)
  
  이 방식에서는 해당 재귀 깊이에서 여전히 활성 상태인 토큰의 KV 쌍만 저장합니다. 따라서 그 단계에서의 어텐션은 캐싱된 토큰들에만 한정되죠. 이렇게 하면 메모리 사용량을 줄일 수 있고, 메모리 입출력(I/O) 과정에서 읽고 쓰는 데이터 양도 줄어듭니다. 말 그대로 필요한 것만 저장하는 방식입니다.

• 재귀적 KV 공유(Recursive KV Sharing)
  
  이 방식에서는 모든 토큰이 최소한 첫 번째 재귀 블록은 반드시 거치게 됩니다. KV 캐싱은 이 첫 번째 단계에서만 일어나고, 이후 단계에서는 이 값을 계속 재사용합니다. 이렇게 하면 생성 시작 속도를 높이고, 중복 계산을 피하고, 메모리까지 더 절약할 수 있습니다.

그리고 중요한 점은, MoR은 각 단계에서 여전히 처리 중인 토큰에 대해서만 연산량이 큰 어텐션(Attention) 연산을 수행한다는 겁니다.

정리하자면:

• 라우팅(Routing)은 MoR을 적응형(Adaptive)으로 만들어 줍니다. → 각각의 토큰이 필요한 만큼만 연산 파워를 받는 겁니다.

• KV 캐싱(KV Caching)은 MoR을 효율적(Efficient)으로 만듭니다. → 덜 저장하고 더 재사용하면서 메모리와 연산을 절약합니다.

또 하나 기억해야 할 점은, MoR이 모든 라우팅 방식과 KV 캐싱 방식을 동시에 쓰는 것이 아니라는 것입니다. 모델의 설정에 따라 하나의 라우팅 전략 + 하나의 KV 캐싱 전략을 조합하는 구조로, 총 네 가지 유효한 조합이 가능합니다.

• 전문가 선택 라우팅(Expert-choice Routing) + 재귀 단위 캐싱(Recursion-wise Caching)

• 전문가 선택 라우팅(Expert-choice Routing) + 재귀적 공유(Recursive Sharing)

• 토큰 선택 라우팅(Token-choice Routing) + 재귀 단위 캐싱(Recursion-wise Caching)

• 토큰 선택 라우팅(Token-choice Routing) + 재귀적 공유(Recursive Sharing)

즉, 속도, 메모리, 단순성 사이의 트레이드오프(Trade-off)에 따라서 어떤 조합을 선택할지가 달라집니다.

실제 성능 결과

테스트 결과, MoR은 우리가 익숙하게 사용해 온 기존 트랜스포머보다 눈에 띄게 향상된 성능을 보여줍니다.

예를 들어서, 전문가 선택 라우팅(Expert-choice Routing)과 2번의 재귀(Recursions)를 적용한 MoR은, 표준 트랜스포머보다 검증 손실(Validation Loss)이 더 낮고, 퓨샷 정확도도 더 높게 나왔습니다.

• MoR: 43.1%

• 표준 트랜스포머: 42.3%

게다가 이런 성능을 달성하면서도 모델 파라미터(Parameter) 수는 약 50% 줄일 수 있었습니다.

재귀(Recursion) 횟수를 3번이나 4번으로 늘렸을 때도, MoR은 풀사이즈 트랜스포머와 비슷한 수준의 성능을 유지했습니다.

• 훈련 시간은 19% 단축

• 최대 메모리 사용량은 25% 절감

이건 계층적 필터링(Hierarchical Filtering)과 재귀 단위 어텐션(Recursion-wise Attention) 덕분이라고 할 수 있을 겁니다. 쉽게 처리할 수 있는 토큰의 시퀀스를 일찍 줄여주기 때문에 이런 효과가 가능한 것이죠.

다만, 3번의 재귀 설정에서는 라우팅 방식에 따라 차이가 나타났습니다.

• 토큰 선택 라우팅(Token-choice Routing): 40.0% 정확도

• 전문가 선택 라우팅(Expert-choice Routing): 42.6% 정확도

즉, 전문가 선택 방식이 더 안정적인 성능을 보여주고 있습니다.

한편, KV 캐시 공유(KV Cache Sharing)는 메모리를 절약하는 효과는 있었지만, 정확도가 약간 떨어졌습니다. 그럼에도 메모리가 매우 중요한 상황이라면 충분히 감수할 수 있는 트레이드오프(Trade-off)라고 하겠습니다.

MoR은 또 하나의 특징으로 연속적 깊이 배치(Continuous Depth-wise Batching)를 활용합니다. 서로 다른 재귀 깊이를 가진 토큰들을 묶어서 한꺼번에 GPU에 태우는 방식인데, 덕분에 GPU 자원을 최대한 활용할 수 있습니다. 큰 배치 상황에서는 최대 2.06배까지 처리량(Throughput)이 향상되는 것을 관찰했다고 합니다.

또 한 가지 흥미로운 점은, MoR은 더 많은 데이터를 먹이는 것보다는, 파라미터를 늘리고 학습 스텝을 줄이는 방식에서 더 큰 이득을 본다는 겁니다.

연구자들은 다양한 설정을 실험한 끝에, 최적의 설계안(Best Design Choice)으로 다음과 같은 세팅을 제안했습니다:

• 미들-사이클 파라미터 공유(Middle-cycle Parameter Sharing)
  → 처음과 마지막 레이어는 고유하게 두고, 중간 레이어들을 공유하는 방식

• 보조 손실(Auxiliary Loss)과 선형 라우터(Linear Router)를 사용한 전문가 선택 라우팅(Expert-choice Routing)

• 재귀 단위 캐싱(Recursion-wise Caching)
  다만, 토큰 선택 라우팅(Token-choice Routing)을 쓰는 경우에는 KV 공유(KV Sharing)가 더 유리할 수 있음

👉 그럼, MoR에 대해서 한 번 정리하고, 이 기법의 장점과 단점을 살펴볼까요?

MoR의 장점

레이어(Layer)를 재사용하면 필요한 파라미터 수를 크게 줄일 수 있습니다. 매번 새로운 레이어를 두는 대신, 같은 레이어 스택을 여러 번 통과시키기 때문이죠.

MoR의 라우팅(Routing) 결정은 무작위로 이뤄지는 게 아닙니다. 토큰의 의미와 난이도에 맞게 정렬되는데, 쉬운 토큰은 적은 횟수만 거치고, 어려운 토큰은 더 많은 단계를 거치도록 하는 식입니다.

또 MoR은 어텐션과 메모리도 더 똑똑하게 사용하죠. 각 단계에서 여전히 “활성 상태”인 토큰에 대해서만 무거운 어텐션 연산을 수행하고, 그 토큰에 대해서만 메모리를 저장합니다.

여기에 더해서, MoR은 ‘재사용(Reusing)’을 효율적으로 하는 전략을 씁니다. 특별한 KV 캐싱(KV Caching) 버전을 적용해서, 첫 번째 패스에서 저장한 데이터를 이후 단계에서도 재사용합니다. 덕분에 속도는 빨라지고 메모리 사용량은 줄어듭니다.

이렇게 연산 효율(Compute Efficiency)이 높아진 덕분에, MoR은 같은 FLOP 예산 안에서도 더 많은 토큰을 처리할 수 있습니다.

또 하나 중요한 장점은 테스트-타임 스케일링(Test-time Scaling)이 가능하다는 점입니다. 즉, 재학습을 하지 않고도 추론(Inference) 시점에서 최대 재귀 깊이를 더 늘릴 수 있다는 것이죠.

MoR의 한계점

하지만, 당연하게도 한계점은 있죠 - MoR의 한계는, 무엇보다도 아직 이에 대한 연구가 충분히 진행되지 않았다는 점에서 비롯되는 것들입니다:

• 훈련 이후 라우팅 용량을 조정하기 어렵다
  보조 손실(Auxiliary Loss)을 사용하는 전문가 선택 라우팅(Expert-choice Routing)에서는, 선택된 토큰과 선택되지 않은 토큰의 출력이 거의 완벽하게 분리됩니다. 그렇기 때문에, 추론(Inference) 시점에서 top-k(즉, 한 단계에서 몇 개 토큰을 통과시킬지) 값을 바꾸기가 힘듭니다. 한 번 학습되고 나면 라우팅 결정이 지나치게 경직(Rigid)되기 때문입니다.

• 라우터의 유연성이 부족하다
  현재 라우터는 토큰의 추론 복잡도(Reasoning Complexity)에 따라서 재귀 깊이를 다이나믹하게 조절하도록 최적화되어 있지 않습니다.

• 실험 규모의 한계
  지금까지 실험된 모델은 최대 17억(1.7B) 파라미터 사이즈 수준입니다. 그 이상, 예컨대 30억(3B) 파라미터 이상의 대규모 MoR 모델이 어떤 성능을 보일지는 아직 알 수 없습니다.

• 토큰별 연산량의 즉각적인 조절 불가
  현재 설계로는, 실행 도중에 각 토큰마다 얼마만큼의 연산을 배분할지 쉽게 바꿀 수 없습니다.

• 희소 연산(Sparse Computation) 기법과의 결합이 미흡
  아직 MoR은 희소 연산 최적화 기법과 통합되지 않았습니다.

• 텍스트 영역 외에는 검증되지 않음
  MoR은 지금까지 텍스트 모델에서만 테스트되었고, 이미지나 멀티모달 같은 다른 영역에서는 아직 실험을 해 본 적이 없습니다.

그렇다면, MoR이 우리에게 던지는 의미는 뭘까요?

맺으며

MoR(Mixture-of-Recursions)기법은, 오늘날 가장 중요한 아이디어 중의 하나를 보여줍니다. 즉, 앞으로의 AI 개발 방향은 단순히 끝없이 레이어(Layer)를 더하는 게 아니라, 정말 필요할 때만 더 깊이 생각할 줄 아는 모델을 만드는 데 있다는 것입니다.

재귀의 깊이(Recursion Depth), 라우팅(Routing), 메모리 사용 방식을 다시 설계해서, MoR은 지금 있는 자원을 훨씬 더 똑똑하게 활용할 수 있는 방법을 제시합니다.

MoR은 유연한 사고의 깊이(Flexible Reasoning Depth)라는 개념의 무대를 열어젖혀 줍니다. 트랜스포머에다가 적응형의 백본(Backbone)을 제공해서, 각 토큰이 필요한 만큼만 적절히 처리되도록 하는 겁니다.

물론 아직 초기 단계고 더 많은 실험이 필요하지만, MoR의 맞춤형 아키텍처는 이 기법을 사실상 범용적인 시스템으로 포지셔닝해 줍니다. 마치 진짜 트랜스포머 장난감처럼요.

이제 남은 것은, 현재의 한계들을 극복하는 것입니다. 그렇게만 된다면, MoR은 고도의 전문성과 효과성을 갖춘 차세대 모델로서 자리 잡을 수 있을 겁니다.

보너스: 참고자료

• Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

• Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs

• Looped Transformers are Better at Learning Learning Algorithms

• Dynamic Early Exit in Reasoning Models

• Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts