Titans: Learning to Memorize at Test Time

Titans: Learning to Memorize at Test Time

배경 및 문제 정의

1. Transformer의 단점:

• Transformer는 Attention 메커니즘으로 동작, 높은 성능을 보이지만 컨텍스트 창 크기가 제한적
• 컨텍스트 창 크기에 따라 시간 및 메모리 복잡도가 이차적으로 증가, 긴 시퀀스 처리하기 어려움
• 이러한 한계는 언어 모델링, 비디오 이해, 장기 시간 시계열 예측과 같은 실제 응용 분야에서 문제2. 효율적인 변형 모델의 시도:
• 최근 연구는 Transformer의 비효율성을 해결하기 위해 Linear Transformer와 같은 대안 제시
• 하지만, 선형 모델은 데이터를 작은 상태(state)로 압축, 이로인해 긴 시퀀스 데이터를 적절히 표현하지 못하는 한계3. 기존 모델의 한계:
• 기존 아키텍처(예: LSTM, Transformer)는 일반화, 길이 외삽, 추론 등의 과제에서 성능이 제한적
• 이러한 모델은 단기 기억이나 장기 기억과 같은 뇌의 다양한 기억 체계 모사에 어려움이 있음

메모리 관점에서 본 문제

1. 기억의 중요성:

• 인간의 학습 과정에서 기억은 필수
• 예를 들어, Transformer의 Attention은 데이터를 그대로 저장, 쿼리와 키의 유사도를 기반으로 검색 > 단기 기억2. 기억 구조 및 업데이트 메커니즘:
• Transformer와 선형 Transformer의 차이는 데이터를 저장하는 메모리 구조와 업데이트 방식에서 옴
• Transformer는 데이터를 압축하지 않고 저장, 선형 Transformer는 데이터를 작은 크기로 압축해 저장3. 효율적인 메모리 설계의 필요성:
• 이상적인 메모리 설계는 다음 질문에 대한 답변이 필수적:
  • Q1: 좋은 메모리 구조는 무엇인가?
  • Q2: 적절한 메모리 업데이트 메커니즘은 무엇인가?
  • Q3: 효과적인 메모리 검색 방법은 무엇인가?
• 인간 기억 체계에서 영감을 얻어 단기, 장기, 작업 기억을 독립적으로 설계하고 상호작용하게 해야 함

Titans 모델의 제안

1. 기본 아이디어:

• Titans는 Transformer의 Attention을 단기 기억으로, 새로운 신경 메모리 모듈을 장기 기억으로 간주, 이를 결합
• 장기 기억은 테스트 시점에서도 데이터를 학습하고 저장할 수 있도록 설계2. 주요 질문 및 기여:
• Titans 모델은 다음과 같은 질문을 기반으로 설계 함:
  • Q4: 효율적인 아키텍처를 설계하려면 메모리 모듈을 어떻게 결합해야 하는가?
  • Q5: 과거 데이터를 효과적으로 저장하려면 깊은 메모리 모듈이 필요한가?
• 이를 통해 모델이 인간의 기억처럼 데이터 중요도에 따라 선택적으로 학습하고 기억을 저장할 수 있도록 함3. 논문의 기여:
• 신경 장기 메모리: 테스트 시점에서 데이터를 학습하고 저장하는 메모리 모듈을 설계
• Titans 아키텍처: 단기 기억과 장기 기억을 결합하여 효율적이고 강력한 학습 구조를 제안
• 실험 결과: Titans는 다양한 작업에서 기존 Transformer 및 선형 모델보다 우수한 성능을 보임

기초 개념 및 배경 지식

1. Attention 메커니즘과 Transformer의 한계

Attention의 정의

• Transformer의 핵심 구성 요소는 Attention 메커니즘
• Attention은 입력 데이터를 쿼리(Query), 키(Key), 값(Value)로 변환, 각 쿼리와 키의 유사도에 따라 가중합을 계산:
• $y_i = \sum_{j=1}^i \frac{\exp(Q_i^T K_j/\sqrt{d})}{\sum_{l=1}^i \exp(Q_i^T K_l/\sqrt{d})} V_j$
• 이 과정은 정확한 의존성 모델링을 가능하게 하지만, 계산 비용이 O(N²)로 증가하여 긴 시퀀스 데이터를 처리하기 어렵습니다.효율적인 Attention 메커니즘
• 메모리 소비를 줄이기 위해 다양한 효율적인 Attention 기법이 제안되었습니다:
  • 희소화(Sparsification): Attention 행렬에서 중요한 요소만 남겨 계산
  • 커널 기반 선형 Attention: Softmax 대신 커널 함수를 사용하여 계산 비용을 선형적으로 줄임
