기술 요약: 연상 재귀 메모리를 통한 LLM 컨텍스트 확장

문제 정의

대규모 언어 모델(LLM)은 기술 보고서 분석, 소프트웨어 개발, 다중 문서 추론과 같은 작업을 위해 수십만 또는 수백만 토큰에 달하는 입력을 처리해야 하는 요구가 점점 커지고 있습니다. 그러나 표준 트랜스포머 아키텍처는 근본적인 병목 현상에 직면해 있습니다: 셀프 어텐션의 계산 및 메모리 비용이 시퀀스 길이($O(N^2)$)에 따라 제곱으로 증가합니다. 또한, 모델의 명목상 윈도우 내에 있더라도 컨텍스트 길이가 길어짐에 따라 성능이 저하되는 경우가 빈번합니다. 반면, 재귀적 아키텍처(예: Mamba, RWKV)는 선형 스케일링을 제공하지만, 일반적으로 처음부터 학습시켜야 하므로 기존의 사전 학습된 LLM을 활용할 수 없으며, 트랜스포머와 비교했을 때 복잡한 알고리즘 작업이나 지시 이행 능력에서 어려움을 겪는 경우가 많습니다.

방법론

저자들은 긴 컨텍스트 길이를 확장하면서도 일정한 메모리 스케일링과 효율성을 유지하기 위한 실질적인 해결책으로 **연상 재귀 메모리 트랜스포머(Associative Recurrent Memory Transformer, ARMT)**를 제안합니다. ARMT는 사전 학습된 베이스 LLM을 감싸는 래퍼(wrapper) 역할을 하며, 입력의 세그먼트별 처리를 가능하게 합니다.

핵심 아키텍처

ARMT는 긴 컨텍스트 입력을 고정된 길이의 겹치지 않는 세그먼트로 나눕니다. 각 세그먼트 내에서 모델은 풀 셀프 어텐션(단기/작업 메모리)을 활용합니다. 결정적으로, 세그먼트 간에 정보를 전파하는 계층별 연상 메모리 모듈(장기 메모리)을 도입합니다. 이 메커니즘은 세 단계로 작동합니다:

1. 메모리 추출(Memory Extraction): 각 트랜스포머 레이어는 입력 세그먼트를 메모리 임베딩으로 압축합니다.
2. 메모리 통합(Memory Consolidation): 이러한 임베딩은 키-값(key-value) 쌍으로서 레이어별 연상 행렬로 통합됩니다.
3. 연상(Association): 후속 세그먼트의 임베딩은 쿼리 벡터로 변환되며, 연상 행렬과 곱해져 이전 세그먼트로부터 관련 정보를 검색합니다.

학습 레시피

본 논문은 사전 학습된 LLM을 ARMT로 적응시키기 위한 종합적인 학습 전략을 설명합니다:

• 지속적 사전 학습(Continued Pre-training): 초기화되지 않은 연상 메모리 파라미터는 효과적인 메모리 전파를 학습하기 위해 긴 컨텍스트(예: FineWeb-Edu의 19B 토큰)에 대한 비지도 언어 모델링을 통해 초기화됩니다.
• 커리큘럼 학습(Curriculum Learning): 장거리 의존성을 처음부터 학습하는 어려움을 해결하기 위해, 모델은 세그먼트 수를 점진적으로 늘려가며(예: 2개에서 4개, 8개로) 학습률을 조정하며 미세 조정됩니다.
• 합성 데이터 생성(Synthetic Data Generation): 긴 컨텍스트 시나리오에서의 데이터 부족 문제를 극복하기 위해, 저자들은 긴 문서의 짧은 구절들을 연결하고 각 구절에 대한 QA 쌍을 생성함으로써 합성 학습 인스턴스를 생성합니다.
• 레이어 프루닝 및 선택(Layer Pruning and Selection): 연구는 연상 메모리가 모든 레이어에 필요한지 조사합니다. 특정 중간 및 최종 레이어(예: 특정 중간 및 최종 레이어)에만 연상 메모리를 유지하는 전략을 제안하여, 성능 저하 없이 학습 가능한 파라미터와 계산 비용을 줄입니다.

주요 기여

1. 도메인 특화 데이터셋: 코드 내 가변 타입 예측을 위한 **ManyTypes-long (MT)**과 긴 문서 질의응답을 위한 **GovReport-long (GR)**이라는 두 가지 새로운 데이터셋을 구축하여, 현실적이고 좁은 도메인의 워크로드를 평가하도록 설계했습니다.
2. 학습 레시피: 지속적 사전 학습, 합성 데이터 생성, 커리큘럼 학습, 선택적 레이어 통합을 결합하여 LLM 컨텍스트를 확장하는 새로운 ARMT 프레임워크를 제시합니다.
3. 경험적 검증: ARMT로 증강된 모델이 다음을 입증하는 광범위한 실험적 연구를 수행했습니다:
   • 원래의 컨텍스트 제한을 훨씬 초과하는 입력(최대 64k 토큰)을 처리하면서도 인리미트(in-limit) 베이스라인 대비 성능 저하 없이 처리합니다.
   • 베이스 모델에 비해 분포 외(OOD) 컨텍스트 길이에 대해 우수한 일반화 능력을 보입니다.
   • 원래 컨텍스트 윈도우 내에서의 베이스라인 성능을 유지하면서도 약 30% 적은 FLOPs를 요구합니다.

