Titans: Learning to Memorize at Test Time

https://arxiv.org/abs/2501.00663

Titans: Learning to Memorize at Test Time

 

댓글에 논문 추천도를 확인하는 것을 추천

 

개요
10년이 넘는 기간 동안 순환 모델(recurrent model)과 어텐션(attention)을 효과적으로 활용하는 방법에 대해 폭넓은 연구가 이루어져 왔다. 순환 모델은 데이터를 고정된 크기의 메모리(은닉 상태)로 압축하는 것을 목표로 하는 반면, 어텐션은 전체 문맥 윈도우에 집중하여 모든 토큰의 직접적인 의존 관계를 포착할 수 있도록 한다. 그러나 이러한 더 정확한 의존 관계 모델링은 이차적(quadratic)인 계산 비용을 발생시켜, 모델이 고정된 길이의 문맥만을 처리하도록 제한한다.

 

본 연구에서는 과거 문맥을 학습해 기억하면서, 어텐션이 장기 과거 정보를 활용함과 동시에 현재 문맥에 집중하도록 돕는 새로운 신경망 기반 장기 메모리 모듈(neural long-term memory module)을 제안한다. 우리는 이 신경 메모리가 빠른 병렬 학습(parallelizable training)이 가능하고 동시에 빠른 추론 속도를 유지한다는 이점을 지닌다는 점을 보인다. 메모리 관점에서, 문맥 길이가 제한적이지만 정확한 의존 관계를 모델링하는 어텐션은 ‘단기 메모리’ 역할을, 데이터를 기억하는 능력으로 인해 신경 메모리는 ‘장기적이고 지속적인 메모리’ 역할을 한다고 주장한다.

 

이 두 모듈을 기반으로 우리는 Titans라 불리는 새로운 아키텍처 계열을 제안하고, 이 아키텍처에 메모리를 효과적으로 통합하기 위한 세 가지 변형 모델(variant)을 소개한다. 언어 모델링, 상식 추론(common-sense reasoning), 유전체학(genomics), 시계열(time series) 과제에서의 실험 결과를 통해 Titans가 트랜스포머(Transformer) 및 최근의 선형 순환 모델(linear recurrent model)보다 더 높은 성능을 보여 줌을 확인했다. 또한, ‘건초 더미에서 바늘 찾기(needle-in-haystack)’ 같은 극단적인 상황에서도 더 높은 정확도를 유지하면서 200만(2M) 이상의 문맥 윈도우 크기까지 효과적으로 확장 가능함을 보였다.

 

1 서론
“진정한 기억의 기술은 주의(attention)의 기술이다!”
— 사무엘 존슨(Samuel Johnson), 1787

 

트랜스포머(Transformers, Vaswani et al. 2017)는 순수 어텐션(attention) 기반 아키텍처로, 문맥 내 학습(in-context learning) 능력과 대규모 확장 학습(Kaplan et al. 2020) 덕분에 시퀀스 모델링 분야에서 사실상의 최신(state-of-the-art) 모델로 자리매김했다. 트랜스포머의 핵심 구성 요소인 어텐션 모듈은 연관 기억(associative memory) 블록처럼 작동하여(Bietti et al. 2024), 쿼리(검색 신호)와 키(문맥)의 쌍wise 유사도를 계산함으로써 키-값(key-value) 연관 관계를 저장하고 검색한다. 이렇게 설계된 트랜스포머의 출력은 현재 문맥 윈도우 내 토큰들의 직접적인 의존 관계에 전적으로 좌우된다. 그러나 이러한 높은 의존 관계 모델링 정확도는 문맥 길이에 대해 시간과 메모리 사용량 모두 이차적(quadratic)으로 증가하는 복잡도를 야기한다. 복잡한 실제 과제—예컨대 언어 모델링(N. F. Liu et al. 2024), 비디오 이해(C.-Y. Wu et al. 2019), 장기 시계열 예측(H. Zhou et al. 2021) 등—에서는 문맥 윈도우가 극도로 커질 수 있으므로, 이러한 다운스트림 과제에서 트랜스포머를 활용하기가 어려워질 수 있다.

 

트랜스포머의 확장성 문제를 해결하기 위해, 최근 연구들은 다양한 선형 트랜스포머(linear Transformers) 변형 모델을 설계하고 있다(Kacham, Mirrokni, and P. Zhong 2024; Katharopoulos et al. 2020; S. Yang, B. Wang, Shen, et al. 2024). 이들은 어텐션에서 소프트맥스(softmax)를 커널 함수로 대체(자세한 내용은 §2.1 참고)하여 메모리 사용량을 크게 줄인다. 그러나 긴 문맥을 처리하는 효율성과 확장 가능성을 확보했음에도, 선형 트랜스포머는 여전히 기존 트랜스포머와 비교했을 때 성능이 떨어지는 경우가 많다. 이는 커널 트릭이 선형 순환 신경망(linear recurrent network)을 형성하게 만들어, 데이터 전체를 행렬 형태(matrix-valued states)로 압축하여 처리하기 때문이다(Katharopoulos et al. 2020). 이로 인해 선형 순환(또는 선형 트랜스포머) 모델에는 모순되는 사실이 존재한다. 한편으로는 (선형 vs. 이차적 복잡도) 확장성과 효율성을 높이기 위해 이러한 선형 모델을 적용하지만, 다른 한편으로는 아주 긴 문맥을 소수의 벡터나 행렬 상태에 제대로 압축하기 어려운 문제가 발생한다(S. Wang 2024).

 

더 나아가, 효율성뿐 아니라, 호프필드 네트워크(Hopfield Networks, Hopfield 1982)부터 LSTM(Jürgen Schmidhuber and Hochreiter 1997), 트랜스포머(Vaswani et al. 2017)에 이르기까지 대부분의 기존 아키텍처들은 일반화(generalization), 길이 외삽(length extrapolation), 추론(reasoning) 과제에서 어려움을 겪는다(Anil et al. 2022; Qin, Y. Zhong, and Deng 2024). 이는 복잡한 실제 과제에서 불가분한 핵심 요소들이다. 비록 이러한 아키텍처들이 인간 뇌에서 영감을 받았다고는 하지만, 각각 다음과 같은 부분이 결여되어 있다: (1) 단기 메모리, 장기 메모리, 메타 메모리, 현재 문맥 집중(attending to current context) 등 학습 과정에 필수적인 구성 요소(Cowan 2008); (2) 이들이 독립적으로 작동하면서도 상호 연결되는 시스템; 그리고/또는 (3) 데이터로부터 능동적으로 학습하고 과거 기록의 추상화를 기억하는 능력. 우리는 인간 뇌와 유사한 효과적인 학습 패러다임에서는 학습에 중요한 역할을 담당하는 서로 다른 모듈이 존재하되, 이 모듈들이 긴밀하게 상호 연결되어야 한다고 주장한다.

 

메모리 관점(Memory Perspective)
메모리는 인간 학습에서 매우 근본적인 정신 과정으로, 떼어놓을 수 없는 구성 요소다 (Terry 2017). 메모리 시스템이 제대로 작동하지 않는다면, 인간과 동물은 기본적인 반사 작용이나 정형화된 행동만 할 수 있을 것이다. 이러한 이유로, 메모리는 기계 학습 분야의 많은 기념비적 연구에 영감을 주었다. 예컨대 호프필드 네트워크(Hopfield Networks, Hopfield 1982), LSTM(Jürgen Schmidhuber and Hochreiter 1997), 트랜스포머(Vaswani et al. 2017) 등이 대표적인 사례다.

 

신경심리학(neuropsychology) 문헌에서 제시하는 일반적인 기억(memory)과 학습(learning)의 정의(Okano, Hirano, and Balaban 2000)에서 영감을 받아, 대부분의 기존 아키텍처들은 메모리를 ‘입력에 의해 야기되는 신경 업데이트(neural update)’로 간주하며, 주어진 목적(objective)에 부합하는 ‘효과적이고 유용한 메모리 획득 과정’을 학습이라고 정의한다. 이런 관점에서, 순환 신경망(Recurrent Neural Networks, RNNs) (Williams and Zipser 1989)은 벡터 형태의 메모리 모듈 𝑀(은닉 상태라고도 함)을 갖는 모델로 볼 수 있으며, 이 모델은 크게 두 단계를 수행한다. 시간 𝑡에서 새 입력 𝑥ₜ가 주어졌을 때, (1) 𝑓(𝑀ₜ₋₁, 𝑥ₜ) 함수를 통해 메모리를 업데이트(압축 포함)하고, (2) 𝑔(𝑀ₜ, 𝑥ₜ) 함수를 통해 해당 입력에 대응하는 메모리를 검색한다(자세한 내용은 §2.1 참조).

 

