[Paper Review] Titans 논문 이해하기

『Titans: Learning to Memorize at Test Time. arXiv preprint. 2024.』

 

2025년 1월 인공지능 학계를 떠들썩하게 만든 2가지 연구 결과가 발표 되었다. 그 중 하나는 대중적으로 이미 많이 알려진 DeepSeek이고 다른 하나는 상대적으로 조용하지만 강한 파급력을 가지고 있는 Titans 모델이다. 그 중 이번에는 Transformer의 대체제라는 이야기가 나오고 있는 Titans 모델에 대해 정리해 보려 한다. 기존 RNN이나 Transformer 모델들에 비해 어떤 장점을 가지는지를 위주로 논문을 정리하였다.

 

NLP 분야가 메인 연구 분야는 아니기에 실험 내용을 자세히 파악하기보다는 새로운 방법론을 파악하는 것을 위주로 논문을 읽어 보았다.

---

0. Abstract

지금까지 10년 이상 RNN과 Attention에 대한 연구가 많이 이루어져 왔지만 아직까지 정확한 의존성을 파악하려하면 할 수록 계산 비용은 Quadratic하게 증가하는 문제가 있어 결국 모델이 다루는 Context 길이는 제한되어 사용된다.

• RNN: 데이터를 고정된 크기의 메모리(hidden state)로 잘 압축하는 것이 목표
• Attention: 전체 Context Window를 참고하여 모든 토큰 간의 의존성을 파악하는 것이 목표

본 논문에서 제안하는 Titans라는 새로운 신경망 아키텍처는 기존 Attention을 활용한 단기 메모리 모듈(Short-term memory module)로, 새롭게 제안하는 장기 메모리 모듈(Long-term memory module)을 기반으로 구성되어 있다.

 

Attention은 짧은 Context 안에서는 정확한 의존성을 모델링하지만 Context의 길이가 제한되어 있기 때문에 장기 기억에 불리하다는 특징을 가지는데, 본 논문에서는 과거의 문맥 정보를 학습하여 기억할 수 있는 새로운 장기 메모리 모듈을 같이 활용하여 Attention이 현재 문맥을 집중적으로 참고하면서도 먼 과거의 정보를 활용할 수 있도록 돕는다. 이 신경망 메모리는 병렬적으로 빠르게 학습할 수 있으며, 동시에 빠른 추론 속도를 유지하는 장점이 있다. 

 

이 두 가지 모듈을 기반으로 메모리를 효과적으로 통합하는 세 가지 변형 모델을 같이 소개하며, 다양한 작업에서 Titans 모델이 기존 Transformer 및 최신 Linear RNN 모델보다 더 효과적인 성능을 보인 것을 확인하였다.

 

참고로 Titans는 Context Window 크기를 200만 이상으로 확장할 수 있다고 한다.

 

1. Introduction

Transformer 모델은 Context 안에서 학습할 수 있는 능력과 대규모 학습이 가능하다는 점 때문에 Sequence Modeling에서 강력한 성능을 보일 수 있었다. Transformer의 핵심 구성 요소인 Attention은 Key-Value의 관계를 학습하고 이를 기반으로 Query-Key의 유사도를 계산하여 정보를 잘 검색하는 것을 목표로 하는데, 이러한 의존성을 정확하게 모델링하는데 있어 Context 길이에 따라 연산 시간과 메모리 요구량이 Quadratic하게 증가하는 문제가 발생한다.

 

Language Modeling, Video Understanding, Long-term Time Series Forcasting 등과 같이 최근 복잡해진 응용 Task에서 큰 Context Window를 필요로 하기 때문에 Transformer의 이러한 문제점이 드러나고 있다. 물론 이러한 Transformer의 Scalability 문제를 해결하기 위해 최근 다양한 Linear Transformer 연구가 이루어지고 있다. Linear Transformer는 기존 Attention에서 사용한 Softmax 연산을 Kernel Function으로 대체하여 메모리 사용량을 크게 줄여 연산 효율성을 줄이고 긴 Context를 처리할 수 있다는 장점을 가진다. 하지만 기존 Transformer만큼의 경쟁력 있는 성능이 나오지 않는다는 한계점이 있어 이또한 완벽한 대체재가 될 수는 없다.

Linear Transformer에서 Transformer만큼 성능이 나오지 못하는 이유
: Kernel Trick을 사용하게 되면 모델이 Linear RNN처럼 동작하게 되는데 이 과정에서 결국 매우 긴 Context를 작은 Vector 형태로 압축하게 되어 정보의 손실이 발생할 가능성이 높다.

 

연산 효율성 문제뿐만 아니라 기존에 사용되었던 모델들(LSTM, Transformer 등)은 보다 근본적인 한계를 가진다. 이들은 인간의 두뇌에서 영감을 받아 설계되었지만 단기 기억, 장기 기억, 메타 기억, 현재 Context에 집중하는 능력이 독립적으로 존재하면서 상호 연결될 수 있는 시스템이 결여되어 있다.

 

즉, 본 논문에서는 기존 모델들이 과거 데이터를 학습해 추상화하여 기억하는 능력이 부족하며, 인간의 두뇌처럼 특정한 기능을 담당하는 개별 모듈들을 가진 채 이를 연결하는 시스템을 구성하는 것이 필요하다고 말하고 있다.

 

Memory Perspective

Memory는 인간의 학습과정에 있어 필수적이고 떼놓을 수 없는 요소로 Hopfield Networks, LSTM, Transformer 등 지금까지 머신러닝 연구에서 메모리 시스템은 중요한 영감을 제공해 왔다. 기존 아키텍처들은 Memory를 "Input에 의해 발생하는 Neural Update"라고 보았고, 학습(Learning)은 "목적에 따라 효과적이고 유용한 Memory를 획득하는 과정"으로 정의하였다. 

 

이러한 관점에서 RNN과 Transformer를 살펴보면,

 

1. RNN

- 벡터 기반 메모리(Vector-valued Memory)

- 새로운 입력 $x_t$를 받으면 이전 Memory $M_{t-1}$을 사용하여 새로운 State를 압축된 형태로 업데이트($f(M_{t-1}, x_t)$)

- 입력에 대한 Memory를 검색하는 과정이 수행되며, 이 때 $g(M_t, x_t)$ 함수를 사용한다.

 

