🤖 machine learning

RKSC: Reasoning-Aware KV Cache Sharing and Confident Early Exit for Multi-Step LLM Inference

RKSC는 어텐션 유사도 기반의 KV 캐시 공유, 신뢰도 게이트형 조기 종료, 그리고 선택적 블록 캐시 관리자를 통해 구조적 중복을 제거함으로써 다단계 LLM 추론을 가속화하는 훈련 불필요 추론 프레임워크로, 다양한 모델과 벤치마크 전반에서 무시할 만한 수준의 오차율을 유지하면서 평균 3배의 속도 향상을 달성한다.

원저자: Anirudh Sekar

게시일 2026-06-10

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Anirudh Sekar

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 복잡한 미스터리를 해결하려는 아주 명석한 탐정이라고 상상해 보십시오. 단순히 하나의 답을 추측하는 대신, 당신은 여덟 명의 서로 다른 탐정(가지/branches)에게 동시에 자신만의 이론을 작성하도록 지시하기로 결정했습니다. 이것이 바로 현대의 "추론(reasoning)" AI 모델이 작동하는 방식입니다. 즉, 최선의 답을 찾기 위해 가능한 여러 경로를 생성하는 것입니다.

이 논문은 이 탐정 팀을 위한 새로운 "매니저"인 RKSC를 소개합니다. 이 매니저의 목표는 단순합니다. 탐정들이 똑같은 일을 반복하며 시간을 낭비하지 않게 하고, 이미 답을 알고 있다면 사건 파일 전체를 읽지 않도록 하는 것입니다.

RKSC가 어떻게 작동하는지 세 가지 간단한 기술로 나누어 설명하겠습니다.

1. "공유 노트" 기술 (ASKS)

문제점:
보통 여덟 명의 탐정이 작업을 시작하면, 그들은 모두 사건 파일의 첫 1,000페이지(접두사/prefix)를 각자 독립적으로 읽습니다. 설령 그들이 완전히 동일한 텍스트를 읽고 있더라도, 컴퓨터는 그 단어들의 의미를 여덟 번이나 따로 계산합니다. 이는 마치 도서관에 있는 여덟 사람이 책의 첫 장을 각자 손으로 베껴 쓰는 것과 같습니다. 공유할 수 있는 복사본이 있는데도 말이죠.

RKSC의 해결책:
RKSC는 **공유 노트(Shared Notebook)**를 도입합니다.

시스템은 사건 파일의 공통 부분을 단 한 번 읽습니다.
그런 다음, 미리 계산된 이 단 하나의 "노트"를 여덟 명의 탐정에게 전달합니다.
마법 같은 점: 기존 시스템은 탐정들이 정확히 같은 단어를 사용할 때만 노트를 공유하지만, RKSC는 비록 표현이 약간 다르더라도 그들이 같은 것을 생각하고 있다면 공유할 수 있을 만큼 영리합니다. RKSC는 단순히 "단어"가 아니라 그들의 "생각"(은닉 상태/hidden states)을 확인합니다.

결과: 팀은 사건 파일의 앞부분을 여덟 번 다시 읽는 일을 멈춤으로써 엄청난 양의 시간을 절약합니다.

2. "확신에 찬 퇴장" 기술 (CGEE)

문제점:
탐정들이 이론을 작성하고 나면, 보통 감독관이 어떤 이론이 맞는지 확인하기 위해 모든 이론을 처음부터 끝까지 검토합니다. 이 "검증(verification)" 단계는 매우 느립니다. 하지만 때로는 어떤 탐정이 너무나 명백하게 확신을 가지고 정답을 말하고 있어서, 나머지 파일을 더 확인하는 것이 시간 낭비인 경우가 있습니다.

RKSC의 해결책:
RKSC는 두 가지 규칙을 가진 영리한 감독관 역할을 합니다.