유사하게, 트랜스포머(Transformers)도 확장 가능한 메모리를 갖춘 아키텍처이며, 이 역시 두 단계로 나눌 수 있다. 즉, 키(key)와 값(value)로 구성된 행렬 쌍이 모델의 메모리 역할을 하며, 모델은 (1) 키와 값을 메모리에 계속 덧붙여가면서(압축 없이) 메모리를 업데이트하고, (2) 쿼리(query)와 키 사이의 유사도를 구함으로써 쿼리 벡터와 대응되는 메모리를 검색한 뒤, 이 유사도를 이용해 값 벡터를 가중하여 출력을 생성한다.

 

이 같은 관점은 기존 패러다임의 핵심 차이를 더욱 잘 이해하고, 보다 효율적인 아키텍처를 설계하는 데 도움을 준다. 예컨대, 트랜스포머(Vaswani et al. 2017)와 선형 트랜스포머(linear Transformers, Katharopoulos et al. 2020)의 가장 큰 차이는 바로 메모리 구조와 메모리 업데이트 단계다. 선형 트랜스포머는 과거 데이터를 고정된 크기의 행렬형 메모리(matrix-valued memory)로 압축하는 데 비해, 트랜스포머는 (문맥 길이가 허용하는 한) 과거 데이터를 전혀 압축 없이 모두 보관한다. 선형 트랜스포머와 선형 RNN(상태공간 모델 포함)은 모두 메모리 업데이트 단계에서 정보를 압축하지만, 두 모델 간 결정적 차이는 ‘메모리 구조’에 있다. 선형 RNN은 (선형 트랜스포머와 달리) 벡터 형태의 메모리를 사용한다(선형 트랜스포머는 행렬 형태). 이러한 사실을 통해 다음과 같은 질문을 제기할 수 있다.

• (Q1) “좋은 메모리 구조란 무엇인가?”
• (Q2) “적절한 메모리 업데이트 메커니즘은 무엇인가?”
• (Q3) “효과적인 메모리 검색 과정은 어떠해야 하는가?”

인간의 기억을 다시 살펴보면, 기억은 단일한 과정이 아니며 단일 기능만 수행하는 것도 아니다(Cowan 2008). 실제로 단기 기억, 작업 기억(working memory), 장기 기억 등 다양한 기억 시스템이 존재하며, 각기 다른 신경 구조를 갖고 서로 독립적으로 작동할 수 있다(Willingham 1997). 이는 또 다른 질문을 이끈다.

• (Q4) “서로 긴밀하게 연결된 여러 메모리 모듈을 효율적으로 포함하는 아키텍처는 어떻게 설계할 수 있을까?”

마지막으로, 기억을 저장하는 것은 과거 정보의 추상화를 부호화하고 저장해야 하는 신경 과정이다. 하나의 벡터나 행렬만으로—그리고 이 안의 매개변수가 선형적 방식으로 데이터를 담고 있다고 가정함으로써—장기 이력을 모두 담기에 충분하다고 보는 것은 과도한 단순화일 수 있다.

• (Q5) “오랜 과거를 효과적으로 저장·기억하기 위해서는 ‘깊은(deep) 메모리 모듈’이 필요한가?”

공헌 및 로드맵
본 논문에서는 효율적으로 장기 기억을 학습하고 테스트 시점에 이를 기억할 수 있는 신경망 기반 장기 메모리 모듈을 설계함으로써, 앞서 제기한 다섯 가지 질문에 답해 보고자 한다. 이를 바탕으로, 이 메모리 모듈을 전체 아키텍처에 어떻게 통합할 수 있는지도 논의한다.

1. 신경 메모리(§3)
   우리는 (딥(deep)) 신경 기반 장기 메모리(neural long-term memory)를 제안한다. 이 모듈은 ‘메타 인-컨텍스트 모델(meta in-context model)’로서, 테스트 시점에 데이터를 자신의 파라미터에 저장·기억하는 법을 학습한다.
   인간 장기 기억 시스템(Mandler 2014)에서 영감을 받아, 놀라움(surprise)을 유발하는(즉, 예상을 벗어나는) 이벤트가 더욱 잘 기억되도록 이 메모리 모듈을 설계했다. 이를 위해, 우리는 연관 기억 손실(associative memory loss)에 대해 입력에 대한 신경망의 기울기(gradient)를 측정함으로써 입력의 ‘놀라움’을 계산한다(자세한 내용은 §3.1 참조).
   또한 제한된 메모리 용량을 더욱 효율적으로 활용하기 위해, 메모리 크기와 데이터의 놀라움 정도를 함께 고려하는 감쇠(decaying) 메커니즘을 제안함으로써 더 나은 메모리 관리가 가능하도록 했다. 우리는 이 감쇠 메커니즘이 최신 순환 모델에서 사용되는 망각(forgetting) 메커니즘(예: Dao and Gu 2024; Gu and Dao 2024; S. Yang, Kautz, and Hatamizadeh 2024)의 일반화 버전임을 보인다. 더 흥미로운 점은, 이 메커니즘이 곧 메타 신경망(meta neural network)을 미니배치 경사 하강(mini-batch gradient descent), 모멘텀, 가중 감쇠(weight decay)로 최적화하는 것과 동일하다는 사실이다.
   마지막으로, 미니배치 경사 하강을 텐서 연산으로 확장(tensorizing)하여 보다 많은 행렬 곱(matmul) 연산을 활용하는 기존 기법(Yu Sun et al. 2024)에 기반해, 우리는 병렬 처리가 가능하면서도 빠른 딥 신경 장기 메모리 학습 알고리즘을 제시한다.
2. Titans 아키텍처(§4)
   장기 신경 메모리를 설계한 뒤에는, 이를 어떻게 깊은 학습(deep learning) 아키텍처에 효과적이면서도 효율적으로 통합할지 결정하는 문제가 남는다. 우리는 이를 위해 Titans라는 새로운 딥 모델 계열을 제안한다. Titans는 세 가지 하이퍼-헤드(hyper-head)로 구성된다.
   1. Core: 이 모듈은 ‘단기 메모리(short-term memory)’를 포함하며, 입력 데이터를 주 흐름(main flow)에서 처리하는 핵심 역할을 담당한다(우리는 제한된 윈도우 크기의 어텐션을 사용한다).
   2. 장기 메모리(Long-term Memory): 이 브랜치는 우리가 제안하는 신경 기반 장기 메모리 모듈로서, 오래전 정보를 저장·기억하는 기능을 담당한다.
   3. 영속 메모리(Persistent Memory): 이는 태스크와 관련된 사전 학습(또는 지식)을 담고 있는, 데이터와는 독립적인 매개변수들의 집합이다.
      마지막으로, 개념 증명을 위해 이 메모리 모듈을 활용하는 세 가지 Titans 변형 모델을 제안한다: (i) 메모리를 맥락(context)으로 사용하는 방식, (ii) 메모리를 층(layer)에 통합하는 방식, (iii) 게이트(gated) 브랜치로 사용하는 방식 등이다.
3. 실험 결과(§5)
   우리는 언어 모델링, 상식 추론, ‘기억 회상’이 중요한 과제(needle in haystack), 시계열 예측, DNA 모델링 등의 실험을 수행하였다. 그 결과, Titan 아키텍처는 다양한 벤치마크 전반에서 최신 순환 모델들과 이들을 슬라이딩 윈도우 어텐션과 결합한 하이브리드 모델들을 모두 능가하였다. 또한, Titans는 같은 문맥 윈도우를 사용하는 트랜스포머보다 높은 성능을 보였으며, 전체 문맥을 활용하는 트랜스포머와도 대등한 성능을 달성했다. 더 나아가 트랜스포머와 달리, Titans는 문맥 윈도우 크기가 200만(2M)을 넘어서는 극단적인 상황에도 확장 가능하다는 점이 특징이다.

