Context Memorization for Efficient Long Context Generation

기술 요약: 효율적인 장문맥 생성을 위한 어텐션-상태 메모리

문제 제기
현대 대규모 언어 모델 (LLM) 응용 프로그램은 추론을 안내하기 위해 점점 더 긴 조건부 접두사 (예: 인-컨텍스트 예시, 검색된 문서, 또는 에이전트 지시사항) 에 의존하고 있습니다. 그러나 접두사 증강 추론은 두 가지 주요 구조적 한계에 직면해 있습니다:

1. 접두사 감쇠 (Prefix Decay): 생성이 진행됨에 따라 모델의 어텐션이 새로 생성된 토큰에 분포하게 되어, 특히 장문맥 시나리오에서 접두사의 영향력이 약화됩니다.
2. 추론 비효율성: 접두사에 대한 어텐션 계산은 그 길이에 비례하여 선형적으로 증가 ($O(L)$) 하므로, 프리필 (prefill) 단계와 모든 디코드 단계에서 상당한 지연 시간과 메모리 오버헤드를 초래합니다. 접두사 캐싱은 프리필 비용을 상쇄하지만 여전히 상당한 메모리 소모를 수반합니다.

기존 해결책은 두 가지 범주로 나뉘며, 둘 다 단점이 있습니다:

• 접두사 내부화 (Prefix Internalization): 모델 파라미터를 미세 조정하거나 하이퍼네트워크를 사용하여 접두사 행동을 내부화하는 방법들입니다. 이러한 방법들은 경사 기반 훈련을 필요로 하므로 자원 집약적이며 업데이트가 느리고 동적인 접두사 변경에는 적합하지 않습니다.
• 접두사 압축 (Prefix Compression): 접두사를 어텐션 메커니즘에 유지하되 압축하는 방법들 (예: 토큰 가지치기 또는 KV 캐시 내보내기) 입니다. 이러한 방법들은 여전히 추론 시 압축된 접두사에 대한 어텐션이 필요하므로, 지연 시간은 압축된 길이에 비례하여 증가하며 접두사의 영향력은 감쇠에 노출된 상태로 남습니다.

방법론: 어텐션-상태 메모리 (ASM)
저자들은 **어텐션-상태 메모리 (Attention-State Memory)**를 제안합니다. 이는 접두사를 경량화된 사전 계산된 어텐션 상태의 조회 기반 메모리로 외부화하는 훈련 없는 접근법입니다. 이 방법은 연결된 키 블록에 대한 어텐션을 서브-블록에 대한 어텐션으로 분해하기 위해 **온라인-소프트맥스 항등식 (online-softmax identity)**을 활용합니다.

핵심 통찰력은 주어진 쿼리에 대해 어텐션 출력과 정규화 항으로 구성된 어텐션 상태 $(a, Z)$가 원래 키와 값에 다시 어텐션하지 않고도 접두사의 어텐션 기여도를 재구성하기에 충분하다는 것입니다.

이 프레임워크는 두 단계로 작동합니다:

1. 오프라인 보정 단계 (메모리 구축):

   • 수집: 모델은 접두사와 연결된 일련의 대표적 쿼리 - 응답 트레이스에 대한 순전파를 수행합니다. 응답 내의 모든 쿼리 토큰에 대해 시스템은 쿼리 벡터와 접두사에 대한 해당 어텐션 상태 $(a, Z)$를 기록합니다.
   • 클러스터링: 수집된 쿼리 벡터는 (예: K-평균을 통해) $K$개의 중심점으로 클러스터링됩니다.
   • 집계: 각 클러스터에 대해 단일 대표 항목이 생성됩니다. 중심점 쿼리는 클러스터 내 쿼리들의 평균입니다. 어텐션 상태 $(\bar{a}, \bar{Z})$는 **어텐션 인식 집계 (attention-aware aggregation)**를 통해 계산되며, 이는 온라인-소프트맥스 업데이트 규칙을 사용하여 클러스터 내 모든 구성원의 상태를 병합합니다. 이는 중심점이 클러스터의 접두사 어텐션에 대한 손실 없는 평균을 대표하도록 보장합니다.
   • 효율성: GPU 메모리를 초과하지 않으면서 매우 긴 접두사 (예: 16K+ 토큰) 를 처리하기 위해 접두사를 청크로 분할하고 독립적으로 인코딩한 후, 온라인-소프트맥스 항등식의 구성 가능성 속성을 사용하여 메모리 내에서 병합할 수 있습니다.