규칙 A (건너뛰기): 만약 한 명의 탐정이 자신의 답에 95% 확신을 갖고 있고 다른 이들은 명확히 불확실한 상태라면, 감독관은 "좋아요, 끝났습니다!"라고 말하며 나머지 검증 과정을 통째로 건너뜁니다.
규칙 B (조기 종료): 만약 감독관이 반드시 검토를 해야 하는 상황이라도, 파일 전체를 다 읽을 필요는 없습니다. 중간까지만 읽고 멈출 수 있습니다. 왜냐하면 이 논문은 탐정이 추론 과정의 특정 지점에 도달하면 그 확신도가 안정화된다는 것을 발견했기 때문입니다. 마지막 30%를 읽는 것은 결론을 바꾸지 못하며, 그저 시간만 낭비할 뿐입니다.

결과: 시스템은 검증 단계를 통째로 건너뛰거나 중간에 짧게 끝냄으로써 엄청난 양의 시간을 절약합니다.

3. "기억력 관리인" (RSBCM)

문제점:
만약 탐정들이 계속해서 더 깊은 이론으로 파고든다면(마치 뿌리가 많은 나무처럼), "공유 노트"가 너무 커져서 컴퓨터의 메모리에 담을 수 없게 될 수 있습니다.

RKSC의 해결책:
RKSC에는 노트를 감시하는 관리인이 있습니다. 만약 노트가 너무 가득 차면, 관리인은 막다른 길(dead-end)에 빠진 이론에 깊이 몰입해 있거나 확신이 없어 보이는 탐정들의 노트를 버립니다. 이를 통해 메모리를 깨끗하게 유지하고, 매우 긴 추론 세션 중에도 시스템이 다운되지 않도록 보장합니다.

무엇을 발견했는가?

저자들은 이 시스템을 다섯 가지 서로 다른 AI 모델(소형부터 중형까지)을 대상으로 네 가지 유형의 어려운 퍼즐(과학, 수학, 논리)에 대해 테스트했습니다.

속도: 이 시스템은 표준 방식과 비교했을 때 AI를 평균 3배 더 빠르게 만들었습니다. 가장 좋은 경우에는 거의 4배까지 빨라졌습니다.
정확도: 정확도는 믿기 힘들 정도로 높았습니다. 1,600번 이상의 검증 체크 중, "조기 종료(Early Exit)" 기술은 단 6번의 실수만을 범했습니다 (오차율 0.37%).
재학습 불필요: 가장 멋진 점은 이것이 AI를 다시 학습시킬 필요가 없다는 것입니다. 이는 자동차를 새로 만드는 것이 아니라, 기존의 자동차에 더 똑똑한 엔진을 장착하는 것과 같습니다. 기존 모델 그대로 즉시 사용할 수 있습니다.

요 요약

RKSC는 AI 탐정 팀에게 똑같은 페이지를 다시 읽지 않도록 공유 노트를 주고, 답이 명백할 때는 검토를 멈추는 영리한 감독관을 배치하며, 작업 공간이 어지러워지지 않도록 관리인을 붙여주는 것과 같습니다. 그 결과, 이 추론 시스템은 예리함을 잃지 않으면서도 훨씬 더 빠르게 생각할 수 있게 되었습니다.

기술 요약: RKSC – 추론 인지형 KV 캐시 공유 및 확신 기반 조기 종료

1. 문제 정의

DeepSeek-R1, o1 스타일의 모델과 같이 복잡한 추론을 위해 설계된 현대적인 대규모 언어 모델(LLM)은 단일 문제에 대해 여러 추론 경로를 생성하는 Tree of Thoughts 또는 self-consistency와 같은 다중 분기 추론 전략을 자주 채택합니다. 본 논문은 현재의 프로덕션 추론 시스템(vLLM 및 SGLang 등)에서 이러한 파이프라인의 효율성을 저하시키는 두 가지 구조적 중복성을 식별했습니다.

교차 분기 KV 낭비 (Cross-Branch KV Waste): 병렬 분기들은 일반적으로 상당 부분의 접두사(문제 문구 및 초기 컨텍스트)를 공유합니다. 그러나 기존 시스템은 키-값(KV) 캐시를 토큰 시퀀스가 바이트 단위로 동일할 때만 재사용합니다. 만약 분기들이 의미적으로는 유사하지만 접두사 토큰이 다르거나, 접두사를 공유하더라도 힌트를 제공하는 표현 방식이 다를 경우, 시스템은 각 분기에 대해 공유된 접두사에 대한 고비용의 $O(n^2)$ 어텐션을 독립적으로 다시 계산하게 됩니다.
과도한 깊이의 검증 (Over-Deep Verification): 프로세스 보상(process-reward) 검증 파이프라인에서 모델은 각 생성된 분기를 점수화하기 위해 전체 깊이의 순방향 패스(forward pass)를 수행합니다. 이는 생성 신뢰도가 이미 결정적이거나 모델의 내부 상태(로짓 분포)가 중간 레이어에서 안정화되었을 때조차 발생합니다. 현재 시스템에는 검증 패스를 완전히 건너뛰거나 per-layer 엔트로피를 기반으로 이를 조기에 종료하는 메커니즘이 부족합니다.