2 사전 지식(Preliminaries)

 

2.1 배경

어텐션(Attention)

 

효율적 어텐션(Efficient Attentions)

긴 시퀀스에서 소프트맥스 어텐션의 메모리 사용량과 처리량 문제를 완화하기 위해, 다양한 연구들이 I/O 처리에 대한 고려(Dao 2024; Dao, D. Fu, et al. 2022), 어텐션 행렬 희소화(sparsifying) (B. Chen et al. 2021; Choromanski et al. 2021; Dai et al. 2019), 소프트맥스 근사(Arora et al. 2024), 커널 기반(선형) 어텐션(Aksenov et al. 2024; Kacham, Mirrokni, and P. Zhong 2024; Schlag, Irie, and Jürgen Schmidhuber 2021; S. Yang, B. Wang, Shen, et al. 2024) 등 다양한 접근을 시도해 왔다.

 

 

최신 선형 모델과 그 메모리 관점
앞서 살펴본 바와 같이, 학습을 ‘효과적이고 유용한 메모리를 획득하기 위한 과정’으로 정의할 수 있다. 이 관점에서 순환 신경망(RNN)의 은닉 상태(hidden state)는 모델이 정보 전체를 압축해 담으려는 메모리 단위로 볼 수 있다. 이에 따라, 일반적인 순환 신경망을 다음과 같은 두 단계(읽기·쓰기 연산)로 나누어 볼 수 있다.

 

 

이를 해결하기 위해, 최근 연구들은 다음 두 가지 방향에서 접근하고 있다.

1. 망각(forgetting) 메커니즘 추가:
   선형 모델에서 필요할 때 메모리를 지우도록 해주는 적응형(데이터 종속형) 망각 게이트를 설계하는 연구들이 진행되어 왔다. 예컨대 GLA (S. Yang, B. Wang, Shen, et al. 2024), LRU (Orvieto et al. 2023), Griffin (De et al. 2024), xLSTM (Beck et al. 2024), Mamba2 (Dao and Gu 2024) 등이며, 이 중 Mamba2는 기존 상태공간 모델(state space models)의 이산화(discretized) 버전과 연결된다(Gu and Dao 2024).
2. 쓰기 연산 개선:
   전통적 순환 모델의 ‘덧셈’ 방식 쓰기 연산의 한계를 극복하기 위해, Widrow and Hoff (1988)는 쓰기 직전에 과거 메모리(즉, 키-값 쌍)를 제거하는 ‘델타 규칙(Delta Rule)’을 제안했다. S. Yang, B. Wang, Yu Zhang, et al. (2024)은 병렬 학습(parallelizable training)과 확장성(scaling)을 높이기 위한 빠른 병렬 알고리즘을 제안했으며, 최근에는 S. Yang, Kautz, and Hatamizadeh (2024)가 델타넷(DeltaNets)에 망각 게이트를 추가해 성능을 개선했다.

메모리 모듈(Memory Modules)

메모리는 신경망 설계에서 언제나 핵심 요소였다(Graves, Wayne, and Danihelka 2014; JH Schmidhuber 1992; Jürgen Schmidhuber and Hochreiter 1997; J. Zhang et al. 2024). 특히 선형 계층(linear layers)을 키-값(연관) 메모리로 바라보는 아이디어는 ‘빠른 가중치 프로그램(fast weight programs)’ 개념까지 거슬러 올라가는데, 이는 순환 신경망 안에 동적으로 기록할 수 있는 빠른 프로그램을 도입해 메모리 역할을 수행하게 하는 방식이다(JH Schmidhuber 1992). 여기에는 Hebbian(Hebb 2005)과 Delta(Prados and Kak 1989)라는 두 가지 주요 학습 규칙이 적용되며, 다양한 연구에서 폭넓게 활용되어 왔다(Irie, Schlag, et al. 2021; Munkhdalai, Sordoni, et al. 2019; Munkhdalai and H. Yu 2017; Schlag, Irie, and Jürgen Schmidhuber 2021; JH Schmidhuber 1992; S. Yang, Kautz, and Hatamizadeh 2024; S. Yang, B. Wang, Yu Zhang, et al. 2024). 그러나 이들 모델은 순간적인 놀라움(momentary surprise) 에만 기반하고 있어(3.1절 참조), 시퀀스 내 토큰의 흐름을 반영하지 못한다는 한계가 있다. 또한 망각 게이트가 없어 메모리 관리가 미흡한 경우가 많다.

 

추가로, 본 논문의 아키텍처와 최신 모델들의 관련성은 부록 C에서 자세히 논의하고, 그 외 관련 연구는 부록 A에서 소개한다.

 

3 테스트 시점에서의 메모리 학습(Learning to Memorize at Test Time)
장기 기억(long-term memory)의 부재를 해결하고, 모델이 정보를 학습·망각·검색(retrieve)할 수 있도록 하기 위해, 본 섹션에서는 테스트 시점에 메타 학습(meta learning)을 통해 데이터를 기억하는 신경 기반 장기 메모리 모듈을 제안한다. 먼저 3.1절에서 신경 메모리 모듈의 설계 동기와 구체적인 메커니즘을 논의하고, 이어서 3.2절에서는 이러한 아키텍처 설계가 빠르고 병렬적인 학습을 통해 어떻게 이점을 얻을 수 있는지 설명한다. 마지막으로 3.3절에서는 태스크에 대한 메타 정보를 학습하기 위한 데이터 비종속적(data-independent) 파라미터로 구성된 영속 메모리(Persistent Memory) 모듈을 결합하는 방법을 소개한다.

 

3.1 장기 메모리(Long-term Memory)

장기 메모리를 설계하려면, 과거 이력을 추상화하여 파라미터로 부호화(encode)할 수 있는 모델이 필요하다. 예를 들어, 대규모 언어 모델(LLMs)은 훈련 데이터 자체를 기억하는 양상을 보이기도 한다(Leybzon and Kervadec 2024; Schwarzschild et al. 2024; Staab et al. 2024). 가장 간단한 생각은 모델을 훈련해 훈련 데이터를 그대로 기억하게 하는 것이다. 하지만 이런 식의 ‘암기(memorization)’는 대부분 바람직하지 않은 현상으로 여겨진다. 그 이유는 (1) 모델의 일반화 성능을 저해한다(Bayat et al. 2024), (2) 사생활 정보 보호 문제(privacy concerns)가 발생할 수 있다(Staab et al. 2024), (3) 테스트 시점에서 분포가 다른 데이터가 들어왔을 때 별 도움이 되지 않을 수 있기 때문이다.

 

따라서 본 연구에서는 온라인 메타 모델(online meta-model) 이 테스트 시점에서 데이터를 어떻게 기억/망각해야 할지를 학습하도록 하는 방식을 택했다. 이 경우, 모델은 데이터를 효과적으로 기억할 수 있는 함수를 학습하지만, 훈련 데이터 자체에 오버피팅(overfitting)되지 않으므로 테스트 시점에서의 일반화 성능이 향상된다.

학습 과정과 놀라움 지표(Surprise Metric)

단, 위 놀라움 지표만 사용하면, 강한 놀라움이 발생한 이후에 들어오는 중요한 정보를 놓칠 위험이 있다. 예를 들어, 연속된 놀라운 이벤트 몇 개 이후 기울기가 매우 작아져서(평탄 지점에 빠져서) 시퀀스 일부 정보를 잃어버릴 수 있다. 인간 기억 관점에서도, 처음 놀라웠던 사건이 오랜 기간 꾸준히 놀라움을 유발하지는 않지만, 한 번 충분히 놀라움을 줬다면 그 후 시점까지의 전체 맥락을 계속 기억하는 경우가 많다.

이 문제를 개선하기 위해, 우리는 놀라움을 두 가지로 나누어 측정한다.

1. 과거 놀라움(past surprise): 바로 직전까지의 놀라움을 측정한다.
2. 순간 놀라움(momentary surprise): 새로 들어오는 데이터의 놀라움을 측정한다.

이를 통해, 메모리 업데이트 공식을 다음과 같이 정의한다.

 

목표(Objective)

 

망각(forgetting) 메커니즘