• 선형 Attention의 일반적인 형태:
• $y_i = \frac{\phi(Q_i)^T \sum_{j=1}^i \phi(K_j)V_j}{\phi(Q_i)^T \sum_{j=1}^i \phi(K_j)}$
• 여기서 ϕ(.)는 커널 함수입니다.선형 Transformer의 한계
• 선형 Attention은 계산 효율성의 이점, 그러나 메모리 압축의 한계로 인해 긴 시퀀스 데이터의 복잡한 의존성을 제대로 학습하지 못함
• 이로 인해 모델 성능이 Transformer에 비해 떨어지는 경우가 많음

2. 현대적인 선형 모델과 메모리 관점

선형 모델과 메모리

• 선형 Transformer와 선형 RNN은 데이터를 고정 크기의 상태(state)로 압축, 저장
• 일반적인 RNN 모델의 메모리 갱신 및 검색 구조:
  • 쓰기: $M_t = f(M_{t-1}, x_t)$
  • 읽기: $y_t = g(M_t), x_t$
• 선형 Transformer는 키-값 쌍을 행렬 메모리로 압축, 이는 긴 시퀀스에서 메모리 오버플로를 유발메모리 관리의 필요성
• 긴 시퀀스 데이터를 처리할 때 중요한 과제는:
  • 망각 메커니즘: 불필요한 데이터를 지우는 능력.
  • 쓰기 메커니즘 개선: 데이터를 효과적으로 저장하는 능력.
• 기존 연구:
  • 망각 게이트: 데이터 의존적인 메커니즘으로 불필요한 데이터를 지움 (예: GLA, LRU, Griffin).
  • 델타 규칙(Delta Rule): 새로운 데이터를 추가하기 전, 이전 값을 제거하는 방식.

3. 메모리 모듈 설계의 통찰

메모리 설계의 문제

• 이상적인 메모리 설계는 다음 질문을 고려해야 함:
  • 메모리 구조는 어떻게 설계해야 하는가?
  • 데이터를 효과적으로 저장하려면 어떤 업데이트 메커니즘이 필요한가?
  • 데이터를 검색하는 방법은 무엇인가?Titans의 메모리 접근
• Transformers와 선형 Transformers의 메모리 차이 기반, Titans는 단기 및 장기 메모리를 구별하여 설계
• 인간 기억 체계를 모방, 각 메모리 모듈이 독립적으로 작동하되, 상호작용할 수 있도록 설계

4. 장기 메모리의 필요성

• 기존 메모리 설계는 데이터를 선형적으로 압축, 이는 복잡한 의존성 학습하는 데에 한계
• Titans는 다층 퍼셉트론(MLP)을 활용한 비선형 장기 메모리를 설계, 더 깊은 데이터 학습
• 이론적으로도 MLP는 선형 모델보다 표현력이 높음, 실험 결과 장기 의존성을 효과적으로 처리할 수 있음

테스트 시점에서 학습 가능한 신경 장기 메모리 모듈 설계

1. 장기 메모리 모듈의 설계

문제 정의

• 기존 메모리 시스템은 학습 단계에서 데이터를 저장, 테스트 시점에는 저장된 데이터 활용
• 그러나 현실 세계의 데이터는 동적인 특징을 가짐, 테스트 시점에서도 새로운 데이터를 학습하고 이를 저장할 수 있는 메모리 시스템 필요설계 원칙
• 기억과 학습의 정의
  • 기억: 입력 데이터가 메모리에 저장되는 과정
  • 학습: 메모리가 데이터의 유용한 특성을 효과적으로 저장하도록 하는 과정
• 중요 설계 요소
  • 데이터를 선택적으로 저장
  • 과거의 데이터를 적절히 망각, 업데이트할 수 있는 기법

2. 서프라이즈 기반 메모리 학습

서프라이즈 메트릭

• 입력 데이터가 기존 데이터와 얼마나 다른지(놀라운지; Surprise)를 측정, 중요한 데이터 학습
• 서프라이즈는 손실 함수의 그래디언트를 이용해 측정
• $Surprise = -\theta_t\nabla l(M_{t-1};x_t)$
    여기서 ϕ(.)는 커널 함수입니다.

