[논문 리뷰 | 정리] Titans: Learning to Memorize at Test Time

Google에서 Transformers라는 구조를 발표한 이후로 약 8년간 인공지능 분야는 Transformers가 지배하고 있었습니다. 잘 알려진 ChatGPT 혹은 Gemini부터, 최근에 갑자기 화두에 오른 DeepSeek-R1까지 Transformers가 사용 안 된 모델을 찾기란 쉽지 않습니다. 이 Transformers는 Attention이라는 기술을 이용해서 입력으로 들어오는 모든 글자 사이의 관계를 파악하여 결론을 도출하는 모델입니다. Transformers는 모든 글자 사이의 관계를 파악하기 때문에, 성능이 올라갔습니다.
최근 들어 사람들은 엄청나게 긴 문장을 처리하는 인공지능을 원하게 됐습니다. 즉, 많은 토큰 (token)을 처리할 수 있는 인공지능을 선호하게 되었다는 거죠. 여기서 기존 Transformers의 문제가 발생합니다. 그 엄청나게 많은 토큰을 모두 처리해야 하니 모델이 커지고 비용이 많이 들게 됩니다. 이 문제는 Transformers의 구조 중, Attention에서 비롯되기 때문에 연구자들은 Linear Attention 혹은 Kernel Attention 등의 저비용 Attention 기술을 연구했지만, 긴 문장을 처리하는데 있어 성능의 하락으로 이어지게 됐습니다.

 
 

---

 

 

Google의 새로운 모델 구조

사실 이 페이퍼는 이해하기 힘들었습니다. 모델 구조도 자세히 안 나와 있고, 표기법들은 혼용되고 있고.... version 1이라서 아직은 완성이 덜 된 것처럼 보이기도 합니다....
 
Titans는 아직 정식으로 학회나 저널에서 출간한 것이 아닌 preprint이기 때문에, 논문이라는 단어 대신 페이퍼라고 부르겠습니다.



 
https://arxiv.org/pdf/2501.00663v1
올해 2025년 1월, Google에서 또 다른 인공지능 구조를 선보였습니다. 바로 Titans입니다. Titans는 사실 Linear Attention 혹은 Kernel Attention의 장점은 가져가면서도 긴 문장을 처리하는데 있어 뛰어난 성능을 보이는 새로운 구조입니다. 즉, Titans의 목표는 "적당한 비용으로 2 백 만 개 정도의 토큰들(단어)을 좋은 성능으로 처리하자"라는 것입니다.
이 목표를 위해 연구진들은 사람의 뇌가 정보를 처리하는 방법을 차용했습니다. 사람의 뇌는 통상적으로, 정보를 받아들이고 -> 단기기억으로 저장하고 -> 단기기억을 장기기억으로 저장하고 -> 일부는 고정관념으로 변환의 과정을 거치게 됩니다. 그리고 여기서 말하는 단기기억, 장기기억, 그리고 고정관념은 서로 각기 다른 부분에 저장이 되고, 필요 시 서로 간의 정보가 연결되도 혼합되어 결론이 나오게 됩니다.
 
 
 
Titans 연구자들은 이런 사람의 뇌 구조에 영감을 받고, 다음과 같은 질문을 하게 됩니다.

Q1. 기억 모듈을 어떻게 잘 설계할 수 있을까?
Q2. 기억을 어떻게 잘 업데이트할 수 있을까?
Q3. 기억을 어떻게 잘 복원할 수 있을까?
Q4. 기억 모듈 간에 연결 방법을 어떻게 잘 설계할 수 있을까?
Q5. 기억 모듈이 효율적으로 긴 정보를 저장할 수 있을까?

 
 
 
사실 위의 질문들은 Linear Attention 혹은 Kernel Attention을 사용하여 적은 비용으로 뛰어난 성능을 발휘하는 인공지능을 만들기 위한 질문들입니다. 사실 전통적인 Attention 구조를 그대로 사용한다면, 뛰어난 성능이 어느정도 나올 것입니다 (2 백 만 개 정도의 정보를 처리하는 게 사실상 불가능해서 그렇지...)

 
 

---

 

 

페이퍼 로드 맵

이 페이퍼는 다음과 같은 큰 세 개의 파트로 구성돼 있습니다.

• Long-term Neural Memory
• Titans Architectures
• Experimental Results

 

Long-term Neural Memory

이 모듈은 Titans에서 중심 역할을 하는 장기 기억(long-term memory)을 담당하는 모듈입니다. 이 모듈의 기능은 다음과 같습니다.