2. 방법론: RKSC 프레임워크

저자들은 세 가지 상호 보완적이고 교체 가능한 메커니즘을 통해 이러한 중복성을 제거하는 RKSC(Reasoning-Aware KV Cache Sharing)를 제안합니다. RKSC는 미세 조정(fine-tuning)이 필요 없고, 아키텍처 변경이 필요 없으며, 커스텀 CUDA 커널을 요구하지 않습니다.

A. ASKS (Attention-Similarity KV Sharing)

ASKS는 토큰 일치 기반의 접두사 캐싱을 의미적 영역으로 일반화합니다.

메커니즘: ASKS는 바이트 단위로 동일한 토큰을 요구하는 대신, 분기 접미사(suffix)의 은닉 상태(hidden states)와 루트(root)의 은닉 상태 사이의 가중 코사인 유사도를 계산합니다.
가중치 부여: 유사도 점수( $\sigma_b$ )는 트랜스포머 레이어가 뒤로 갈수록 더 많은 작업 특정 정보를 담고 있고 의미적 발산을 더 잘 감지한다는 점을 고려하여, 레이어에 지수적으로 증가하는 가중치( $w_l = \exp(\alpha l/L)$ )를 적용합니다.
동작: 만약 $\sigma_b \geq \tau$ 라면, 해당 분기는 접두사와 의미적으로 일치하는 것으로 간주되어 zero-copy repeat_interleave 연산을 통해 공유 KV 캐시를 전달받습니다. 임계값 미만의 분기는 독립적인 계산 방식으로 회귀합니다.
일반화: 이는 토큰 일치형 캐싱(토큰이 동일할 때 $\sigma_b=1$ )을 엄격히 일반화하며, 표면적인 표현은 다르지만 의미적 표현은 유사한 분기들에 대해서도 재사용을 가능하게 합니다.

B. CGEE (Confidence-Gated Early Exit)

CGEE는 이중 레벨 종료 전략을 통해 검증 비용을 줄입니다.

검증 스킵 (Verify-Skip, 레벨 1): 검증 패스를 실행하기 전, 시스템은 생성 신뢰도를 확인합니다. 승리한 분기가 높은 절대 신뢰도( $p(b) \geq \tau_{conf}$ )를 가지고 경쟁자들을 명확한 상대적 차이로 압도하는 경우, 전체 검증 순방향 패스를 건너뜁니다.
레이어 수준 엔트로피 종료 (Layer-Level Entropy Exit, 레벨 2): 검증이 진행되는 동안, 경량 순방향 훅(forward hooks)이 각 레이어의 로짓 엔트로피를 모니터링합니다. 만약 엔트로피가 임계값( $\theta$ ) 아래로 떨어지고 중간 레이어 $l^*$ 에서 안정화(변화량 $< \epsilon$ )되면, 순방향 패스가 조기에 종료됩니다. 부분 출력값이 검증 결과로 사용되며, 나머지 $L - l^*$ 레이어는 건너뜁니다.

C. RSBCM (Reasoning-Selective Block Cache Manager)

심층 트리 탐색 중 발생하는 무제한적인 메모리 증가를 방지하기 위해 RSBCM은 공유 캐시를 관리합니다.

교체 정책 (Eviction Policy): 블록은 중요도 점수 $\omega = \text{branch\_score} / (\text{depth} + 1)$ 에 따라 교체됩니다. 이는 얕은 깊이의 고득점 분기 블록을 우선적으로 유지하고, 깊은 곳의 저신뢰도 궤적 블록을 먼저 제거하도록 합니다.

3. 주요 기여