매우 긴 시퀀스(예: 수백만 개의 토큰)를 다룰 때, 어떤 과거 정보는 더 이상 필요치 않을 수 있으므로 이를 적절히 망각(제거)하는 것이 중요하다. 심지어 딥(deep)하거나 매우 큰 행렬형 메모리를 쓰더라도, 유한한 메모리 용량은 결국 제한이 있기 때문이다.

 

이를 위해, 우리는 적응형(adaptive) 망각 기법을 도입한다. 이는 필요한 경우 메모리를 지우고, 그렇지 않다면 유지함으로써 메모리를 더 효과적으로 관리할 수 있도록 해준다. 구체적으로, 다음과 같이 업데이트 규칙을 수정한다.

이후 본 절에서 보이겠지만, 이런 가중 감쇠(weight decay) 형태의 망각 메커니즘은 최신 순환 신경망(RNN)의 게이트 구조(Dao and Gu 2024; Orvieto et al. 2023)와 밀접한 관련이 있다.

 

메모리 아키텍처(Memory Architecture)

이 식은 온라인 선형 회귀(online linear regression) 문제에 해당하므로, 결국 과거 데이터의 종속 관계가 선형적이라고 가정하는 해가 최적이라는 뜻이다.

 

 

---

 

메모리 아키텍처 비교 선형 메모리 (Matrix) 입력 데이터 W ∈ ℝᵈⁱⁿˣᵈⁱⁿ 출력 딥 메모리 (MLP) 입력 데이터 Hidden Layer 1 Hidden Layer 2 출력

1. 왼쪽: 선형 메모리 (Matrix)
   • 단순한 행렬(W)을 사용하여 메모리를 구현합니다
   • 입력 데이터가 선형 변환만 거치게 됩니다
   • 데이터의 선형적 관계만 포착할 수 있습니다
   • 표현력이 제한적입니다
2. 오른쪽: 딥 메모리 (MLP)
   • 2개 이상의 층을 가진 신경망을 사용합니다
   • 비선형 활성화 함수를 통해 복잡한 패턴을 학습할 수 있습니다
   • 더 풍부한 표현력을 가집니다
   • 복잡한 데이터 관계도 포착할 수 있습니다

논문에서 주장하는 핵심은, 단순한 선형 변환(행렬)보다는 여러 층을 가진 신경망(MLP)을 사용하면 더 효과적으로 정보를 기억하고 처리할 수 있다는 것입니다. 실제 실험 결과(5.5절)에서도 이러한 딥 메모리 구조가 더 좋은 성능을 보였다고 합니다.

---

메모리 검색(Retrieving a Memory)

앞서 설명했듯이, 테스트 시점에 데이터를 기억하도록 학습된 장기 메모리 모듈을 어떻게 설계·훈련하는지 논의했다. 하지만 핵심적인 질문은, 메모리에 저장된 정보를 어떻게 검색(retrieve)할 것인가? 하는 것이다.

 

 

---

그냥 transformer는 쿼리 형식으로 메모리 검색이지만, 여기서는 mlp 2개이상으로 순전파로 검색하는것으로 생각하면됨.

메모리 검색 방식 비교 Transformer 본 논문의 Deep Memory Query Key 1 Value 1 Key 2 Value 2 유사도 계산 가중합 출력 입력 Hidden Layer 1 (with non-linear activation) Hidden Layer 2 (with non-linear activation) 출력 단순 순전파 비선형 변환 Transformer: Query-Key 유사도 기반 Value 검색Deep Memory: 비선형 변환을 통한 직접 매핑

 

---

 

그림 1: 매트럴을 사용하여 신경 기억 훈련을 병렬로 수행할 수 있는 방법을 보여주는 그림입니다.

 

3.2 장기 메모리 학습을 병렬화하는 방법

청크 단위로 파라미터를 정의하기

---

순차적 vs 병렬 메모리 학습 순차적 처리 (느림) x₁ x₂ x₃ 각 시점 t마다 순차적으로 계산: Mt = (1 - αt)Mt-1 - θt∇ℓ(Mt-1; xt) 병렬 처리 (빠름) 배치 입력 [x₁, x₂, ..., xb] GPU 행렬 연산 한 번에 처리 배치 출력 [M₁, M₂, ..., Mb] 배치 단위 행렬 연산으로 변환: ∑(θi/βi)βb∇ℓ(W0; xi) = ΘbBb(W0X - X)X^T

 

---

 

그림 2: Context(맥락)으로서의 메모리(MAC) 아키텍처
이 아키텍처에는 (1) 코어(core), (2) 문맥적(장기) 메모리, (3) 영속(persistent) 메모리의 세 가지 브랜치가 있다. 코어 브랜치는 입력 시퀀스에 해당하는 토큰들과 그에 대응하는 장기·영속 메모리를 연결(concatenate)한다. 이후 어텐션이 시퀀스에 적용되어, 어떤 정보를 장기 메모리에 저장할지 결정한다. 테스트 시점에는 컨텍스트 메모리에 해당하는 파라미터는 계속 학습되고, 코어 브랜치의 파라미터는 인-컨텍스트 학습(in-context learning)을 담당하며, 영속 메모리 파라미터는 태스크 관련 지식을 저장하는 역할을 하므로 고정된다.

 

3.3 영속 메모리(Persistent Memory)

우리의 장기 메모리는 ‘문맥적 메모리(contextual memory)’로 볼 수 있는데, 이는 모든 파라미터가 입력에 따라 동적으로 변화한다는 의미다. 하지만 효과적인 메모리 시스템은 입력과 무관하게(input-independent) 태스크 지식을 저장할 수 있는 파라미터도 필요하다. 예컨대, 태스크를 숙달하기 위해서는 "태스크를 어떻게 수행해야 하는지"라는 지식을 기억해야 하는데, 이러한 정보를 저장·관리하는 파라미터가 별도로 필요한 것이다.

 

 

 

 

그림 3: Titans의 다른 변형들에서 사용되는 어텐션 마스크 예시

4 메모리를 어떻게 통합할 것인가?

여기서 남아 있는 중요한 질문은 다음과 같다. “설계된 신경 메모리를 딥러닝 아키텍처에 어떻게 효과적이고 효율적으로 결합할 수 있을까?” 앞서 논의했듯이, 메모리 관점에서 트랜스포머의 K, V 행렬 쌍은 연관 기억 블록(associative memory block)으로 해석할 수 있다. 트랜스포머가 토큰 간 의존 관계를 매우 정확하게 모델링하는 대신 문맥 윈도우 크기가 제한된다는 점을 감안하면, 이를 ‘단기 메모리(short-term memory) 모듈’ 로 볼 수 있다.

 

반면, 우리 신경 메모리는 지속적으로 데이터를 학습하고 그 가중치에 저장할 수 있는 능력이 있으므로, ‘장기 메모리(long-term memory)’ 의 역할을 수행할 수 있다. 본 절에서는 세 가지 다른 Titans 변형 모델을 제안하여 이 질문에 답해 보고자 한다. 이후 실험을 통해 각 변형 모델이 저마다의 장단점을 갖고 있으며, 매우 긴 문맥을 다룰 때 효율성과 효과성 간의 상충관계(trade-off)를 보인다는 점을 보일 것이다.

 

4.1 Context(맥락)으로서의 메모리

 

이 아키텍처의 주요 장점

 

그림 4: 게이팅으로서의 메모리(MAG) 아키텍처
이 아키텍처 역시 (1) 코어, (2) 문맥적(장기) 메모리, (3) 영속 메모리 세 가지 브랜치를 가진다. 다만 영속 메모리만을 컨텍스트에 직접 통합하며, 장기 메모리는 코어 브랜치와 **게이팅(gating)**을 통해 결합한다. 테스트 시점에서의 동작 방식은 그림 2와 동일하다.

 

4.2 게이트 메모리(Gated Memory)
두 번째 변형 모델(그림 4 참조)에서는, 한쪽 브랜치(branch)에서 입력 데이터를 장기 메모리에 직접 업데이트하고, 다른 브랜치에서는 슬라이딩 윈도우 어텐션(SWA)을 사용한다. 구체적으로,