2. Transformer

- 확장 가능한 메모리(Growing Memory)

- Transformer의 Memory는 Key와 Value 쌍으로 표현되며 다음의 단계를 따른다.

(1) 새로운 Key와 Value를 압축없이 Memory에 추가하면서 Memory를 업데이트

(2) Query 벡터를 사용하여 Key 벡터들과의 유사도를 계산하고, 그 결과를 바탕으로 찾은 Value 벡터에 Weight를 주어 최종 출력을 생성

 

3. Linear Transformer

- 고정된 크기의 행렬 메모리(Matrix-valued Memory)

- 과거 데이터를 압축하여 고정된 크기의 행렬에 저장

- 압축 방식 때문에 Transformer와 달리 정보에 손실이 발생할 수 있음(Scalability는 높지만 정보 손실 가능성이 있다는 것)

 

결국 논문의 저자들은 효과적인 메모리 구조를 설계하기 위해 다음과 같은 중요한 질문들을 제시한다:

 

Q1: 좋은 메모리 구조란 무엇인가?

(What constitutes a good structure for the memory?)

 

Q2: 적절한 메모리 업데이트 매커니즘은 무엇인가?

(What is a proper memory update mechanism?)

 

Q3: 효율적인 메모리 검색 방법은 무엇인가?

(What is a good memory retrieval process?)

 

Q4: 서로 연결된 메모리 모듈들을 효과적으로 결합한 아키텍처를 어떻게 설계할 것인가?

(How to design and efficient architecture that incorporates different interconnected memory modules?)

 

Q5: 과거의 정보를 효과적으로 저장하고 기억하기 위해 더 깊은 메모리 모듈이 필요할까?

(Is a deep memory module needed to effectively store/remember long past?)

 

Contributions and Roadmap

본 논문에서는 앞서 제시한 5가지 핵심 질문에 답하기 위해 Long-term neural memory module을 설계하고, 이를 효율적으로 학습하여 Test time에 memory를 형성하는 법을 연구하였다. 

 

이 장기 기억 모듈을 설계하고, 아키텍처에 통합하는 과정은 다음과 같다.

 

Neural Memory

본 논문에서 제안하는 장기 기억 모듈은 신경망을 기반으로 하여 Test시 데이터를 Parameter에 저장하고 기억하는 방법을 학습한다(마치 메타러닝의 Meta in-context 모델처럼).

 

이 모델은 인간의 장기 기억 시스템에서 영감을 받았으며, 사람이 예측하기 어려운 정보(Surprising Event)가 더 잘 기억된다는 점을 반영하여 설계하였다. 이에 Associative memory Loss를 사용하여 입력에 대한 Surprise 정도를 판단하였으며, 그 정도는 입력에 대한 신경망의 Gradient를 통해 측정하였다.

 

또한 Decay 매커니즘을 도입하여 Memory 크기와 Surprise 정도를 고려한 동적 업데이트를 수행하였다. 이는 RNN 모델에서 사용되는 Forgetting 매커니즘을 일반화한 것으로 볼 수 있으며 Meta Neural Network의 최적화 과정과 동일하게 일어난다(ex. Mini-batch gradient, Momentum, Weight Decay). 

 

Titans Architectures

장기 신경망 메모리를 설계하였으니 이제 남은 질문은 "어떻게 이 메모리를 딥러닝 아키텍처에 효과적으로 통합할 것인가?"에 대한 것이다. 이에 대해 본 논문에서는 Titans라는 새로운 딥러닝 모델을 제안하며, 다음의 3가지 핵심 모듈로 구성된다.

1. Core(핵심 모듈)
   • 단기 기억 역할
   • 제한된 크기의 Attention을 사용하여 데이터의 주요 흐름을 처리
2. Long-term Memory(장기 기억 모듈)
   • 장기 기억을 담당
   • 과거의 데이터를 저장하고 기억
3. Persistent Memory(영속적 기억 모듈)
   • 특정 Task와 관련된 고정된 파라미터를 저장
   • 기존의 기억 모듈과 달리 특정 데이터에 의존하지 않고 지식을 저장

이러한 구조를 가지는 Titans 모델의 Memory 모듈을 활용하는 3가지 변형 모델 또한 같이 제안하였다.

• Memory를 Context로 활용하는 모델
• Memory를 Network의 Layer로 추가하는 모델
• Memory를 Gated branch(게이트) 방식으로 활용하는 모델

Experimental Results

Language Modeling, Commonsense Reasoning, Recall-Intensive Tasks, Needle in Haystack Tasks, Time Series Forecasting, DNA Modeling Task에서 실험을 수행하였고, RNN 모델 대비 좋은 성능, Transformer-동일한 Context Window 크기 대비 더 좋은 성능, Transformer-전체 Context에서도 경쟁력 있는 성능을 보였다. 

 

Titans는 또한 200만 이상의 Context Window 크기로 확장할 수 있다는 장점을 가진다.

여기까지 읽고 정리해보자면, 결국 Titans라는 모델은 기존 Transformer가 계산량의 문제로 전체 Context에 대한 의존성을 잡아내기 어렵기 때문에 장기 기억 모듈을 새롭게 설계하고 단기 기억에 특화된 Attention과 같이 활용하여 좋은 성능을 내는 새로운 아키텍처를 구성하였다고 볼 수 있다.

즉, 장기 기억 모듈을 어떻게 설계하였는지에 집중해서 보는 것이 중요!

 

2. Preliminaries

Preliminary는 기초 개념을 나타내는 말로, 이 Section에서는 논문 전반에서 사용할 Notation과 Background Concept을 정리한다. 사용할 Notation은 다음과 같이 정리할 수 있으며,

• $x \in \mathbb{R}^{N \times d_{in}}$ : 입력 데이터
• $\mathcal{M}$: 신경망 또는 신경망 기반 Memory 모듈
• $Q, K, V$: Attention 매커니즘에서 사용되는 Query, Key, Value
• $M$: Attention Mask
• $S^{(i)}_j$: Sequence의 $i$번째 Segment, $j$번째 토큰
• $\mathcal{N(x)}$: Weight를 조정한 상태에서 수행된 순전파
• $\mathcal{N^*(x)}$: Weight 조정 없이 수행된 순전파
• $\mathcal{N}^{(k)}$: 신경망의 $k$번째 layer

