https://arxiv.org/abs/2501.00663
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
arxiv.org
0. Abstract
순환 모델은 데이터를 고정된 크기의 메모리(hidden state)로 압축하는 것을 목표로 하는 반면, 어텐션은 전체 컨텍스트 창을 참조하여 모든 토큰 간의 직접적인 의존성을 포착할 수 있도록 한다. 그러나 이러한 더 정확한 의존성 모델링은 계산 비용이 제곱적으로 증가하여 모델이 고정된 길이의 컨텍스트에만 제한되는 단점이 있다.
어텐션은 제한된 컨텍스트 내에서 정확한 의존성 모델링이 가능하기 때문에 단기 메모리로 작용하며, 반면 신경망 메모리(neural memory)는 데이터를 장기적으로 저장할 수 있어 더 지속적인 장기 메모리 역할을 한다. 우리는 이러한 두 가지 모듈을 기반으로 새로운 아키텍처 계열인 Titans를 소개한다.
Titans는 언어 모델링, 상식 추론, 유전체학(Genomics), 시계열 분석 등의 다양한 작업에서 트랜스포머 및 최신 선형 순환 모델보다 더 높은 성능을 보였다. 또한, Titans는 200만 개 이상의 컨텍스트 창을 처리할 수 있다.
1. Introduction
1.1. 트랜스포머(Transformers)의 강점과 한계
- 트랜스포머는 최고 성능의 시퀀스 모델로 자리 잡았으며,
- 어텐션이 연관 기억처럼 동작하여 맥락 내 학습(in-context learning) 및 확장성 을 제공한다.
- 하지만 컨텍스트 창에 의존하여 컨텍스트 창(문맥 길이)이 커질수록 시간 및 메모리 복잡도가 O(n²)로 증가 한다는 문제점이 있음.
1.2. 선형 트랜스포머(Linear Transformers) 대안(LRN)
- 기존 트랜스포머의 비효율성을 해결하기 위해 소프트맥스를 커널 함수로 대체한 선형 트랜스포머가 등장.
- 메모리 사용량 감소, 연산 효율성 증가(O(n))
- 그러나 성능이 기존 트랜스포머보다 낮음 → 데이터 압축 과정에서 정보 손실 발생
- 결국 데이터가 행렬(matrix) 상태로 압축 되어 정보 손실이 발생할 수 있다.
- 너무 긴 컨텍스트를 작은 벡터/행렬 상태로 압축할 수 없음
1.3. 기존 신경망 모델의 한계
- 일반화(generalization), 길이 확장(length extrapolation), 추론(reasoning)에서 어려움
- 인간 두뇌에서 영감을 받았지만, 단기/장기 기억 부족
- 모듈 간 독립적이면서 연결된 시스템 부족(단기, 장기, 메타 기억 등)
- 과거 정보를 추상적으로 저장하는 능력 부족
1.4. 핵심 주장
- 인간 두뇌처럼 독립적이면서 상호 연결된 학습 모듈이 필요
- 기존 신경망과 차별화된 새로운 학습 패러다임 제안 필요
1.5. 메모리 관점
- 신경심리학에서 정의하는 기억과 학습의 개념에 따르면, 대부분의 기존 신경망 아키텍처는 메모리를 입력으로 인해 업데이트되는 신경 상태 변화로 간주하며, 학습을 주어진 목표 내에서 효과적이고 유용한 기억을 획득하는 과정으로 정의한자.
- 이러한 관점에서 순환신경망은 벡터 값의 메모리 모듈 M(은닉 상태, hidden state)을 포함한 모델로 볼 수 있다. 이는 두 가지 주요 단계를 거친다.
1. 새로운 입력 xt가 주어졌을때, 모델이 메모리를 업데이트 한다.(압축 포함)
2. 입력과 관련된 메모리를 검색한다.
- 유사하게, 트랜스포머(Transformers) 는 점진적으로 증가하는(growing) 메모리 구조를 가진 모델 로 해석할 수 있다. 이때, 키(key)와 값(value) 행렬이 모델의 메모리 역할을 수행 하며, 다음 두 단계를 따른다
1. 새로운 키-값 쌍을 메모리에 추가하여 업데이트 (압축 없음).
2. 쿼리(query)와 키(key) 간의 유사도를 찾아 해당 값을 가중 평균하여 출력으로 생성.
- 이러한 메모리 관점은 기존 모델들의 차이점을 명확히 하고, 더 효과적인 아키텍처 설계를 가능
- 이러한 차이점 등을 통해 다음과 같은 질문이 도출된다.
- 기억 구조 설계에서 제기되는 5가지 질문
1. Q1: 좋은 메모리 구조란 무엇인가?
2. Q2: 적절한 메모리 업데이트 방식은?
3. Q3: 효율적인 메모리 검색 방법은?
- 인간의 기억은 단일한 과정이 아니며, 단일한 기능을 수행하지 않는다. 기억은 여러 시스템(예: 단기 기억, 작업 기억, 장기 기억 등)으로 구성되며, 각각의 기능과 신경 구조가 다르고 독립적으로 작동할 수 있다. 이에 따라 기억의 여러 시스템에 이러한 질문이 나온다.
4. Q4: 서로 다른 메모리 모듈을 효과적으로 통합하려면?
- 이 질문에 대한 대답은 현 논문에서 다루게 되는데, 위에서 본 모델의 기억을 저장하는 과정은 단순한 벡터(RNN) 또는 행렬(linear Transformer)에 선형적으로 데이터를 저장하는 방식만으로 충분하지 않을 수 있다. 이에 따라 아래의 질문이 나오게 되는데
5. Q5: 오래된 정보를 효과적으로 저장하려면 더 깊은 기억 모듈이 필요한가?
- 본 논문에서는 효율적이고 효과적인 장기 신경 기억 모듈(long-term neural memory module)을 설계하여 앞서 제기한 5가지 질문에 답하고자 한다.
1. 신경 기억 모듈(Neural memory)
- 테스트 시점에서 데이터를 기억하도록 학습
- 인간의 장기 기억 시스템에서 영감을 받아 예상과 다르게 놀라운(surprising) 사건일수록 기억에 더 오래 남도록 설계
- 놀라움 정도(surprise)를 측정(연관 기억 손실(associative memory loss)의 그래디언트(gradient)를 활용 하여 이를 평가)
- 메모리 용량이 제한적인 문제를 해결하기 위해 감쇠 메커니즘을 도입
- 메모리 크기 대비 데이터의 놀라움 정도를 고려하여 더 나은 기억 관리를 가능하게 한다.
- 현대 순환 모델(recurrent models)에서의 망각 메커니즘을 일반화한 것과 같다는 것을 보인다.
- 미니배치 그래디언트 디센트, 모멘텀, 가중 감쇠(weight decay)를 활용한 메타 신경망(meta neural network) 최적화 과정과 동등함을 발견하였다. 이를 기반으로, 우리는 미니배치 그래디언트 디센트의 행렬 연산을 최적화한 병렬 알고리즘을 제안하여, 더 빠르고 병렬적으로 학습 가능한 딥 신경 기반 장기 기억 모델을 개발한다.
2. Titans 아키텍처
- 장기 신경 기억 모듈을 설계한 후, 이를 딥러닝 아키텍처에 어떻게 효과적이고 효율적으로 통합할 것인가?
- 이를 해결하기 위해, Titans 라는 새로운 딥러닝 모델 패밀리를 제안하며, 세 가지 핵심 하위 모듈(hyper-heads)로 구성 된다.
- Titans 3가지 기억 하위 모듈
1. Core
- 단기 기억(short-term memory) 모듈
- 데이터의 주된 흐름을 처리하며, 제한된 윈도우 크기(limited window size)를 가진 어텐션(attention) 기법을 사용
2. Long-term Memory
- 본 연구에서 개발한 신경 장기 기억 모듈
- 오래된 과거 데이터를 저장 및 검색하는 역할 수행
3. Persistent Memory
- 특정 데이터와 독립적이며, 태스크(task)와 관련된 지식을 학습하는 파라미터 집합
- 또한, Titans의 세 가지 변형(variants)을 제시하며, 기억 모듈을 다음과 같은 방식으로 통합한다.(MLP에서 어텐션을 적용하는 방식)
1. 컨텍스트(context)로 활용
2. 레이어(layer)로 활용
3. 게이트(gated branch) 형태로 활용
2. Preliminaries
1. Titans
- 메모리 사용량을 절감하고 처리 속도를 개선하는 신경 기억 모듈 개발
- Titans 모델을 통해 장기 기억을 효과적으로 처리하는 방법 제안
2. Background
1. Attention
- 주어진 입력 x ∈ R^(n × din) 에 대해, 인과적 어텐션(causal attention)은 입력 의존적인 키(K), 값(V), 쿼리(Q) 행렬을 사용 하여 출력을 계산
- 여기서 W_Q, W_K, W_V 는 학습 가능한 파라미터
- 메모리 사용량이 크고, 긴 시퀀스에서는 처리량(throughput)이 낮아지는 문제가 있다.
2. Linear Attention
- 전통적인 소프트맥스 기반 어텐션을 대체하기 위해 커널 함수 𝜙(𝑥, 𝑦) 사용
- 기존 방식보다 높은 처리량(throughput) 제공
- 반복적으로 계산되는 항목 $k_{j}$, $k_{l}$을 재사용하여 연산량 절감
- 특정 커널을 항등 행렬(identity matrix)로 선택하면, 순환 구조(recurrent format)로 변형 가능
- 선형 어텐션을 순환 모델(RNN)처럼 동작하도록 만들어, 매우 효율적인 추론 가능
3. Memory
- 학습(learning)이란 효과적이고 유용한 메모리를 획득하는 과정으로 정의할 수 있다.
- 순환 신경망(RNN)의 은닉 상태(hidden state)를 일종의 메모리(memory) 유닛 으로 볼 수 있으며, 모델은 이 메모리 유닛에 정보를 압축하려고 한다.
- 순환 신경망(RNN) 에서, 은닉 상태는 메모리 역할을 하며, 이를 읽기(read) 및 쓰기(write) 연산으로 분리 할 수 있다.
4. 선형 트랜스포머의 메모리 문제
1. 문제점
- 긴 시퀀스(long context data) 를 다룰 때, 이 덧셈 기반 저장 방식(계속된 덧셈 연산)이 메모리 과부하(memory overflow)를 초래하여 모델 성능 저하를 유발
2. 해결 방식
1. 망각 메커니즘
- 기억을 필요할 때 삭제할 수 있도록 데이터 기반 망각(forging) 게이트를 추가하는 방식
- 아래에서 자세히 설명
2. 쓰기 연산 개선
- 전통적인 RNN의 메모리 쓰기 방식은 단순히 새로운 정보를 기존 메모리에 더하는 방식이기에 불필요한 정보가 계속 축적되면서 메모리 효율성이 떨어짐
- 델타 룰이라는 새로운 정보를 추가하기 전에, 기존 정보를 먼저 제거하는 방식을 통해 진행
3. Learning to Memorize at Test Time
3.1 신경망 메모리 설계의 동기 및 개념 논의
1. 장기 메모리 (Long-term Memory)
- 신경망 기반의 장기 기억 모듈을 설계하려면, 과거의 정보를 모델의 파라미터에 추상적으로 인코딩할 수 있어야 함.
- 단순 암기, 즉 기억은 일반적으로 부정적으로 취급되는데 그 이우는 일반화를 제한하고, 프라이버시 문제를 유발하기 때문이다.
- 테스트 시 데이터를 학습하면서 적절히 기억/망각하는 온라인 메타 모델(online meta-model)이 필요 함.
- 단순 암기가 아니라, 필요한 정보만 학습하여 테스트 시 더 나은 일반화 성능을 제공하는 모델
2. 학습 과정과 서프라이즈 메트릭 (Learning Process and Surprise Metric)
- 장기 기억 모델을 훈련하는 핵심 아이디어는 온라인 학습 문제로 다루는 것으로, 과거 정보 $x_{1}, ..., x_{t_1}$ 를 메모리 모듈($M_{t}$)의 파라미터에 압축하는 방식으로 진행된다.
1. 서프라이즈 매트릭
- 인간은 예상과 다른 사건(예: 놀라운 사건)을 더 잘 기억하는 경향이 있는데, 이를 모델에 적용하여, 입력 데이터에 대한 그래디언트 크기를 서프라이즈 지표(surprise metric)로 사용
- 서프라이즈(Surprise) 개념을 활용하여, 새로운 데이터가 기존 데이터와 얼마나 다른지 측정
- 그래디언트가 클수록 이전 데이터와 차이가 크며, 즉 더 놀라운 데이터이다.
$$M_{t}=M_{t-1}-\Theta_{t}\bigtriangledown l (M_{t-1};x_{t})$$
- 너무 놀라운 데이터가 연속해서 등장하면, 이후의 정보는 상대적으로 덜 중요하게 여겨질 위험
- 그래디언트가 일정 수준 이하로 줄어들면(flat area), 중요한 정보도 놓칠 가능성
- 인간 기억에서는, 한 번 놀라운 사건이 발생하면 이후에도 지속적으로 관련 정보를 기억하는 경향
2. 개선된 서프라이즈
- 서프라이즈를 두 가지 요소로 분리
1. 과거 서프라이즈(Past Surprise): 최근 과거의 놀라움 정도를 측정
2. 순간 서프라이즈(Momentary Surprise): 새롭게 들어오는 데이터의 놀라움 정도를 측정
$$M_{t}=M_{t-1}+S_{t}$$
$$S_{t}=\eta_{t}S_{t-1}(Past Suprise)-\Theta_{t}\bigtriangledown l( M_{t-1};x_{t})(Momentary Surprise)$$
- 이 수식은 모멘텀을 적용한 그래디언트 디센트(gradient descent with momentum)와 유사한데, 이전 서프라이즈가 새로운 데이터의 서프라이즈에 영향을 줄 수도 있고, 아닐 수도 있다.
- 즉, 위에서 제기된 서프라이즈 매트릭에서의 이전 서프라이즈 데이터가 계속 새로운 서프라이즈 데이터의 영향을 주는 것을 방지하고, 선택적으로 가져갈 수 있게 한다.
3. 목적 함수
- 서프라이즈 메트릭은 손실 함수 ℓ에 기반하며, 이는 테스트 시 메모리가 학습해야 할 목표(objective) 다. 모리 모듈은 손실 함수 l를 기반으로 동작하는 메타 모델(meta model) 이다.
4. 연관 기억 (Associative Memory) 모델
- 연관 기억(associative memory) 을 사용하여 과거 데이터를 키(key)와 값(value)의 쌍 으로 저장
- 입력 xt를 선형 변환하여 키(key)와 값(value)을 생성
- 메모리 모듈이 키와 값의 관계를 학습 하도록 설계되며, 이를 위해 손실 함수를 정의
$$l(M_{t-1};x_{t})=||M_{t-1}(k_{t}-v_{t})||^{2}_{2}$$
- 이 손실 함수를 최소화함으로써, 모델은 test-time에 key와 value 사이의 매핑을 기억하는 방법을 학습
- 메모리 내부에서 학습하여 메타 러닝과 유사하며, 메모리 외부에서는 전체 아키텍쳐의 다른 파라미터가 최적화된다.
5. 망각 매커니즘
- 매우 긴 시퀀스를 다룰 때, 불필요한 정보를 잊어버리는 것이 중요
- 메모리 용량이 제한적이므로, 필요 없는 정보는 삭제해야 한다.
- 메모리 유지/망각을 조정하는 게이트 변수를 통해 망각을 조절한다.
$$M_{t}=(1-\alpha_{t})M_{t-1}+S_{t}$$
$$S_{t}=\eta_{t}S_{t-1}-\Theta_{t}\bigtriangledown l(M_{t-1};x_{t})$$
- $\alpha_{t}\in[0,1]$ 는 메모리를 유연하게 제어하는 gating 메커니즘으로, 얼마나 많은 정보를 잊어야 하는지 결정한다. 이 weight decay 메커니즘은 RNN의 gating 메커니즘과 밀접한 관련이 있다.
6. 메모리 아키텍쳐
- 다층 퍼셉트론(MLP) 을 사용하여 장기 기억을 모델링(가중치가 기억을 저장하는 용도로 사용)
- 최소 한 개 이상의 레이어(LM ≥ 1)를 포함하는 단순한 MLP 구조 사용( 선형 변환 + 활성화 함수(ReLU 등)를 적용하여 비선형성을 도입)
- 인 딥 메모리 모듈은 선형 모델보다 표현력이 뛰어나다.
- 일반적인 신경망 모델에서는 과거의 정보가 점차 소멸되지만, 딥 메모리 모듈은 적절한 학습 기법을 사용하여 장기적인 정보 저장이 가능
7. 메모리 검색 (Retrieving a Memory)
- 위에서는 장기 기억 모듈을 학습하고 테스트에서 활용하는 방법에 대해서 제안했다.
- 여기서는 "메모리에서 정보를 어떻게 검색(retrieve)할 것인가?" 에 대해 설명한다.
1. 추론(inference) 시 가중치를 업데이트하지 않고(forward pass만 사용) 메모리를 검색
2. 주어진 입력 $x_{t}$에 대해, 선현 변환 행렬 W_{Q}를 사용하여 쿼리 벡터 생성
$$q_{t}=x_{t}W_{Q}$$
3. 쿼리를 메모리에 입력하여 해당하는 정보를 검색
$$y_{t}^{*}=M_{t}(q_{t})$$
3.2 빠르고 병렬화 가능한 학습 설계 논의
- long-term memory 모듈의 설계는 모멘텀과 weight decay를 사용한 gradient descent로 loss function 최적화하여 메타 모델을 학습시키는 것과 동일
- Long-term memory 모듈의 학습에는
- 학습 프로세스를 병렬화하고 하드웨어 가속기(TPU, GPU)를 최대한 활용하려면 프로세스를 텐서화(tensorize)하고 더 많은 matmul을 사용해야 한다.
- mini-batch gradient descent, 데이터에 따라 다른 learning rate, weight decay를 사용하여 내부 루프에서 가중치를 계산하는 것을 matmuls와 sum만 사용하도록 재구성
3.3 지속적인 메모리 모듈(persistent memory module) 추가
- 우리의 장기 기억은 문맥적 기억(contextual memory)으로 볼 수도 있으며, 이는 출력이 전적으로 문맥(context)에 의존함을 의미한다. 따라서, 장기 기억뿐만 아니라, 학습 가능한 입력과 독립적인 매개변수 집합을 사용하여 작업(task) 관련 기억으로 작용하도록 한다. 이러한 유형의 기억은 지속적 기억(persistent memory) 또는 메타 기억(meta-memory)으로 불려왔다.
- 지속적 기억의 필요성
1. 기억(Memory) 관점
- 우리의 신경망 장기 기억(neural long-term memory)은 문맥적 기억으로, 모든 매개변수가 입력에 의존적이다. 그러나 효과적인 기억 시스템은 입력과 독립적인 매개변수도 필요하며, 이는 작업 지식을 추상적으로 저장하는 역할을 한다. 즉, 어떤 작업을 숙달하기 위해서는 "작업을 수행하는 방법"에 대한 지식을 기억해야 하며, 이러한 매개변수가 그러한 지식을 저장하는 역할을 한다.
2. 전달망(Feedforward Network) 관점
- 트랜스포머(Transformer) 아키텍처에서는 어텐션(attention) 모듈 이후에 완전 연결층(fully connected layer)이 존재하며, 이는 데이터와 독립적인 매개변수를 가진 어텐션 가중치와 유사한 것으로 나타났다.
- 완전 연결층의 ReLU를 Softmax로 대체하면 어텐션과 유사한 가중치를 형성할 수 있으며, 이때 가중치는 데이터와 독립적이다.
$$FFN(x)=W_{V}⋅Softmax(W_{K}x)$$
- W_{K}와 W_{V}는 입력과 독립적일 때 어텐션 모듈의 K와 V 행렬과 유사하게 동작한다. 지속적 기억의 가중치는 동일한 기능을 수행할 것으로 기대되며, 시퀀스의 첫 번째 부분에 이를 사용하면 입력과 독립적인 어텐션 가중치를 형성할 수 있다.
3. 기술적(Technical) 관점
- 인과적 마스크(causal mask)를 적용한 어텐션은 시퀀스의 초기 토큰에 대한 편향을 내재적으로 가지며, 이로 인해 초기 토큰의 어텐션 가중치가 지나치게 높아지는 문제가 발생한다. 이러한 문제는 모델 성능을 저하시킬 수 있다. 기술적인 관점에서, 시퀀스 시작 부분에 학습 가능한 매개변수를 추가하면 어텐션 가중치를 보다 효과적으로 재분배하여 이러한 효과를 완화할 수 있다.
4. How to Incorporate Memory?(메모리를 어떻게 통합할 것인가?)
- 기존 트랜스포머 구조에서는 Key, value의 행렬 쌍 구조는 연관기억으로 볼 수 있지만, 트랜스포머 특유의 제한된 컨텍스트 윈도우 크기 때문에 장기 기억이 아닌 단기 기억 모듈로 간주될 수 있다.
- 이러한 단기 기억의 한계를 극복하기 위해 지속적으로 학습 및 저장할 수 있는 구조(Neural memory)를 도입하여, 이전 데이터를 가중치에 저장하여 장기적인 정보 유지 가능하게 해준다.
- 따라서 긴 컨텍스트(long context)에서도 효과적으로 동작한다.
- 아래는 뉴럴 메모리를 딥러닝 모델에 통합하는 세 가지 방법(Titans 변형 모델) 제안
4.1. MAC, 메모리를 컨텍스트(Context)로 활용하기
- 기존 MLP 네트워크에 컨텍스트 모듈(Context Module)을 추가하여,과거 학습된 장기 기억을 현재 입력과 결합.
- 첫 번째 아키텍처 설계에서는 메모리를 현재 정보의 컨텍스트로 사용
- 그 이전 세그먼트들은 히스토리(과거 정보)로 간주
- 이전 메모리 상태 $M_{t−1}$를 활용하여 현재 컨텍스트 S(t) 에 맞는 정보를 검색
1. 진행 방식
0. 쿼리 생성
- 현재 세그먼트 S(t) 를 선형 변환 하여 쿼리 벡터 생성
$$q_{t}=S(t)W_{Q}$$
1. Contextual Memory (Neural Memory)에서 과거 정보 검색 (Retrieval)
- 현재 시퀀스(Core Sequence)가 주어지면, 이를 쿼리로 사용하여 Neural Memory에서 관련된 정보를 검색
- 이전 메모리 M_{t-1} 에 쿼리를 입력하여, 과거 관련 정보 h_{t} 를 검색
- 좌측 상단의 Neural Memory는 이전 시퀀스 정보를 저장하는 역할
$$h_{t}=M_{t-1}(q_{t})$$
2. Core Sequence와 Persistent Memory 결합
- Persistent Memory는 학습 가능한 데이터 독립적인 가중치로 Task-specific 정보를 포함한다.
- Persistent Memory는 Test-time에도 변하지 않으며(Fixed), Attention 모듈에서 활용된다.
- Core Sequence + Persistent Memory 결합
3. Attention을 통한 정보 조합
- Retrieval된 과거 정보, Core Sequence, Persistent Memory가 Attention 모듈을 통해 결합된다.
$$\widetilde{S}(t)=[p_{1},p_{2},...p_{N_{p}}||h_{t}||S(t)]$$
- Attention 모듈은 각 요소의 중요도를 결정하여 최종 출력으로 전달
$$y_{t}=Attn( \widetilde{S}(t)$$
- Attention을 통해 유용한 정보 선택 및 조합
4. Neural Memory 업데이트 (Test-time Learning)
- Test-time에도 Neural Memory는 지속적으로 업데이트(Update)된다.
$$M_{t}=M_{t-1}(y_{t})$$
- 반면, Persistent Memory는 고정(Fixed)된 채로 남아 Task-specific 정보를 유지
- Test-time에도 Neural Memory 업데이트 가능 (적응적 학습)
- 최종 출력 계산
$$o_{t}=y_{t} ⊗ M_{t}^{*}(y_{t})$$
2. 특징
1. 어텐션을 통한 선택적 기억 활용
- 어텐션(attention) 모듈이 현재 데이터가 장기 기억을 필요로 하는지 판단하여 불필요한 정보를 걸러낸다.
2. 메모리 효율성 증가
- 어텐션을 통해 중요한 정보만 보관
3. 테스트 단계에서도 메모리 업데이트 가능
4.2. MAG, Gated Memory
- 두 번째 방식에서는, 입력 데이터를 활용하여 직접 장기 기억을 업데이트하는 브랜치 와 슬라이딩 윈도우 어텐션(SWA, Sliding Window Attention) 브랜치 를 동시에 사용한다.
- LTM을 선택적으로 사용할 수 있도록 게이트(Gating Mechanism)를 적용하는 방식입니다.
1. 진행 방식
1. 입력 데이터 변환
- 입력 데이터 x에 프리픽스(prefix) 추가
- 즉, 입력 데이터 앞에 추가적인 정보(프리픽스)를 결합
$$\widetilde{x}=[p_{1}, p_{2}, ...,p_{N_{p}}||x]$$$
2. 슬라이딩 윈도우 어텐션(SWA) 적용
- SWA는 짧은 범위 내에서 정보를 집중적으로 처리하는 단기 기억 역할
$$y=SW-Atten^{*}(\widetilde{x} )$$
3. 메모리와 게이팅(gating) 적용
- 메모리 모듈 M(x) 에 $\widetilde{x}$를 입력하여 최종 출력 생성
$$o=y\otimes M(\widetilde{x})$$
2. 특징
1. 슬라이딩 윈도우 어텐션(SWA)으로 단기 기억 강화
2. 뉴럴 메모리는 장기 기억 역할 수행
3. 게이팅 메커니즘으로 학습 조절
4.3. Memory as a Layer (계층으로서의 메모리, MAL)- 뉴럴 메모리를 딥러닝 모델의 하나의 계층(layer)로 사용하는 방식
- 순환 신경망(RNN)과 어텐션을 결합하는 하이브리드 모델 형태
1. 입력 데이터 변환
- 프리픽스(prefix) 추가
2. Persistent Memory 추가 (Learnable Data-Independent Weights)
- Persistent Memory는 Task-specific 정보를 저장하고, Attention에 영향을 준다.
$$\widetilde{x}=[p_{1}, p_{2}, ...,p_{N_{p}}||x]$$$
3. 뉴럴 메모리를 활용하여 특징 추출
- Contextual Memory 영역에 있는 Neural Memory(과거 정보를 저장하는 메모리) 에서 관련 정보를 검색(Retrieval)한다.
- Core Sequence는 Retrieval을 위한 쿼리 역할을 하며, 필요한 정보를 Neural Memory에서 가져온다.
- Neural Memory에서 과거 정보를 검색(Retrieval)하여 활용한다.
$$y=M( \widetilde{x} )$$
4. 슬라이딩 윈도우 어텐션(SWA) 적용
- Persistent Memory, Core Sequence, Retrieval된 Neural Memory 정보가 Attention 모듈을 통해 결합한다.
- 각 정보의 중요도를 조절하며, 최종 출력이 생성한다.
- SWA를 사용하여 최종 출력을 생성한다.
$$o=SW-Atten(y)$$
5. Test-time Learning (In-context Learning)
- Neural Memory는 Test-time에도 업데이트 가능하며, 새로운 정보를 학습할 수 있다
- Neural Memory는 Test-time Learning 가능, Persistent Memory는 변하지 않는다.
5. Experiment
5.1. Experimental Setup
- 이 실험 섹션에서는 Titans 모델과 그 변형들이 언어 모델링, 상식 추론, "needle in haystack", DNA 모델링, 시간 시계열 예측 등 여러 작업에서 어떻게 성능을 발휘하는지 평가
- 모델: Titans의 세 가지 변형(MAC, MAG, MAL)과 별도의 신경망 메모리 모듈을 사용. 이들은 170M, 340M, 400M, 760M 파라미터 규모로 실험을 진행
- 데이터셋: FineWeb-Edu 데이터셋에서 15B(170M, 340M, 400M 모델) 및 30B 토큰(760M 모델)을 학습
- 기준 모델: 다양한 최신 모델들(Transformer++, RetNet, GLA, Mamba, GPT4, Llama3, Mistral 등)과 비교
- LLama 2 토크나이저, 치 크기는 50만 토큰, 학습률은 4e-4, 가중치 감소는 0.1로 설정하여 AdamW 옵티마이저를 사용
5.2. Language Modeling
- 언어 모델링과 상식 추론 작업에서의 Perplexity와 정확도 성능을 비교합니다. 실험은 Titans의 변형들과 여러 기준 모델들을 사용하여 진행
- Neural Memory Module은 언어 모델링에서 가장 우수한 성능을 보이며, 가중치 감소와 모멘텀의 중요성이 강조
- Titans의 변형(MAC, MAG, MAL)은 기존 하이브리드 모델들보다 더 뛰어난 성능을 보이며, 특히 MAC은 긴 의존성 처리에 강점을 가진다.
5.3. Needle in a Haystack (NIAH)
- 긴 시퀀스를 처리하는 모델의 성능을 평가하기 위해 Needle in a Haystack (NIAH) 작업이 설계되었다.
- 이 작업은 긴 산만한 텍스트(“haystack”)에서 정보 조각(“needle”)을 검색하는 능력을 측정한다.
- Neural Memory Module이 모든 시퀀스 길이에서 기존 모델들보다 뛰어난 성능을 보였다.
- 뛰어난 성능을 Titans가 기존 시퀀스 모델들과 차별화되는 세 가지 중요한 요소 덕분으로 설명할 수 있다.
1. Neural Memory는 momentum과 가중치 감소(weight decay)를 사용하여 메모리 용량을 더 잘 관리할 수 있습니다. Mamba2는 메모리를 지울 수 없기 때문에 시퀀스 길이가 길어질수록 성능이 급격히 떨어진다.
2. Mamba2와 비교했을 때, Titans는 깊고 비선형적인 메모리를 사용하여 더 나은 메모리 관리를 제공,
3. DeltaNet은 델타 규칙을 사용해 메모리를 지울 수 있지만, 완전한 잊기 메커니즘이 없어 메모리를 완전히 지울 수 없습니다. 반면 Titans는 Neural Memory와 잊기 메커니즘을 결합하여 더 뛰어난 성능을 보인다.
- Titans 변형들 중에서는 MAC가 가장 우수한 결과를 보여준다.
5.4. BABILong Benchmark
- BABILong 벤치마크를 사용한 더 어려운 작업을 다룹니다. 이 작업에서는 모델이 매우 긴 문서에서 분산된 사실들을 추론해야한다.
- BABILong은 더 긴 문서에서 복잡한 사실을 추론하는 작업으로, 모델의 성능을 더 잘 평가할 수 있다.
- Few-shot 설정: 대형 사전 학습된 모델을 사용하여 성능을 평가.
- Fine-tuning 설정: MAC 변형의 Titans를 미세 조정하여 다른 미세 조정된 모델들과 성능을 비교.
5.5. The Effect of Deep Memory
- 깊은 메모리(deep memory)훈련 시간과 모델 성능에 미치는 영향을 평가
- 메모리 깊이 (LM)을 1, 2, 3, 4로 설정하고, 성능 비교를 위해 Mamba 모델을 기준으로 사용
- 모든 모델은 동일한 훈련 절차를 사용하여 Pile 데이터셋의 일부로 훈련
- 메모리 깊이가 증가하면 Perplexity가 개선되고, 특히 파라미터 수가 적은 모델에서 더 효과적
- 깊은 메모리는 모델이 더 나은 성능을 보이게 하며, 특히 시퀀스 길이가 긴 경우 유리합니다.
- 그러나 메모리 깊이가 깊어질수록 훈련 속도가 느려지므로, 효율성과 효과성 간에 균형을 맞추는 것이 중요
5.6. Time Series Forecasting(시계열 예측)
- Titans의 메모리 모듈의 효과를 더 넓은 범위의 작업에서 보여주기 위해, 시계열 예측 작업에 대한 성능을 평가
- 평가에 사용된 시계열 예측 벤치마크 데이터셋은 ETT, ECL, Traffic, Weather
- Neural Memory는 시계열 예측 작업에서도 뛰어난 성능을 발휘하며, 다양한 아키텍처와 비교하여 우수한 성과를 기록
5.7. DNA Modeling
- 자연어를 넘어서 Titans의 능력을 이해하기 위해, DNA 모델링 작업에서 Neural Memory Module의 성능을 추가로 평가
- Titans의 Neural Memory Module은 DNA 모델링 작업에서도 뛰어난 성능을 보여주며, 자연어 처리 외에도 생물학적 데이터에 대한 우수한 처리 능력을 입증
5.8. Efficiency
- Neural Memory와 Titans의 효율성을 최신 시퀀스 모델들과 비교
- 모델 훈련 처리량(training throughput)은 시퀀스 길이와 배치 크기에 따라 진행
- Titans (MAL) 모델은 기준 모델들과 메모리 모듈보다 더 빠른 성능을 보였습니다. 이 성능 향상의 주요 이유는 Flash-Attention(Dao 2024)의 고도로 최적화된 커널 덕분
- Flash-Attention은 SWA와 전체 어텐션 모듈을 구현하는 데 사용되어 훈련 처리량을 개선하는 데 중요한 역할한다.
5.9. Ablation Study
- Neural Memory Module을 기본 모델로 삼고, 한 번에 하나씩 구성 요소를 변경하여 그 효과를 평가
- 변경된 요소는 다음과 같다.
1. 깊은 메모리를 선형 메모리로 교체.
2. 컨볼루션 제거.
3. 서프라이즈 측정에서의 모멘텀 제거.
4. 가중치 감소 (혹은 잊기 메커니즘) 제거.
5. 영속적인 메모리 제거.
- 가장 큰 기여는 가중치 감소(weight decay), 모멘텀(momentum), 컨볼루션(convolution), 영속적인 메모리(persistent memory) 순으로 나타난다.
- MAC와 MAG는 언어 모델링과 상식 추론에서는 유사한 성능을 보였으나, 긴 컨텍스트 작업에서는 MAC가 더 좋은 성능을 보였다.
6. Conclusion
- 이 논문에서는 테스트 시 학습하는 메타 인컨텍스트 학습자로서 신경망 장기 기억(Neural Long-Term Memory)을 제시했다.
- Neural Memory Module은 본질적으로 순환 모델로, 더 놀라운 (surprising) 또는 놀라운 토큰에 근접한 토큰을 기억하도록 적응적으로 학습한다. 현대의 순환 모델들과 비교했을 때, 더 표현력 있는 기억 갱신(메모리 갱신)과 저장 메커니즘을 제공한다.
- 이 기억을 사용하여, Titans 아키텍처와 그 세 가지 변형을 제시했습니다. 우리는 이 메모리 모듈을 컨텍스트, 게이팅, 레이어로 통합할 것을 제안
- 다양한 작업에 대한 실험적 평가 결과, Titans는 Transformers 및 최신 선형 순환 모델보다 효과적이며, 특히 긴 컨텍스트에서 좋은 성능을 보인다.
참고 목록
https://kimjy99.github.io/%EB%85%BC%EB%AC%B8%EB%A6%AC%EB%B7%B0/titans/
[논문리뷰] Titans: Learning to Memorize at Test Time
Titans 논문 리뷰
kimjy99.github.io
Question.
1. key, value에 loss 최소화시 왜 loss의 최소화가 좋은 영향을 미치는 건지?
- Loss의 최소화가 아닌 적절한 차이를 유지하면서 줄이는 것, key와 value의 관계를 왜곡하지 않으면서도 최적의 매핑을 학습하는 것이 좋은 영향을 미틴다.
- $M_{t-1}은 key와 value 사이의 차이를 특정 방식으로 조정하는 역할을 한다.$
2. test time에 value 값이 없는데 학습이 어떻게 가능한지?
- 모델이 학습 중에 key와 value 간의 관계를 충분히 학습했기 때문에, 모델이 연관 학습을 통해 추론한 value를 신뢰할 수 있다는 판단 가정하게 Test시에도 그 값을 사용해 진행
3. MLP가 따로 2개로 구성되어 돌아가는 것인지
- Titan의 어떤 구조를 사용하느냐에 따라 다르다.
- LTM과 Titan의 모델에서 사용하는 MLP가 각각 독립적일 수도 있고, 때때로 합쳐서 사용하는 경우도 있다.
- 예를 들어, Titan의 MLP가 key와 value를 매핑하는 데 사용되는데, 이때 LTM의 MLP가 Titan 모델의 일부로 합류하여 기억된 정보를 활용할 수 있는 방식으로 사용할 수 있어.