KV 캐시 축출은 정교할 필요가 없을까: Random Attention의 재검토

긴 추론 과정은 대규모 언어 모델이 복잡한 문제를 해결하도록 돕지만, 생성 토큰이 늘어날수록 Key와 Value를 저장하는 KV 캐시가 심각한 메모리 병목이 된다. 기존 KV 캐시 압축 기법은 대체로 같은 방식으로 작동한다. 각 캐시 토큰이 나중에 얼마나 중요할지 추정하고, 점수가 높은 토큰을 남기는 방식이다. Random Attention은 이 중요도 신호가 실제로 큰 가치를 제공하는지에 의문을 제기한다.

제안 방식은 간단하다. 프롬프트는 그대로 보존하고, 그 밖의 캐시 항목은 각 어텐션 헤드 안에서 균등하게 무작위 축출한다. 토큰별 점수 계산이나 정렬은 필요하지 않다. 네 개 모델과 여섯 개 추론 과제에서 평가한 결과, 이 방식은 가장 강력한 기존 축출기와 비슷한 성능을 보였다. vLLM 배포에서는 해당 비교 방식보다 32~43% 높은 처리량도 달성했다.

핵심 내용

• 프롬프트 보존이 우선이다. 통제 실험에 따르면 선택기 사이의 성능 차이 상당 부분은 선택 신호가 프롬프트 관련 정보를 우연히 보존했는지에서 발생한다.
• 추론 텍스트에는 중복이 있다. 모델은 추론을 이어가면서 여전히 필요한 정보를 다시 표현하거나 요약한다. 앞선 토큰이 사라져도 뒤의 텍스트가 대체 정보를 제공할 수 있다.
• 어텐션 헤드 사이에도 중복이 존재한다. 각 헤드는 추론 궤적을 자체적으로 표현한다. 가장 중요한 단 하나의 토큰을 정확히 찾지 못하더라도 관련 정보의 복사본이 충분히 남으면 추론을 이어갈 수 있다.
• 점수 계산 비용을 줄인다. 중요도 추정과 정렬을 제거하면 캐시 압축 과정의 계산량과 구현 복잡도를 낮출 수 있다.

이 결과가 모든 토큰이 동일하게 중요하거나 무작위 축출이 모든 상황에서 안전하다는 뜻은 아니다. 더 정확한 해석은 긴 추론 작업에서 생성된 궤적 자체가 상당한 중복성을 가지며, 정교한 선택기가 추가하는 품질 향상이 그 비용에 미치지 못할 수 있다는 것이다.

추론 서비스 관점에서 이 연구는 캐시 전체를 세밀하게 순위화하기보다, 손실에 취약한 프롬프트를 먼저 보호하고 나머지는 텍스트 반복과 다중 헤드 구조의 중복성에 맡기는 설계를 제안한다. 더 긴 컨텍스트, 다른 유형의 작업, 더 엄격한 품질 조건에서도 같은 결과가 나오는지는 추가 검증이 필요하다. 그럼에도 복잡한 캐시 정책을 도입하기 전에 단순한 무작위 기준선이 어느 정도 성능을 내는지 측정하게 만드는 의미 있는 출발점이다.

출처: Hugging Face Daily Papers