이 Section에서 다루는 Background Concept는 다음과 같다.

• Attention과 그 효율적인 변형 모델들(Efficient Variants)에 대해
• 최신 Linear RNN 모델에 대해
• 위 아키텍처들을 Memory 관점에서 분석하고, Titans 모델을 어떻게 설계하게 되었는지에 대해

 

2.1. Backgrounds

Attention

Causual Attention은 입력 기반의 Key, Value, Query 행렬을 사용하여 출력 $y \in \mathbb{R}^{N \times d_{in}}$을 계산한다. 이 때 Query, Key, Value를 다음과 같이 정의할 수 있다.

$$Q = xW_Q, K = xW_K, V = xW_v$$

이 때 모든 $W$는 학습 가능한 가중치를 의미하며, 출력 y는 softmax 기반의 weighted sum을 통해 계산된다.

$$y_i = \sum_{j=1}^{i}\frac{exp(Q_i^TK_j/\sqrt{}d_{in})V_j}{\sum_{\ell=1}^{i}exp(Q_i^TK_{\ell}/\sqrt{d_{in}})}$$

이를 활용한 Transformer 모델은 뛰어난 정보 검색과 표현 학습 능력을 가지지만 출력을 계산하기 위해 최소 $N \times d$개의 연산이 필요하기 때문에, 메모리 사용량이 크고 긴 시퀀스를 처리할 때 처리 속도가 낮아지는 문제 또한 가진다.

 

Efficient Attention

각 Sequence에서 Softmax 기반의 Attention의 메모리 사용량을 줄이고 처리 속도를 개선하기 위해 다양한 Attention이 연구되었으며, 본 논문에서는 그 중 Kernel 기반의 선형 Attention에 초점을 맞추고 있다. 일반적으로 사용하던 Softmax 함수 대신 Kernel Function인 $\phi(.,.)$를 사용하는 것으로 기존 Softmax 공식 대신 Kernel Function을 사용한 Attention 수식은 다음과 같다.

$$y = \sum_{j=1}^{i}\frac{\phi(Q_i^TK_j)V_j}{\sum_{\ell=1}^{i}\phi(Q_i^TK_\ell)}$$

이 방식은 중간 계산 결과를 재사용할 수 있기 때문에 연산량을 줄이고 더 높은 처리량을 달성할 수 있다. 특히 Kernel Function으로 Identity Matrix를 사용하면 위 공식을 Recurrent하게 변환할 수 있어 기존 Attention보다 더 효율적으로 추론을 수행할 수 있다.

$$\mathcal{M}_t=f(\mathcal{M}_{t-1}, x_t),\quad y_t=g(\mathcal{M}_t, x_t)$$

 

Modern Linear Models and Their Memory Perspective

학습은 효과적이고 유용한 Memory를 획득하는 과정으로 볼 수 있다. 비슷한 관점에서 RNN의 Hidden State는 일종의 Memory로 볼 수 있으며, 모델은 입력 정보를 이 Memory에 잘 압축할 수 있도록 학습된다. 이러한 과정을 Read 연산과 Write 연산의 구성으로 생각할 수 있고, 결국 모델은 효율적인 Memory의 Read/Write 매커니즘을 학습한다고 볼 수 있다.

$$M_t = f(M_{t-1}, x_t)$$

Write Operation: 메모리를 Update 하는 함수

$$y_t = g(M_t, x_t)$$

Read Operation: 메모리에서 정보를 Retrieve 하는 함수

 

같은 관점으로 Linear Transformer의 수식을 다시 보면 Key와 Value Matrix를 Memory에 추가적으로 Write하는 과정을 확인할 수 있다. 이는 단순 덧셈으로 이루어져 긴 Context를 처리할 때 Memory Overflow문제를 초래하고, 과거 정보가 계속해서 추가되며 모델 성능이 떨어지는 현상이 발생한다.

 

이러한 문제를 극복하기 위해 크게 2가지 방법론이 제안되었다.

(1) Adding Forget Mechanism

- 기억을 지속적으로 추가함에 따라 불필요한 정보를 지우는 기능을 추가

ex. GLA, LRU, Griffin, xLSTM, Mamba2

 

(2) Improving Write Operation

- 기존 방법론은 단순히 기존 Memory에 새로운 정보를 추가하는 방식을 사용했기에 기존 정보를 먼저 제거하는 Delta Rule을 적용

두 방식 다 정보를 지우는 내용이지만 방법론 1은 과거의 정보를 선택적으로 유지하는 것에 초점을 맞추고, 방법론 2는 이전 정보를 아예 제거한 뒤 새로운 정보를 추가한 방식으로 보인다(Gradient Descent 관련 Delta Rule과 다른 맥락).

 

Memory Module

Fast Weight Programs에는 2가지 학습 규칙이 존재한다.

(1) Hebbian Learning

- 자극을 반복적으로 받은 정보는 더욱 강화된다

(2) Delta Learning

- Gradient Descent의 기초 형태로 Error를 기반으로 가중치를 조정하며 학습

 

하지만 위의 Memory 모델은 몇가지 문제점을 가진다.

• Momentary Surprise만 반영 → 장기적인 Context 고려하는데 한계
• Forgetting Gate 부재 → 필요없는 정보까지 무조건 저장하게 되어 메모리 관리가 비효율적

따라서 논문에서는 이러한 문제를 해결하는 새로운 아키텍처를 제안하였다.

여기까지 정리해보면 본 논문에서는 기존에 사용하던 Linear Attention이 장기적인 Context에서 성능이 떨어진다고 말하며 새로운 방법론을 제안하려 한다. 또한 그 이유를 Memory 관점에서 설명하고 있다.

 

3. Learning to Memorize at Test Time

이 섹션에서는 모델이 어떻게 Test time에 정보를 기억하고 활용할 수 있는지에 대한 새로운 Memory 구조를 제안한다.

 

3.1. Long-term Memory

Neural Long-Term Memory Module을 설계하기 위해서는 과거 정보를 추상화하여 모델 파라미터에 저장할 수 있는 구조가 필요하다. 단순하게는 신경망을 학습시켜 학습 데이터를 암기하도록 유도하면 된다고 생각할 수 있지만, 이 경우 모델의 일반화 성능을 제한하거나 개인정보 보호 문제, 테스트시의 성능 저하 등의 문제가 있다. 

 