• Test Time (Inference Time)에서도 정보를 어떻게 저장할지를 학습합니다.
• 인간의 지능에서 영감을 받아 처음 보는 데이터를 더욱 잘 기억하도록 합니다. 즉, training 단계에서는 보지 못한, 인공지능에게 놀라움 (surprise)를 선사할 수 있는 데이터를 선별하고, 이 데이터를 더욱 잘 기억하도록 합니다.
• 필요 없는 지식을 망각합니다. 한정된 모델 크기를 효율적으로 사용하기 위해, 긴 텍스트에서 필요 없는 부분을 잊도록 하는 decay mechanism을 보입니다.

 

Titans Architectures

위에서 Long-term Neural Memory를 설계한 후, 남은 작업은 이 모듈을 어떻게 효율적이고 효과적으로 depp learning architecture에 적용시키냐 입니다. Titans는 세 개의 모듈로 구성돼 있습니다.

• Core: 단기기억을 담당하며, 정답을 도출하는데 직접적인 영향을 주고, 현재 입력 받은 토큰에 영향을 많이 받음. Test Time 때도 학습이 됨.
• Contextual Memory: 장기기억 (Long-term memory)을 담당하며, 이전 토큰들의 정보를 저장하거나 복원하는 것을 담당함. Test Time 때도 학습이 됨. 입력 토큰에 영향을 받음.
• Persistent Memory: Task 자체에 대한 정보를 처리함. (이 부분은 이해하지 못함.) Tes Time에는 학습이 안 됨.

페이퍼에서는 위의 세 가지를 변형하여 Titans의 변형 세 가지를 더 소개합니다. 즉, 총 네 개의 모델을 보이는 것이죠.

 

Experimental Results

실험은 크게 {언어 | 시계열 | DNA} 분야에서 진행되었으며, 실험을 통해 Titan 구조가 다른 최신 recurrent models를 뛰어 넘었다는 것을 알 수 있습니다. 추가로 같은 context window 크기를 가지는 Transformers보다 더 뛰어난 성능을 보이고, 전체 context를 한 번에 받는 Transformers와 비슷하다고 합니다. 즉, Transformers의 비효율성을 해결하고 2 백 만 개 정도의 토큰을 처리할 수 있는 확장성을 기대할 수 있습니다.
 
 
 

---

 

 

 

Titans의 장기 메모리 (Contextual Memory) 이해해 보기

 Titans 페이퍼는 개인적으로 느꼈을 때, 이해하기 난해한 페이퍼였습니다. 그래도 최대한 저자의 의식 흐름을 따라가면서 한 번 파헤쳐 보겠습니다.
 
 
 
핸드폰에서는 수식이 깨져서 보입니다.

 

 

Classical Attention

우선, 전통적인 Attention 수식을 보겠습니다. Attention을 수식으로 표현하는 방법에는 행렬 표현 방법과 벡터 표현 방법이 있습니다.
 
행렬 표현 방식의 Attention: $ Y=softmax(Q^\top K)V $
벡터 표현 방식의 Attention: $ y_t=\sum_{j=1}^i\dfrac{exp(Q_i^\top K_j)}{\sum_{l=1}^i exp(Q_i ^ \top K_l)}V_j $
 
여기서 $ Q = xW_Q $이고, 같은 방식으로 $ K $와 $ V $도 정의됩니다. 위의 방식은 많은 연산을 요구하게 됩니다. 그래서 연구자들은 빠르고 적절한 성능을 위해 위의 수식을 $ \phi() $ 라는 임의의 함수를 도입하여 변경하게 됩니다 (논문마다 이 $ \phi() $는 다릅니다). 이 $ \phi() $라는 함수를 kernel이라고 부르며, 이를 활용한 Attention을 Kernel Attention이라고 부릅니다.
 

Kernel Attention

 Kernel Attention에서 사용하는 $ \phi() $는 다음과 같은 성질을 만족합니다. 쉽게 생각하면 선형성을 만족한다고 여길 수도 있을 것 같습니다.
 
$$ \phi(xy) = \phi(x) \phi(y) $$

 
이런 특성을 가지는 kernel $ \phi() $를 Classical Attention의 $exp()$ 대신에 넣는다면, 벡터 표현 방식의 Attention은 다음과 같이 쓸 수 있습니다.
 
$$ y_t=\sum_{j=1}^i\dfrac{\phi(Q_i^\top k_j)}{\sum_{l=1}^i \phi (Q_i ^ \top K_l)}V_j=\sum_{j=1}^i\dfrac{\phi(Q_i)^\top \phi(K_j)}{\sum_{l=1}^i \phi (Q_i) ^ \top \phi(K_l)}V_j= \dfrac{\phi(Q_i)^\top \sum_{j=1}^i \phi(K_j) }{\phi (Q_i) ^\top \sum_{l=1}^i  \phi(K_l)}V_j $$
 