모멘텀과 감쇠를 결합한 업데이트

• 서프라이즈 메트릭에 과거 서프라이즈와 순간 서프라이즈를 결합하여 메모리 업데이트
• $S_t = \eta_t S_{t-1} -\theta_t\nabla l(M_{t-1};x_t)$
  • $S_t$: 현재 서프라이즈 값
  • $\eta_t$: 과거 서프라이즈의 중요도를 조정하는 감쇠 계수
  • $\theta_t$: 현재 서프라이즈를 반영하는 학습율

적응형 망각 메커니즘

• 메모리 용량이 제한적이므로 불필요한 정보를 잊도록 설계
• $M_t=(1-\alpha_t)M_{t-1}+S_t$
  • $\alpha_t$: 망각 비율을 조정하는 가중치, 중요하지 않은 데이터를 잊게 함
  • $(1-\alpha_t): 기존 메모리 상태를 유지하는 정도$

3. 메모리 모듈의 구조

기본 메커니즘

• 신경 장기 메모리는 입력 데이터를 키-값 쌍으로 변화하여 저장:
• $k_t = x_tW_K$,  $v_t=x_tW_V$
  • $k_t$: 입력 데이터를 키로 변환
  • $W_K$: 키를 생성하는 선형 변환 행렬
  • $v_t$: 입력 데이터를 값으로 변환
  • $W_V$: 값을 생성하는 선형 변환 행렬
• 손실 함수:
• $l(M_{t-1};x_t)=||M_{t-1}(k_t)-v_t||_2^2$
  • 이전 메모리 $M_{t-1}$와 현재 입력 데이터 $x_t$의 예상 값 사이의 차이를 최소화메모리 검색
• 메모리에 저장된 데이터를 검색할 때, 입력 데이터 $q_t = x_tW_Q$를 쿼리로 사용:
• $y_t=M*(q_t)$
  • 여기서 $M*(q_t)$는 업데이터 없이 메모리를 조회한 결과, 메모리에서 $q_t$에 해당하는 값을 반환하는 함수

4. 학습 효율화

병렬 학습

• 긴 시퀀스를 처리하기 위해 병렬 연산을 사용하여 학습 속도를 높임
• 미니배치 기반 그래디언트 계산을 통해 효율적으로 학습:
• $M_t=\beta_tM_0-\sum_{i=1}^t\frac{\theta_i}{\beta_t}\nabla l(M_{t^\prime};x_i)$
  • $t^\prime$: 현재 배치의 시작점
  • $\beta_t$: 학습율과 감쇠 계수를 포함한 가중치추가적인 최적화
• 모델의 학습 속도를 높이기 위해 각 매개변수를 시퀀스의 청크(chunk) 단위로 설정
• 청크 내에서 일정한 학습률과 감쇠 계수를 사용하여 효율적인 계산 수행

5. 지속적인 메모리

• 컨텍스트 의존적 메모리와는 별도로, 데이터 독립적인 지속적 메모리를 추가 설계
• 지속적 메모리는 작업(task) 관련 정보를 저장하며, 입력 데이터와 무관하게 고정된 상태로 유지

6. 장기 메모리 모듈의 특성

1. 데이터 감쇠와 망각
   • 적응형 망각 메커니즘으로 메모리 용량을 최적화
2. 깊이 있는 학습
   • 선형 메모리보다 다츤 퍼셉트론(MLP)을 활용하여 더 복잡한 패턴 학습 가능
3. 테스트 시점 학습
   • 기존 모델과 달리, 테스트 단계에서도 메모리를 업데이트하녀 성능 향상

How to Incorporate Memory? (메모리를 통합하는 방법)

1. 메모리 통합의 중요성

• 문제 정의: 설계괸 신경 장기 메모리 모듈을 어떻게 딥러닝 아키텍쳐에 통합할지 결정하는 것은 모델 성등과 효율성을 좌우하는 중요한 문제입니다.
• 아이디어: Transformer의 Attention 메커니즘을 단기 메모리로 간주하고, 신경 장기 메모리를 장기 메모리로 통합하여 상호 보완적인 작동을 유도합니다.2. 세가지 Titans 아키텍쳐
• Memory as a Context (MAC)
• 기본 아이디어:
  • 장기 메모리를 컨텍스트로 사용하여 Attention 모듈에 제공
  • 장기 메모리가 현재 입력 데이터와 관련된 정보를 검색하고, 이를 Attention 입력으로 결합합니다.
