[OUTTA Alpha팀 논문 리뷰 요약] Part 10-8. Titans: Learning to Memorize at Test Time

논문 링크: https://arxiv.org/abs/2501.00663

 

Titans: Learning to Memorize at Test Time

 

OUTTA 논문 리뷰 링크: [2025-1] Titans: Learning to Memorize at Test Time

[2025-1] Titans: Learning to Memorize at Test Time

→ 매우 자세하게 나와있음! 요약 한 번 읽고 읽으면 더 이해 잘 될듯~

---

Abstracts

• 기존 Transformer 기반 모델들은 긴 문맥을 다룰 때 연산량 증가, 메모리 한계, 정보 손실 문제를 겪음.
• Titans 모델은 Neural Memory Module을 활용하여 Test Time에서도 학습할 수 있는 새로운 메모리 구조를 제안함.
• 기존 Transformer, Linear Transformer, Modern Linear RNN (Mamba, DeltaNet, RetNet 등)과 비교하여 성능이 우수함을 입증.
• 핵심 기여점
  1. Neural Memory 설계 → 인간의 장기 기억 시스템을 모방하여 모델이 학습한 정보를 효과적으로 저장/검색.
  2. Test Time에서의 학습 가능 → 기존 Transformer는 학습된 가중치로 추론만 수행하지만, Titans는 Inference 중에도 기억을 최적화함.
  3. 기존 모델보다 효율적인 구조 → 연산량 감소, 긴 컨텍스트 처리 가능, 일반화 성능 향상.

---

1. Introduction

1.1 기존 Transformer의 한계

• Transformer의 Self-Attention 메커니즘은 문맥(Context)을 기반으로 정보를 검색하고,
  • 연관 기억(Associative Memory)처럼 작동함.
  • Query-Key 유사도를 활용하여 적절한 Value를 선택함.
• 문제점
  • 메모리 사용량이 O(N²)으로 증가 → 긴 문맥을 다룰 때 비효율적.
  • 고정된 컨텍스트 윈도우 → 이전 정보를 저장할 수 없음.
  • Long-term Memory 부재 → 과거 정보를 유지하는 기법이 없음.

1.2 기존 해결책과 한계

1. Linear Transformers
   • Attention의 softmax를 Kernel Function으로 대체하여 메모리 사용량을 줄임.
   • 결과적으로 선형 연산(O(N))으로 처리 가능.
   • 하지만 Transformer보다 성능이 낮음.
2. Linear RNN vs. Linear Transformer
   • Linear RNN: 벡터 기반 메모리 사용 → 과거 정보를 압축하여 유지.
   • Linear Transformer: Key-Value 기반 행렬 형태로 데이터를 저장.
   • 하지만 기억을 효과적으로 업데이트하는 메커니즘이 부족함.

---

2. Preliminaries

2.1 Transformer의 Self-Attention 수식

Transformer의 Self-Attention 연산:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$​ $$A = \text{softmax} \left(\frac{Q K^T}{\sqrt{d_k}} \right) V$$

• 문제점:
  • 모든 토큰이 상호작용하므로 연산량 O(N²).
  • 컨텍스트 윈도우가 제한됨 → 긴 문맥에서 정보 손실 발생.

2.2 Linear Attention

• Attention을 Kernel Function $\phi(x)$으로 근사하여 선형 연산으로 변환: $$\phi(Q^T K) = \phi(Q) \phi(K)$$
• 이를 적용하면 연산량이 O(N)로 감소하지만,
  • 정보 압축 과정에서 손실 발생
  • 완전한 메모리 유지가 어려움

2.3 RNN 기반 Memory 구조

• 기본적인 메모리 업데이트$$M_t = f(M_{t-1}, x_t)$$
• Transformer vs. Linear Transformer
  • Transformer는 모든 과거 데이터를 보존하지만,
  • Linear Transformer는 고정된 크기의 행렬로 데이터를 압축하여 저장.
    • 기억의 유연성이 부족함.

---

3. Titans: Test Time에서 학습하는 Neural Memory

3.1 Neural Memory 구조

• Test Time에서도 지속적으로 학습, 업데이트할 수 있는 메모리 설계.
• Surprise Metric을 도입하여 기억할 정보를 선택적으로 업데이트.

3.1.1 Surprise Metric (예상치 못한 정보 강화)

• 사람이 예상 밖의 사건을 더 잘 기억하는 것처럼, 모델도 "서프라이즈"를 기반으로 학습.
• 메모리 업데이트 수식: $$M_t = M_{t-1} - \theta_t \nabla \ell(M_{t-1}; x_t)$$
  • $\theta_t$ = Learning Rate 역할
  • $\nabla \ell(M_{t-1}; x_t)$ = Surprise Score
  • Surprise가 클수록 기억을 더 강하게 유지