위의 수식은 Titans 페이퍼에 나와 있는 3 번째 수식과 동일합니다. 여기서 페이퍼와 동일하게 $ \phi() $를 identity matrix(단위 행렬)로 설정한다면, ($ \phi(x) = Ix = x $), 다음과 같이 수식을 쓸 수 있습니다. (보기 쉽게 $ v_j $를 분자로 올렸습니다.)
 
$$ y_i=\dfrac{Q_i^\top \sum_{j=1}^i K_j^\top V_j}{Q_i ^\top \sum_{l=1}^i  K_l} $$
 
여기서 분모는 normalization을 하는 역할을 수행합니다. 보통은 수식에서 이를 생략한다고 하니, 아래와 같이 재귀적인 수식이 나오겠네요.
 
$$ y_i=Q_i^\top \sum_{j=1}^i K_j^\top V_j $$
 
이를 중간 변수인 $\mathcal{M}_t$를 도입하여 점화식 형태로 바꿀 수 있습니다. ($ \mathcal{M}_0 = 0 $)
 
$$ \mathcal{M} _t = \mathcal{M} _{t-1} + K_t^\top V_t $$
$$ y_t = Q_t \mathcal{M}_t $$
 
이렇게 하면, 페이퍼의 4번과 5번 수식이 나오게 됩니다. 수식을 해석하면, $ t $는 입력 토큰의 위치로써, $ y_t $는 $ t $번째 토큰에 대한 출력 값입니다. $ \mathcal{M}_t $는 히든 스테이트 혹은 압축된 메모리라고 할 수 있겠네요.  실제로 Titans 페이퍼에서는 $ \mathcal{M}_t $가 RNNs의 hidden state와 유사하며, 메모리 유닛이라고 지칭합니다.
 
 

---

 
 

Kernel Attention에서 확장해보자

 Titans 페이퍼에서는 Kernel Attention의 형태만 차용할 뿐, 수식적으로는 다른 형태를 취하게 됩니다. Titans 페이퍼에서는 Kernel Attention의 점화식 형태에 다음과 같은 의미를 부여합니다. 아래는 일반화된 Kernel Attention의 형태입니다. $ x_t $가 갑자기 등장한 것처럼 보이기도 하지만, $ Q = x W_Q $라는 걸 생각하면, 사실 놀라운 건 아닙니다.
 
Memory Write Operation: $ \mathcal{M}_t = f( \mathcal{M}_{t-1}, x_t) $
Memory Read Operation : $ y_t = g( \mathcal{M}_t, x_t) $
 
위의 수식에서 $ \mathcal{M}_t $는 $ t $ 시점의 토큰까지의 모델의 기억이 되고, 그 기억은 업데이트하는 방식으로 $ f(,x_t) $라는 함수를 취했습니다. 위에서 언급한 "Contextual Memory가 입력 변수 $ x $에 영향을 받"는다는 것이 여기서 수식적으로 나오게 되네요.
 

Titans의 장기 기억 방식

 Titans의 장기기억은 위에서 언급한 $ \mathcal{M}_t = f( \mathcal{M}_{t-1}, x_t) $ 수식을 그대로 가져가되, $ f(,) $ 함수를 재구성했습니다. 여기서 surprise라는 용어가 등장합니다. Titans 페이퍼는 다음과 같이 생각합니다. (직접적으로 언급된 말은 아닙니다.)
 

기억은 놀라움(surprise)를 형상화 한 것

이 생각을 기반으로, $ f(, ) $ 함수에 surprise라는 값을 도입합니다. 다음과 같이 말이죠
 
$$ f( \mathcal{M}_{t-1}, x_t) = (1 - \alpha _t ) \mathcal{M}_{t - 1} + S_t $$
$$ S_t = \eta_t S_{t-1} - \theta \nabla \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $$
 
여기서 $ S_t $가 $ t $ 번째 토큰을 만났을 때, 장기기억 모듈이 얼마나 놀랐는지를 측정하는 surprise에 해당하는 값입니다. 그리고 그 값을 기반으로 $ M $을 업데이트하죠. 즉, surprise로 기억을 업데이트합니다.
 
$ S_t $ 외에도 수식을 보면 처음 등장한 변수들이 아래와 같이 있습니다. 
 