그림 3b에서 볼 수 있듯이, 슬라이딩 윈도우 어텐션은 정교한 단기 메모리 역할을 수행하며, 신경 메모리 모듈은 서서히 잊혀지는(fading) 메모리처럼 동작한다. 한편, 이 아키텍처는 헤드(head) 구조가 서로 다른 멀티헤드 아키텍처(X. Dong et al. 2024)로도 볼 수 있다.

 

4.3 계층(Layer)로서의 메모리

마지막 변형 모델은 깊은 신경망의 한 층(layer) 으로 신경 메모리를 사용하는 MAL(Memory As a Layer) 구조(그림 5 참조)다. 이는 문헌에서도 흔히 볼 수 있는 방식으로, RNN 계열 모델과 완전 혹은 슬라이딩 윈도우 어텐션을 쌓아올려 하이브리드 형태로 구성하는 것이다. 구체적으로, 입력 x에 대해

의 과정을 거친다(SW-Attn은 슬라이딩 윈도우 어텐션). 이 설계의 주요 단점은, 각 층의 표현력에 의해 전체 모델 성능이 제한될 수 있다는 점이다. 즉, 어텐션과 신경 메모리를 협업적으로(complementary) 결합하지 못하고, 계층적으로 단순 연결하는 데 그칠 수 있다.

 

본 논문 실험에서는, H3 (D. Y. Fu et al. 2023) 아키텍처와 유사하게 LMM이라는 신경 메모리 모듈을 시퀀스 모델로 교체하여 사용한다.

어텐션 없이 메모리만 사용하는 경우

위 설명에서는 MAL을 “LMM + 어텐션” 순차적 조합으로 소개했다. 하지만 이를 어텐션 없이 순차 모델처럼 쓸 수도 있다. 메모리 관점에서(1절 참조), 각 메모리 시스템은 다른 컴포넌트와 무관하게 독립적으로 작동할 수 있어야 한다. 따라서 장기 메모리 모듈은 단기 메모리(어텐션)가 없어도 충분히 강력한 모델이길 기대할 수 있다. 이 변형 모델을 실험에서는 LMM 또는 **Titans (LMM)**이라 부른다. 또한 Titans와 최신 순환 모델들 간의 연관성은 부록 C에서 추가로 논의한다.

 

그림 5: 계층(Layer)으로서의 메모리(MAL) 아키텍처
이 구조에서는 메모리 레이어가 어텐션 모듈 이전에 과거 및 현재 컨텍스트 정보를 압축해 전달하는 역할을 맡는다.

 

---

 

 

• 그림 2
  • 논문에서 제안하는 전체 아키텍처(“Memory as a Context” 방식)를 개념적으로 그린 ‘큰 그림’에 가깝습니다.
  • 그림을 보면, 세 가지 ‘메모리 영역’(Persistent / Contextual / Core)을 구분해 놓았고, Contextual Memory(즉, Neural Memory 모듈)가 테스트 시점에도 학습(Learning)을 계속하며(오른쪽 상단), Attention을 통해 Core 쪽의 Sequence와 데이터를 주고받는 흐름을 묘사하고 있죠.
  • Persistent Memory는 고정된 파라미터로서(오른쪽 하단 ‘Fixed’), 계속 유지되는 지식(태스크 정보 등)을 담고 있다고 보시면 됩니다.
• 그림 3a, 3b
  • 이 둘은 **“Titans의 변형(variant)마다 어텐션이 어떻게 적용되는지”**를 좀 더 구체적으로 보여주는 그림입니다.
  • 예컨대,
    • 그림 3a(MAC: Memory as a Context) 는 세그먼트를 나눈 뒤, 해당 세그먼트와 (Persistent + Long-term) 메모리를 합쳐서(Concatenate) Attention을 수행하는 식의 어텐션 마스크 구조를 그려둡니다.
    • 그림 3b(MAG: Memory as Gating) 는 슬라이딩 윈도우 어텐션을 쓰는 코어 모듈과, 별도의 메모리 모듈을 게이팅으로 결합하는 방식을 그림으로 나타낸 것입니다.

 

 

“복잡하게 해서 얻는 이점이 뭔가?”에 대해

• 장기 메모리(Long-term Memory)를 모델 내부에서 계속 학습(업데이트)할 수 있다는 점이 가장 큰 포인트입니다.
  • 전통적인 트랜스포머는 ‘고정된 윈도우(context window) 안의 토큰들’에 대해서만 직접적인 연관을 학습해요(짧은 문맥).
  • 반면, 이 논문에서는 Neural Memory 모듈이 “테스트 시점에도” 새로운 입력을 받아서 메모리를 갱신하고, 오래된 정보를 (망각 게이트 등을 통해) 정리해나가며, 필요하다면 과거 정보도 다시 ‘회상(retrieve)’ 하는 구조를 만들고자 합니다.
• 이렇게 되면 매우 긴 시퀀스(수백만 토큰 등)를 전부 Attention 으로 처리해야 하는 부담이 줄어들 수 있다고 주장합니다.
  • Transformer는 문맥 길이가 늘어나면 연산량이 O(N^2)로 늘어나지만, 여기서는 필요할 때만 메모리에 요약/업데이트하고, Attention은 제한된 윈도우 범위에서 수행하므로, “긴 시퀀스 전체를 단번에” 보지 않고도 과거 맥락을 어느 정도 효율적으로 저장·검색할 수 있다는 취지죠.
• 또한 Attention이 모든 토큰을 다 고려하지 않고, “이 세그먼트에서 정말 유용한 정보만” 골라서 메모리에 반영하게끔 할 수 있으므로, “길이가 매우 긴 문맥”에서도 메모리를 좀 더 효율적으로 쓴다는 아이디어가 있습니다.
• 다만 말씀대로, 구조가 복잡해지고(메모리 모듈 + Attention + 게이팅 등), 트랜스포머만 쓰는 것보다 초기 구현·학습 측면에서는 오버헤드가 있는 게 사실입니다.
  • 논문 측 주장은, “하지만 (1) 아주 긴 시퀀스에서 O(N^2)를 줄이고, (2) 필요한 정보만 메모리에 요약해서 저장·검색함으로써, 궁극적으로 확장성(Scalability) 측면에서 이득을 볼 수 있다”는 점입니다.
  • 즉, “학습은 조금 더 복잡해지겠지만, 훨씬 더 긴 문맥까지 모델을 돌릴 수 있다”가 저자들이 말하는 장점이라 보시면 됩니다.

---

 

4.4 아키텍처 세부 사항(Architectural Details)
단순성과 가독성을 위해, 본문에서는 구현 관련 세부 사항(예: 잔차 연결(residual connection), 선형 계층으로 게이팅(gating), 정규화(normalization) 활용 등)을 생략했다. 실제 구현에서는 모든 블록에 잔차 연결을 사용하며, 쿼리·키·값을 계산할 때 비선형 활성 함수로는 SiLU(·) (Elfwing, Uchibe, and Doya 2018)을 사용한다. 또한 쿼리와 키는 ℓ_2​-노름을 이용해 정규화한다.

 

컨볼루션(Convolution)

최근의 선형 순환 모델(Gu and Dao 2024; S. Yang, Kautz, and Hatamizadeh 2024)들을 참고하여, 쿼리·키·값 투영(projection) 이후에 1D 딥스와이즈-세퍼러블 컨볼루션(Depthwise-Separable Convolution) 층을 도입했다. 성능에 극적인 변화를 주지는 않으나, 이러한 1D 컨볼루션은 성능 개선에 기여할 뿐 아니라 계산 효율도 높다는 점이 보고되었다.

 

게이팅(Gating)

최종 출력 투영 전에 정규화와 선형 계층을 활용한 게이팅을 적용하는 최근 아키텍처(Mehta et al. 2023)를 참고하였다.

 

 

5 실험(Experiments)
이제 우리는 Titans 및 그 변형 모델들의 성능을 언어 모델링, 상식 추론, ‘건초 더미에서 바늘 찾기’(needle in haystack), DNA 모델링, 시계열 예측 과제에서 평가한다[^1]. 구체적으로, 본 절에서는 다음과 같은 실증적(empirical) 질문에 답한다.

