💬 NLP

MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

원저자: Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Sau

게시일 2026-06-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Saurous, Guillaume Lajoie, Charlotte Frenkel, Razvan Pascanu, Blaise Agüera y Arcas, João Sacramento

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

개요: 현재 AI의 문제점

매우 똑똑한 사서(현재 AI의 표준인 트랜스포머)가 당신의 질문에 답하기 위해 책을 읽고 있다고 상상해 보세요.

장점: 이 사서는 책의 처음부터 끝까지의 세부 사항을 기억하는 데 매우 뛰어납니다.
단점: 이를 위해 사서는 책상 위에 점점 높아지는 인덱스 카드 더미를 쌓아두어야 합니다. 책이 길어질수록 이 더미는 더 높아집니다. 결국 이 더미가 너무 거대해져서 방 전체를 차지하게 되면, 사서는 답을 찾기 위해 그 모든 카드를 뒤적거리느라 압도당하고 맙니다. 이 때문에 긴 텍스트를 처리할 때 속도가 느려지고 비용이 많이 듭니다.

최근 과학자들은 다른 종류의 사서(Mamba나 DeltaNet 같은 RNN)를 구축하려고 시도했습니다. 이 사서들은 작은 고정 크기의 메모장만을 가지고 있습니다. 이들은 빠르고 넓은 공간을 필요로 하지 않습니다. 하지만 이들은 아주 긴 이야기의 맨 앞부분을 기억하는 데 어려움을 겪거나, 이야기가 너무 복잡해지면 실수를 하기도 합니다.

해결책: MesaNet의 등장

이 논문의 저자들은 새로운 사서인 MesaNet을 만들었습니다. 그들은 작은 메모장을 가진 사서의 속도와 큰 카드 더미를 가진 사서의 기억력을 모두 원했습니다.

이를 위해 그들은 특별한 도구인 Mesa 레이어를 도입했습니다.

비유: "즉석 전문가" vs "느린 학습자"

당신이 방금 본 숫자 목록을 바탕으로 수학 문제를 풀려고 한다고 가정해 봅시다.

기존 RNN (느린 학습자): 새로운 숫자가 나타날 때마다, 사서는 아주 작은 추측을 하고, 자신이 틀렸는지 확인한 뒤, 메모장을 약간씩 조정합니다. 이 과정을 단계별로 수행합니다. 효율적이긴 하지만, 단순히 "추측하고 조정"하는 방식이기 때문에 즉각적으로 완벽한 답을 내놓지는 못할 수도 있습니다.
MesaNet (즉석 전문가): 새로운 숫자가 나타나면, MesaNet 사서는 단순히 추측만 하지 않습니다. 그들은 지금까지 본 모든 숫자를 바탕으로 메모장을 조정할 가장 완벽한 방법을 찾아내기 위해 초고속 정신적 계산을 즉시 실행합니다. 현재의 맥락에 가장 잘 맞는 답을 보장하기 위해 한 번에 전체 최적화 문제를 해결합니다.

작동 원리 ("테스트 타임 트레이닝")

이 논문에서는 이를 **"테스트 타임 트레이닝(Test-Time Training)"**이라고 부릅니다.

보통 AI 모델은 공장에서 한 번 훈련된 후 그대로 실행됩니다. 즉, 실제로 당신과 대화할 때는 새로운 것을 배우지 않습니다.

MesaNet의 비결: MesaNet은 새로운 단어를 읽을 때마다, 그 순간을 위해 내부 메모리를 "재학습"하기 위해 아주 짧은 순간 멈춥니다. 그리고 이렇게 묻습니다: "지금 이 순간까지 읽은 모든 내용을 고려했을 때, 이 정보를 저장하는 가장 완벽한 방법은 무엇인가?"
비용: 이 "재학습" 과정은 기존 방식보다 약간 더 많은 컴퓨팅 파워(빠른 수학 솔버를 실행하는 것과 같은)를 소모합니다.
이점: 매번 최적의 답을 구하기 때문에, 기존의 "느린 학습자" RNN들보다 길고 복잡한 이야기를 훨씬 더 잘 이해합니다.

"청크(Chunk)" 혁신

이 아이디어의 원형(이전 논문)은 한 페이지씩 읽는 것처럼 계산을 하나씩 수행해야 했기 때문에 훈련 속도가 너무 느렸습니다.