$ \alpha_t $: Forget coefficient. 필요 없는 예전 기억을 잊게 해주는 값. $ 0 \sim 1 $.
$ \eta_t $: Past Surprise coefficient. 이전 surprise를 얼만큼 반영할지. $ 0 \sim 1$.
$ \theta_t $: Momentary Surprise coefficient. 순간적인 surprise를 얼만큼 반영할지.
$ S_{t-1} $ : Past Surprise. 과거부터 현재까지의 surprise의 누적 합
$ \nabla \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $: Momentary Surprise. 새로 들어오는 surprise에 대한 값.
$ \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $: Momentary Surprise를 위한 Loss 함수. Test Time 때 모델을 update하기 위해 사용됨.
 
여기서 마지막의 $ \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $를 보면, Test Time 때 모델을 update하기 위해 사용된다고 합니다. 
 

Test Time Learning?

 Test Time 때 Titans 모델이 학습한다는 이유는 바로 $ \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $ 덕분입니다. 이 Loss 함수는 다음과 같이 Key와 Value로 구성돼 있습니다. 여기서 $ h(, ) $는 페이퍼에서 $ \mathcal{M}_{t-1}(k_t) $로 나와 있고, 임의의 어떤 연산(MLP 연산)인 것을 확인할 수 있습니다. (페이퍼에서 수식 11 ~ 12)
 
$$ \mathcal{L}( \mathcal{M}_{t-1}; x_t) = \lvert\lvert h( \mathcal{M}_{t-1}, x_t W_k) - x_t W_v \rvert\rvert_2^2 $$
 
위의 방식에서 변수들 모두 모델 내부 값에 해당되기 때문에, test time 때 최적화가 가능하며, 이 때 최적화 대상은 간접적으로 $ \mathcal{M}_t $가 되겠네요. 이 최적화 과정을 inner loop라고 페이퍼에서는 표현합니다. 그 외로 training 단계에서는, inner loop 최적화인 $ \mathcal{M}_t $를 최적화하는 과정과, outer loop인 $ W_k, W_v $, 그리고 그 외의 모델의 모든 부분을 최적화 하는 단계가 따로 있는 것 같습니다. 
 

Output

 

주어진 수식을 종합하면, long-term memory는 다음과 같이 구성됩니다.

Memory Write Operation: $ \mathcal{M}_t = (1 - \alpha _t ) \mathcal{M}_{t - 1} + S_t $
Surprise Update: $ S_t = \eta_t S_{t-1} - \theta \nabla \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $
Momentary Surprise: $ \nabla \mathcal{L}( \mathcal{M}_{t-1}; x_t) = \nabla \lvert\lvert h( \mathcal{M}_{t-1}, x_t W_k) - x_t W_v \rvert\rvert_2^2$
Memory Read Operation : $ y_t = g( \mathcal{M}_t, x_t) $
 
페이퍼에서는 다음과 같은 표기법을 사용하여 표현합니다.

Memory Write Operation: $ \mathcal{M}_t = (1 - \alpha _t ) \mathcal{M}_{t - 1} + S_t $
Surprise Update: $ S_t = \eta_t S_{t-1} - \theta \nabla \mathcal{L} ( \mathcal{M}_{t-1}; x_t) $
Momentary Surprise: $ \nabla \mathcal{L}( \mathcal{M}_{t-1}; x_t) = \nabla \lvert\lvert \mathcal{M}_{t-1}(x_t W_k) - x_t W_v \rvert\rvert_2^2$
Memory Read Operation : $ y_t = \mathcal{M}^*(x_t .W_Q) $
 
 

---

 

Persistent Memory

 앞에서 Titans는 Core, Contextual, 그리고 Persistent memory가 있다고 했습니다. 이 세 메모리 중에서 유일하게 Persistent memory만 test time 때 학습이 안 되는 것으로 나와 있습니다.
 페이퍼에 따르면, 이 Persistent memory는 task 자체에 대한 정보를 담고 있고, 일종의 prompt tuning 방식으로 동작되는 것 같습니다. prefix 방식으로 입력 tokens의 앞단에 부착되며, Core 모듈이 이 tokens를 처리하는 방식으로 진행되죠.
 
 

---

 

마무리

  Titans 페이퍼에서 중요한 개념만 정리하고, 해석했습니다. Titans 페이퍼는 모델에 대한 구체적인 설명이 부족하여 (페이퍼에서도 그렇게 밝힘) 이해하기 난해했지만, 굵직한 부분은 이해하고 잘 전달된 것 같습니다. 이 Titans 모델은 Transformer나 Mamba 같은 모델들의 성능을 대부분의 영역에서 뛰어 넘었습니다. 아마도 이 구조가 앞으로는 대세가 될 날이 오지 않을까 기대합니다.