따라서 본 논문은 Meta-Learning 기반의 Memory 모델이 필요하다고 주장하며, 때문에 모델은 테스트시 실시간으로 데이터를 학습하고 필요할 때 잊어버리는 매커니즘을 학습해야 한다고 말한다.

Meta Learning의 Learn-to-Learn을 살짝 변형한 Learn-to-Memorize라고 보여진다.

 

Learning Process and Surprise Metric

Human Memory 분야에서 우리의 예상과 다르게 발생한 사건은 더 오래 기억된다는 연구가 존재한다. 본 논문은 이러한 개념을 신경망에 차용하여 특정 입력을 받았을 때 놀람(Surprise) 정도를 정의하였다. 모델에서 Surprise는 입력에 대한 Gradient 크기로 정의하였으며, 이에 Gradient가 클수록 새로 들어온 입력이 기존 Memory와 다르다는 의미가 된다. 이를 Memory 갱신 방식으로 나타내면 다음과 같다.

$$M_t = M_{t-1} - \theta_t \nabla\ell(M_{t-1};x_t)$$

• $\nabla\ell(M_{t-1};x_t)$: 입력 $x_t$에 대한 Loss 함수의 Gradient
• $\theta_t$: Surprise score를 반영하는 Learning rate

수식을 보면 알겠지만, Surprise 값이 클수록 모델이 이를 강하게 반영하여 Memory를 업데이트 한다. 하지만 단순히 위와 같이 수식을 설정하면 다음과 같은 문제가 발생할 수 있다.

(1) 큰 Surprise가 발생한 후 중요한 정보를 놓칠 가능성

- Surprise input이 여러번 반복되면 이후 Input의 Gradient 값이 작아질 수 있음 → 모델이 특정 순간에만 집중(Local Minima)

(2) 인간의 기억 관점에서 정말 놀라운 사건은 시간이 지나도 기억날 수 있음

- 위 공식은 Momentary Surprise만을 반영하기 때문에 장기 기억의 특성을 반영하지 못함

 

이러한 이유로 Surprise Metric을 개선하기 위해 본 논문에서는 Surprise를 2가지 요소로 분리하였다.

(1) Past Surprise: 매우 가까운 과거에 발생한 Surprise의 크기를 측정

(2) Momentary Surprise: 현재 들어오는 Input에 대한 Surprise 크기를 측정

$$M_t = M_{t-1} + S_t$$

$$S_t=\eta_tS_{t-1}-\theta_t\nabla\ell(M_{t-1};x_t)$$

• $S_t$: Surprise를 저장하는 상태 값
• $\eta S_{t-1}$: Past Surprise를 특정 비율만큼 반영
• $\theta_t\nabla\ell(M_{t-1};x_t)$: Momentary Surprise

위 공식은 Gradient Descent + Momentum과 유사한 형태를 가진다. $S_t$는 Momentum 역할을, $\eta_t$는 Surprise Decay를 결정하는 역할을, $\theta_t$는 Momentary Surprise를 얼마나 반영할지를 결정하는 역할을 수행한다.

 

Objective

본 논문이 제안하는 메모리 모듈은 Loss Function $\ell(.;.)$을 기반으로 학습하는 Meta-model이다. 위에서 설명한 Surprise Metric을 보면 손실 함수를 기반으로 정의되며, 모델은 이 Objective(목표)를 가지고 Test time에 어떻게 작동할지 학습한다.

 

본 연구에서 주로 다루는 Memory는 Associative Memory로 과거 데이터를 Key-Value 쌍으로 저장하는 방식을 사용한다. 예를 들어 입력 $x_t$가 주어졌을 때, 2개의 Linear Layer를 사용하여 $x_t$를 Key와 Value로 변환한다.

$$k_t=x_tW_K, \quad v_t=x_tW_V$$

정보가 들어왔을 때 Neural Network를 거치며 하나의 정보를 Key와 Value 형태로 저장한다고 생각

 

이후 메모리 모듈은 Key와 Value간의 관계를 학습해야한다. 이를 위해 본 논문에서는 손실함수를 다음과 같이 정의하였다.

$$\ell(\mathcal{M}_{t-1};x_t)=\left\| \mathcal{M}_{t-1}(k_t)-v_t\right\|^2_2$$

• $\mathcal{M}_{t-1}(k_t)$: 현재 Memory state에서 Key $k_t$에 대응하는 Value
• $v_t$: 실제 Value

이 Loss Function을 Memory module의 Inner-loop에서 최적화하면, 모델은 Test time에서의 Key와 Value 간의 Mapping을 학습할 수 있다. Inner-loop에서는 Memory module의 Weight를 최적화하고, Outer-loop에서는 전체 아키텍처의 나머지 파라미터($W_K, W_V$)를 최적화 한다. 이 파라미터들은 Loss Function 내에서 하이퍼파라미터 역할을 수행한다.

메타러닝 방식으로 학습되며 Inner-loop에서는 Memory를 최적화하고 Outer-loop에서는 전체 아키텍처를 최적화한다. Memory의 최적화는 Key-Value간의 관계(Associative Memory)를 학습하는 식으로 진행된다.

 

Forgetting Mechanism

아무리 Deep Memory나 대규모의 행렬 기반 메모리(matrix-valued memory)를 사용하더라도 긴 Sequence를 처리할 때, 모델이 어떤 과거 정보를 잊어야 할지 효율적으로 관리하는 것이 중요하다. 이를 해결하기 위해 본 논문에서는 Adaptive Forgetting Mechanism을 사용하였다.

$$\mathcal{M}_t = (1-\alpha_t)\mathcal{M}_{t-1}+S_t, \quad S_t=\eta_tS_{t-1}-\theta_t\nabla\ell(M_{t-1};x_t)$$

• $\alpha_t\in[0,1]$: Adaptive Forgetting Mechanism, Memory를 얼마나 지울지 결정
• 수식을 보면 알겠지만, 0일 때 기존 Memory를 그대로 유지하고, 1이면 전체 Memory를 초기화

LSTM이나 GRU같은 RNN 모델의 Gate Mechanism과 비슷한 개념이지만, 기존 RNN의 Forget Gate보다 더 유연하고 Data Adaptive한 방식으로 작동한다.