새로운 전환: 저자들은 이러한 계산을 청크(덩어리) 단위로 수행하는 방법을 찾아냈습니다. 한 페이지씩 읽는 대신, 사서가 한 장의 챕터를 통째로 집어 들고 강력한 컴퓨터 칩을 사용하여 챕터 전체에 대한 수학 문제를 한꺼번에 해결한 뒤 다음으로 넘어가는 것을 상상해 보세요. 이 방식 덕분에 방대한 양의 데이터를 빠르게 훈련할 수 있게 되었습니다.

연구 결과

연구팀은 거대한 데이터셋(수십억 개의 단어)으로 MesaNet을 테스트하여 트랜스포머 및 다른 빠른 RNN들과 비교했습니다.

시작 부분에서의 우수성: MesaNet은 문장이나 이야기의 시작 부분을 이해하는 데 탁월합니다. 처음 몇 백 단어 구간에서는 거대한 트랜스포머조차 능가하는 성능을 자주 보여줍니다.
긴 문맥에서의 우수성: 다른 빠른 RNN들이 이야기가 길어짐에 따라 잊어버리거나 혼란을 겪는 반면, MesaNet은 훨씬 더 잘 버텨냅니다. 긴 문서를 다른 모델들보다 더 정확하게 이해할 수 있습니다.
트레이드오프(Trade-off): MesaNet은 "공짜"가 아닙니다. 수학 문제를 풀기 위해 읽기 과정 중에 더 많은 컴퓨팅 파워(FLOPs)를 사용합니다. 하지만 저자들은 이 추가적인 노력이 특히 긴 문맥의 깊은 이해를 요구하는 작업에서 더 높은 정확도로 보답한다는 것을 발견했습니다.
동적 속도: 이 시스템은 언제 열심히 일해야 할지 스스로 알 수 있을 만큼 똑똑합니다. 문장이 단순하면 수학 단계를 적게 수행하고, 문장이 복잡하면 더 많은 단계를 수행합니다. 과업의 난이도에 따라 노력을 동적으로 할당합니다.

요 요약

MesaNet은 이야기를 읽으면서 동시에 완벽하게 기억하는 방법을 끊임없이 재계산하는 사서와 같은 새로운 유형의 AI 아키텍처입니다. "최선의 메모리"를 찾기 위해 매 단계마다 복잡한 수학 문제를 해결함으로써, MesaNet은 특히 길고 복잡한 텍스트를 다룰 때 다른 빠른 메모리 효율적 모델들보다 뛰어난 수준의 이해력을 달합니다. 이는 더 높은 지능을 얻기 위해 약간의 추가 컴퓨팅 파워를 기꺼이 사용하는 방식입니다.

기술 요약: MESANET: 국소적 최적 테스트 타임 트레이닝을 통한 시퀀스 모델링

문제 정의

시퀀스 모델링은 현재 소프트맥스 셀프 어텐션(softmax self-attention)을 사용하는 인과적 트랜스포머(causal transformer) 아키텍처가 지배하고 있습니다. 이러한 모델은 효과적이긴 하지만, 추론 시 시퀀스 길이에 따라 메모리와 계산 요구량이 선형적으로 증가하는 문제를 겪습니다. 최근의 연구들은 소프트맥스 연산을 선형화하여, 일정한 메모리와 계산 비용을 갖는 순환 신경망(RNN) 모델(예: DeltaNet, Mamba, xLSTM)을 만드는 방식으로 이를 해결하고자 노력해 왔습니다. 이러한 모델들은 대개 그들의 순환 역학(recurrent dynamics)이 인컨텍스트 회귀 목적 함수(in-context regression objective)로부터 도출되며, 온라인 학습 규칙(예: 경사 하강법 단계)을 통해 근사적으로 최적화된다는 관점으로 통일됩니다. 그러나 기존의 접근 방식들은 일반적으로 각 타임 스텝에서 최적의 가중치를 구하는 대신 점진적인 업데이트(1차 방법론)에 의존하며, 이는 새로운 증거에 빠르게 적응하면서 과거의 정보를 유지하는 능력을 제한할 수 있습니다. 또한, 이전 세대의 최적 테스트 타임 트레이닝 레이어(예: 원래의 Mesa 레이어)는 포겟 게이트(forget gate) 사용 시의 수치적 불안정성과 병렬화의 부재로 인해 확장 가능한 학습이 불가능하다는 한계가 있었습니다.

