Learning to (Learn at Test Time): RNNs with Expressive Hidden States

Self-attention performs well in long context but has quadratic complexity. Existing RNN layers have linear complexity, but their performance in long context is limited by the expressive power of their hidden states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. The key idea is to make the hidden state a machine learning model itself, and the update rule a step of self-supervised learning. Since the hidden state is updated by training even on test sequences, our layers are called Test-Time Training (TTT) layers. We consider two instantiations: TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B parameters, comparing with a strong Transformer and Mamba, a modern RNN. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research.

한 줄 요약

TTT 레이어는 테스트 시점에 학습하는 RNN 구조로, 긴 문맥에서 성능을 유지하면서도 선형 복잡도를 달성한다.

핵심 기여도

TTT 레이어는 히든 상태를 자체 학습 모델로 구현하여, Mamba와 같은 기존 RNN의 긴 문맥 성능 한계를 극복.
TTT-Linear와 TTT-MLP 두 가지 인스턴스를 제안, TTT-MLP는 16k 이상의 긴 문맥에서 성능을 지속적으로 개선.
125M에서 1.3B 파라미터 규모로 실험, Mamba와 Transformer와 비교하여 선형 복잡도를 유지하면서도 성능 경쟁력 보임.
Mamba는 16k 문맥 이후 성능이 더 이상 개선되지 않는 반면, TTT-MLP는 계속 퍼플렉시티가 감소함.

핵심 아이디어

기존 RNN은 고정 크기의 히든 상태로 문맥을 압축하기 때문에 긴 시퀀스에서 정보 손실이 발생한다. 이에 반해, TTT 레이어는 히든 상태 자체를 머신러닝 모델로 구현하고, 테스트 시점에서도 자기 감독 학습을 통해 업데이트하도록 설계했다. 이는 히든 상태가 시퀀스의 구조와 관계를 학습적으로 압축할 수 있게 해준다. TTT-Linear는 선형 모델, TTT-MLP는 2층 MLP를 히든 상태로 사용하며, 이는 각각 다른 표현력과 성능을 보인다. 핵심 아이디어는 "학습을 (테스트 시점에) 학습하는 것"으로, 이는 기존 네트워크 아키텍처 개념을 재정의하는 새로운 접근법이다.

기술적 접근법

**TTT 레이어**: 히든 상태를 자체 학습 모델로 구현, 업데이트 규칙은 자기 감독 학습 단계.
**TTT-Linear**: 히든 상태는 선형 모델.
**TTT-MLP**: 히든 상태는 2층 MLP.
**실험 환경**: EasyLM 기반, 125M~1.3B 파라미터 규모.
**비교 대상**: Transformer (TF finetune), Mamba.
**FLOPs 일치 조건**: TTT-Linear와 TTT-MLP는 Mamba와 동일한 FLOPs로 훈련 및 추론 수행.
**Transformer의 FLOPs**: TTT-Linear/MLP 대비 2.8× 더 높음.

주요 결과

**16k 이상 문맥에서**: TTT-MLP는 퍼플렉시티가 계속 감소하지만, Mamba는 감소하지 않음.
**32k 문맥에서**: TTT-MLP는 Transformer (TF finetune)보다 낮은 퍼플렉시티를 보임.
**FLOPs 일치 조건**: TTT-Linear와 TTT-MLP는 Mamba와 동일한 FLOPs로 훈련됨.
**TTT-MLP**: 16k 이상에서 Mamba 대비 +10% 이상의 퍼플렉시티 개선.

의의 및 한계

TTT 레이어는 RNN의 선형 복잡도 장점을 유지하면서, 긴 문맥에서의 표현력 한계를 해결하는 새로운 방향을 제시한다. 특히 TTT-MLP는 MLP 기반 히든 상태로 더 복잡한 관계를 학습할 수 있어, Mamba와 같은 기존 RNN보다 긴 문맥에서 우수한 성능을 보인다. 그러나 TTT-MLP는 메모리 I/O 문제를 여전히 겪고 있어, 효율적인 구현이 필요하다. 또한, TTT-MLP는 짧은 문맥에서는 TTT-Linear보다 성능이 떨어지는 경향이 있어, 모델 복잡도와 문맥 길이 간의 균형이 중요하다.

실용적 활용

TTT 레이어는 대규모 언어 모델에서 긴 문맥 처리가 필요한 상황, 예를 들어 문서 요약, 대화 시스템, 코드 생성 등에 적용 가능하다. 특히, Mamba와 같은 기존 RNN이 한계를 보이는 16k 이상의 긴 문맥에서 유용하며, 선형 복잡도를 유지하면서도 Transformer의 성능에 근접할 수 있어, 실용적이고 효율적인 대안이 될 수 있다.