Maglev: Sliding Recurrent Memory

We introduce , a recurrent Transformer architecture with fixed-size memory that generalizes sliding-window attention while remaining parallelizable during training. consists of two coupled models: a prefiller Q, which leverages full attentionIn practice, we use interleaved full and sliding-window attention for Q, as this yields stronger performance. The essential requirement is that Q be more expressive than P, with access to the full history. to produce memory targets m'_t, and a decoder P, which uses only sliding-window attention and recurrent K/V injection to produce decoder memories m_t for next-token prediction. We train with a memory consistency loss that aligns m_t with m'_t, allowing inference to use P alone. Empirically, improves validation loss and downstream pretraining benchmarks over sliding-window and latent recurrent transformer baselines. Moreover, sharing parameters between P and Q reduces parameter memory while preserving most of the gains.

한 줄 요약

Maglev는 슬라이딩 윈도우 어텐션을 일반화하면서 추론 시 고정 메모리로 작동하는 반복 트랜스포머 구조를 제안한다.

핵심 기여도

Maglev는 Q(_prefiller)와 P(decoder)의 두 모델로 구성된 반복 트랜스포머 구조를 제안.
Q는 전체 어텐션과 슬라이딩 윈도우 어텐션을 혼합하여 더 강력한 성능을 달성.
P는 슬라이딩 윈도우 어텐션과 반복 K/V 주입을 사용하여 메모리 m_t를 생성.
메모리 일관성 손실(memory consistency loss)을 통해 m_t와 m'_t를 정렬하여 추론 시 P만 사용 가능.
P와 Q 간 파라미터 공유를 통해 메모리 절감하면서 대부분의 성능 향상 유지.

핵심 아이디어

Maglev는 기존 트랜스포머의 고정 메모리 문제와 반복 모델의 순차성 제약을 동시에 해결하려는 시도이다. 기존 슬라이딩 윈도우 어텐션은 먼 정보를 버리고, LSTMs는 작은 상태만 순차적으로 전파하여 병렬성에 제약이 있다. Maglev는 Q와 P의 두 모델을 통해 이러한 한계를 극복한다. Q는 전체 어텐션과 슬라이딩 윈도우 어텐션을 혼합하여 더 표현력 있는 메모리 m'_t를 생성하고, P는 슬라이딩 윈도우 어텐션과 반복 K/V 주입을 통해 m_t를 생성한다. 이 두 메모리 간의 일관성을 위해 memory consistency loss를 사용하여 P가 Q의 역할을 학습하게 한다. 이는 추론 시 Q를 제거하고 P만 사용할 수 있게 만든다.

기술적 접근법

**Q(prefiller)**: 전체 어텐션과 슬라이딩 윈도우 어텐션을 혼합하여 m'_t 생성.
**P(decoder)**: 슬라이딩 윈도우 어텐션 + 반복 K/V 주입을 사용하여 m_t 생성.
**Memory consistency loss**: m_t와 m'_t를 정렬하는 손실 함수.
**Parameter sharing**: P와 Q 간 파라미터 공유로 메모리 절감.
**Inference**: Q 제거, P만 사용하여 반복 모델로 작동.

주요 결과

**nanochat d20 설정**에서 Maglev는 슬라이딩 윈도우 및 LRT 기반 모델 대비 **검증 BPB 개선**.
**다운스트림 프리트레이닝 벤치마크**에서 슬라이딩 윈도우 및 LRT 기반 모델 대비 성능 향상.
P와 Q 간 파라미터 공유로 **메모리 절감**하면서 대부분의 성능 향상 유지.

의의 및 한계

Maglev는 고정 메모리로 작동하면서도 트랜스포머의 비선형 표현력을 유지하는 반복 모델을 제안하여, 대규모 언어 모델의 추론 효율성을 향상시킬 수 있다. 특히, P와 Q 간 파라미터 공유를 통해 모델 크기를 줄이면서 성능을 유지하는 점이 학술적·실용적 가치를 지닌다. 그러나 본 연구는 계산 자원 제약으로 인해 **초기 탐색 수준**이며, **P와 Q의 파라미터 공유 패턴**이나 **다양한 메모리 주입 방식**에 대한 연구는 여전히 개방적이다. 또한, Q가 P와 함께 학습되는 방식이 아닌, 사전 학습된 모델에서 메모리 타겟을 제공하는 방식도 제안되어 있어, 이에 대한 실험은 아직 수행되지 않았다.

실용적 활용

Maglev는 대규모 언어 모델의 추론 효율성을 향상시키는 데 활용될 수 있으며, 특히 **메모리 제약이 있는 장치**나 **실시간 추론이 필요한 시스템**에서 유용할 수 있다. 또한, **다운스트림 태스크**에서의 프리트레이닝 성능 향상도 기대할 수 있다.