Transactional Attention: Semantic Sponsorship for KV-Cache Retention

논문 개요

본 논문은 대규모 언어 모델 (LLM) 의 긴 컨텍스트 처리 시 발생하는 '잠재 토큰 (Dormant Tokens)' 문제를 해결하기 위해 거래적 주의 (Transactional Attention, TA) 라는 새로운 메커니즘을 제안합니다. 기존 KV 캐시 압축 방법론들이 주의 (Attention) 점수나 재구성 손실에 의존할 때, API 키, 비밀번호 등 통계적 신호는 거의 없으나 생성 시점에 필수적인 정보를 유실하는 한계를 극복하고, 구조적 패턴을 기반으로 이러한 토큰을 보호하는 '스폰서십 (Sponsorship)' 방식을 도입했습니다.

---

1. 문제 정의: 잠재 토큰 (Dormant Tokens) 의 실패

• 배경: 멀티턴 대화나 에이전트 워크플로우에서 모델은 초기에 받은 API 키, 비밀번호, 설정 값 등을 수천 토큰 뒤에 다시 참조해야 하는 경우가 많습니다.
• 문제점: 이러한 정보 (잠재 토큰) 는 입력 단계에서 거의 0 에 가까운 주의 (Attention) 점수를 받습니다.
• 기존 방법의 한계:
  • H2O, TOVA, SnapKV 등: 누적 주의 점수, 슬라이딩 윈도우, 재구성 손실 (Reconstruction Loss) 등을 기반으로 토큰을 평가합니다.
  • 결과: 잠재 토큰은 통계적으로 중요도가 낮아 보이기 때문에 KV 캐시 압축 과정에서 반드시 삭제 (Eviction) 됩니다.
  • 실험 결과: 4,000 토큰 컨텍스트에서 KV 캐시 예산 (K) 이 16 토큰 (전체의 0.4%) 일 때, 기존 6 가지 방법론 (H2O, TOVA, SnapKV 등) 은 **신호 (Credential) 검색 정확도가 0%**였습니다.

2. 제안 방법: 거래적 주의 (Transactional Attention, TA)

TA 는 통계적 신호 대신 구조적 패턴 (Structural Patterns) 을 활용하여 토큰을 보호하는 스폰서십 메커니즘을 도입합니다.

핵심 메커니즘

1. 앵커 감지 (Anchor Detection):
   • key:, password:, Authorization: 과 같이 필수 정보를 예고하는 구조적 패턴 (앵커) 을 탐지합니다.
   • 구현 방식: 정규식 (Regex), 의미적 동사 ("remember", "said" 등), 임베딩 기반 유사도 (언어 무관) 등.
2. 스폰서십 배정 (Sponsorship Allocation):
   • 앵커가 탐지되면, 그 뒤에 이어지는 값 (Value-bearing spans, 예: API 키 문자열) 을 보호하기 위해 보호 신호 (Voucher) 를 주입합니다.
   • 지수 감쇠 (Exponential Decay): 앵커 바로 다음 토큰부터 $L$개의 토큰까지 보호 예산을 할당하며, 거리가 멀어질수록 보호 강도가 감쇠합니다 ($V_j = B \cdot 0.8^{j-i}$).
3. 유틸리티 점수 재조정:
   • 기존 주의 점수 ($A_i$) 에 스폰서십 보너스 ($V_i$) 를 더한 새로운 유틸리티 점수를 계산합니다.
   • 공식: $u_i = \alpha A_i + \beta R_i + \gamma S_i + \delta P_i - \lambda F_i + V_i$
   • 이를 통해 통계적으로 무시당하던 토큰도 Top-K 선택 기준을 통과하여 KV 캐시에 유지됩니다.

TA-Fast (경량화 버전)

• 주의 (Attention) 계산 자체를 제거한 변형입니다.
• SDPA 및 FlashAttention과 호환되어 메모리 오버헤드를 52% 감소시킵니다.
• 앵커 감지와 보너스 배정만 수행하므로 지연 시간 (Latency) 증가가 1% 미만입니다.

---

3. 주요 기여 (Contributions)

1. 100% 검색 성공률: KV 캐시 예산이 16 토큰 (K=16) 일 때, 기존 6 가지 베이스라인이 0% 였던 반면 TA 는 **100%**의 자격증명 검색 정확도를 달성했습니다.
2. TA-Fast: 메모리 효율성을 극대화하고 SDPA/FlashAttention과 호환되는 주의 없는 변형을 제안했습니다.
3. 학습된 앵커 탐지기: F1 점수 0.946 을 기록하며, 재학습 없이 다양한 모델 규모 (Llama, Mistral 등) 로 전이 학습이 가능합니다.
4. 직교성 (Orthogonality): TA 는 기존 압축 방법론을 대체하지 않고 점수 함수에 추가 신호를 주입하므로, 기존 방법론과 결합 가능하며 일반적인 긴 컨텍스트 성능을 저하시키지 않습니다.

---

4. 실험 결과

• Needle-in-a-Haystack (Haystack 내 바늘 찾기):
  • 4K 컨텍스트, K=16 조건에서 TA 는 100% 정확도, 기존 방법론들은 0% 정확도.
  • K=128 이상이어야 기존 방법론들이 100% 에 도달했으나, TA 는 모든 예산에서 100% 유지.
• 기능 호출 에이전트 (Function-Calling Agents):
  • 200 번의 멀티턴 기능 호출 테스트에서 TA-Fast-Semantic 은 100% 성공, TOVA 는 0% 성공.
• 모델 확장성:
  • Llama-3.2-1B, Llama-3.1-8B 에서 K=16 시 100% 성공.
  • Mistral-7B 는 토크나이저 특성상 토큰 수가 더 길어 K=24 에서 100% 성공 (K=16 은 5%).
• 메모리 및 성능:
  • 7K 컨텍스트에서 TOVA 는 OOM(메모리 부족) 발생 (>50GB), TA-Fast 는 6.8GB 로 100% 성공.
  • 지연 시간 증가: 0.358s $\to$ 0.361s (약 1% 미만).

---

5. 의의 및 결론

• 패러다임 전환: KV 캐시 압축이 단순히 "통계적으로 중요한 토큰"을 찾는 것을 넘어, **"구조적으로 필수적인 토큰"**을 사전에 보호하는 선제적 (Proactive) 접근법을 제시했습니다.
• 실무 적용성: 에이전트 기반 LLM 이 도구를 호출하고 상태를 유지해야 하는 실제 환경에서, API 키나 세션 토큰 같은 중요한 정보가 유실되는 치명적인 오류를 방지합니다.
• 한계: 앵커 패턴이 명확하지 않은 비구조화된 문맥에서는 정확도가 떨어질 수 있으며, 악의적인 앵커 주입 (Adversarial Injection) 으로 스폰서 예산을 고갈시킬 수 있는 위험이 존재합니다 (Allowlist 필터링으로 대응 가능).

요약하자면, Transactional Attention 은 LLM 이 긴 문맥에서 '잊혀진' 중요한 정보를 잃지 않도록, 구조적 패턴을 기반으로 KV 캐시를 지능적으로 보호하는 혁신적인 솔루션입니다.