Maglev: Sliding Recurrent Memory

기술 요약: Maglev – 슬라이딩 재귀적 메모리 (Sliding Recurrent Memory)

문제 정의
현재의 트랜스포머(Transformer) 아키텍처는 메모리 용량과 계산 효율성 사이의 근본적인 트레이드오프에 직면해 있다. 전체 인과적 셀프 어텐션(Full causal self-attention)은 강력한 비선형 검색과 인컨텍스트 연산을 가능하게 하지만, 컨텍스트 길이에 따라 비용이 선형적으로 증가한다. 반면, 슬라이딩 윈도우 어텐션(Sliding-window attention)은 어텐션 및 KV-캐시 비용을 제한하지만, 먼 거리의 정보를 버리게 되어 지속적인 장기 기억을 제공하는 데 실패한다. 고전적인 재귀 모델(예: LSTM)은 이력을 압축하지만, 작은 상태(state)를 순차적으로 전파하므로 용량이 제한된다. 메모리 트랜스포머와 선형 어텐션 모델들이 이 간극을 메우려 시도하고 있으나, 이들은 종종 인위적인 세그먼트 경계를 부과하거나, 구조화된 선형 업데이트에 의존하거나, 훈련 시 순차적 언롤링(sequential unrolling)을 요구한다. 핵심 과제는 병렬 프리트레이닝(pretraining)이 가능하면서도 추론 비용이 유계(bounded)인, 트랜스포머를 위한 실용적인 토큰 단위 비선형 재귀 메모리 메커니즘을 개발하는 것이다.

방법론: Maglev 아키텍처
Maglev는 훈련 중의 메모리 구축과 추론 중의 메모리 전파를 분리함으로써 이 문제를 해결한다. 이 아키텍처는 두 개의 결합된 모델로 구성된다:

1. 프리필러 (Prefiller, $Q$): 관찰된 시퀀스를 병렬로 처리하는 인과적 트랜스포머로, 더 넓은 수용 영역(예: 인터리브된 풀 및 슬라이딩 윈도우 어텐션)을 사용한다. 이는 모든 토큰 위치에 대해 타겟 메모리 벡터 $m'_t$의 궤적(trajectory)을 생성한다.
2. 디코더 (Decoder, $P$): 다음 토큰을 예측하고 자체적인 메모리 벡터 $m_t$를 생성하는 슬라이딩 윈도우 트랜스포머이다. 결정적으로, $P$는 전체 이력에 접근할 수 없다. $P$는 오직 최근의 토큰들과 이동된(shifted) 메모리 벡터들의 유계된 윈도우만을 본다.

훈련 메커니즘: 리프티드 병렬 훈련 (Lifted Parallel Training)
순차적 언롤링(이는 병렬화를 방해함) 대신, Maglev는 2단계 훈련 전략을 채택한다:

• 패스 1 (Pass 1): 프리필러 $Q$가 전체 타겟 메모리 궤적 $m'_{1:T}$를 병렬로 생성한다.
• 패스 2 (Pass 2): 디코더 $P$는 입력 토큰과 함께 이동된 타겟 궤적 $m'_{0:T-1}$을 소비하여 다음 토큰을 예측하고 자신의 메모리 $m_{1:T}$를 생성한다.
• 목적 함수 (Objective): 모델은 다음과 같은 복합 손실 함수를 사용하여 최적화된다:
  • 다음 토큰 예측 손실 ($L_{CE}$): 언어 모델링을 위한 표준 크로스 엔트로피 손실.
  • 메모리 일관성 손실 ($L_{cons}$): 디코더의 출력 메모리 $m_t$를 프리필러의 타겟 $m'_t$와 정렬하는 항. 이는 디코더가 순차적 언롤링 없이도 다음 단계를 위해 필요한 메모리 상태를 재현하도록 학습시킨다.

추론 메커니즘
추론 시에는 프리필러 $Q$가 폐기된다. 디코더 $P$는 재귀적으로 실행되며, 이전에 생성된 자신의 메모리 $m_{t-1}$을 시스템에 다시 주입한다.