3.2 Forgetting Mechanism (망각 메커니즘)

• 불필요한 정보를 제거하여 메모리를 최적화
• 망각 수식: $$M_t = (1 - \alpha_t) M_{t-1} + S_t$$​
  • $\alpha_t \to 1$ → 메모리 삭제
  • $\alpha_t \to 0$ → 기존 정보 유지

3.3 Persistent Memory (고정 메모리 저장)

• Task-specific 정보를 저장하는 고정된 메모리 공간 추가
• Attention에서 Key-Value 행렬처럼 작동: $$x_{\text{new}} = [p_1, p_2, \dots, p_{N_p}] \| x$$

---

4. Titans Architecture

Titans는 Attention과 Neural Memory를 결합한 새로운 아키텍처로, 세 가지 변형(variants)을 제안.

4.1 Memory as Context (MAC)

• 현재 컨텍스트 + 과거 정보를 Attention에 입력.

$$h_t = M_{t-1}(q_t)$$ $$\tilde{S}(t) = [p_1, \dots, p_{N_p}] \| h_t \| S(t)$$ $$y_t = \text{Attn}(\tilde{S}(t))$$

4.2 Memory as Gating (MAG)

• Sliding Window Attention (SWA)과 Neural Memory를 결합.

$$y = \text{SWA}^*(x)$$ $$o = y \otimes M(x)$$

4.3 Memory as a Layer (MAL)

• Neural Memory를 모델의 하나의 Layer로 사용.

$$y = M(x)$$ $$o = \text{SWA}(y)$$

---

5. Experiments

5.1 Language Modeling

• Titans (MAC, MAG, MAL) → Transformer 및 기존 순환 모델보다 성능 우수

5.2 Needle in a Haystack (NIAH)

• 긴 문서에서 특정 정보를 검색하는 태스크
• Titans는 2M+ 토큰 문맥에서도 성능 유지

5.3 BABILong Benchmark

• GPT-4, Llama3.1-8B보다 성능 우수
• Titans는 Retrieval-Augmented Generation (RAG) 적용 모델보다도 뛰어남.

---

6. Conclusions

• Titans는 Test Time에서도 학습 가능한 Neural Memory를 제안.
• 긴 문맥을 다루는 새로운 접근법으로 Transformer 대비 효율성 증가.
• MAC이 가장 높은 성능을 보이며, Forget Mechanism이 핵심 역할.

---

Titans 논문 Q&A

---

Q1. Titans 모델은 긴 문맥 처리에서 어떻게 성능을 개선했는가?

• 기존 Transformer 한계
  • Attention 연산량 O(N²) → 긴 문맥에서 비효율적
  • 고정된 컨텍스트 윈도우 → 과거 정보 유지 어려움
• Titans의 성능 개선 요소
  1. Neural Memory 결합 → Transformer와 달리 이전 정보를 유지하며 학습
  2. Long-term Memory 적용 → 2M+ 토큰을 유지하면서 성능 유지 가능
  3. Test Time에서 학습 가능 → Inference 중에도 정보 업데이트
  4. Adaptive Forgetting 도입 → 불필요한 정보 삭제하여 효율적인 기억 관리
• 실제 성능 개선 사례
  • Needle in a Haystack (NIAH) Task:
    • 긴 문서에서 특정 정보를 검색하는 문제
    • 기존 Transformer는 컨텍스트 길이가 4K+ 토큰을 넘으면 성능 저하
    • Titans는 16K+ 토큰에서도 성능 유지
  • BABILong Benchmark
    • GPT-4, Llama3-8B보다 긴 컨텍스트에서 정확도 높음
    • Retrieval-Augmented Generation (RAG) 모델보다 효율적

---

Q2. ‘놀라움(Surprise)’ 지표가 메모리 효율성에 어떻게 기여하는가?

• 기존 Transformer 문제
  • 모든 데이터를 동일한 중요도로 저장 → 불필요한 정보까지 유지
  • 긴 문맥을 처리할 때 메모리 낭비
• Surprise Metric 도입 효과
  1. 중요한 정보만 선택적으로 저장
     • 기본 수식: $$S_t = \eta_t S_{t-1} - \theta_t \nabla \ell(M_{t-1}; x_t)$$
     • Surprise Score가 높을수록 더 강하게 기억
  2. 메모리 사용량 최적화
     • 중요하지 않은 정보는 빠르게 제거
     • 장기 기억(Long-term Memory) 유지 비용 감소
  3. Inference 중에도 지속적인 정보 업데이트 가능
     • 새로운 정보가 등장하면 기존 기억을 수정