1. Titans는 다운스트림 과제에서 다른 모델 대비 어떤 성능을 보이는가? (5.2, 5.6, 5.7절)
2. Titans의 실제 컨텍스트 길이(context length)는 어느 정도인가? (5.3, 5.4절)
3. Titans는 컨텍스트 길이가 늘어날 때 어떻게 확장(scale)되는가? (5.8절)
4. 메모리의 깊이(depth)는 성능과 효율에 어떤 영향을 미치는가? (5.5절)
5. Titans의 각 구성 요소는 전체 성능에 어떻게 기여하는가? (5.9절)

 

5.1 실험 환경(Experimental Setup)

모델(Models).
본 실험에서는 Titans의 세 가지 변형 모델인

1. MAC (Memory as a Context)
2. MAG (Memory as a Gate)
3. MAL (Memory as a Layer)
   에 더해,
4. 신경 메모리 모듈 단독
   을 함께 평가한다.

장기 메모리를 별도 모듈로 분리한 이유는, 1절에서 언급한 “학습을 ‘효과적이고 유용한 메모리를 획득하는 과정’으로 정의”하는 관점 때문이다. 즉, 주어진 데이터로부터 주어진 목적을 충족하도록 학습하기 위해서는, 어텐션 없이도 장기 메모리가 스스로 데이터를 잘 학습할 수 있으리라 기대한다.

이 모델 각각에 대해, 아래 네 가지 규모로 나누어 실험한다.

1. 170M 파라미터
2. 340M 파라미터
3. 400M 파라미터
4. 760M 파라미터

처음 세 가지(170M, 340M, 400M)는 FineWeb-Edu 데이터셋(Penedo et al. 2024)에서 샘플링한 150억(15B) 토큰으로, 마지막 모델(760M)은 같은 데이터셋에서 300억(30B) 토큰으로 학습했다.

 

비교 대상(Baselines).
우리는 다음과 같은 최신 선형 순환 모델, 트랜스포머, 그리고 ‘순환 + 어텐션’ 하이브리드 모델들과 성능을 비교한다:

• 언어 과제:
  • Transformer++ (Touvron et al. 2023)
  • RetNet (Yutao Sun et al. 2023)
  • Gated Linear Attention (GLA) (S. Yang, B. Wang, Shen, et al. 2024)
  • Mamba (Gu and Dao 2024)
  • Mamba2 (Dao and Gu 2024)
  • DeltaNet (S. Yang, B. Wang, Yu Zhang, et al. 2024)
  • TTT (Yu Sun et al. 2024)
  • Gated DeltaNet (S. Yang, Kautz, and Hatamizadeh 2024)
• ‘건초 더미에서 바늘 찾기’(needle in haystack) 과제:
  • GPT4 (Achiam et al. 2023)
  • Llama3 with RAG (Touvron et al. 2023)
  • RecurrentGemma2-9B (Botev et al. 2024)
  • Mistral (Jiang et al. 2023)
  • 모두 (Yuri Kuratov et al. 2024)의 벤치마크에서 제공된 모델들임.
• 시계열(time series) 과제:
  • Mamba 기반 모델 (Behrouz, Santacatterina, and Zabih 2024)
  • 트랜스포머 기반 모델 (Y. Liu et al. 2023; Nie et al. 2022; Yunhao Zhang and Yan 2023)
  • 선형(linear) 모델 (Das et al. 2023; Z. Li et al. 2023; H. Wu et al. 2023; Zeng et al. 2023)

 

학습(Training).
학습 절차는 S. Yang, Kautz, and Hatamizadeh (2024)를 따르며, 어휘집 크기 32K의 LLaMA 2 토크나이저를 사용하고, 학습 시 시퀀스 길이를 4K 토큰으로 설정한다. 옵티마이저는 AdamW를 쓰며, 학습률(learning rate)은 4×10^−4, 코사인 감소 스케줄(cosine annealing schedule), 미니배치 크기는 50만(0.5M) 토큰, 가중 감쇠(weight decay)는 0.1로 설정한다.

Footnotes

1. 본 논문에서 제시된 모든 실험 코드는 별도의 부록 또는 저장소에서 제공될 것으로 추정된다. ↩

5.2 언어 모델링(Language Modeling)
우선 언어 모델링과 상식 추론 과제에서 퍼플렉시티(perplexity)와 정확도(accuracy)를 측정했다. 세 가지 규모(340M, 400M, 760M 파라미터)에 대해, Titans 변형 모델들과 기존 베이스라인의 결과를 표 1에 정리했다.

• 비하이브리드(Non-hybrid) 모델 비교
  Transformer++를 포함한 비하이브리드 모델들 중에서는, 우리의 신경 메모리 모듈이 퍼플렉시티와 정확도 양 측면에서 가장 우수한 성능을 보였다.
  TTT 역시 기울기 기반의 순환 모델이므로, TTT와의 성능 비교는 **가중 감쇠(weight decay)**와 **모멘텀(momentum)**의 효과를 잘 보여준다. 앞서 논의했듯이, 가중 감쇠는 필요 시 과거 데이터를 망각(forgetting)하는 게이팅(gating) 메커니즘으로 해석될 수 있고, 모멘텀은 놀라움(surprise) 지표를 좀 더 유연하게 관리하여 메모리를 개선한다.
  일부 베이스라인 모델(Mamba, Mamba2, Gated DeltaNet 등)도 게이팅 메커니즘을 사용하지만, 우리 신경 메모리 모듈이 더 높은 성능을 보인다는 점은 **놀라움 지표(surprise mechanism)**의 도입과 깊고 비선형적인(deep, non-linear) 메모리 구조가 결합된 결과로 해석할 수 있다. 깊은 메모리의 효과는 5.5절에서 더 자세히 논의한다.
• 하이브리드(Hybrid) 모델 비교
  Samba(Mamba + 어텐션)와 Gated DeltaNet-H2(Gated DeltaNet + 어텐션) 같은 하이브리드 모델들과 비교한 결과, Titans의 세 변형(MAC, MAG, MAL) 모두 우수한 성능을 나타냈다.
  • Titans (MAL)이 우수한 이유는, 설계와 사용된 어텐션은 동일하지만, 핵심 신경 메모리 모듈의 성능 덕분이라고 볼 수 있다.
  • Titans (MAG)와 (MAC)을 비교해 보면, 두 모델 모두 성능은 유사하나, MAC이 긴 의존 관계(long dependencies)를 처리할 때 조금 더 우수한 경향을 보였다.
  • 한편, MAG와 MAC 모두 MAL 변형을 능가했는데, 이는 세 변형에서 사용하는 구성 요소는 같으나 아키텍처 설계 차이 때문으로 추정된다. 이는 **기존 문헌의 하이브리드 모델들(특히 MAL 방식 사용)**과 다른 접근의 중요성을 시사한다(단, Hymba (X. Dong et al. 2024) 같은 예외도 존재).

표 1: 언어 모델링 및 상식 추론 과제에서 Titans와 순환/트랜스포머 기반 베이스라인들의 성능 비교. 하이브리드 모델은 ∗로 표시. 단순 모델과 하이브리드 모델 각각에서 최고의 결과를 볼드 처리.

 

5.3 Needle in a Haystack(건초 더미에서 바늘 찾기)
모델을 긴 문맥 길이로 확장한다고 해서, 실제로 아주 긴 시퀀스를 효과적으로 처리할 수 있다는 보장이 있는 것은 아니다(Hsieh et al. 2024). Needle-in-a-haystack(NIAH) 과제는 모델의 실질적인 유효 문맥 길이(effective context length)를 평가하기 위해 고안되었다.

본 섹션에서는 **RULER 벤치마크(Hsieh et al. 2024)**의 단일 NIAH(S-NIAH) 태스크를 활용해, 길이 2K, 4K, 8K, 16K의 시퀀스에서 Titans와 베이스라인 모델들을 평가했다(결과는 표 2 참조).