• 재귀적 K/V 주입: 별도의 메모리 토큰을 사용하는 대신, Maglev는 학습된 프로젝션과 게이팅 메커니즘을 통해 이동된 메모리 상태를 키(Key)/값(Value) 특징에 주입한다.
• 유계된 복잡도: 디코더는 고정된 크기의 슬라이딩 윈도우 토큰과 메모리를 유지한다. 재귀적 상태는 K/V 캐시를 통해 앞으로 전달되며, 이를 통해 추론 비용(연산 및 캐시 크기)이 전체 시퀀스 길이에 독립적으로 유계되도록 보장하며, 이는 표준 슬라이딩 윈도우 어텐션의 프로파일과 일치한다.

주요 기여

1. 토큰 단위 비선형 재귀: Maglev는 모든 토큰의 디코더 상태가 메모리 역할을 수행하며, 추가적인 메모리 토큰이나 인위적인 세그먼트 경계 없이 트랜스포머의 전체 비선형 깊이를 통과하는 메커니즘을 도입한다.
2. 재귀 모델을 위한 병렬 훈련: 프리필러를 사용하여 타겟을 병렬로 생성함으로써, Maglev는 전통적인 RNN 훈련의 순차적 병목 현상을 피하고 표준 시퀀스 병렬 처리량을 가진 비선형 재귀 모델의 병렬 프리트레이닝을 가능하게 한다.
3. 파라미터 공유: 저자들은 프리필러 $Q$와 디코더 $P$가 대부분의 파라미터(특히 트랜스포머 블록)를 공유할 수 있으며, 성능을 크게 저하시키지 않고도 훈련 설정의 파라미터 메모리 오버헤드를 줄일 수 있음을 입증했다.
4. 유계된 추론: 배포된 모델은 지속적인 메모리 능력을 갖춘 재귀적 K/V 주입을 통해 슬라이딩 윈도우 어텐션의 효율성을 유지하면서도 유계된 컨텍스트를 유지한다.

실험 결과
nanochat d20 아키텍처(435M 파라미터, 43.52B 토큰으로 훈련됨)에서 Maglev를 인터리브된 풀/슬라이딩 윈도우 트랜스포머, 순수 슬라이딩 윈도우 모델, 그리고 잠재 재귀 트랜스포머(Latent Recurrent Transformer, LRT) 베이스라인과 비교 평가하였다.

• 성능: 가장 우수한 Maglev 변형 모델(별도 파라미터, $\lambda=1$)은 FineWeb-Edu 검증 bits-per-byte(BPB)를 슬라이딩 윈도우 베이스라인의 0.7413에서 0.7251로 개선했다.
• 다운스트림 태스크: 평균 다운스트림 정확도 56.4%를 달성하여, 매칭된 슬라이딩 윈도우 트랜스포머(54.1%)와 매칭된 LRT 베이스라인을 모두 능가했다.
• 일관성 가중치: 연구에서는 최적의 일관성 가중치($\lambda$)가 파라미터 공유 전략에 따라 달라짐을 언급했다. 높은 $\lambda$는 별도 파라미터 모델에는 유리했으나, 공유 파라미터 모델에는 제약이 될 수 있었다.

의의 및 주장
본 논문은 Maglev가 매 토큰마다 다시 쓰일 수 있는 지속적인 토큰 단위 메모리를 갖춘 비선형 트랜스포머를 구축하는 실용적인 경로를 제공한다고 주장한다. Maglev는 풀 어텐션의 표현력과 유계된 컨텍스트 모델의 효율성 사이의 간극을 성공적으로 메웠다. 훈련 중에 메모리 생성과 메모리 사용을 분리함으로써, Maglev는 이전에는 비선형 재귀에 대해 어려웠던 작업인 비선형 재귀 아키텍처를 위한 병렬 프리트레이닝을 달성했다. 저자들은 Maglev를 이 설계 공간에 대한 예비 조사로 규정하며, 이 접근 방식을 확장하고, 배포를 위한 커널 효율성을 최적화하며, 대안적인 메모리 주입 메커니즘(예: 잔차 스트림 주입 또는 크로스 어텐션)을 탐구하기 위한 추가 작업이 필요하다고 명시했다.