• 작동 방식:

1. 시퀀스를 고정된 크기의 세그먼트 $S(i)$로 분할
2. 이전 메모리 $M_{t-1}$를 사용하여 현재 세그면트 $S(t)$와 관련된 정보를 검색
   • $h_t = M^*_{t-1}(q_t)$,  $q_t=S(t)W_Q$
     • $h_t$: 장기 메모리에서 검색된 과거 정보
     • $W_Q$: 쿼리 행렬
3. 검색된 정보 $h_t$와 지속적 메모리를 현재 세그먼트와 결합
   • $\hat{S}(t)=[p_1,p_2,\dots,p_{N_p}]||h_t||S(t)$
     • $[p_1,p_2,\dots,p_{N_p}]$: 지속적 메모리.
     • $||$: 벡터 연결 연산
4. Attention 모듈이 결합된 데이터를 처리하여 최종 출력 생성
   • $y_t = Attn(\hat{S}(t))$

• 장점:
  • Attention이 현재와 과거 데이터를 통합하여 유의미한 정보만을 선택적으로 저장
  • 메모리 오버플로를 방지하며, 테스트 단계에서도 지속적으로 학습 가능

아래 그림은 MAC 아키텍쳐를 시각화한 것이다. 입력 시퀀스에 입력 데이터를 기반으로 검색된 관련 Contextual 메모리 정보와 입력에 독립적이고 태스크와 관련된 데이터인 Persistant 메모리 데이터를 통합해 Attention 모듈에서 입력하는 과정을 보여준다.

Memory a a Gating (MAG)

• 기본 아이디어:
  • 단기 메모리와 장기 메모리를 게이팅 메커니즘으로 결합하여 데이터를 처리
  • 슬라이딩 윈도우 Attention을 단기 메모리로 사용

슬라이딩 윈도우 어텐션?
어텐션의 입력 시퀀스를 작은 윈도우 단위로 나누어 어텐션 연산을 수행하는 방법
메모리 효율적인 모델을 구현하는 데에 사용
긴 컨텍스트를 처리할 때 한번에 하지않고
윈도우 단위로 나누어 처리해 메모리 사용량을 줄이고 계산 비용 최적화 가능

• 작동 방식:
  1. 입력 데이터에 지속적 메모리를 추가:
     • $x=[p_1,p_2,\dots,p_{N_p}]||x$
  2. 슬라이딩 윈도우 Attention(SW-Attn)을 단기 메모리로 사용하여 현재 데이터를 처리:
     • $y=SW-Attn^*(x)$
  3. 장기 메모리 모듈이 입력 데이터를 처리한 결과 $M(x)$와 결합:
     • $o=y\otimes M(x)$
       • $\otimes$: 게이팅 연산 (예: 점별 곱셈과 비선형 활성화)
       • $M(x)$: 장기 메모리가 입력 데이터를 처리한 결과
• 장점:
  • 장기 메모리를 추가로 활용하여 단기 메모리의 한계를 보완
  • 게이팅 메커니즘으로 중요한 정보만 선택적으로 활성화

아래 그림은 MAG 아키텍쳐를 시각화한 것이다. 입력에 Persitant 메모리만 결합하고 슬라이딩 윈도우 attention과 Contextual 메모리에서 결합된 입력 데이터를 처리한 결과를 게이팅 연산을 통해 최종 출력을 생성하는 과정이다.

Memory as a Layer (MAL)

• 기본 아이디어:
  • 장기 메모리를 네트워크의 독립적인 레이어로 설계하여 다른 모듈과 병렬로 처리
  • 슬라이딩 윈도우 Attention(SW-Attn)과 결합
• 작동 방식:
  1. 입력 데이터에 지속적 메모리를 추가
     • $x=[p_1,p_2,\dots,p_{N_p}]||x$
  2. 장기 메모리를 레이어로 처리하여 데이터 압축
     • $y=M(x)$
  3. 슬라이딩 윈도우 Attention이 이를 처리하여 최종 출력 생성
     • $o=SW-Attn(y)$
• 장점:
  • 구조적으로 단순하며, 기존 하이브리드 모델들과 유사한 설계를 제공
  • 장기 메모리를 독립적으로 평가 가능

3. 각 아키텍쳐의 비교

