Titans: Learning to Memorize at Test Time
Digest (CISELQ)
- Context: Transformer의 self-attention은 전 문맥에 대한 정확한 의존성 모델링을 허용하지만 문맥 길이에 대한 2차 비용이 따르며, 선형 RNN/SSM(Mamba, RetNet, RWKV 등)은 고정 크기 hidden state에 과거를 압축해야 하므로 긴 문맥에서 세부 정보를 잃어버린다.
- Issue: 장기 문맥(100K~2M 토큰)에서 정확도와 효율성을 동시에 확보할 수 있는 범용 아키텍처가 부재하며, 특히 추론(test-time)에서도 새로운 정보를 기억하고 갱신하는 메커니즘이 부족하다.
- Solution: 저자들은 “learning to memorize at test time”이라는 관점으로 Neural Long-Term Memory Module을 제안한다. 입력 토큰이 흘러들어오면 해당 모듈의 파라미터를 test-time에 직접 gradient 업데이트하여 “놀라움(surprise)” 신호를 기반으로 기억을 저장하고, 망각 게이트로 용량을 관리한다.
- Evaluation: 언어 모델링, needle-in-a-haystack, BABILong, 상식 추론, 시계열 예측, DNA 모델링 등 다양한 영역에서 Transformer, Mamba, RetNet, DeltaNet 등과 비교한다.
- Limitations: Test-time 학습에 따른 연산·메모리 오버헤드, online 업데이트의 안정성(발산 가능성), 공개된 구현/체크포인트가 제한적이라 독립 재현이 쉽지 않다.
- Questions / Takeaways: Attention을 단기 기억, neural memory를 장기 기억으로 분리하는 이원화가 Transformer+SSM 하이브리드의 새로운 설계 원리로 자리잡을 수 있는가?
섹션별 요약
Introduction
논문은 “기억(memory)“을 기계학습 시스템의 일차 관심사로 재정의한다. Attention은 전체 문맥을 직접 조회할 수 있으나 quadratic cost로 인해 실용적 상한이 존재하고, 선형 RNN 계열은 고정 state에 정보를 압축하므로 표현력이 한계가 있다. 저자들은 인간 기억의 단기/장기 구분에서 영감을 얻어, attention을 “정확하지만 좁은 단기 기억”으로, 새로 도입하는 neural memory module을 “장기·연속 기억”으로 설정하는 이원화를 제안한다.
Methods
핵심은 세 가지 구성 요소다. (1) Neural Long-Term Memory M: 작은 MLP 형태의 파라미터 집합으로, 스트리밍 토큰에 대해 test-time에 online gradient 업데이트를 수행한다. 업데이트 크기는 associative-memory loss의 surprise signal(예측 오차의 gradient norm)로 결정되며, momentum-like term으로 과거 surprise를 누적한다. (2) Forgetting Gate: 데이터 의존적 decay 계수 α_t가 과거 파라미터를 감쇠시켜 고정 용량을 유지한다. (3) Persistent Memory: 입력에 독립적인 학습 가능한 토큰(태스크 지식)을 prefix로 붙여 전역 지식을 보존한다. 이 구성요소를 attention과 결합하는 방식으로 Titans는 세 가지 변형을 제시한다 — MAC(Memory as Context): 메모리 출력이 attention의 추가 컨텍스트로 주입됨; MAG(Memory as Gate): attention과 memory 경로를 게이트로 결합; MAL(Memory as Layer): memory를 별도 레이어로 쌓음. 전체 학습은 chunk 단위로 parallelizable하게 구현된다.
Results
저자들은 주요 벤치마크에서 Transformer, Mamba/Mamba2, RetNet, DeltaNet, GLA, Hyena 등과 비교한다. 언어 모델링(perplexity)과 상식 추론(PIQA/HellaSwag 등)에서 동급 파라미터의 Transformer++·Mamba를 상회하는 보고가 있으며, needle-in-a-haystack 실험에서는 2M 토큰 이상 문맥에서도 검색 정확도를 유지한다. BABILong(long-context reasoning)에서는 GPT-4-class 모델을 능가하는 사례도 보고된다. 시계열(ETT, traffic) 및 DNA(Genomic Benchmarks)에서도 경쟁적 결과를 제시한다.
| 영역 | 경쟁 모델 | 핵심 결과 |
|---|---|---|
| Language Modeling | Transformer++, Mamba2, DeltaNet | 동급 파라미터에서 낮은 perplexity |
| Needle-in-Haystack | GPT-4, Mamba, RecurrentGemma | 2M+ 문맥에서도 높은 검색 정확도 |
| BABILong | GPT-4, Llama3, Mamba | 장기 추론에서 SOTA급 |
| Time Series | Linear, PatchTST, iTransformer | 장기 예측에서 경쟁적 MSE |
| DNA Modeling | HyenaDNA, Mamba | 상위권 정확도 |
Discussion
저자들은 Titans가 “parametric memory” 관점에서 LSTM의 cell state, Transformer의 KV cache, SSM의 hidden state를 통합하는 프레임이라고 주장한다. Test-time 업데이트는 continual learning과 in-context learning의 중간 지점으로 해석된다. 한계로는 online 업데이트의 학습 안정성, hyperparameter 민감도, forgetting gate 튜닝이 언급된다.
Insights
- Attention을 단기 기억, 별도 neural module을 장기 기억으로 분리하는 설계는 Transformer↔SSM 이분법을 뛰어넘는 통합적 관점을 제공한다.
- “Surprise-driven” 업데이트 규칙은 Kalman gain과 meta-learning의 fast-weight 전통을 잇는 현대적 재해석이다.
- MAC/MAG/MAL 세 가지 결합 방식은 memory를 “입력·게이트·레이어” 어느 위치에 두느냐에 따라 표현력과 효율이 달라진다는 설계 축을 제시한다.
Discussion Points
- Test-time에 파라미터를 갱신한다는 것은 inference와 training의 경계가 모호해짐을 의미한다. 서빙 파이프라인에 어떤 함의가 있는가?
- Forgetting gate가 catastrophic forgetting 또는 adversarial memory injection에 얼마나 강건한가?
- Persistent memory 토큰은 soft prompt/prefix tuning과 어떤 점에서 다른가?
메타데이터
| 항목 | 값 |
|---|---|
| 제목 | Titans: Learning to Memorize at Test Time |
| 저자 | Ali Behrouz, Peilin Zhong, Vahab Mirrokni |
| 소속 | Google Research |
| arXiv | 2501.00663 |
| 공개 | 2024-12 / 2025-01 개정 |
| 분야 | Architecture, Long-Context LLM, Memory |
왜 이 연구를 하는가?
장기 문맥 처리는 RAG, 에이전트, 코드베이스 이해, 장편 스토리텔링 등 실서비스에서 점점 중요해지고 있다. 그러나 Transformer의 O(L²) 비용은 수백만 토큰 수준에서는 비현실적이고, 기존 linear-time 대안(Mamba, RetNet, Linear Attention)은 세부 검색(needle retrieval)에서 Transformer에 밀리는 경향이 있다. Titans는 “정확한 단기 조회”와 “압축된 장기 기억”이라는 상보적 장점을 한 아키텍처 안에 결합해, 양쪽의 한계를 극복하려는 시도다. 또한 test-time에 메모리를 학습한다는 발상은 continual/online learning 커뮤니티의 연구 흐름과 LLM 서빙을 연결한다.
방법 (Method)
flowchart LR X[입력 토큰 스트림] --> SA[Short-Term Memory: Attention] X --> LM[Long-Term Neural Memory M_t] LM -- surprise = grad L --> UP[파라미터 online 업데이트] UP -- forgetting gate alpha_t --> LM P[Persistent Memory prefix] --> SA SA --> COMB{결합 방식} LM --> COMB COMB -->|MAC: context로 주입| OUT[Titans 출력] COMB -->|MAG: gated sum| OUT COMB -->|MAL: stacked layer| OUT
핵심 업데이트 규칙은 도식적으로 M_t = (1 - alpha_t) * M_{t-1} - eta_t * grad_M L(M_{t-1}; k_t, v_t) + beta_t * S_{t-1} 형태이며, 여기서 L은 associative-memory loss, S_{t-1}은 momentum-like surprise 누적 항, alpha_t는 data-dependent forgetting, eta_t/beta_t는 learning-rate·모멘텀 계수다. 이 업데이트를 chunk 단위로 풀어 parallel training을 가능하게 한다.
발견
| # | 발견 | 근거 | 함의 |
|---|---|---|---|
| F1 | 2M+ 토큰 needle-in-haystack에서도 높은 검색 정확도 | 초장문맥 retrieval 실험 | 장기 문맥 에이전트의 실현 가능성 |
| F2 | 동급 파라미터에서 Mamba2/Transformer++ 대비 perplexity 개선 | LM 벤치마크 | 메모리 모듈이 표현력 보완 |
| F3 | BABILong에서 GPT-4급 장기 추론 성능 | 다단계 추론 벤치마크 | 단순 압축이 아닌 구조적 기억 효과 |
| F4 | 시계열·DNA 등 비언어 도메인에서도 경쟁적 | ETT, Genomic Benchmarks | 범용 sequence backbone 가능성 |
| F5 | MAC/MAG/MAL 변형 간 trade-off 존재 | 아키텍처 ablation | 결합 위치가 설계 자유도 축 |
이론적 의의
Titans는 “fast weights”(Schmidhuber 1992), meta-learning의 inner-loop 업데이트, 그리고 associative memory(Hopfield, Modern Hopfield Networks)의 계보를 현대적 대규모 LM 스택에 통합한다. 또한 SSM/Linear Attention이 제시한 “state = 압축된 과거”라는 관점에 “state는 test-time에도 학습된다”는 차원을 덧붙여, 모델의 parametric memory와 contextual memory 사이의 경계를 재정의한다. 이는 continual learning, RAG, in-context learning을 하나의 스펙트럼에서 보게 만드는 이론적 통합을 시도한다.
재현성 및 신뢰도 평가
| 항목 | 평가 | 비고 |
|---|---|---|
| 코드 공개 | 부분 | 공식 구현 공개 범위 제한적 |
| 데이터 | 공개 | 표준 벤치마크(LM, BABILong, ETT 등) |
| 하이퍼파라미터 | 부분 기재 | chunk size, eta·alpha 스케줄 세부 미흡 |
| 독립 재현 | 미확인 | 커뮤니티 재현 보고 초기 단계 |
| Evidence | B | 다영역 강한 실험, 외부 검증은 진행 중 |
| Reproducibility | C | 체크포인트·학습 스크립트 제한 |
관련 연구
- Linear / Sub-quadratic Attention: Linear Attention - Transformers are RNNs, Performer, Reformer, Longformer, BigBird
- State Space Models: Mamba - Linear-Time Sequence Modeling with Selective State Spaces, S4, H3, Hyena
- Retention / RNN 재해석: RetNet - Retentive Network - A Successor to Transformer for LLMs, RWKV - Reinventing RNNs for the Transformer Era, DeltaNet, GLA
- Fast Weights / Meta-Learning: Schmidhuber 1992, Ba et al. 2016, MAML
- Associative Memory: Modern Hopfield Networks (Ramsauer et al. 2021)
- Long-Context Benchmark: BABILong, RULER, Needle-in-a-Haystack
원자적 인사이트
- 단기/장기 이원화: attention을 “정확한 단기 조회”, 별도 neural module을 “압축된 장기 기억”으로 배치하면 quadratic 비용 없이 초장문맥에서도 세부 검색을 보존할 수 있다.
- Surprise-driven online update: test-time에 gradient 기반으로 memory를 갱신하는 규칙은 Kalman 갱신·fast weights의 현대적 재해석이며, inference와 learning의 경계를 재정의한다.
- 결합 위치가 설계 축: MAC/MAG/MAL은 memory를 “컨텍스트/게이트/레이어” 중 어디에 붙이느냐에 따라 표현력과 효율이 달라짐을 보여주는 체계적 축이다.
- Forgetting gate의 중심성: data-dependent alpha_t가 없으면 online 업데이트는 발산하거나 포화되므로, 망각은 장기 기억 시스템의 필수 구성요소다.
핵심 용어 정리
- Neural Long-Term Memory (M_t): 입력 키-값을 재구성하도록 test-time에 gradient 업데이트되는 소형 파라미터 모듈.
- Surprise Signal: associative-memory loss의 gradient norm으로, 예측과 실제가 크게 다를 때 memory를 강하게 갱신하도록 trigger.
- Forgetting Gate alpha_t: 과거 파라미터를 감쇠시키는 데이터 의존적 계수로, 용량 유지와 안정성을 담당.
- Persistent Memory: 입력에 독립적인 학습 가능한 prefix 토큰으로, 태스크 레벨 지식을 보존.
- MAC / MAG / MAL: 각각 Memory-as-Context, Memory-as-Gate, Memory-as-Layer 결합 변형.
- Chunk-Parallel Training: 온라인 업데이트를 chunk 단위로 풀어 병렬 학습을 가능케 하는 전략.
- Needle-in-a-Haystack: 초장문맥 안에 숨겨진 소량 정보를 정확히 검색하는 벤치마크.
태그
Architecture LongContext Attention NeuralMemory TestTimeLearning Titans Transformer SSM MetaLearning AssociativeMemory