방법론

저자들은 새롭고 수치적으로 안정적이며 청크 단위로 병렬화 가능한 Mesa 레이어를 중심으로 하는 언어 모델 아키텍처인 MesaNet을 소개합니다. 이 레이어는 매 타임 포인트에서 **국소적 최적 테스트 타임 트레이닝(locally optimal test-time training)**을 수행하도록 설계되었습니다.

Mesa 레이어

단일 경사 하강법 단계(예: DeltaNet)를 통해 상태를 업데이트하는 표준 선형 RNN과 달리, Mesa 레이어는 $\Phi_t$ 를 테스트 타임 최적화 문제의 해로 정의합니다. 구체적으로, 이는 지금까지 관찰된 모든 데이터( $e_{1:t}$ )에 대해 누적된 정규화된 제곱 오차 손실을 최소화합니다:
$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$
여기 여기서 $\zeta_{tt'}$ 는 누적 포겟 계수(cumulative forget factors)를 나타내며, $\Lambda$ 는 대칭 양의 정규화 항입니다.

출력은 $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$ 로 계산됩니다. 결정적으로, 저자들은 (포겟 게이트 사용 시 수치적 불안정성을 초래하고 병렬화가 불가능한) 재귀적 최소 제곱법(recursive least squares) 대신, 공액 경사법(Conjugate Gradient, CG) 방법을 사용하여 최적의 쿼리 응답 $q^*_t = (H_t + \Lambda)^{-1} q_t$ 를 구합니다.

병렬화: 저자들은 여러 타임 스텝에 대한 CG 방법과 게이티드 리니어 어텐션(Gated Linear Attention, GLA) 사이의 등가성을 활용합니다. 상태 변수 $G_t$ 와 $H_t$ 를 유지하고 시퀀스의 청크 단위로 선형 시스템을 병렬로 해결함으로써, 이 레이어는 훈련 시 $O(T)$ 복잡도와 추론 시 $O(1)$ 복잡도를 달성합니다.
동적 계산량(Dynamic Compute): 정지 기준(stopping criterion)에 도달하는 데 필요한 CG 반복 횟수( $k$ )는 시퀀스와 토큰마다 다릅니다. 이를 통해 모델은 추론 비용과 성능 사이의 절충안을 선택하며 테스트 타임 계산량을 동적으로 할당할 수 있습니다.

아키텍처 (MesaNet)

MesaNet은 표준 디코더 전용 트랜스포머 백본(Llama 2와 유사)을 채택하되, 시퀀스 믹싱 컴포넌트(일반적으로 Multi-Head Attention)를 Mesa 레이어로 교체합니다. 아키텍처는 다음을 포함합니다:

잔차 블록(Residual Blocks): 시퀀스 믹싱 레이어(Mesa)와 채널 믹싱 레이어(SwiGLU MLP)를 포함하는 $N$ 개의 블록을 쌓습니다.
전처리(Preprocessing): 키(Keys)와 쿼리(Queries)는 정규화되며 Mesa 레이어 이전에 작은 윈도우와 컨볼루션됩니다.
하이브리드 변형: 단기 시퀀스 모델링과 장기 시퀀스 모델링의 강점을 결합하기 위해 LRU(Hawk) 레이어와 Mesa 레이어를 교대로 사용하는 "Hawk-Mesa" 모델도 탐구되었습니다.

주요 기여

병렬화 가능하고 안정적인 Mesa 레이어: 본 논문은 원래의 공식이 가진 순차적 병목 현상과 수치적 불안정성을 극복한 버전의 Mesa 레이어를 제시합니다. 이는 효율적인 훈련을 위해 CG 솔버를 활용하며, 이를 통해 현대적 하드웨어(GPU/TPU)에서 최적의 테스트 타임 회귀라는 이론적 이점을 유지하면서도 효율적인 훈련을 가능하게 합니다.
동적 테스트 타임 계산량 할당: 이 레이어는 주어진 시퀀스의 난이도에 따라 반복 횟수가 달라지는 외부 솔버를 명시적으로 호출합니다. 이는 입력 복잡도에 따라 추론 계산량을 변화시킬 수 있는 원칙적인 메커니즘을 제공합니다.
광범-한 경험적 평가: 저자들은 SlimPajama 데이터셋을 사용하여 140M, 440M, 1B 파라미터 규모의 MesaNet을 훈련했습니다. 이들은 아키텍처, 토크나이저, 하이퍼파라미터를 통제하여 Mamba2, xLSTM, DeltaNet, Gated DeltaNet 및 Transformer와 엄격한 "애플 투 애플(apples-to-apples)" 비교를 수행했습니다.

결과

언어 모델링 성능: MesaNet은 모든 규모에서 다른 RNN 베이스라인(Mamba2, xLSTM, DeltaNet)보다 낮은 검증 퍼플렉시티(PPL)를 달eric합니다. 1B 파라미터 규모에서 MesaNet은 동일한 아키텍처를 가진 Transformer 베이스라인과 대등하거나 약간 상회하는 성능을 보입니다.
토큰 위치 분석: 토큰 위치에 따른 음의 로그 가능도(NLL)에 대한 상세 분석은 다음과 같은 뚜렷한 패턴을 보여줍니다:
- 초기 토큰: MesaNet을 포함한 대부분의 RNN은 시퀀스의 초기 토큰에서 Transformer보다 우수한 성능을 보입니다.
- 긴 컨텍스트: Transformer가 후기 토큰에서 탁월한 성능을 보이는 반면, MesaNet과 Hawk-Mesa는 다른 RNN들에 비해 (512 토큰 이상까지) 그 우위를 훨씬 더 길게 확장합니다. 다만, 매우 긴 깊이에서는 결국 Transformer에 뒤처지게 됩니다.
다운스트림 벤치마크:
- 글로벌 추론(Global Reasoning): MesaNet은 장거리 의존성이 필요한 벤치마크(예: Lambada, HellaSwag, RACE)에서 모든 다른 RNN을 능가하지만, 여전히 Transformer 베이스라인에는 미치지 못합니다.
- 인컨텍스트 회상(In-Context Recall): MesaNet은 다른 RNN들뿐만 아니라 1024 토큰 윈도우를 가진 Sliding Window Attention(SWA) 모델보다도 뛰어난 성능을 보이지만, 전체 어텐션을 사용하는 Transformer보다는 뒤처집니다.
- 퓨샷 학습(Few-Shot Learning): MesaNet은 단어 섞기(word scrambling) 및 번역 작업에서 강력한 성능을 입증하며, 종종 섞기 작업에서 Transformer를 능가합니다.
효율성: 추론 시 CG 솔버로 인한 추가 FLOPs가 필요함에도 불구하고, MesaNet은 적절한 CG 단계(예: 10-30회)를 사용할 때 H100 GPU 및 TPUv5에서 경쟁력 있는 토큰 처리량(throughput)을 유지합니다.

의의 및 주장

본 논문은 MesaNet을 신경망 내에서의 최적 테스트 타임 트레이닝을 향한 중요한 단계로 규정합니다. 저자들은 추론 시 순차적 최적화 문제(구체적으로, 최적의 선형 연상 메모리 가중치를 찾는 문제)를 해결함으로써, 모델이 단순한 경사 업데이트를 사용하는 모델보다 새로운 증거에 더 효과적으로 적응할 수 있다고 주장합니다.

이 연구는 하나의 트레이드오프를 강조합니다: 퍼플렉시티와 다운스트림 추론 성능의 향상은 추론 계산량(FLOPs)의 증가라는 비용을 수반합니다. 저자들은 이것이 성능 향상을 위해 테스트 타임 계산량을 늘리는 최근의 추세와 일치한다고 제안합니다. 그들은 Mesa 레이어가 이 계산량을 동적으로 할당할 수 있는 유연한 메커니즘을 제공하여, 네트워크 내부에서 복잡한 순차적 문제를 잠재적으로 해결할 수 있다고 주장합니다.

결론적으로, MesaNet이 현재 계산 비용 측면에서 단점이 있음에도 불구하고, 긴 컨텍스트 이해 및 글로벌 추론 작업에서 다른 RNN들을 능가하는 능력은 국소적 최적 테스트 타임 트레이닝이 효율적인 시퀀스 모델링을 위한 유망한 방향이며, 잠재적으로 RNN의 효율성과 Transformer의 성능 사이의 간극을 메울 수 있음을 시사합니다. 저자들은 현재의 구현 방식이 가진 계산 비용 문제를 인정하며, 이를 완화하기 위해 웜 스타트(warm-start) 최적화나 하이브리드 추론 전략을 활용하는 후속 연구를 제안합니다.