구글, 메모리 절반으로 추론 속도 2배 높인 MoR 발표

구글 딥마인드와 한국과학기술원(KAIST) 연구진이 최근 새로운 언어 모델 아키텍처인 'Mixture-of-Recursions' (MoR)를 발표했다. 이 모델은 트랜스포머 아키텍처의 성능을 유지하면서 추론 속도를 두 배로 높이고, 학습 계산량을 줄이며, 약 50%의 KV 캐시 메모리 사용을 감소시키는 것으로 알려져 있다. 논문이 발표되자마자 소셜 미디어에서 큰 관심을 받았으며, 일부는 이를 '트랜스포머 살인자'라고 평가하기도 했다. 트랜스포머 아키텍처는 2017년 등장한 이후 대형 언어 모델의 기술적 기반으로 자리잡았다. 현재 대부분의 선두 주자 모델들이 이 아키텍처 위에 구축되어 있다. 그러나 모델의 규모가 커지면서 트랜스포머 아키텍처는 계산과 메모리 자원에 대한 요구가 점점 늘어나고, 학습과 배포 비용이 매우 높아졌다. 과거 효율성 최적화 방법들은 일반적으로 단일 방향에 집중되었는데, 예를 들어 파라미터 공유를 통해 모델 크기를 줄이는 것, 또는 적응형 계산을 통해 필요한 만큼의 계산력을 할당하는 것 등이 있다. 그러나 여러 효율 목표를 동시에 최적화하는 것은 거의 찾아볼 수 없었다. MoR 아키텍처의 핵심 혁신은 재귀 계산과 동적 라우팅 메커니즘을 결합하여 다중 효율 문제를 하나의 통합 프레임워크 내에서 해결하는 것이다. 표준 트랜스포머 모델에서는 입력 텍스트의 각 토큰이 동일한 수의 계산층을 거치지만, MoR는 이를 변경하여 각 토큰이 자신의 복잡도에 따라 다른 깊이의 처리를 받도록 허용한다. MoR는 공유 파라미터 블록을 사용하여 파라미터 효율성을 높이며, 경량화된 '라우터'를 통해 각 토큰이 얼마나 많은 재귀 계산을 거쳐야 하는지를 결정한다. 연구팀은 '엑스퍼트 선택'과 '토큰 선택' 등의 다양한 라우팅 전략을 테스트하여 계산 부하를 균형 있게 분배하고 정보 처리 과정에서 논리적 문제를 피하는 방법을 찾았다. 파라미터 공유 면에서 'Mid-Cycle' 전략이 가장 우수한 성능을 보였다. 이 전략은 모델의 첫 번째 계층과 마지막 계층이 독립적인 파라미터를 사용하도록 하면서 중간 계층 사이에서만 가중치를 공유하여, 파라미터 효율성과 모델의 표현 능력 사이에서 좋은 균형을 이рус었다. 메모리 관리는 MoR의 또 다른 중요한 개선 사항이다. 파라미터가 공유되었더라도, 전통적인 재귀 모델은 각 재귀 단계마다 독립적인 KV 캐시를 생성하여 메모리 사용량이 여전히 높았다. MoR는 이 문제를 해결하기 위해 두 가지 새로운 KV 캐시 전략을 제안했다. 첫 번째는 '재귀식 캐시'로, 특정 재귀 단계로 라우팅된 토큰만 KV 데이터를 저장하며, 이 데이터에서의 주의 계산을 제한하여 KV 캐시의 메모리 사용량과 데이터 읽기/쓰기 횟수를 효과적으로 줄였다. 두 번째 '재귀 공유' 전략은 모든 토큰이 첫 번째 재귀 블록을 거치는 특성을 활용하여, 첫 번째 단계에서만 KV 데이터를 캐싱하여 후속 모든 재귀 단계에서 재사용함으로써 메모리를 최대한 절약한다. 연구팀은 1억 3500만에서 17억 개의 파라미터를 가진 다양한 모델 규모에서 테스트를 진행했다. 결과는 같은 학습 계산 예산 하에서 MoR 아키텍처를 사용한 모델이 벤치마크 트랜스포머 모델보다 약 50% 적은 파라미터를 가졌음에도 불구하고, 여러 샘플 학습 작업에서 평균 정확도 43.1%를 달성해 벤치마크 모델의 42.3%를 상회했다. 더욱 중요한 것은 MoR의 높은 계산 효율성이 같은 계산 예산 내에서 더 많은 학습 데이터를 처리할 수 있게 해주어, 이는 다시 모델의 최종 성능을 향상시키는 것으로 나타났다. 고정 학습 데이터 양을 비교한 실험에서는 MoR 설정이 벤치마크 모델보다 25% 적은 학습 계산량을 사용했음에도 불구하고 성능이 뛰어났으며, 학습 시간은 19% 줄고, 최대 메모리 사용량은 25% 감소했다. 추론 성능 면에서도 MoR의 장점이 두드러진다. MoR는 연속 깊이 배치 처리 기술을 사용하여 같은 계산 단계에 있는 토큰들을 하나의 배치로 결합해 처리한다. 이는 같은 파라미터 블록을 사용하기 때문에 가능하다. 이 기술은 모델의 조기 종료 메커니즘과 함께 작동하여 처리 흐름을 크게 개선했다. 3억 6000만 규모의 모델 테스트에서 MoR-4 설정은 특정 조건에서 최대 2.06배의 추론 속도 향상을 실현했다. 연구팀은 또한 MoR 모델이 다양한 유형의 토큰을 처리할 때 의미적인 중요성과 관련된 패턴을 보이는 것을 발견했다. 내용이 풍부한 토큰, 예를 들어 "People"나 "defensively confident"는 더 많은 재귀 횟수(세 번)를 받았으며, 기능적인 단어인 "and"와 같은 토큰은 적은 재귀 횟수를 필요로 했다. 이는 모델이 중요한 정보를 처리하는데 더 많은 계산 자원을 사용하도록 학습되었다는 것을 시사한다. 실제로 MoR의 등장은 구글의 이전 연구를 바탕으로 이루어졌다. 예를 들어 Mixture-of-Depths (MoD) 기술에서는 이미 계산 자원을 동적으로 할당하는 방법을 탐색했으며, 재귀 트랜스포머는 파라미터 공유 기술의 이론적 기초를 제공했다. MoR는 이러한 연구를 이어받아, 단일 차원의 최적화에서 파라미터, 계산, 메모리 등 여러 차원의 협력적인 최적화로 발전시켰다. 이는 대형 언어 모델의 배포와 사용 비용을 줄이는 데 큰 실질적 의미를 가진다. 전반적으로 아직 MoR가 트랜스포머를 완전히 대체할 수 있을지는 예단하기 어렵다. 그러나 MoR은 확실히 미래의 언어 모델 아키텍처 설계에서 성능과 효율성을 모두 크게 개선할 수 있는 잠재력을 지닌 방향성을 제시하고 있다. 업계 전문가들은 MoR이 대형 모델의 성능을 유지하면서도 효율성을 크게 높여줄 수 있을 것으로 기대하고 있으며, 구글 딥마인드의 지속적인 연구 노력은 AI 모델의 효율성 향상에 중요한 역할을 하고 있다. 이 연구는 AI 모델의 실제 적용 가능성과 비용 효율성을 크게 향상시킬 것으로 보인다.