• 신경 메모리 모듈은 모든 태스크에서 베이스라인 대비 최고의 성능을 달성했다. 이는 기존 시퀀스 모델들과 비교했을 때, Titans가 다음과 같은 장점을 지닌 덕분으로 해석된다.
  1. TTT와 비교: 모멘텀과 망각(가중 감쇠) 메커니즘을 통해 메모리 용량을 보다 유연하게 관리할 수 있으므로, 시퀀스 길이가 증가해도 성능 하락 없이 안정적으로 유지된다.
  2. Mamba2와 비교: Mamba2도 게이팅(망각) 메커니즘을 갖고 있으나, Titans는 **깊은 비선형 메모리(deep non-linear memory)**를 사용하여 더 나은 메모리 관리를 구현한다. 또한 Mamba2는 특정 메모리를 완전히 제거할 수 없기 때문에 시퀀스가 길어질수록 성능 저하가 크게 나타난다.
  3. DeltaNet과 비교: DeltaNet도 ‘델타 규칙(delta rule)’을 통해 메모리를 제거할 수 있지만, ‘망각(forgetting)’ 자체는 지원하지 않으므로 시퀀스가 길어질수록 한계가 드러난다.
• Titans의 변형 모델들도 대체로 MAC 변형이 가장 좋은 결과를 보였거나 그에 상응하는 수준을 기록했다.

표 2: RULER 벤치마크의 S-NIAH 태스크에서 Titans와 베이스라인들의 성능 비교. 단순 모델과 하이브리드 모델 각각에서 최고의 결과를 볼드 처리.

 

(a) Few-shot Setup

(b) Fine-Tuning Setup

 

그림 6: BABILong 벤치마크에서의 Titans와 베이스라인 성능. Titans (MAC)은 GPT4 등 매우 큰 모델도 포함한 모든 베이스라인보다 우수한 결과를 보임.

 

5.4 BABILong 벤치마크(BABILong Benchmark)
앞서 단일 NIAH 태스크(하나의 ‘바늘’을 찾는 과제) 결과에서는 Titans가 경쟁 모델 대비 우수한 성능을 보였으나, 이로는 매우 긴 시퀀스에서의 진정한 강점을 충분히 드러내지 못했다. 따라서 본 절에서는 BABILong 벤치마크(Yuri Kuratov et al. 2024) 중 더 까다로운 태스크를 활용한다. 이 과제는 상당히 긴 문서에 흩어져 있는 사실들을 종합적으로 추론해야 하므로, 모델이 매우 긴 문맥을 실제로 처리·추론할 수 있는 능력을 평가하기에 적합하다. 실험 세팅과 학습 과정은 벤치마크의 오리지널 구성과 동일하게 따른다.

BABILong에는 크게 두 가지 설정이 있다.

1. Few-shot 설정: 사전 학습(Pre-trained)이 충분히 된 대형 모델을 활용.
2. Fine-tuning 설정: Titans(MAC 변형)을 미세 조정(fine-tune)한 뒤, 다른 미세 조정된 모델들과 비교.

Few-shot 설정 결과

그림 6a는 Few-shot 설정에서의 결과다. 여기서 Titans(MAC)은 Mamba2.8B(Gu and Dao 2024), RWKV-6-7B(Peng, Goldstein, et al. 2024), RecurrentGemma-9B(Botev et al. 2024), Gemma-9B(Team et al. 2024), Llama3.1-8B(Touvron et al. 2023), GPT-4, GPT4o-mini(Achiam et al. 2023) 등 기존 모델들을 모두 능가한다. 주목할 점은, Titans(MAC)이 이들 대비 훨씬 적은 파라미터로도 더 높은 성능을 달성한다는 것이다.

 

Fine-tuning 설정 결과

fine-tuning 설정에서는, 크기가 작은 Titans(MAC) 모델을 미세 조정한 뒤 다음과 비교했다.

1. 소형 모델(Titans과 파라미터 수가 거의 비슷)
   • Mamba(Gu and Dao 2024)
   • RMT(Bulatov, Yury Kuratov, and Burtsev 2022)
2. 대형 모델(Retrieval-Augmented Generation, RAG)
   • Llama3.1-8B(Touvron et al. 2023) + RAG (P. Lewis et al. 2020)
3. 초거대 모델
   • GPT-4(Achiam et al. 2023)
   • GPT4o-mini
   • Qwen2.5-72B(A. Yang et al. 2024)
   • Llama3.1-70B(Touvron et al. 2023)

(Yuri Kuratov et al. 2024)에서 보고된 베이스라인 결과와, Titans(MAC)의 결과를 그림 6b에 나타냈다. 그 결과, Titans는 GPT-4 같은 초거대 모델들을 포함해 모든 모델보다 더 높은 성능을 보였다. 또한 RMT처럼 트랜스포머 기반에 메모리를 추가한 모델보다도 우수했는데, 이는 Titans가 보다 강력한 메모리 구조를 갖춘 덕분이다. 예컨대, RMT는 과거 데이터를 길이 16짜리 벡터로만 압축해 저장하는 반면, Titans는 온라인 메모리 학습(in-context online memory learner) 방식을 통해 모델 파라미터 자체에 과거 정보를 더 풍부하게 부호화할 수 있다.

 

흥미로운 점은, Llama3.1-8B에 RAG를 결합한 모델 역시 약 70배나 많은 파라미터를 쓰면서도 Titans보다 성능이 낮게 나타났다는 사실이다.

 

그림 7: 메모리 깊이(memory depth)가 퍼플렉시티에 미치는 영향. 더 깊은 장기 메모리는 긴 시퀀스에서 더욱 안정적으로 성능이 향상되는 경향을 보인다.

 

표 3: 장기 예측(long-term forecasting) 성능 비교. 가장 우수한 결과를 볼드 처리.

 

5.5 깊은 메모리(Deep Memory)의 효과

그림 8: 메모리 깊이가 학습 처리량에 미치는 영향

 

5.6 시계열 예측(Time Series Forecasting)
우리 메모리 모듈의 범용성을 입증하기 위해, 시계열 예측 과제에서도 성능을 평가했다. 구체적으로, Simba 프레임워크(Patro and Agneeswaran 2024)에 내장된 Mamba 모듈을 우리 신경 메모리로 대체하여 테스트했다.
ETT, ECL, Traffic, Weather (H. Zhou et al. 2021) 등 대표적인 시계열 예측 벤치마크 데이터셋에서의 결과를 표 3에 요약했다. 그 결과, 신경 메모리 모듈은 Mamba 기반, 선형(Linear) 기반, 트랜스포머(Transformer) 기반 아키텍처를 모두 능가하는 성능을 보였다.

 

표 4: GenomicsBenchmarks (Grešová et al. 2023)에서 사전 학습된 DNA 모델들의 다운스트림 결과. Top-1 분류 정확도(%)를 보고.

 

5.7 DNA 모델링(DNA Modeling)
자연어 영역을 넘어 Titans의 능력을 평가하기 위해, 우리의 신경 메모리 모듈을 DNA 모델링 과제에 적용했다. 구체적으로 사전 학습된(pre-trained) 모델들을 GenomicsBenchmarks(Grešová et al. 2023)의 다운스트림 태스크에서 평가했다. 실험 설정은 Nguyen et al. (2024)과 동일하게 구성했으며, Arora et al. (2024)에서 보고된 기존 베이스라인 성능과 비교했다.

 

표 4는 Titans(LMM)과 베이스라인 모델들의 성능(Top-1 분류 정확도)을 보여 준다. 결과적으로, LMM은 다양한 게놈(genomics) 태스크에서도 최신 아키텍처와 견줄 만한 경쟁력을 보였다.

 

5.8 효율성(Efficiency)
본 섹션에서는 우리의 신경 메모리 모듈과 Titans가 최신 시퀀스 모델들과 비교하여 얼마나 효율적인지 살펴본다. 그림 9는 다양한 시퀀스 길이 × 배치 크기에 대한 모델들의 학습 처리량(throughput)을 보여 준다.

• 순환 모델 간 비교
  우리의 신경 메모리 모듈이 Mamba2, Gated DeltaNet보다 소폭 느린 이유는 크게 두 가지다:
  1. 더 깊은 메모리(deep memory)로 인한 표현력 높은 전이(transition) 과정(메모리 업데이트)의 존재
  2. Mamba2 구현에서 고도로 최적화된 커널(kernel) 이 사용됨
• Titans (MAL)과 비교
  반면, Titans(MAL)은 베이스라인 모델들과 신경 메모리 모듈보다도 학습 처리량이 높았다. 이는 Flash-Attention(Dao 2024)을 통해 구현된 SWA 및 풀 어텐션 모듈이 고도로 최적화되어 있기 때문이다.

그림 9: Titans와 베이스라인 모델들의 학습 처리량 비교

 

