Associative Recurrent Memory Transformer

https://arxiv.org/abs/2407.04841

Associative Recurrent Memory Transformer

 

 

ICML 2024 Next Generation of Sequence Modeling Architectures Workshop 제출 논문인데 숏이네요 

짧습니다.

 

긴 입력 전체를 attention으로 보지 말고 입력을 segment 단위로 처리하며 각 layer 마다 associative key-value memory를 축적하는 구조를 제안해 16k 토큰 만으로 50M 토큰까지 QA를 수행할 수 있음을 보여줌 

접근	장점	한계
Long-context Transformer 확장	기존 Transformer 성능 유지	attention cost, context window 한계
RMT 계열 recurrent memory	segment 단위 처리 가능	memory token 수가 작아 저장 용량 제한
Mamba/RWKV/SSM	긴 시퀀스 효율적 처리	key-value recall, copying, 사후 질문형 memory task에 약할 수 있음
RAG	외부 검색으로 긴 문맥 우회	여러 근거를 조합해야 하는 reasoning task에서 실패 가능

기존 RMT는 Segmen마다 memory token로 넘김

ARMTsms 각 layer에서 memory token을 단순 전달하는 것이 아닌 이를 key-value association martix에 저장함

현재 segment 토큰들이 위 matrix에 query를 날려 과거 segment에서 축적된 정보를 가져와 입력에 활용 

각 memory token을 key-value pair로 바꾼 뒤 기존 memory matrix에 같은 key에 저장된 old value를 지우고, 새로운 value를 다시 쓰는 과정 

ARMT가 RMT보다 훨씬 많은 key-value pair를 저장할 수 있었음 

RMT와 Mamba는 학습 길이를 넘어가면 점진적으로 성능이 하락함 

ARMT 는 긴 문맥을 버티는게 아닌 key-value 형태로 저장, 갱신, 검색에 강하다.

문제의식	Transformer는 긴 입력을 처리할 때 self-attention 비용이 커지고, RMT류 recurrent memory는 segment-level 처리는 가능하지만 memory token 수가 제한되어 장기 정보 저장 용량에 한계가 있음. Mamba/RWKV 같은 recurrent sequence model도 효율적이지만 key-value recall, 복사, 과거 정보 검색형 task에서 약점이 있을 수 있음.
핵심 목표	새로운 segment를 처리할 때 시간·공간 복잡도를 일정하게 유지하면서, 수백만~수천만 토큰 규모의 long-context에서 필요한 정보를 저장하고 회수할 수 있는 구조를 만드는 것.
제안 방법	ARMT는 RMT에 layer-wise associative memory를 추가한 구조.  현재 segment는 Transformer self-attention으로 처리하고, 과거 segment의 정보는 각 layer의 associative memory matrix에 key-value 형태로 저장함.
기존 RMT와 차이	RMT는 memory token을 다음 segment로 그대로 넘기는 방식에 가깝지만, ARMT는 memory token을 (k_i, v_i)로 변환해 associative matrix A_s^l에 저장함. 즉, memory token을 단순 전달하지 않고 검색 가능한 key-value memory로 구조화함.
Associative Block 역할	각 memory token m_i에서 key k_i, value v_i, 저장 강도 β_i를 만들고, 이전 memory에서 같은 key의 old value \bar{v}_i를 읽은 뒤, v_i-\bar{v}_i를 통해 기존 값을 새 값으로 갱신함.
중요한 수식 직관	\bar{v}_i는 이전 memory에 저장된 old value이고, v_i-\bar{v}_i는 old value를 지우고 new value로 바꾸기 위한 delta임. β_i는 저장 강도, γ_i는 normalization vector z_s^l가 같은 key를 과도하게 중복 누적하지 않도록 막는 correction term임.
핵심 기술적 포인트	ARMT는 단순히 정보를 누적하는 것이 아니라 같은 key가 다시 등장하면 최신 value로 rewrite할 수 있음. 이 때문에 긴 sequence에서 entity 상태, 위치, 속성처럼 시간에 따라 바뀌는 정보를 추적하는 데 유리함.
실험 1: Associative Retrieval	Remember task와 Rewrite task를 사용. Remember는 unique key-value pair를 기억하는 능력을 평가하고, Rewrite는 같은 key가 여러 번 등장할 때 최신 value를 기억하는 능력을 평가함.
Associative Retrieval 결과	ARMT는 RMT보다 훨씬 많은 key-value pair를 저장했고, PRMT ablation이 큰 개선을 보이지 않아 성능 향상의 핵심이 단순 layer-wise memory가 아니라 associative memory matrix임을 보임. Rewrite task에서는 50 pair로 학습했지만 500 update에서도 정확한 recall을 유지해 약 10배 길이 일반화를 보임.
실험 2: BABILong	BABILong은 긴 context 안의 fact와 distractor sentence 중 필요한 정보를 찾아 QA를 수행하는 benchmark. QA1은 single supporting fact, QA2~QA5는 여러 supporting facts 또는 relation reasoning을 요구함.
BABILong 주요 결과	GPT-2 137M 기반에 ARMT를 적용한 145M 모델이 16K tokens로 학습했음에도 QA1에서 50M tokens까지 평가됨. best model 기준 50M tokens에서 79.9% accuracy를 기록했고, QA2~QA5에서도 10M tokens까지 강한 성능을 보임.
비교 모델 대비 성능	ARMT는 RMT, Mamba, GPT-4 few-shot, GPT-4+RAG 대비 BABILong의 장문 QA에서 대체로 우수한 성능을 보임. 특히 500K~10M token 구간에서 다섯 개 QA task 전반에 걸쳐 가장 안정적인 성능을 보인 구조로 제시됨.
주요 기여	1) RMT에 associative memory를 결합한 ARMT 제안, 2) key-value memory capacity 향상, 3) memory rewrite operation에 강한 구조 제시, 4) 16K 학습 후 최대 50M token까지 extrapolation, 5) BABILong에서 long-context QA 성능 기록 제시.
한계점	Segment를 순차 처리해야 하므로 병렬화가 제한적임. 300K 이하 short/medium context에서는 Mamba/RWKV보다 느릴 수 있음. 또한 Wikitext-103 language modeling 실험에서는 ARMT가 RMT와 유사한 수준에 머물러, 일반 LM 성능 개선 구조로는 아직 충분히 검증되지 않음. 실험도 137M~145M 규모라 대형 LLM scaling 검증이 필요함.
연구적 의미	이 논문은 long-context 문제를 단순히 context window 확장으로 해결하려는 것이 아니라, task-relevant 정보를 recurrent associative memory에 저장하고 필요할 때 회수하는 방향을 제안함. 즉, “긴 attention”보다 “구조화된 장기 memory”가 중요하다는 관점을 보여줌.
최종 평가	ARMT는 일반 언어모델링 대체재라기보다, 초장문 context에서 sparse fact를 저장·갱신·회수하는 memory-augmented Transformer 구조로 보는 것이 타당함. 장기 기억, entity tracking, long-context QA, multi-segment evidence retrieval 연구에 중요한 참고점이 있음.