Titans: Learning to Memorize at Test Time

Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.

한 줄 요약

Titans는 2M 이상의 컨텍스트 길이를 처리하면서 정확도를 향상시키는 신경망 장기 기억 모듈 기반의 새로운 아키텍처 패밀리이다.

핵심 기여도

새로운 신경망 장기 기억 모듈(Neural Memory Module)을 제안하여, 2M 이상의 컨텍스트 길이를 처리 가능하게 함.
Titans 아키텍처 3가지 변형(MAC, MAG, MAL)을 제시하며, MAC이 긴 컨텍스트 NIAH 작업에서 가장 뛰어난 성능을 보임.
기존 Transformer와 최근 선형 RNN 대비, 긴 컨텍스트 작업에서 정확도 향상과 빠른 추론 속도를 동시에 달성.
MAC 모델은 가중치 감소, 모멘텀, 컨볼루션, 지속 메모리가 성능에 긍정적 영향을 미침 (Table 5 참조).

핵심 아이디어

기존 Transformer는 정확한 의존성 모델링을 제공하지만, 이는 컨텍스트 길이에 대해 제곱 비용을 유발하여 확장성에 한계가 있다. 반면, 선형 RNN은 효율적이지만 데이터 압축 과정에서 정보 손실이 발생한다. 이 논문은 이러한 문제를 해결하기 위해, **장기 기억**(Neural Memory)과 **단기 기억**(Attention)을 병합한 새로운 학습 패러다임을 제안한다. Neural Memory는 테스트 시점에서 학습하며, 놀라운 토큰 또는 놀라움에 가까운 토큰을 적응적으로 기억한다. 이는 RNN과 달리, 메모리 업데이트와 저장 메커니즘이 더 표현력이 높다. 이 두 모듈을 기반으로 Titans 아키텍처를 제안하며, 이는 컨텍스트, 게이팅, 레이어의 세 가지 방식으로 메모리 모듈을 통합할 수 있다.

기술적 접근법

**Neural Memory Module**: 토큰을 적응적으로 기억하며, 놀라움 기반의 메모리 업데이트를 수행.
**Titans 아키텍처**: MAC (Memory as Context), MAG (Memory as Gating), MAL (Memory as Layer)의 세 가지 변형.
**가속 기법**: 컨볼루션, 모멘텀, 가중치 감소(weight decay), 지속 메모리(persistent memory)를 활용.
**하이퍼파라미터**: MAC 모델에서 가중치 감소, 모멘텀, 컨볼루션, 지속 메모리가 성능에 긍정적 영향 (Table 5 참조).

주요 결과

**Needle-in-haystack (NIAH) 작업**: MAC 모델이 2M 이상의 컨텍스트 길이에서 기존 기반 모델 대비 정확도가 높음.
**언어 모델링**: MAC과 MAG가 MAL보다 더 높은 성능을 보임.
**추론 속도**: Neural Memory 모듈은 빠른 추론 속도를 유지하면서도 병렬화된 학습이 가능.
**확장성**: Titans는 2M 이상의 컨텍스트 길이를 처리 가능하며, 이는 기존 Transformer와 선형 RNN보다 우수.

의의 및 한계

Titans는 단기 기억(attention)과 장기 기억(neural memory)을 결합한 새로운 학습 패러다임을 제시하며, 긴 컨텍스트 작업에서 기존 모델을 능가하는 성능을 보인다. 특히, MAC 모델은 NIAH 작업에서 뛰어난 정확도를 달성하며, 병렬화된 학습과 빠른 추론을 동시에 제공한다. 그러나, 현재는 대규모 모델의 실험 결과는 아직 보고되지 않았으며, 향후 연구가 필요하다. 또한, 메모리 업데이트 메커니즘의 복잡성은 모델의 해석성과 안정성에 영향을 줄 수 있다.

실용적 활용

Titans는 긴 텍스트, 유전자 분석, 시간 시리즈 예측 등 긴 컨텍스트가 필요한 다양한 분야에 적용 가능하다. 특히, NIAH 작업과 같은 정밀한 정보 추출 작업에서 뛰어난 성능을 보여주며, 대규모 언어 모델과 생물정보학 분야에서 실용적 활용이 기대된다.