Titans + MIRAS: Helping AI have long-term memory

Titans와 MIRAS 프레임워크는 AI 모델이 실행 중에 코어 메모리를 업데이트하여 훨씬 빠르게 작동하고 방대한 Context를 처리할 수 있도록 돕는 새로운 접근 방식을 제시합니다. 기존 Transformer는 Attention 메커니즘으로 시퀀스 모델링을 혁신했지만, 시퀀스 길이가 길어질수록 계산 비용이 급증하여 전체 문서 이해나 유전체 분석과 같은 극도로 긴 Context에 적용하기 어렵다는 한계가 있었습니다. Mamba-2와 같은 효율적인 RNN 및 SSM은 Context를 고정된 크기로 압축하여 빠른 선형 스케일링을 제공했지만, 이 고정 크기 압축은 매우 긴 시퀀스의 풍부한 정보를 충분히 포착하지 못했습니다.

Titans 아키텍처: 실시간 Context 학습

Titans는 RNN의 속도와 Transformer의 정확성을 결합한 특정 아키텍처입니다. 이는 인간 두뇌의 단기 및 장기 기억 분리와 유사하게, 명확하지만 상호 연결된 메모리 모듈을 필요로 하는 효과적인 학습 시스템을 구현합니다. Attention 메커니즘이 정확한 단기 기억에 뛰어나다면, Titans는 새로운 신경 장기 기억 모듈을 도입합니다. 전통적인 RNN의 고정 크기 벡터 또는 행렬 메모리와 달리, 이 모듈은 Deep neural network, 특히 Multi-layer perceptron (MLP)으로 작동합니다. 이를 통해 훨씬 높은 표현력을 제공하여 중요한 Context를 잃지 않고 방대한 정보를 요약할 수 있습니다.

Titans는 단순히 데이터를 수동적으로 저장하는 것이 아니라, "Surprise metric"이라는 핵심 개념을 통해 입력 전체에 걸쳐 토큰을 연결하는 중요한 관계와 개념적 주제를 인식하고 유지하는 방법을 능동적으로 학습합니다. 이 "Surprise metric"은 현재 모델이 기억하는 내용과 새로운 입력이 전달하는 내용 사이의 큰 차이를 감지하는 것입니다.
* Low surprise: 예를 들어, 새로운 단어가 "cat"이고 모델의 메모리 상태가 이미 동물 단어를 예상하고 있다면, Gradient는 낮습니다. 이 단어를 영구적인 장기 상태에 저장할 필요가 없습니다.
* High surprise: 모델의 메모리 상태가 심각한 금융 보고서를 요약하고 있는데, 새로운 입력이 "banana peel" 사진(예상치 못한 사건)이라면, Gradient는 매우 높을 것입니다. 이는 새로운 입력이 중요하거나 비정상적임을 신호하며, 장기 기억 모듈에 영구적으로 저장될 우선순위를 부여해야 합니다.

모델은 이 내부 오류 신호(Gradient)를 수학적으로 "이것은 예상치 못했고 중요하다!"라고 말하는 것과 동일하게 사용합니다. 이를 통해 Titans 아키텍처는 가장 새롭고 Context를 깨는 정보만을 선택적으로 장기 기억에 업데이트하여 전체 프로세스를 빠르고 효율적으로 유지합니다. Titans는 이 메커니즘을 두 가지 핵심 요소로 개선합니다.
* Momentum: 모델은 "순간적인 놀라움(Momentary surprise)"(현재 입력)과 "과거의 놀라움(Past surprise)"(최근 Context 흐름)을 모두 고려합니다. 이는 개별적으로 놀랍지 않은 토큰이라도 관련성 있는 후속 정보가 포착되도록 합니다.
* Forgetting (Weight decay): 매우 긴 시퀀스를 처리할 때 메모리의 유한한 용량을 관리하기 위해, Titans는 Adaptive Weight decay 메커니즘을 사용합니다. 이는 더 이상 필요하지 않은 정보를 버리는 망각 게이트(Forgetting gate) 역할을 합니다.

MIRAS 프레임워크: 시퀀스 모델링의 통합적 관점

MIRAS는 Titans와 같은 접근 방식을 일반화하기 위한 이론적 프레임워크입니다. 현대 Transformer부터 새로운 고속 Linear RNN에 이르기까지 모든 주요 시퀀스 모델링 혁신은 본질적으로 동일한 것, 즉 고도로 복잡한 Associative memory 모듈이라는 통일된 관점을 제공합니다. MIRAS는 AI 모델링을 바라보는 방식을 다양성 대신 동일한 문제(새로운 정보를 이전 기억과 효율적으로 결합하면서 핵심 개념을 잊지 않는 것)를 해결하는 다른 방법으로 봅니다.

MIRAS는 시퀀스 모델을 네 가지 핵심 설계 선택을 통해 정의합니다.

Memory architecture: 정보를 저장하는 구조(예: Vector, Matrix, 또는 Titans와 같은 Deep Multi-layer perceptron).

Attentional bias: 모델이 최적화하는 내부 학습 목표로, 무엇을 우선순위로 둘지 결정합니다.

Retention gate: 메모리 Regularizer입니다. MIRAS는 "Forgetting mechanisms"을 새로운 학습과 과거 지식 유지를 균형 맞추는 특정 형태의 Regularization으로 재해석합니다.

Memory algorithm: 메모리 업데이트에 사용되는 Optimization algorithm입니다.