여기서 Adaptive라는 의미는 고정된 값이 아니라 입력 데이터에 따라 동적으로 변화하는 방식을 의미한다. 예를 들어 위 수식에서는$\alpha_t$뿐만 아니라  $\eta_t, \theta_t$가 있가 Adaptive 하며 이들이 특정 규칙을 강요하는 것이 아니라 데이터에 따라 변하기 때문에 Adaptive하다고 볼 수 있다.

반대로 RNN의 경우 한번 학습이 끝나고 나면 Forget gate에서 고정된 파라미터만큼 Forgetting을 수행한다.

 

Memory Architecture

기존에 많이 사용하던 Vector-valued Memory와 Matrix-valued Memory는 결국 과거 데이터를 선형 형태로 압축하는 형태를 띈다. 이는 결국 과거 데이터가 선형 관계를 따른다고 가정하는 것으로 복잡한 데이터 패턴을 학습하는 데에 한계가 존재한다.

 

따라서 본 논문에서는 Long-term Memory 모듈의 아키텍처로 MLP를 사용하였으며 최소 1개 이상의 Layer를 포함한다. MLP 구조를 선택한 이유는 "더 효과적인 장기 기억 설계 방법을 탐구하기 위해서"이며 이 부분에 있어서는 추후 더 성능 좋은 모듈을 적용시켜도 좋을 것이라고 말한다.

 

또한, Layer가 2개 이상이여야 단순 선형 관계 이상의 복잡한 패턴을 효과적으로 기억할 수 있다고 말하고 있다.

 

Retrieving a Memory

앞서 우리는 테스트 시 기억을 학습하는 Long-term Memory Module을 설계하고 학습하였다. 이제는 어떻게 메모리에서 정보를 검색(Retrieve)할지를 결정하면 된다. 본 논문에서는 Memory 검색에 있어 Weight를 업데이트하지 않고 단순 Forward Pass를 사용하였다.

 

메모리 검색은 다음의 순서로 이루어진다.

1. Query Vector 생성

$$q_t = x_tW_Q$$

입력 $x_t$를 Query vector $q_t$로 변환하여 Memory에서 정보를 검색할 준비를 한다. 여기서 $W_Q$는 학습이 완료된 Matrix.

 

2. Memory Retrieval

$$y_t = \mathcal{M}^*(q_t)$$

• $\mathcal{M}^*$: Inference mode memory module(가중치 업데이트 X)
• $y_t$: 쿼리 $q_t$와 가장 관련 있는 정보

즉, Memory에서 Retrieve를 실행할 때에는 Weight 업데이트를 진행하지 않고 단순 Inference를 사용한다고 정리할 수 있다.

 

3.2. How to Parallelize the Long-term Memory Training

본 논문에서 제안하는 Long-term Memory Module의 학습에 사용되는 Associative Memory Loss는 결국 Gradient Descent + Momentum + Weight Decay를 최적화 하는 과정과 동일하다. 이론적으로 이 방식은 O(N) FLOPs가 필요하기 때문에(N은 Sequence), 본 논문에서는 Tensorization, Matmul 연산을 활용해 학습을 병렬화 하였다.

 

본 논문에서는 Mini-batch Gradient Descent에서 Forward Pass를 Matmul 연산으로 변환하였다. 

기존 신경망 학습 과정에서 Weight를 업데이트하는 방식은 Sequential 연산으로 매 t단계마다 업데이트가 이루어진다. 하지만 이를 행렬 곱셈 연산으로 변환 시 GPU에서 한 번에 여러 개의 Mini-batch를 받아 처리할 수 있도록 병렬화가 가능해진다(한 번에 여러 개의 Gradient를 동시에 계산한다고 생각).

 

Parameters as the Function of Chunks

기존에는 $\alpha_t$, $\theta_t$, $\eta_t$와 같은 파라미터들이 입력 데이터 $x_t$에 따라 변화하는 Input-dependent한 방식으로 모델의 표현력을 증가시키지만, 계산량이 많다는 단점이 있다. 

 

여러 토큰의 집합인 Chunk 내에서는 모든 토큰이 동일한 파라미터 값을 공유하여 모델의 표현력이 조금 줄어드는 대신 계산량 또한 줄어 학습 속도를 크게 개선할 수 있다(Chunk마다 Weight Decay 행렬과 Learning Rate 행렬을 하나의 Scalar로 저장할 수 있기 때문에). 또한 LTI(Linear Time-Invariant) System으로 변환하여 Global  Convolution과 같은 최적화 기법을 활용할 수 있다.

 

본 논문에서는 Token 기반의 방식으로 연구를 진행하였지만, 추후 Chunk 기반 방식도 고려할만 하다고 말하고 있다.

아무래도 성능 자체는 떨어질 것으로 예상되기 때문이지 않을까..?

 

3.3. Persistent Memory

본 논문에서는 효과적인 Memory System이 구축되기 위해서는 Input에 의존하는 파라미터 뿐만 아니라 Input과 무관한 파라미터 또한 포함해야 한다고 말한다. 따라서 입력과 관계 없이 특정 Task에 대한 정보를 저장하는 Persistent Memory의 필요성을 주장한다.

 

Persistent Memory

$$P= [p1, p2, ... p_{N_p}],\quad x_{new} = [p_1, p_2, ..., p_{N_p}] || x$$

Persistent Memory P는 학습 가능한 파라미터 세트로 정의할 수 있으며, 새로운 입력을 생성할 시 입력 시퀀스 맨 앞에 이 파라미터 들을 추가한다(||=연결). 즉, 입력 데이터를 그대로 사용하면서 추가적인 Task 관련 정보를 포함한 새로운 입력을 생성하는 방식이다.

 

Persistent Memory를 도입하는 이유는 3가지 관점으로 설명할 수 있다.

(1) Memory Perspective

기존 Long-term Memory는 Input에 따라 다르게 동작하는 Contextual Memory였지만, 효과적인 Memory System이라면 Task 자체에 대한 지식이 있어야 함.

 → Persistent Memory를 통해 Task를 수행하는 방법에 대한 정보를 저장

 

(2) Feedforward Network Perspective

Sukhbaatar의 연구에 따르면 Fully Connected Layer의 ReLU 활성화 함수를 Softmax로 변경하면, 해당 Layer가 Attention 처럼 작동한다고 한다.

