Tensor Cache: Eviction-conditioned Associative Memory for Transformers

기술 요약: 텐서 캐시 (Tensor Cache)

문제 제기
자기회귀형 (Autoregressive) 트랜스포머 추론은 접두사를 재계산하지 않도록 키 - 밸류 (KV) 쌍을 캐싱하는 데 의존합니다. 그러나 유지되는 KV 상태는 컨텍스트 길이, 깊이, 동시성에 따라 선형적으로 증가하여 메모리 병목 현상을 초래합니다. 슬라이딩 윈도우 캐싱은 가장 최근 $W$개의 토큰만 유지함으로써 메모리 사용량을 제한하지만, 이 접근법은 '경직된 망각 (hard forgetting)'으로 고통받습니다. 즉, 토큰이 윈도우에서 퇴출되면 그 정보가 완전히 폐기되어 윈도우 외부의 관련 증거에 접근할 수 없게 됩니다. 기존 유계 상태 (bounded-state) 방법들은 퇴출된 토큰을 완전히 폐기하거나, 모든 토큰을 실행 평균 (running average) 으로 압축하려는 시도를 하는데, 이는 노이즈를 유발하거나 최근과 원거리 컨텍스트를 효과적으로 구분하지 못할 수 있습니다.

방법론: 텐서 캐시 (Tensor Cache, TC)
텐서 캐시는 KV 퇴출을 삭제 사건이 아닌 연관 기억 (associative-memory) 쓰기 연산으로 변환하는 2 단계 유계 상태 아키텍처를 도입합니다.

1. 2 단계 아키텍처:

   • L1 (1 단계 캐시): 가장 최근 $W$개의 토큰을 보유하는 고정 크기의 KV 링 버퍼입니다. 이는 로컬 컨텍스트에 대한 정확한 소프트맥스 어텐션을 지원합니다.
   • L2 (2 단계 캐시): L1 윈도우에서 퇴출된 토큰의 정보만을 저장하는 고정 크기의 외적 행렬 $A$(패스트-가중치 메모리) 입니다.
   • 쓰기 메커니즘: 링 버퍼가 엔트리를 덮어쓸 때, 이동된 $(k, v)$ 쌍은 폐기되지 않습니다. 대신 외적 업데이트를 통해 L2 행렬 $A$에 기록됩니다: $A \leftarrow \lambda A + \eta (k \otimes v)$. 여기서 $\lambda$(감쇠) 와 $\eta$(쓰기 속도) 는 헤드 당 학습된 스칼라입니다.
   • 읽기 메커니즘: 미래의 쿼리 $q_t$는 $r_t = q_t A$를 계산함으로써 압축된 과거에 접근합니다. 선형 어텐션 항등식 $q(k \otimes v) = \langle q, k \rangle v$를 활용하여, 이 연산은 개별적으로 저장하지 않고도 모든 퇴출된 토큰의 정보를 효과적으로 검색합니다.
   • 퓨전: 학습된 스칼라 게이트 $g$가 로컬 어텐션 (L1) 의 출력과 메모리 읽기 (L2) 의 출력을 융합합니다: $y_t = y_{local} + \sigma(g) m_t$.

2. 학습 최적화 (병렬 스캔):
   이 논문은 표준 청크화 (chunked) 학습 단축키의 중요한 문제를 식별합니다. 단순한 접근법은 평균 ($\bar{k} \otimes \bar{v}$) 으로 청크를 요약하는데, 이는 스트리밍 추론 중에는 존재하지 않는 $C^2 - C$개의 가짜 교차 토큰 외적 ($i \neq j$인 $k_i \otimes v_j$) 을 도입합니다.

   • 해결책: 저자들은 순차적 토큰당 쓰기와 수학적으로 동등한 **병렬 가중합 스캔 (parallel weighted-sum scan)**을 구현합니다. 이는 float32 오차 범위 내 (상대 오차 $< 10^{-7}$) 에서 학습과 추론 간의 격차를 해소하여, 모델이 교차 토큰 간섭의 노이즈 없이 올바른 연관 역학을 학습하도록 보장합니다.

