본문으로 건너뛰기
🐨코알라 오딧세이
🐨
이다운
코알라 오딧세이

새 글 알림 받기

GitHub

© 2026 이다운

글 목록
AI/LLM

Titans: Learning to Memorize at Test Time

트랜스포머에 신경 기억 레이어를 추가하여 장단기 기억을 모두 효율적으로 처리하는 Google의 Titans 아키텍처를 살펴봅니다.

Titans: Learning to Memorize at Test Time

AI/LLM 시리즈 : Titans: Learning to Memorize at Test Time

GPT, Claude 등 현재의 LLM들은 많은 장점을 가지고 있지만, 완벽하다고 보기는 아직 어렵습니다. LLM의 대표적인 단점 중 하나는 입력 Context의 길이가 증가할수록 계산 및 메모리 비용이 기하급수적으로 증가한다는 점입니다.

이를 해결하기 위해 Mamba 같은 모델이 주목받기도 했지만, 문맥을 압축하는 과정에서 중요한 세부 정보를 놓친다는 지적도 있습니다. 최근 Google은 이 문제를 근본적으로 해결하기 위해 트랜스포머 기반 LLM에 신경 기억(neural memory) 레이어를 추가한 새로운 아키텍처 Titans를 제안했습니다.

기존 모델의 한계

논문에서는 현재 LLM의 문제를 두 가지로 정리합니다.

순환 모델(Recurrent Models) 은 데이터를 고정 크기의 히든 상태로 압축하여 저장하려 하지만, 모든 데이터를 효과적으로 기억하기 어렵습니다.

주의 메커니즘(Attention) 은 전체 문맥 창을 참조하며 토큰 간 직접적인 의존성을 잘 모델링하지만, 이차적 계산 복잡도로 인해 문맥 창 크기를 고정해야 하는 한계가 있습니다.

softmax 대신 커널 함수를 적용한 Linear Transformer 변형들이 메모리 사용량을 줄이려 했으나, 데이터를 행렬 형태로 압축하는 선형 순환 네트워크가 되어 기존 트랜스포머보다 성능이 뒤처지는 문제가 남았습니다. 즉, 확장성과 성능 사이의 트레이드오프가 여전히 해결되지 않은 것입니다.

Titans가 해결하고자 한 다섯 가지 질문

  1. 효과적인 기억 구조는 무엇인가?
  2. 적절한 기억 업데이트 메커니즘은 무엇인가?
  3. 효율적인 기억 검색 방법은 무엇인가?
  4. 서로 연결된 다양한 기억 모듈을 어떻게 효과적으로 통합할 것인가?
  5. 긴 과거를 저장하기 위해 깊은(딥) 기억 모듈이 필요한가?

이를 해결하기 위해 Titans는 단기 기억, 장기 기억, 지속 기억(Persistent Memory) 세 가지 기억 체계를 통합한 아키텍처를 제안합니다.

핵심 구성 요소

1. 장기 기억 메모리 (Long-term Neural Memory)

단순히 학습 데이터를 암기하는 것이 아니라, 테스트 시점에서 새로운 데이터에 대해 스스로 기억하고 불필요한 정보를 잊는 메타-모델을 구축하는 것이 핵심입니다.

(1) 서프라이즈 메트릭

인간은 예상과 다른 사건을 더 강하게 기억합니다. 이에 착안하여 입력에 대한 기울기(gradient)를 서프라이즈 지표로 정의합니다. 기울기가 클수록 이전 데이터와 크게 다르므로 기억할 가치가 높다고 판단합니다. 과거 서프라이즈(Past Surprise)와 순간 서프라이즈(Momentary Surprise)를 모멘텀 개념과 결합하여 업데이트합니다.

  • ηt\eta_t: 이전 서프라이즈의 영향을 얼마나 유지할지 조절하는 감쇠(decay) 계수
  • θt\theta_t: 현재 입력의 서프라이즈를 얼마나 반영할지 결정하는 계수

(2) 목표 함수 및 키-값 연관 학습

입력 xtx_t에 대해 두 개의 선형 층으로 키 kt=xtWKk_t = x_t W_K와 값 vt=xtWVv_t = x_t W_V를 생성합니다. 손실 함수는 메모리 Mt−1M_{t-1}이 키 ktk_t에 대해 예측한 값과 실제 값 vtv_t 간의 L2 노름으로 측정됩니다.

ℓ(Mt−1;xt)=∥Mt−1(kt)−vt∥22\ell(M_{t-1}; x_t) = \| M_{t-1}(k_t) - v_t \|_2^2

이 inner-loop 최적화를 통해 메모리 모듈은 테스트 시점에서 키와 값 사이의 연관성을 학습합니다.

(3) 망각(Forgetting) 메커니즘

수백만 토큰에 달하는 긴 시퀀스에서 모든 정보를 유지하면 메모리 한계와 성능 저하가 발생합니다. 이를 방지하기 위해 forgetting gate 역할을 하는 계수 αt\alpha_t를 도입합니다.

Mt=(1−αt)Mt−1+StM_t = (1 - \alpha_t) M_{t-1} + S_t

αt\alpha_t 값에 따라 이전 메모리의 일부 혹은 전체를 선택적으로 잊을 수 있습니다.

(4) 메모리 아키텍처 및 정보 검색