$$FFN(x) = W_vSoftmax(W_kx)$$

위의 식에서 $W_k$, $W_v$가 Attention의 Key, Value 행렬과 같은 역할을 수행하는데, 이 때 이 파라미터 들은 Input 데이터와 무관한 Persistent Memory 역할을 수행할 수도 있다. 즉, Persistent Memory는 모델이 특정 Task에서 Input과 독립적인 형태로 Attention을 조절하도록 돕는다.

일반적으로 Transformer에서는 $W_k$, $W_v$는 입력 데이터에 영향을 받지만, 논문에서는 해당 파라미터들을 입력 데이터와 무관한 값으로 고정한다면~ 이런식으로 서술하는 것 같다.

또한 ReLU 대신 Softmax를 사용하였을 때 Attention 효과를 얻는다는 부분은, Softmax가 입력 값들을 확률분포로 변환하여 특정 차원에 Attention 하게 만들기 때문에 FFN의 내부 연산이 $W_k$, $W_v$를 Key, Value로 인식하는 것 같다.

 

(3) Technical Perspective

Transformer의 Causal Masked Attention은 초기 토큰들에 높은 가중치를 할당하는 경향이 있다. 이 경우 앞부분 토큰들이 높은 Attention을 가지기 때문에 후반부 토큰 정보가 충분히 반영되기 어렵다.

→ Persistent Memory를 사용해 Attention Weight을 균등하게 분배하도록 해야한다.

 

4. How to Incorporate Memory?

Section 3에서 설명한 것처럼, Transformer의 Key와 Value Matrix는 Associative Memory 역할을 수행할 수 있다. 하지만 이들은 현재 Context Window 내부에서만 작동하는 Short-term Memory로, 지속적으로 데이터를 학습하고 이를 Weight로 저장할 수 있는 별다른 Long-term Memory가 필요하다.

 

4.1. Memory as a Context

Long-term Memory 역할을 수행하는 첫번째 방법은 Memory를 Context로써 사용하는 것이다. 즉, 현재 정보를 처리할 때 과거 정보를 Context로 사용하는 것으로 다음의 순서에 따라 계산된다.

 

1. Retrieving Memory

긴 시퀀스 $x\in \mathbb{R}^{N\times d_{in}}$가 주어졌을 때, 이를 고정된 Chunk 크기의 세그먼트 $\mathsf{S}^{(i)}$로 분할한다. 현재 $\mathsf{S}^{(t)}$가 들어오면 이전 $\mathsf{S}^{(t-1)}$를 과거 정보로 사용하며, 이전 메모리 $M_{t-1}$에 Query를 입력하여 해당 세그먼트와 관련된 과거 정보를 검색한다.

$$h_t=\mathcal{M}^*_{t-1}(q_t), \quad q_t=\mathsf{S}^{(t)}W_Q$$

• $q_t$: 현재 세그먼트 $\mathsf{S}^{(t)}$를 선형 변환하여 Query 벡터 생성
• $\mathcal{M}^*_{t-1}(q_t)$: 이전 메모리에서 $q_t$에 해당하는 과거 정보를 검색

2. Memory를 활용하여 Attention 수행

1단계에서 검색을 통해 찾은 $h_t$를 현재 입력과 함께 Attention Module에 전달한다(RNN과 비슷). 이 때 Persistent Memory $P$도 함께 포함하여 새로운 입력 시퀀스를 구성한다.

$$\tilde{S}^{(t)}=[p1, p2, ... p_{N_p}] \left\| h_t\right\| \mathsf{S}^{(t)}$$

이렇게 구성한 새로운 입력 시퀀스를 Attention 모듈에 전달하여 최종 출력을 얻는다.

$$y_t=Attention(\tilde{S}^{(t)})$$

적용되는 Attention Map은 Figure 3(a)에 나타나있다.

 

3. Memory Update

2단계를 통해 얻은 Attention 모듈의 출력을 사용하여 Long-term Memory를 업데이트한다.

$$\mathcal{M}_t=\mathcal{M}_{t-1}(y_t)$$

최종 출력 $o_t$는 Attention의 결과로 생성된 $y_t$와 업데이트된 Memory $\mathcal{M}^*_{t}$를 결합하여 생성된다.

$$o_t=y_t\otimes\mathcal{M}^*_t(y_t)$$

 

위와 같은 방식을 사용하면 다음 3가지의 장점을 가진다.

(1) Attention이 과거 정보가 필요한지 여부를 스스로 결정할 수 있다.

- 과거 Memory를 무조건 사용하는 것이 아니라 필요할 때에만 활용할 수 있다.

(2) Attention이 Memory에 저장할 정보의 중요도를 조절한다.

- 모든 데이터를 기억하지 않고 중요한 정보만 저장할 수 있어 Memory Overflow를 방지한다.

(3) Test-time에서도 Long-term Memory가 계속해서 학습할 수 있다.

- Test-time에서 Persistent Memory는 고정되어 있고, Attention 모듈은 새로운 입력을 바탕으로 최적의 정보를 계속해서 선택한다. 때문에 Long-term Memory는 계속 업데이트가 되어 Test-time에서도 새로운 데이터를 학습할 수 있도록 설계된걸 알 수 있다.

 

4.2. Gated Memory

앞서서 Memory를 Context로써 사용하는 방법을 소개했다면, 이번에는 Sliding Window Attention(SWA)와 Long-term Memory를 결합하는 아키텍처(Figure 4)에 대해 소개한다.

 

이번 방식은 4.1 에 소개한 아키텍처와 달리 입력을 Chunk로 나누지 않고, 전체 Sequence를 그대로 사용한다. SWA가 Short-term Memory 역할을 수행하고, Neural Memory는 Long-term Memory 역할을 수행한다. 이 아키텍처에서는 입력 데이터를 2개의 Branch로 나누어 처리한다.

 

[1st branch] Long-term Memory

$$\tilde{x}=[p_1, p_2, ... p_{N_p}] || x$$

Input x에 Persistent Memory를 추가하여 새로운 입력 $\tilde{x}$를 생성

 

[2nd branch] Sliding Window Attention(SWA)

$$y=SW-Attention^*(\tilde{x})$$

SWA를 사용하여 현재 Context에서 필요한 정보를 학습