실험 결과

저자들은 Gemma-3-1B-IT와 SmolLM-2-360M-IT 백본을 사용하여 MT 및 GR 데이터셋에 대해 ARMT를 평가했습니다.

• 성능: ARMT 모델은 MT 및 GR 데이터셋에서 최대 65k 토큰의 컨텍스트 길이까지 안정적인 성능을 유지했습니다. 반면, 베이스 모델(미세 조정된 경우 포함)은 고유 컨텍스트 윈도우(예: 8k 또는 32k)를 넘어서면 급격한 성능 저하를 보였습니다. ARMT는 Long-OOD 영역(32k–65k)에서 베이스 모델보다 유의미하게 우수한 성능을 보였습니다.
• 효율성: ARMT는 컨텍스트 길이에 관계없이 일정한 GPU 메모리 사용량을 보여준 반면, 베이스 모델의 메모리 사용량은 선형적으로 증가했습니다. 32k 토큰 시퀀스의 경우, ARMT는 동일한 메모리 예산 하에서 베이스 모델(8)보다 4배 높은 배치 크기(32)를 허용했습니다.
• FLOPs 감소: 이론적 분석과 경험적 추론 시간 측정을 통해 전체 어텐션 FLOPs가 $T/S$(전체 시퀀스 길이를 세그먼트 크기로 나눈 값) 비율만큼 감소함을 확인했으며, 이는 풀 어텐션 모델 대비 총 FLOPs의 약 30% 감소로 이어졌습니다.
• 절제 연구(Ablation Studies):
  • 레이어 프루닝: 약 20%의 레이어(특히 사전에 선택된 중간 및 최종 레이어)에만 연상 메모리를 배치한 모델이 전체 ARMT 모델과 대등하거나 더 나은 성능을 달-성했습니다.
  • 사전 학습: 연상 메모리를 초기화하는 데 있어 지속적 사전 학습이 필수적임을 발견했으며, 이는 인도메인 및 OOD 성능을 모두 크게 향상시켰습니다.
  • 베이스라인: ARMT는 Mamba-2, DeltaNet, xLSTM을 포함한 다른 롱 컨텍스트 베이스라인보다 특히 긴 컨텍스트 일반화 측면에서 뛰어난 성능을 보였으며, 처음부터 학습된 모델들보다 덜 광범위한 사전 학습을 필요로 했습니다.

의의 및 주장

본 논문은 ARMT를 소형 및 중형 규모(최대 1B 파라미터)의 LLM에서 긴 컨텍스트 처리를 가능하게 하는 실용적이고 계산 효율적인 접근 방식으로 규정합니다. 저자들은 이 방식이 트랜스포머의 강력한 단기 컨텍스트 성능과 재귀 모델의 선형 스케일링 사이의 간극을 메운다고 주장합니다.

주요 의의에 대한 클레임은 다음과 같습니다:

• 프라이버시 및 로컬 배포: 긴 컨텍스트 처리를 효율적으로 가능하게 함으로써, 원격의 대규모 API 기반 LLM에 의존하지 않는 프라이버시 보호 애플리케이션을 용이하게 합니다.
• 확장성: 일정한 메모리로 임의의 길이의 컨텍스트를 처리할 수 있게 하며, 이는 실제 문서 분석에 있어 중요한 요구사항입니다.
• 효율성: 일정한 메모리 스케일링과 감소된 FLOPs의 결면은 ARMT를 자원이 제한된 환경에서 실행 가능한 대안으로 만듭니다.
• 일반화: 제안된 학습 레시피는 표준 트랜스포머가 긴 시퀀스를 다룰 때 나타나는 "중간에서의 손실(lost in the middle)" 현상과 성능 저하 문제를 효과적으로 해결합니다.

저자들은 실험이 1B 파라미터 이하의 모델과 특정 작업 세트(코드 및 문서 QA)로 제한되었다는 점을 언급하며 한계를 인정합니다. 또한 연상 메모리와 트랜스포머 표현 간의 상호작용에 대한 근본적인 메커니즘이 부분적으로 이해되지 않은 상태라는 점도 밝힙니다. 그러나 그들은 소형 모델에서 컨텍스트를 확장하는 입증된 능력은 실질적인 로컬 롱 컨텍스트 AI를 향한 중요한 진전이라고 주장합니다.