거의 모든 성공적인 기존 시퀀스 모델은 Attentional bias와 Retention 모두에 Mean squared error (MSE) 또는 Dot-product similarity에 의존합니다. 이러한 의존성은 모델을 Outlier에 민감하게 만들고 표현력을 제한할 수 있습니다. MIRAS는 최적화 및 통계 문헌에서 얻은 정보를 바탕으로 더 풍부한 설계 공간을 탐색하는 생성적 프레임워크를 제공하여 이러한 한계를 뛰어넘습니다. 이를 통해 Non-Euclidean objectives 및 Regularization을 가진 새로운 아키텍처 생성이 가능해집니다.

MIRAS를 사용하여 세 가지 특정 Attention-free 모델이 개발되었습니다.
* YAAD: 이 MIRAS 변형은 주요 오류 또는 "Outlier"(예: 대규모 문서의 단일 오타)에 덜 민감하도록 설계되었습니다. 실수를 처리하기 위해 더 부드러운 수학적 페널티인 Huber loss를 사용하여 단발성 문제에 과민 반응하지 않습니다. 이로 인해 입력 데이터가 지저분하거나 일관성이 없을 때 모델이 더 강건해집니다. Huber loss는 $L_{\delta}(x) = \begin{cases} \frac{1}{2}x^2 & \text{if } |x| \le \delta \\ \delta(|x| - \frac{1}{2}\delta) & \text{otherwise} \end{cases}$ 와 같이 정의될 수 있으며, 작은 오차에는 Quadratic, 큰 오차에는 Linear 페널티를 적용합니다.
* MONETA: 이 모델은 더 복잡하고 엄격한 수학적 페널티인 Generalized norms ($L_p$-norm)의 사용을 탐색합니다. 이는 모델이 Attention을 기울이는 것과 잊는 것 모두에 대해 이러한 더 엄격한 규칙을 사용하는 것이 더 강력하고 안정적인 장기 기억 시스템으로 이어질 수 있는지 조사합니다.
* MEMORA: 이 모델은 메모리가 엄격한 확률 맵처럼 작동하도록 강제하여 가능한 최고의 메모리 안정성을 달성하는 데 중점을 둡니다. 이 제약 조건을 사용하여 메모리 상태가 업데이트될 때마다 변경 사항이 제어되고 균형을 이루도록 보장합니다. 이는 새로운 정보를 통합하는 깨끗하고 안정적인 프로세스를 보장합니다.

실험 및 결과

Titans와 MIRAS 변형(YAAD, MONETA, MEMORA)은 Transformer++, Mamba-2, Gated DeltaNet을 포함한 선도적인 아키텍처와 엄격하게 비교되었습니다. 또한 유전체 모델링(DNA) 및 시계열 예측에 대한 Titans 테스트를 통해 아키텍처의 다재다능함이 검증되었습니다.

표준 언어 모델링 데이터셋(C4, WikiText)과 Zero-shot reasoning task(HellaSwag, PIQA) 모두에서 모델은 일관되게 더 높은 정확도와 낮은 Perplexity(LLM이 텍스트를 보고 얼마나 놀랐는지 측정)를 시연했습니다. Ablation study는 메모리 아키텍처의 깊이가 결정적임을 명확히 보여줍니다. 동일한 크기이지만 다른 깊이의 장기 기억 모듈을 비교했을 때, 더 깊은 메모리를 가진 모듈은 언어 모델링에서 일관되게 더 낮은 Perplexity를 달성했으며, 시퀀스 길이가 크게 증가해도 성능을 유지하는 더 나은 스케일링 특성을 보였습니다.

언어 모델링 및 상식 추론 작업에서 Titans 아키텍처는 유사한 크기의 최신 Linear recurrent 모델(Mamba-2, Gated DeltaNet) 및 Transformer++ Baseline을 능가합니다. 새로운 MIRAS 변형(MONETA, YAAD, MEMORA) 또한 이러한 Baseline에 비해 향상된 성능을 달성하여, 강력한 Non-MSE 최적화 메커니즘 탐색의 이점을 입증했습니다. 중요한 것은 이러한 모델이 효율적이고 병렬화 가능한 훈련과 빠른 선형 추론 속도를 유지한다는 것입니다.

가장 중요한 장점은 이러한 새로운 아키텍처가 극도로 긴 Context를 처리할 수 있다는 것입니다. 이는 극도로 긴 문서에 분산된 사실에 대한 추론을 요구하는 BABILong 벤치마크에서 강조됩니다. 이 어려운 설정에서 Titans는 GPT-4와 같은 매우 큰 모델보다 훨씬 적은 매개변수를 가짐에도 불구하고 모든 Baseline을 능가합니다. Titans는 200만 토큰 이상의 Context window 크기로 효과적으로 확장할 수 있는 능력을 추가로 보여줍니다.

결론적으로, Titans와 MIRAS 프레임워크의 도입은 시퀀스 모델링에서 중요한 발전을 의미합니다. 데이터를 받아들이면서 기억하는 방법을 학습하는 Deep neural network를 메모리 모듈로 사용하여, 이 접근 방식은 고정 크기 순환 상태의 한계를 극복합니다. 또한 MIRAS는 온라인 최적화, Associative memory 및 아키텍처 설계 간의 연결을 밝히는 강력한 이론적 통합을 제공합니다. 표준 Euclidean paradigm을 넘어서, 이 연구는 RNN의 효율성과 긴 Context AI 시대에 필요한 표현력을 결합한 차세대 시퀀스 모델의 길을 엽니다.