$$o=y\otimes\mathcal{M}(\tilde{x})$$

이렇게 만들어진 두 결과를 결합하여 최종 출력을 생성한다. 여기서 결합연산 $\otimes$는 Non-linear Gating을 의미한다.

단기기억으로 SWA, 장기기억으로 Neural Memory를 구성하며 독립적인 Branch를 사용하므로 Multi-head Attention과 유사한 효과를 얻을 수 있다.

Multi-head Attention은 각 Head가 다른 형태의 Attention을 수행
→ 여기서는 SWA와 Neural Memory가 서로 다른 방식으로 정보 처리

 

4.3. Memory as a Layer

세번째 아이디어는 Neural Memory를 Neural Network의 독립적인 Layer로 쌓는 것이다. 이 접근 방식은 기존 Hybrid 모델 구조와 유사한데, 여기서 Hybrid 모델이라 함은 RNN과 SWA를 같이 사용하던 이전 방식을 말한다.

$$\tilde{x}=[p_1, p_2, ... p_{N_p}] || x$$

$$y=\mathcal{\tilde{x}}$$

$$o=SW-Attention(y)$$

이 방법은 Persistent Memory 추가 → Input을 Neural Memory에 통과 → SWA를 적용하여 최종 출력을 생성한다. 식에서 볼 수 있듯이 Neural Memory를 순차적으로 연결하고(이전 방식들은 병렬), 독립적인 Layer처럼 사용한다.

 

하지만, 이 방식은 Attention과 Neural Memory의 장점을 모두 활용하기 어렵다. 모델의 성능이 각 Layer의 표현력에 의존하기 때문에 각 Layer가 강력하지 않으면 전체 모델의 성능이 제한되며, Attention과 Neural Memory간의 상호작용이 부족하다.

 

Memory Without Attention

논문에서는 또 다른 변형 아키텍처로 Attention 없이 Neural Memory만 사용하는 버전을 같이 소개한다. 이 모델은 MAL에서 Attention을 제거하고 Neural Memory만 사용하는 것으로 LMM 또는 Titans (LMM) 이라고 부른다.

 

논문에서 이 모델을 주요하게 여기는 이유는 "진정한 Long-term Memory라면 Short-term Memory 없이도 독립적으로 강력한 성능을 내야한다."고 생각하기 떄문이다.

 

그 결과가 어떻게 나왔는지는 추후 실험 결과를 통해 확인해보자.

 

4.4. Architectural Details

앞선 아키텍처 소개에서는 간단하게 개념을 전달하기 위해 디테일한 세부 구현 내용을 생략하였다. 따라서 논문에서는 이 섹션을 통해 Titans 모델을 효과적으로 학습시키기 위해 사용된 핵심 기술들을 소개한다.

 

(1) Residual Connection

- 모든 Block에서 Residual Connection을 사용하였다.

 

(2) SiLU 활성화 함수 사용

- SiLU(Sigmoid-Weighted Linear Unit)은 ReLU보다 더 Smooth하고 Gradient 소실 문제가 적다.

- 최근 Transformer와 같은 대형 모델에서 좋은 성능을 보임

 

(3) L2 Norm 사용

- Query, Key의 크기를 일정하게 유지

 

(4) 1D Depthwise-Separable Convolution 추가

- 최신 연구 중 Convolution이 RNN 성능을 향상시킨다는 결과가 존재

- Query, Key, Value 계산 후 1D Convolution을 적용하면 성능 향상

 

(5) Gating Mechanism

- 출력 Projection 전에 정규화 + Gating 적용

 

(6) Theorem 4.1.

Transformer, Linear RNN, DeltaNet 같은 모델들은 $TC^0$문제에 같혀있어 복잡한 문제를 해결하기 어렵다.

→ Titans는 더 표현력이 풍부하고 복잡한 Sequence를 가진 문제도 해결할 수 있다!

Section 4의 1st 정리이기에 4.1.로 이름 붙인 것 같다.

https://openreview.net/forum?id=QZgo9JZpLq

The Illusion of State in State-Space Models

 

5. Experiments

논문에서는 Titans 모델 성능을 검증하기 위해 총 5가지 질문에 대한 실험을 진행하였다.

(1) Titans 모델은 기존 모델들과 비교했을 때 성능이 더 좋은가?

- How do Titans perform compared to baselines in downstream tasks?

- Section 5.2, 5.6. 5.7

 

(2) Titans의 실제 Context 길이는 얼마인가? 

- What is the actual context length of Titans?

- Section 5.3, 5.4

 

(3) Titans는 Context 길이에 따라 어떻게 크기가 바뀌는가?

- How do Titans scale with respect to context length?

- Section 5.8

 

(4) 메모리의 깊이가 성능과 효율성에 영향을 미치는가?

- How the depth of memory can affect both performance and efficiency?

- Section 5.5

 

(5) Titans 모델의 각 구성 요소가 성능에 어떻게 기여하는가?

- What is the contribution of each Titans' component in its performance?

- Section 5.9

 

5.1. Experimental Setup

Titans 모델에는 MAC(Memory as a Context), MAG(Memory as a Gate), MAL(Memory as a Layer), Neural Memory(No Attention)로 총 4가지 변형 모델이 존재한다. 모델 별 파라미터 수, 비교 Baseline 모델 리스트, 구체적인 학습 방법은 논문을 참고.

 

5.2. Language Modeling

Language Modeling과 CommonSense Reasoning Task에서의 Titans 모델 성능을 평가하였다. 여기서 말하는 Non-hybrid 모델이란 Attention과 RNN이 결합되지 않은 순수 모델을 의미한다(Samba, Gated DeltaNet-H2, MAC, MAG, MAL이 Hybrid 모델).

 

vs TTT: TTT 모델 또한 Gradient-based RNN 모델이지만 Titans(LMM)의 성능이 더 좋음

vs Mamba: Mamba, Mamba2, Gated DeltaNet도 Gating Mechanism을 사용하지만, Titans가 더 높은 성능을 기록

 

MAC, MAG 방식이 MAL 방법보다 전체적인 성능이 높았고, 이에 기존 Hybrid 모델들이 대부분 MAL 방식을 따르고 있는데, Titans의 MAC, MAG 방식이 더 우수한 성능을 보인다는 것을 발견한 것이 주요한 발견이라고 논문에서는 말하고 있다.

 