ASKS: 다양한 표현 시나리오에서 기존의 토큰 일치형 매칭보다 28.6% 추가적인 KV 재사용을 달성하는 의미적 유사도 게이트를 도입하여 기존 접두사 캐싱을 엄격히 일반화했습니다.
CGEE: 신뢰도가 결정적일 때 검증을 스킵하거나, per-layer 엔트로피가 안정화될 때 패스를 중단하는 이중 레벨 종료 메커니즘을 제안하여 재학습 없이 검증 지연 시간을 줄였습니다.
RSBCM: 심층 탐색 트리 내에서 캐시 안정성을 유지하기 위한 깊이 우선 교체 전략을 제공합니다.
제로 오버헤드 설계: 학습 가능한 파라미터를 추가하지 않으며 아키텍처 수정도 필요하지 않은, 기존 추론 파이프라인에 즉시 적용 가능한 드롭인(drop-in) 래퍼 역할을 합니다.

4. 실험 결과

프레임워크는 1,000개의 문제를 대상으로 5개의 모델 제품군(Qwen2.5, Mistral, Falcon3, Llama-3를 포함한 7B–10B 파라미터 규모)과 4개의 벤치마크(GPQA Diamond, MMLU-STEM, ARC-Challenge, GSM8K)에서 평가되었습니다.

속도 향상: RKSC는 접두사를 재계산하는 No-KV 베이스라인 대비 평균 3.008배의 속도 향상을 보였으며, vLLM의 토큰 일치형 접두사 캐싱 대비 1.66배 개선된 성능을 보였습니다. 관찰된 최대 속도 향상은 3.990배였습니다 (Llama-3-8B on MMLU-STEM).
정확도: RKSC는 CGEE로 인한 에러율이 단 0.37%(1,616회의 검증 호출 중 6개 오류)에 불과할 정도로 높은 정확도를 유지했습니다. 오류는 생성 유창성이 가끔 정답 신호와 어긋나는 다지선다형 벤치마크에 집중되었으며, 자유 형식 생성 작업에서는 오류가 거의 나타나지 않았습니다.
구성 요소별 분석:
- KV 공유만으로 1.343배의 속도 향상(25.5% 지연 시간 감소)을 제공했습니다.
- 결합된 RKSC 시스템(KV + CGEE)은 처리량 분해에서 1.622배의 속도 향상을, 확장 평가에서 3.008배의 속도 향상을 제공했습니다.
- CGEE의 검증 스킵 비율은 어려운 GPQA 문제(~~42%)부터 쉬운 MMLU-STEM 문제(~~92%)까지 작업 난이도에 따라 적응적으로 나타났습니다.
아키텍처 불가지론 (Architecture Agnosticism): KV 공유 속도 향상은 모든 5개 모델에서 일관되게 나타났으며(표준 편차 0.02×), 이는 이 이득이 특정 모델 아키텍처가 아닌 추론 워크로드의 구조(접두사 길이 및 분기 계수)에 의해 유도됨을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 RKSC가 현재의 프로덕션 시스템이 간과하고 있는 다단계 LLM 추론의 근본적인 구조적 비효율성을 해결한다고 주장합니다. KV 재사용을 렉시컬(lexical) 정체성으로부터 분리하고(ASKS), 트랜스포머 레이어 후반부의 "엔트로피 붕괴(entropy collapse)" 현상을 활용함으로써(CGEE), RKSC는 추론 시점의 스케일링에 드는 에너지와 시간 비용을 크게 낮춥니다.

저자들은 이러한 이득이 훈련의 계산적 오버헤드나 스펙큘레이티브 디코딩(speculative decoding)과 같은 아키텍처적 복잡성 없이 달성되었다는 점을 강조합니다. 이 프레임워크는 다중 분기 파이프라인에서 중복된 계산을 줄임으로써 최첨단 추론을 제한된 하드웨어에서도 더 쉽게 접근할 수 있게 만드는 실용적인 드롭인 최적화 도구로 제시됩니다. 결과적으로 다단계 추론 작업에 있어, 상당 부분의 검증 및 프리필(prefill) 계산은 구조적으로 중복되며 의미적 유사성과 신뢰도 지표를 통해 안전하게 제거하거나 단축될 수 있음을 시사합니다.