주요 기여

• 퇴출 조건부 연관 기억: 모든 토큰을 섭취하는 기존 외적 메모리 (예: Infini-attention, mLSTM) 와 달리, TC 는 외적 행렬을 슬라이딩 윈도우 퇴출에 의해서만 공급되는 L2 캐시로 엄격하게 사용합니다. 이는 토큰을 'L1 에 거주'와 'L2 에 압축'으로 깔끔하게 분할합니다.
• 학습 아티팩트 수정: 이 논문은 가짜 교차 토큰 제품을 도입하는 이전에 간과되었던 학습 단축키 (청크화 평균) 를 식별하고 해결합니다. 제안된 병렬 가중합 스캔은 동일한 계산 비용으로 이 불일치를 제거합니다.
• 유계 상태 효율성: TC 는 개별 퇴출 엔트리를 유지하지 않으면서도 이전 컨텍스트가 예측에 영향을 미칠 수 있는 메커니즘을 제공하여, 전체 KV 유지에 대한 실용적인 대안을 제시합니다.

실험 결과
저자들은 시스템 확장, 합성 연관 회상, 그리고 실제 텍스트 장문 컨텍스트 언어 모델링 (OpenWebText 및 셰익스피어) 을 통해 TC 를 평가했습니다.

• 합성 연관 회상: 슬라이딩 윈도우 너머의 정보 검색이 필요한 작업에서 TC 는 100% 정확도를 달성하여 전체 KV 성능과 일치했습니다. 전체 KV 대비 유지되는 추론 상태를 72–84% 줄였으며, 간격이 증가함에 따라 성능이 저하된 Infini-attention 과 우연 수준으로 수행된 슬라이딩 윈도우 베이스라인을 크게 능가했습니다.
• 장문 컨텍스트 언어 모델링: OpenWebText 에서 TC 는 1,024 에서 32,768 토큰까지의 모든 평가 컨텍스트 길이에서 **최저 평균 음의 로그 가능도 (NLL)**를 달성했습니다. 단, $L=2,048$에서는 Infini-attention 과 통계적으로 구별되지 않는 유일한 예외가 있었습니다. $L=32,768$에서 TC 는 NLL 5.14 를 기록하여 전체 KV(6.00) 와 Infini-attention(5.42) 을 능가했으며, 전체 KV 대비 피크 GPU 메모리를 2.4$\times$ 적게 사용했습니다.
• 처리량: TC 는 순수 슬라이딩 윈도우 방법에 비해 L2 읽기/업데이트로 인해 약 30% 의 처리량 오버헤드가 발생했지만, 테스트된 모든 컨텍스트 길이에서 Infini-attention 보다 일관되게 빨랐습니다.
• 용량 진단: 원시 연관 메모리 용량에 대한 실험은 더 많은 관련 없는 연관성이 고정된 상태로 압축됨에 따라 검색 성능이 저하됨을 보여주어, 메모리의 유계 특성을 확인했습니다.

의의 및 주장
이 논문은 텐서 캐시가 슬라이딩 윈도우 어텐션의 메모리 효율성과 전체 컨텍스트 어텐션의 회상 능력 사이의 격차를 성공적으로 연결한다고 주장합니다. KV 퇴출을 파괴적인 삭제가 아닌 구조화된 메모리 연산으로 취급함으로써, TC 는 전체 KV 캐시의 선형 메모리 비용 없이도 트랜스포머가 즉각적인 윈도우 외부의 관련 증거에 접근할 수 있게 합니다. 저자들은 주요 기여가 외적 메모리 자체의 발명이 아니라, 퇴출 조건부 L2 캐시로서의 구체적인 적용과 이전에 이러한 방법을 방해했던 학습 아티팩트의 엄격한 수정이라고 강조합니다. TC 는 장문 컨텍스트 추론을 위한 메모리 - 품질 프런티어를 개선하는 실용적인 유계 상태 대안으로 제시됩니다.