아키텍쳐	장점	단점
MAC	긴 의존성을 효과적으로 처리하며, 유연한 메모리 관리 가능	계산 비용이 상대적으로 높음
MAG	단기 및 장기 메모리를 게이팅으로 결합하여 효율적	설계가 비교적 복잡
MAL	간단한 구조로 장기 메모리 성능 평가 가능	단기와 장기 메모리의 상호작용이 제한적

실험 결과

실험 목표

Titans 모델의 성능을 검증, 기존 모델 대비 성능을 비교:

1. Titans는 다양한 작업에서 기존 모델보다 우수한 성능을 보이는가?
2. Titans의 컨텍스트 길이는 실제로 얼마나 긴가?
3. 메모리의 깊이가 모델 성능과 효율성에 미치는 영향은 무엇인가?

실험 환경

모델 및 비교군

• Titans 모델:
  • MAC, MAG, MAL
  • 단독 메모리 모듈(LMM; Long-term Memory Module)
• 모델 크기 및 훈련
  • 170M, 340M, 400M, 760M 파라미터로 테스트
  • 데이터: FineWeb-Edu 데이터셋에서 15B~30B 토큰으로 학습.
• 평가 데이터
  • 언어 모델링: 위키피디아, LAMBADA, HellaSeag 등에서 perplexity(PPL)과 정확도(ACC) 측정
  • Needle-in-a-Haystack (NIAH): 긴 시퀀스에서 특정 정보를 추출하는 작업
  • BABILong 벤치마크: 긴 문서에서 논리적 추론을 요구하는 작업
  • 시계열 예측: 장기적 의존성을 평가
• 비교군
  • Transformer 기반 모델 (Transformer++, RetNet 등)
  • 선형 Attention 모델(GLA, DeltaNet, 등)
  • 하이브리드 모델(Sambe, Gated DeltaNet-H2 등)실험 결과
• 논문에 나와있는 표를 보시길..
• 실험 결과상으로를 제안 기법의 성능이 뛰어납니다.

결론

주요 기여

1. 신경 장기 메모리 설계:
   • Titans는 테스트 시점에서 학습 가능한 신경 장기 메모리(NLTM)를 제안
   • 서프라이즈 메트릭과 적응형 망각 메커니즘을 통해, 테스트 단계에서 새로운 데이터를 학습하고 저장 가능
2. 아키텍처 혁신:
   • 단기 메모리(Attention)와 장기 메모리를 결합한 세 가지 아키텍처(MAC, MAG, MAL)를 제안, 다양한 작업에 적합한 유연성 제공
3. 실험 결과:
   • Titans는 언어 모델링, 긴 시퀀스 검색(Needle in a Haystack), 상식 추론, 시간 시계열 예측 등 다양한 작업에서 기존 모델 대비 우수한 성능을 보임
   • 특히 MAC 변형은 긴 시퀀스에 대해 압도적인 성능을 보임, 초대형 모델(GPT-4 등)을 능가하기도 함

연구 의의

• 긴 시퀀스 처리의 혁신
  • Titans는 컨텍스트 창 크기를 2M 토큰 이상으로 확장할 수 있어, 기존 Transformer 모델의 한계 극복
• 인간 기억에서 영감
  • 인간의 단기 기억과 장기 기억 구조를 모델에 통합, 데이터의 중요도에 따라 선택적으로 학습 및 저장하도록 설계

향후 연구 방향

1. 메모리 구조의 개선:
   • Titans의 신경 장기 메모리 구조를 더 정교하게 설계하여 효율성을 높이는 연구
   • 특히, 더 깊은 메모리 구조(MLP 개선)와 병렬 학습 최적화에 초점
2. 대규모 모델로의 확장:
   • Titans를 초대형 모델로 확장하여 더 다양한 작업에서 성능 검증
   • 대규모 데이터셋에서 Titans의 효율성을 평가하는 추가 연구
3. 응용 분야 확장:
   • 언어 모델링 외에도, 비디오 이해, 생물학적 데이터 분석(예: 게놈 데이터), 시계열 데이터 처리 등 다양한 분야에 Titans 적용

결론

• Titans는 테스트 시점 학습(Test-Time Learning)을 가능하게 하는 새로운 방향성 제시
• 단기 및 장기 메모리 통합을 통해 기존 Transformer 및 순환 모델의 한계 극복, 다양한 작업에서 탁월한 성능 입증
• Titans는 긴 시퀀스 처리와 메모리 관리라는 두 가지 측면에서 중요한 혁신을 이룸, 향후 연구에서 중요한 기반 제공