---

Q3. Titans의 메모리 모듈 간 상호작용과 역할은?

• 메모리 구조
  • Core (Short-term Memory): 현재 컨텍스트 내에서 정보 저장
  • Long-term Memory: 과거 정보 저장 및 검색
  • Persistent Memory: 학습된 지식을 영구적으로 유지
• 모듈 간 상호작용
  1. 단기 기억 (Core) → 장기 기억 (Long-term Memory)
     • 최근 정보를 임시 저장 → Surprise Score 기반으로 장기 기억으로 이동
  2. 장기 기억 (Long-term Memory) ↔ Persistent Memory
     • 자주 사용되는 정보는 Persistent Memory에 저장
     • 특정 태스크에서 공통적으로 활용할 수 있는 지식 보존
• 작업별 메모리 기여 예시
  • 언어 모델링
    • Core Memory → 현재 문맥 유지
    • Long-term Memory → 과거 문맥 활용하여 일관성 유지
  • 질문 응답 (QA)
    • Long-term Memory → 긴 문서에서 답변 검색
    • Persistent Memory → 일반적인 사실(Fact) 저장
  • 시계열 예측 (Time Series Forecasting)
    • Core Memory → 최근 데이터 반영
    • Long-term Memory → 과거 패턴 학습

---

Q4. Titans 모델은 긴 문맥 처리 시 계산 복잡도와 메모리 사용량을 어떻게 최적화하는가?

• 기존 Transformer 문제점
  • Self-Attention 연산량 O(N²) → 시퀀스 길이가 길어질수록 비효율적
  • 모든 데이터를 메모리에 저장 → 메모리 사용량 급증
• Titans의 최적화 기법
  1. Neural Memory 적용 → 연산량 감소
     • Transformer는 모든 토큰이 Attention 연산 수행
     • Titans는 핵심 정보만 저장하고 검색 → O(N) 연산 가능
  2. Adaptive Forgetting 도입 → 메모리 최적화
     • 불필요한 정보를 삭제하여 메모리 소비 절약
  3. Long-term Memory + Persistent Memory 구조 적용
     • Transformer는 모든 데이터를 유지하지만, Titans는 필요한 정보만 유지
  4. 실제 메모리 사용량 비교
     • 기존 Transformer (GPT-4) 16K+ 토큰 처리 시 80GB+ VRAM 필요
     • Titans는 2M+ 토큰을 50GB VRAM으로 처리 가능

---

Q5. 적응형 망각(Adaptive Forgetting)이 모델 학습 및 추론에서 수행하는 역할과 한계는?

• 기존 Transformer 문제점
  • 모든 정보를 동일한 중요도로 저장 → 비효율적
  • 불필요한 정보까지 유지하여 메모리 낭비
• Adaptive Forgetting 역할
  1. 중요한 정보만 유지하고 불필요한 정보 삭제
     • Forget Mechanism 수식: $$M_t = (1 - \alpha_t) M_{t-1} + S_t$$​
       • $\alpha_t \to 1$ → 과거 기억 삭제
       • $\alpha_t \to 0$ → 기존 정보 유지
  2. Test Time에서도 지속적으로 학습 가능
     • 기존 모델은 추론(Inference) 시 기억을 업데이트할 수 없음
     • Titans는 새로운 정보를 반영 가능
  3. 메모리 사용량 최적화
     • Transformer는 모든 데이터를 저장하지만, Titans는 필요한 정보만 유지
• Adaptive Forgetting의 한계점
  1. 잘못된 정보 삭제 가능성
     • 중요한 정보가 낮은 Surprise Score를 가질 경우 삭제될 위험
  2. 추론 속도 저하 가능성
     • 과거 정보를 계속 업데이트해야 하므로, 모델이 추론 시 추가 연산 수행
  3. 최적의 Forgetting Rate 설정이 필요
     • $\alpha_t$ 값이 과도하게 높으면 정보 손실 발생
     • $\alpha_t$ 값이 낮으면 불필요한 정보 유지

---

최종 요약

Titans는 Neural Memory를 결합하여 Transformer보다 긴 문맥을 효과적으로 처리함
Surprise Score 기반 선택적 기억 유지로 메모리 사용량 최적화
Core, Long-term, Persistent Memory 간의 상호작용으로 성능 향상
Self-Attention의 O(N²) 연산을 O(N) 수준으로 줄이는 최적화 구조 적용
Adaptive Forgetting을 통해 Test Time에서도 학습 가능하지만, 정보 손실 가능성이 존재