2. 온라인 추론 단계:

   • 검색: 들어오는 쿼리에 대해 시스템은 코사인 유사도 (선택적으로 화이트닝 변환 및 RoPE 처리 포함) 를 사용하여 가장 가까운 메모리 중심점을 식별합니다.
   • 병합: 검색된 사전 계산된 어텐션 상태 $(\bar{a}, \bar{Z})$는 온라인-소프트맥스 항등식에 의해 정의된 병합 연산자를 사용하여 쿼리의 자기 어텐션 출력 (접두사 토큰이 아닌 것들에 대해 계산됨) 과 병합됩니다. 이는 접두사 토큰에 어텐션하지 않고도 접두사가 존재했을 때 생성되었을 정확한 어텐션 출력을 복원합니다.
   • 복잡도: 계층적 인덱싱을 통한 메모리 크기에 대한 조회 비용은 로그적으로 ($O(\log K)$) 증가하며, 원래 접두사 길이와 무관합니다.

주요 기여

1. 훈련 없는 외부화: 이 논문은 경사 기반 훈련이나 보조 모델 없이 긴 접두사를 컴팩트한 메모리 사전으로 외부화하는 방법을 제안합니다.
2. 온라인-소프트맥스 확장: 저자들은 온라인-소프트맥스 항등식을 효율적인 어텐션 계산에서 교차 쿼리 접두사 재사용 영역으로 확장하여, 검색을 통한 접두사 어텐션의 손실 없는 재구성을 가능하게 합니다.
3. 성능 및 효율성: 실험은 ASM 이 전체 어텐션 성능을 일치시키거나 초과하면서 계산 비용을 크게 줄임을 보여줍니다.

실험 결과
이 방법은 ManyICLBench(인-컨텍스트 학습) 및 RuleArena/NBA(검색 증강 생성) 에서 LLaMA-3.1-8B를 사용하여 평가되었습니다.

• 인-컨텍스트 학습 (ICL): ManyICLBench 에서 ASM 은 1K 에서 8K 항목에 이르는 메모리 예산 전반에 걸쳐 표준 인-컨텍스트 학습보다 정확도를 향상시켰습니다. 8K 메모리 예산에서 전체 어텐션 대비 어텐션 지연 시간을 1.36 배 단축했습니다. ASM 은 종종 레이블 토큰과 같은 중요한 정보를 붕괴시키는 KV 캐시 압축 베이스라인 (KVzip) 보다 우수한 성능을 발휘했습니다.
• 검색 증강 생성 (RAG): NBA 벤치마크 (약 20K 토큰의 규칙 포함) 에서 ASM 은 메모리 발자국의 20% 만 사용하여 전체 어텐션 RAG 성능을 능가했습니다. 최적의 성능은 중간 코드북 크기 (4K 항목) 에서 관찰되었으며, 이는 항목 수가 단순히 극대화되어야 하는 것이 아니라 작업별 하이퍼파라미터임을 시사합니다.
• 지연 시간: 계층적 조회는 매우 작은 메모리 예산에서 더 높은 오버헤드를 발생시키지만, 지연 시간 증가는 선형 미만입니다. ASM 은 약 4K 메모리 항목에서 전체 어텐션보다 빨라지며 16K 항목에서 1.8 배의 속도 향상을 달성합니다.

의의 및 주장
이 논문은 어텐션-상태 메모리가 접두사 지식을 쿼리별 어텐션 계산과 분리함으로써 장문맥 생성의 구조적 비용을 해결한다고 주장합니다. 접두사를 사전 계산된 조회 기반 메모리로 외부화함으로써, 이 방법은 접두사 길이에 따른 추론 지연 시간의 선형적 증가를 제거하고 표준 어텐션 메커니즘과 관련된 접두사 감쇠 문제를 없앱니다.

저자들은 이 작업을 접두사 프롬프트와 모델 파라미터를 넘어 존재하는 컴팩트하고 재사용 가능한 표현으로 LLM 지식을 외부화하는 더 넓은 연구 방향을 향한 한 걸음으로 위치시킵니다. 그들은 이 접근법이 정확성과 효율성 간의 유리한 균형을 제공하며, 파라미터 내부화에 대한 훈련 없는 대안이자 접두사 압축보다 더 효율적인 대안을 제공한다고 강조합니다.