메모리 구조로는 간단한 MLP(다층 퍼셉트론)를 사용하며, 깊은 메모리 모듈(2층 이상)이 선형 모델보다 이론적으로 표현력이 뛰어납니다. 정보 검색 시에는 추가적인 가중치 업데이트 없이 순전파만으로 관련 정보를 추출합니다.

yt=M∗(qt)y_t = M^*(q_t)

2. 장기 메모리 학습의 병렬화

장기 메모리 모듈은 이론상 전체 시퀀스 길이 NN에 대해 O(N)O(N) FLOPs가 필요합니다. GPU/TPU의 병렬 처리 능력을 활용하기 위해 다음과 같은 기법을 적용합니다.

  • 미니배치 기반 업데이트: 시퀀스를 크기 bb의 청크로 나누어 행렬 곱셈과 덧셈만으로 업데이트를 구현
  • 병렬 연관 스캔: 모멘텀 항을 선형 순환 형태로 나타내고, parallel associative scan으로 각 청크 내 모든 StS_t를 동시에 계산
  • 청크 단위 파라미터 설정: 토큰별 파라미터를 청크 단위 상수로 설정하면 저장 비용이 줄고, 전역 합성곱(global convolution)으로도 계산 가능

이를 종합한 결과가 MAC(Memory as a Context) 아키텍처입니다. MAC는 코어(Core) 브랜치와 지속적(Persistent) 메모리 브랜치를 포함하여 인컨텍스트 학습과 작업별 지식 저장 기능을 함께 제공합니다.

3. Persistent Memory (지속 기억)

장기 기억이 입력에 의존하는 반면, Persistent Memory는 입력과 무관하게 학습 가능한 파라미터 집합으로 작업 관련 메타 지식을 저장합니다. 모델이 현재 문맥뿐 아니라 과거에 축적된 추상적 지식을 활용할 수 있도록 합니다.

학습 가능한 파라미터 P=[p1,p2,…,pNp]P = [p_1, p_2, \ldots, p_{N_p}]를 시퀀스 앞에 연결(Concatenation)하여 입력을 확장합니다.

xnew=[p1,p2,…,pNp]∥xx_{\text{new}} = [p_1, p_2, \ldots, p_{N_p}] \parallel x

신경 기억 모듈의 통합 방식

설계된 장기 기억 모듈을 아키텍처에 통합하는 세 가지 방식을 제안합니다.

(1) Memory as a Context (MAC)

긴 시퀀스를 세그먼트로 분할하고, 이전 장기 메모리에서 관련 정보를 검색하여 현재 세그먼트에 연결한 후 어텐션 모듈에 전달합니다.

S~(t)=[p1,…,pNp]∥ht∥S(t)\tilde{S}^{(t)} = [p_1, \ldots, p_{N_p}] \parallel h_t \parallel S^{(t)}

테스트 시 지속 메모리는 고정되고, 어텐션은 인컨텍스트 학습을 수행하며, 장기 메모리는 계속 업데이트됩니다.

(2) Gated Memory (MAG)

두 개의 분기를 사용합니다. 한 분기에서는 장기 메모리를 업데이트하고, 다른 분기에서는 Sliding Window Attention(SWA)을 수행합니다. 두 결과를 비선형 게이팅으로 결합하여, SWA가 단기 기억을, 신경 기억 모듈이 장기 기억을 각각 담당합니다.

(3) Memory as a Layer (MAL)

입력에 지속 메모리 파라미터를 연결한 후 신경 기억 모듈에 전달하고, 그 출력을 Sliding Window Attention에 전달하여 최종 출력을 생성합니다. 각 층 성능에 의존하기 때문에 어텐션과 기억 모듈의 상호 보완적 처리 능력을 온전히 활용하지 못할 수 있다는 단점이 있습니다.

구현 세부사항 및 이론적 우수성

  • 모든 블록에 Residual Connection 사용
  • SiLU 활성화 함수와 ℓ2\ell_2-노름으로 쿼리·키 정규화
  • Query, Key, Value 투영 후 1D Depthwise-Separable Convolution 적용
  • 최종 출력 전 Linear Layer를 이용한 게이팅 및 정규화

Theorem 4.1: 기존 Transformer, 대각선 선형 순환 모델, DeltaNet 등이 TC0(상수 깊이 회로)로 표현력이 제한되는 반면, Titans는 TC0를 넘어선 문제 해결이 가능하여 이론적으로 더 높은 표현력을 가짐을 증명합니다.

성능 및 결론

Titans는 언어 모델링, 상식 추론, 시계열 예측, DNA 모델링 등 다양한 작업에서 기존 Transformer와 순환 모델을 능가하는 성능을 보였습니다. 특히 2M 토큰 이상의 초장기 문맥에서도 안정적인 정확도와 효율적인 메모리 관리를 달성했습니다.

단기 기억(어텐션), 장기 기억(신경 기억 모듈), 지속 기억(Persistent Memory)이라는 세 가지 기억 체계를 유기적으로 통합한 Titans는, 긴 문맥 처리라는 LLM의 고질적인 한계를 근본적으로 해결하려는 시도라는 점에서 앞으로의 연구 방향에 큰 영향을 줄 것으로 기대됩니다.

관련 글

댓글

0 / 1000