5.9 소 Ablation Study(Ablation Study)
마지막으로, Titans 아키텍처에서 서로 다른 구성 요소를 제거·변경하며 성능 변화를 살펴보았다. 구체적으로, 우리의 신경 메모리 모듈을 기준(base) 모델로 두고,

1. 깊은 메모리를 선형 메모리로 교체,
2. 컨볼루션 제거,
3. 놀라움 지표(surprise measure)에서 모멘텀(momentum) 제거,
4. 가중 감쇠(weight decay, 망각 기법) 제거,
5. 영속 메모리(persistent memory) 제거,
   등을 각각 수행했다. 결과는 표 5에 나타나 있다.

결론적으로, 신경 메모리의 모든 구성 요소가 성능에 긍정적 기여를 했으며, 특히 가중 감쇠와 모멘텀, 컨볼루션, 영속 메모리 순으로 큰 영향을 미쳤다.

 

아키텍처 설계의 영향(The Effect of Architectural Design)

아울러, Titans의 세 가지 변형(MAC, MAG, MAL)을 다음 세 과제 측면에서 비교했다:

1. 언어 모델링, 2) 상식 추론, 3) 긴 문맥 길이 NIAH(BABILong).
   표 5에 요약된 결과에 따르면, MAC과 MAG는 언어 모델링 및 상식 추론에서는 유사한 성능을 보이지만, 긴 문맥 NIAH에서는 MAC이 크게 앞선다. 또한 두 모델 모두 MAL 변형보다 성능이 우수하다. 한편, 그림 9에서 보이듯이 학습 속도 측면에서는 MAL 쪽이 더 빠를 수 있기에, 빠른 학습과 높은 표현력 사이의 트레이드오프가 존재함을 시사한다.

 

표 5: Titans에 대한 소아블레이션(Ablation) 결과. 각 구성 요소가 성능에 모두 긍정적인 영향을 미침.

 

6 결론(Conclusion)

본 논문에서는 테스트 시점에서도 데이터를 기억하도록 학습하는 신경 장기 메모리(neural long-term memory) 모듈을 제안했다. 제안 모델은 메타 인-컨텍스트 학습자(meta in-context learner)로서, 놀라움을 유발하는 토큰(또는 그 주변 토큰)을 적응적으로 기억하는 순환 구조를 갖춘다. 최신 순환 모델과 비교했을 때, 더 풍부한 메모리 업데이트·저장 메커니즘을 지닌 점이 특징이다.

 

이 메모리를 바탕으로, 우리는 Titans 아키텍처와 세 가지 변형(MAC, MAG, MAL)을 소개했다. 즉, 메모리 모듈을 (1) 컨텍스트로 사용하거나, (2) 게이팅(Gating)을 통해 코어와 결합하거나, (3) 계층(Layer) 형태로 추가하는 방식이다. 다양한 과제—언어 모델링, 긴 문맥(수백만 토큰), 상식 추론, DNA 모델링, 시계열 예측 등—에서의 실험 결과, Titans가 트랜스포머 및 최신 선형 순환 모델보다 더욱 효과적임을 확인했다. 특히, Titans는 200만(2M) 토큰을 초과하는 거대 컨텍스트 길이에도 확장 가능하면서, 베이스라인보다 정확도가 더 높다.

 

Titans는 PyTorch 및 JAX로 구현되었으며, 조만간 모델 훈련 및 평가에 사용한 코드를 공개할 계획이다.

 

---

이상하고 이상하고 이상한 점들

 

 

• 규모 대비 실험 부족
  • 논문에서 주장하는 방향성이나 기술적인 복잡도를 생각해보면, 사실 실험 규모가 그리 크지 않아 보여요.
  • 특히 “수백만(2M) 길이를 커버한다”라든가 “GPT-4보다 성능이 좋다”고까지 얘기하면서, 정작 대규모 벤치마크나 정말 극단적인 스케일에서의 검증은 제한적이죠.
  • 저자들이 구글 혹은 딥마인드 수준의 거대 리소스를 활용했다면, 더 큰 데이터나 더 많은 파라미터 규모로 확실한 실험을 해볼 수도 있었을 것 같은데, 논문에서는 상대적으로 소규모(수억 ~ 수십억 토큰 정도)에서만 결과를 제시한 것처럼 보입니다.
• ‘장기 메모리’의 실제 이점 증명 부족
  • 논문 전반에서, “테스트 시점에도 메모리가 업데이트된다”는 게 핵심적인 세일즈 포인트입니다.
  • 그런데 이것이 정말로 “길이가 몇 백만인 문맥을 Transformer 없이 효율적으로 다룰 수 있다”라는 식의 의미 있는 장점을, 구체적인 벤치마크 전부에서 확실히 보여주는지는 조금 애매합니다.
  • 가령 실제로 2M 토큰 길이를 한번에 처리해 볼 수 있는 대규모 벤치마크나, 온라인·스트리밍 형태의 정말 긴 입력(예: 비디오 스트림, 실시간 로그 등)에서의 성능을 대대적으로 보여줬다면 납득이 쉬웠을 텐데, 논문이 제시하는 태스크는 대부분 “텍스트”나 “DNA 시퀀스” 같은 것이고, 그마저도 비교적 작은 규모로만 실험한 인상을 줍니다.
• 메모리 모듈의 자원 소모와 학습 난이도
  • 말씀하신 것처럼, ‘consistent memory’(테스트 시점에서 지속적으로 업데이트되는 메모리)는 자칫하면 GPU나 TPU 메모리 부담이 생각보다 클 수 있습니다.
  • 게다가 “깊은(deep) 메모리”를 쓰면서 모멘텀·게이팅·감쇠(Weight Decay) 등 온갖 기법을 다 붙이면, 실제 트레이닝 과정이 복잡해지고(오버헤드 증가), 디버깅도 어려워지죠.
  • 논문에서 “수백만 토큰 문맥도 OK”라고 주장하는데, 정작 실제로 그걸 훈련하는 과정에서 얼마나 리소스를 썼고, 어느 정도 학습 시간이 걸렸는지, 그리고 그때의 메모리 사용량 등은 자세히 나오지 않은 듯합니다.
• 어텐션 없이 ‘깊은 메모리’만으로도 충분?
  • 논문에서 “단일 메모리 모듈(LMM)”만으로도 꽤 높은 성능을 낸다고 주장하잖아요.
  • 그렇다면 굳이 이렇게 MAC, MAG, MAL 같은 복잡한 하이브리드 아키텍처가 필요할까?
  • 논문 말미에 “트레이드오프가 있다”는 식으로 언급은 하지만, 독자가 보기엔 “굳이 어텐션과 결합하지 않고, 메모리 모듈만 잘 최적화해도 되지 않을까?”라는 의문이 들 수 있습니다.
  • 물론 “어텐션의 장점 + 메모리의 장점”을 살리고 싶었을 거지만, 실제론 “둘이 섞어서 이득을 보는 상황”이 충분히 시연된 건지 좀 궁금합니다.
• 실험 결과의 해석 / 주장
  • 논문에서 “GPT-4보다도 좋다”라고 언급되는 부분(needle in haystack, BABILong 결과 등)이 꽤 자극적이긴 한데, 사실 GPT-4가 그 태스크에 얼마나 최적화되어 있는지, 프롬프트나 few-shot 설정은 어땠는지 등이 제대로 공개된 건 아니지 않습니까.
  • “GPT-4는 정식 API로 접근해서 간단한 prompt로만 돌려본 수준”이라면, 사실 진지한 fine-tuning(혹은 세심한 프롬프트 엔지니어링)과는 조건이 달라서, 공정 비교라고 보기 어려울 수도 있어요.
• 이론적 표현력 vs. 실제 응용
  • 논문 후반부에 “Titans는 Transformer나 DeltaNet보다 이론적으로 더 표현력이 높다(TC0를 넘어선다)”는 식의 문장이 있는데, 이 부분이 일종의 이론적 과대포장처럼 보이기도 합니다.
  • 이론적 계산 모델에서의 표현력과, 실제 대규모 태스크에서 보여주는 퍼포먼스는 언제나 괴리가 있을 수 있는데, 그걸 얼마나 잘 연결했는지가 약간 미흡해 보입니다.