5.3. Needle in a Haystack

다음 실험은 NIAH(Needle in a Haystack)이다. 이 실험에서는 모델이 긴 Sequence에서 중요한 정보(Needle)을 잘 찾아낼 수 있는지를 평가한다.

Titans 모델들이 기존 모델들보다 더 좋은 성능을 보였으며, 특히 MAC 모델이 가장 좋은 성능을 보였다.

 

vs TTT: Momentum(Surprise Metric)과 Weight Decay를 사용하여 Memory 관리가 뛰어나 더 좋은 성능을 보임

vs Mamba2: Mamba2는 기존 정보를 지울 수는 있지만, 완전히 제거할 수 없다. 이에 시퀀스 길이가 길어짐에 따라 Titans의 성능이 더 높아짐

vs DeltaNet: DeltaNet은 Forget Mechanism이 부족하여 메모리를 효과적으로 지우지 못한다.

 

5.4. BABILong Benchmark

BABILong Benchmark 실험을 통해 극단적으로 긴 문서에서의 Reasoning 능력을 평가(NIAH는 단순 검색). 여기서는 Few-shot Learning과 Fine-Tuning Setting으로 실험을 진행하였으며, Few-shot에서는 Pretrained 모델을 그대로 평가, Fine-Tuning에서는 MAC 모델을 조금 Fine-Tuning한 후 평가를 진행하였다.

 

Figure 6에서 볼 수 있듯이 다른 모델들 보다 좋은 성능을 보였으며, In-context Online Memory Learner가 과거 데이터를 효과적으로 저장하는데 큰 역할을 수행하였다고 말한다.

 

5.5. The Effect of Deep Memory

이번 섹션에서는 Neural Memory의 깊이가 모델 성능과 학습 속도에 미치는 영향을 평가한다. Neural Memory 모듈의 깊이를 1~4로 변경하면서 실험을 진행하였으며, 깊은 Memory일 수록 성능이 증가하는 것을 확인하였다.

 

또한 Sequence Length에 따라 모든 모델의 학습 속도가 선형적으로 증가하였다.

 

5.6. Time Series Forecasting

시계열 예측 평가에서도 Long-term Dependency, Weight Decay, Surprise Metric 덕분에 좋은 성능을 낼 수 있었고,

 

5.7. DNA Modeling

NLP Task에서 뿐만 아니라 DNA 서열 분석 Task에서도 효과적인 성능을 보였다.

 

5.8. Efficiency

학습 속도의 경우에는 MAL 모델이 Titans 모델들 중에서는 가장 빨랐다. 나머지는 조금 오래 걸림.. 가장 성능이 좋지 않았던 MAL 방식이 아무래도 기존 Hybrid 모델들과 비슷한 방법론을 차용해서이지 않을까?

본인 모델들 끼리는 색을 좀 비슷한 계열로 통일해주지;;

 

5.9. Ablation Study

Neural Memory의 각 구성요소에 대해 Ablation Study를 진행하였다. 결론적으로 모든 구성요소가 성능에 기여하는 것을 확인하였으며, Weight Decay를 뺐을 때 성능이 가장 크게 감소하였다(메모리 효과적 관리 용도). 

 

6. Conclusion

본 논문에서는 Neural Long-term Memory를 활용한 새로운 Sequence 모델 Titans를 제안하였다. Transformer와 RNN 모델들의 한계를 뛰어넘는 성능을 보여주며, Memory as a Context, Memory as a Gate, Memory as a Layer 3가지 변형 모델들을 제안하며 장기 기억 메모리를 어떻게 효과적으로 관리할 수 있는지 그 방법에 대해 탐색하였다.

 

긴 Context에서 특히 좋은 성능을 보였으며, Surprise Metric을 활용해 중요한 정보를 더 효과적으로 저장하도록 설계하였다.

길었던 Titans 논문 리뷰 작성이 마무리 되었다.
NLP 분야는 요즘 논문을 안찾아보다보니 중간중간 모르는 개념들을 개인적으로 확인해보느라 더 오래 걸렸던 것 같은데(핑계...), 그래도 최근 이슈가 되었던 논문을 한번 훑어보았다는 것에 의미를 두려 한다. 사실 이전 Mamba도 제대로 이해하지 못한 상태에서 Titans 논문을 읽었는데, 이번 논문을 읽다보니 얼른 Mamba 모델 내용도 정리해 두어야겠다는 생각이 들었다.

Titans 모델에서 결국 핵심적으로 작용하는 것은 Surprise Metric과 Long-term Memory+Attention(Short-term Memory) Fusion이라고 생각하며, Persistent Memory와 3가지 변형 모델을 설계하는데까지 논문의 저자가 정말 많은 사전 조사를 했고, 정말 많은 이론들을 알고있구나라는 느낌을 받았다(당연하겠지만...).

Titans 모델은 좋은 성능을 보였지만, 한 가지 아쉬운 점은 논문을 처음부터 끝까지 읽으면서 완전 색다른 방법론이라고 느껴졌던 부분이 없었다는 점이다. Surprise도 결국엔 Momentum이고, Persistent Memory도 사전정보를 넣는 것에, RNN과 Attention을 합쳐 Long-term Memory로 작동하는 것들이 다 어디선가 본 애들이 잘 짜집기 되었다는 느낌을 받았다. 결국 아직은 Attention이라는 대괄호 속에 갇혀있는 것 같고, 이정도로 Transformer의 대체 모델이라고 하기에는 어렵지 않나..? 라는게 지금의 생각이다.

물론 Titans 연구를 기반으로 앞으로 나올 후속 연구들에 대해 기대하고 있다.

이제 Titans에 이미지를 넣어보자... 하하..

---

세 줄 요약

1. Neural Long-term Memory를 활용하여 기존 Transformer와 RNN 모델보다 더 효과적으로 긴 Context를 학습할 수 있는 Titans 아키텍처를 제안
2.  Surprise Metric을 활용해 중요한 정보를 선택적으로 기억하며 Persistent Memory, Weight Decay 방법 등을 활용한 MAC, MAG, MAL 3가지 변형 모델을 제안
3. Language Modeling, Time Series 등 다양한 Task에서 좋은 성능을 보였으며 2M 이상의 Context Window를 처리할 수 있는 확장성을 가짐