🤖 AI

KV-PRM: Efficient Process Reward Modeling via KV-Cache Transfer for Multi-Agent Test-Time Scaling

KV-PRM은 사전 계산된 KV 캐시를 직접 활용하여 텍text 기반 재인코딩의 계산 병목 현상을 제거함으로써 스코어링 복잡도를 O(L²)에서 O(L)로 줄여, 기존 방식들과 대등하거나 이를 능가하는 성능을 여러 추론 벤치마크에서 달성하는 동시에 속도와 메모리 측면에서 막대한 이득을 얻는 효율적인 프로세스 보상 모델입니다.

원저자: Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

게시일 2026-07-13

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

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

당신이 거대한 규모의 고액 판돈이 걸린 게임 쇼를 운영하고 있다고 상상해 보세요. 여기서 '멀티 에이전트 스쿼드(Multi-Agent Squad)'라고 불리는 AI 탐정 팀이 초고난도 수학 미스터리를 풀려고 노력 중입니다. 그들은 단순히 답을 추측하는 것이 아니라, 단서를 주고받으며 문제를 아주 작은 단계로 쪼개어 해결합니다. 이들이 잘못된 길로 빠지지 않도록 하기 위해, '판사(Judge)'(프로세스 보상 모델, PRM)가 그들이 적어 내려가는 모든 단서 하나하나를 검토합니다.

과거의 방식: 철저한 재독자 (The Exhaustive Re-Reader)
과거에는 판사가 단서를 확인할 때마다 믿기 힘들 정도로 지치는 작업을 수행해야 했습니다. 탐정들이 5,000단어짜리 이야기를 썼다고 가정해 봅시다. 이를 채점하기 위해 판사는 이야기를 처음부터 끝까지 다시 읽어야 했습니다. 이야기를 처음부터 마지막 단어까지 다시 읽어야만 점수를 줄 수 있었던 것입니다. 만약 이야기가 길어지면, 읽는 데 걸리는 시간은 단순히 조금 늘어나는 수준이 아니라 폭발적으로 증가합니다. 이야기가 두 배 길어지면 읽는 시간은 네 배가 됩니다. 이것이 논문에서 말하는 $O(L^2)$ 비용입니다. 이는 마치 건초더미 속에서 바늘을 찾기 위해 매번 전체 건초더미를 새로 만드는 것과 같습니다. 논문은 이것이 엄청난 에너지와 시간 낭비이며, 특히 탐정들이 길고 복잡한 이야기를 쓸 때 더욱 그렇다고 주장합니다.

새로운 방식: KV-PRM ("마법의 기억"을 가진 독자)
이 논문의 저자인 펑 쿠앙(Peng Kuang)과 그의 팀은 탐정들이 이미 공짜로 힘든 일을 해내고 있다는 사실을 깨달았습니다! 탐정들이 이야기를 쓰는 동안, 그들의 뇌(AI의 내부 'KV 캐시')는 자연스럽게 지금까지 생각했던 모든 단어에 대한 고해상도의 매우 상세한 기억을 저장합니다. 이것은 단순한 텍스트가 아니라, 이야기의 영혼을 담은 완벽하고 연속적인 기록과 같습니다.

이 논문은 KV-PRM이라는 새로운 종류의 판사를 제안합니다. 이 판사는 텍스트를 다시 읽지 않습니다. 대신, 그 "마법의 기억" 기록을 살짝 훑어봅니다.

작동 원리: 판사는 단 하나의 작은 "검증 토큰(verify token)"(생각해 보면 마법의 물음표 "?"와 같습니다)을 가져와서 메모리에 묻습니다. "지금까지 저장된 내용을 바탕으로 볼 때, 이 경로는 괜찮은가?"
결과: 판사가 텍스트를 다시 읽을 필요가 없기 때문에, 비용은 거대한 폭발에서 단순한 선형적 이동으로 줄어듭니다. 논문은 이 메모리가 텍스트 자체보다 엄격하게 더 많은 정보를 담고 있음을 수학적으로 증명했습니다. 이는 평면적인 종이 한 장을 보는 것과 달리, 이야기의 3D 홀로그램을 보는 것과 같습니다. 홀로그램은 더 적은 공간에 더 많은 세부 정보를 담고 있습니다.

숫자로 보는 성능: 얼마나 빠른가?
연구팀은 다양한 AI 크기(0.6B, 4B, 8B 파라미터)를 사용하여 가장 어려운 수학 퍼즐들(MATH, GSM8K, AIME 등)로 테스트를 진행했습니다. 측정 결과는 다음과 같습니다:

속도: KV-PRM은 실제 시간 기준으로 최대 37배 더 빠릅니다. 긴 이야기의 경우, 기존 방식의 판사가 확인하는 데 172.0 밀리초가 걸렸다면, 새로운 판사는 단 4.6 밀리초밖에 걸리지 않았습니다.
에너지: 확인당 사용하는 연산량(FLOPs)이 최대 5,000배 적습니다.
메모리: 작업을 수행하는 데 필요한 컴퓨터 메모리가 34.2배 적습니다.
정확도: 이렇게 훨씬 빠름에도 불구하고, 단순히 속도만 따라가는 것이 아니라 기존의 느린 판사들보다 오히려 더 높은 점수를 기록하는 경우가 많았습니다.

그들이 명시적으로 제외한 것들
이 논문은 무엇이 효과가 없거나 정답이 아닌지를 매우 명확히 밝히고 있습니다:

판사를 작게 만드는 것: 연구팀은 비용을 아끼기 위해 더 작은 AI(0.6B 또는 4B)를 기존 방식의 텍스트 읽기 판사로 사용해 보았습니다. 그 결과 속도는 빨라졌지만 지능은 떨어졌습니다. 심지어 아주 작은 8B KV-PRM이 거대한 8B 텍스트 읽기 판사보다 훨씬 더 뛰어난 성과를 냈습니다. 논문은 단순히 모델을 줄이는 것이 해결책이 아니라, 모델이 읽는 방식을 바꾸는 것이 핵심이라고 주장합니다.
더 많은 토큰을 사용하는 것: 연구팀은 "이야지를 확인하기 위해 물음표를 하나 이상 사용하면 어떨까?"라는 의문을 가졌습니다. 그들의 수학적 증명(정리 2)과 실험에 따르면, 첫 번째 물음표가 거의 모든 유용한 정보를 포착합니다. 두 번째나 세 번째 물음표를 추가해도 이득은 거의 없지만 비용은 더 들게 됩니다. 따라서 단 하나의 검증 토큰을 사용하는 것이 최적의 지점입니다.

보너스 기술: "KV 스티어링(KV Steering)"
새로운 판사는 텍스트(불연속적이고 끊기는 신호)가 아니라 "마법의 기억"(매끄럽고 연속적인 신호)을 보기 때문에, 연구팀은 멋진 부수 효과를 발견했습니다. 수학적 그래디언트를 사용하여 탐정들이 생각하는 도중에 그들의 생각을 더 나은 답으로 유도할 수 있는 것입니다. 이를 KV 스티어링이라 부릅니다. 논문은 이것이 개념 증명으로서 작동함을 보여주었으며, 검색(search) 과정을 거치지 않고도 AIME 퍼즐의 정확도를 최대 3.33 퍼센트 포인트 향상시켰습니다. 논문은 기존의 텍스트 기반 판사로는 이것이 구조적으로 불가능하다고 언급합니다. 왜냐로 종이 조각을 제어하는 것과 달리, 메모리 신호를 제어하는 것은 가능하기 때문입니다.

결론
이 논문은 단순히 이 방식이 작동할 수도 있다고 제안하는 것이 아니라, 여러 데이터셋과 모델 크기에 걸쳐 측정하고 수학적으로 증명했습니다. 연구팀은 AI의 "마법의 기억"을 재사용함으로써 텍스트를 다시 읽는 대신, 훨씬 더 빠르고, 저렴하며, 종종 더 정확하게 복잡한 문제를 해결할 수 있다는 것을 발견했습니다. 이는 "책 전체를 다시 읽는 것"에서 "저자의 완벽한 노트를 확인하는 것"으로의 전환입니다.

기술 요약: KV-PRM

문제 정의

본 논문은 대규모 언어 모델(LLM)의 테스트 시간 스케일링(Test-Time Scaling, TTS), 특히 멀티 에이전트 시스템(Multi-Agent Systems, MAS) 및 추론 파이프라인 내에서 발생하는 결정적인 계산 병목 현상을 다룹니댜. 프로세스 보상 모델(Process Reward Models, PRM)은 중간 추론 단계를 평가하여 탐색 알고리즘(예: Beam Search, MCTS)을 가이드하는 데 필수적이지만, 기존의 PRM은 텍스트 기반입니다. 즉, 점수를 매길 때마다 전체 궤적(trajectory) 텍스트를 처음부터 다시 인코딩해야 합니다.

자기 주의 집중(Self-attention) 메커니즘은 시퀀스 길이( $L$ )에 따라 제곱으로 증가하므로, 단일 PRM 스코어링 호출의 계산 비용은 $O(L^2)$ 가 됩니다. 궤적이 수천 개의 토큰에 달하고 탐색 과정에서 수십 개의 후보를 검토해야 하는 롱 컨텍스트(long-context) MAS 시나리오에서, 이러한 중복된 재인코딩은 생성 자체에 필요한 연산량(FLOPs)과 맞먹거나 이를 초과하는 심각한 병목 현상을 초래합니다. 이는 롱 컨텍스트 추론에서 PRM의 실질적인 적용을 제한합니다.

방법론: KV-PRM

저자들은 LLM의 자기회귀 생성 단계에서 자연스럽게 생성되는 키-값(Key-Value, KV) 캐시를 직접 활용함으로써 텍스트 재인코딩을 제거하는 매우 효율적인 프로세스 보상 모델인 KV-PRM을 제안합니다.

핵심 메커니즘

KV 캐시 전송: 텍est 토큰을 처리하는 대신, KV-PRM은 시퀀스 끝에 단 하나의 "검증 토큰"(verify token) (예: "?")을 추가합니다. 그 후, 베이스 모델의 기존 KV 캐시에 조건화된 경량 LoRA 어댑터를 통해 이 토큰을 처리합니다.
복잡도 감소: 사전 계산된 캐시를 읽음으로써, 스코어링 비용은 텍스트 재인코딩 방식의 $O(L^2)$ 에서 단일 토큰을 $L$ 개의 캐시된 위치에 대해 처리하는 $O(L)$ 로 감소합니다.
아키텍처: 베이스 모델( $f_\theta$ )이 궤적과 KV 캐시를 생성합니다. 보상 모델( $f_{\theta+\Delta\theta}$ )은 동일한 베이스 가중치를 사용하되 활성화된 LoRA 어댑터를 사용하여 캐시를 해석하고 검증 토큰의 로짓(logits)을 통해 점수를 출력합니다. 추론 중에 베이스 모델을 통한 그래디언트 흐름은 없으며, 훈련 중에만 LoRA 파라미터가 업데이트됩니다.

이론적 토대

본 논문은 이 접근 방식을 정당화하는 이론적 프레임워크를 제공합니다:

정보 용량 (Proposition 1): KV 캐시는 디코딩된 텍스트보다 엄격하게 더 큰 정보 용량을 가집니다. 길이 $L$ 인 KV 캐시의 정보를 손실 없이 텍스트로 표현하려면 $\Omega(d \cdot L / \log |V|)$ 개의 토큰이 필요하며, 이는 캐시가 훨씬 더 정보 밀도가 높음을 의미합니다.
오차 분해 (Theorem 1): 검증 오차는 베이즈 최적 하한(Bayes-optimal floor)과 근사 격차(approximation gap)로 분해됩니다. KV 캐시 상의 베이즈 최적 예측기는 더 풍부한 표현 덕분에 텍스트 상의 예측기와 같거나 더 낮은 오차를 달성합니다.
수익 체감 (Theorem 2): 추가적인 검증 토큰( $k$ )을 더함에 따른 한계 정보 이득은 지수적으로 감소합니다. 단일 검증 토큰( $k=1$ )이 추출 가능한 보상 정보의 대부분을 포착하므로, $O(L)$ 설계는 거의 최적에 가깝습니다.

주요 기여

이론적 프레임워크: 검증을 위한 KV 캐시의 표현적 우위를 공식화하고, 근사 격차가 읽기 깊이(readout depth)에 따라 지수적으로 감소함을 증명하여 단일 토큰 읽기의 효율성을 정당화했습니다.
알고리즘 (KV-PRM): 생성 KV 캐시를 재사용하여 스코어링 복잡도를 $O(L^2)$ 에서 $O(L)$ 로 줄이고, 텍스트 기반 PRM의 아키텍처적 병목을 제거하는 새로운 PRM 아키텍처를 도입했습니다.
경험적 검증: MATH, GSM8K, AIME 벤치마크에서 다양한 모델 규모(Qwen3-0.6B, 4B, 8B)와 탐색 알고리즘(Beam Search, MCTS, Weighted Voting)에 걸쳐 KV-PRM이 텍스트 기반 PRM과 일치하거나 이를 능가함을 입증했습니다.
KV 스티어링 (KV Steering): 잠재적 에이전트 간 메시지를 그래디언트 기반으로 최적화할 수 있는 KV 캐시 스코어링의 미분 가능성을 탐구했습니다. 이는 이산적인 텍스트 기반 PRM으로는 구조적으로 불가능한 기능입니다.

실험 결과

MATH, GSM8K, AIME 2024, AIME 2025에서 평가되었습니다:

정확도: KV-PRM은 텍스트 기반 PRM과 일치하거나 엄격히 능가합니다. 예를 들어, MATH 데이터셋의 MCTS와 Qwen3-8B 환경에서 KV-PRM은 **69.57%**의 정확도를 기록하여 Text-PRM의 **68.92%**를 앞질렀습니다.
효율성 (FLOPs): 호출당 스코어링 FLOP스를 최대 5,000배 감소시켰습니다 (예: 8B 모델의 MATH 작업에서 Text-PRM 대비 $1/2782$ 수준의 비용).
지연 시간 (Latency): 실제 실행 시간(wall-clock) 스코어링 지연 시간을 15배에서 37배까지 단축했습니다. 8B 파라미터와 4,096의 시퀀스 길이에서, KV-PRM은 4.6 ms 만에 스코어링을 완료한 반면 Text-PRM은 172.0 ms가 소요되었습니다.
메모리: $O(L^2)$ 어텐션 맵의 구체화를 피함으로써 시퀀스당 메모리 풋프린트를 34.2배 줄였습니다.
읽기 깊이 (Readout Depth): 검증 토큰 수를 늘리는 것( $k > 1$ )이 수익 체감을 가져온다는 실험 결과를 통해, $k=1$ 이 충분하다는 이론적 주장을 확인했습니다.

의의 및 주장

본 논문은 KV-PRM이 중복된 텍스트 재인코딩에서 벗어나 **표현 재사용(representation reuse)**으로 나아가는 검증기 설계의 근본적인 변화를 나타낸다고 주장합니다.

확장성: 스코어링 비용을 생성 비용의 무시할 만한 수준으로 낮춤으로써, KV-PRM은 기존에 계산적으로 불가능하다고 여겨졌던 롱 컨텍스트 및 멀티 에이전트 추론 시나리오에서 PRM의 실질적인 적용을 가능하게 합니다.
새로운 능력: KV 캐시 스코어링의 미분 가능성은 **추론 시간 최적화(inference-time optimization)**인 KV 스티어링을 위한 새로운 방향을 열어줍니다. 이는 탐색 없이도 그래디언트 기반의 에이전트 상호작용 정교화를 가능하게 하며, 이는 이산적인 텍스트 출력으로는 불가능한 일입니다.
제한 사항: 저자들은 이 방법이 생성기와 검증기가 동일한 아키텍처를 공유해야 하며, 이론적 경계에 선형 표현 가설(Linear Representation Hypothesis)에 의존한다는 점, 그리고 KV 스티어링이 현재 개념 증명 단계라는 점을 언급했습니다.

요약하자면, KV-PRM은 생성 중에 이미 계산된 고충실도 잠재 표현을 활용함으로써 테스트 시간 스케일링의 계산 병목을 해결하며, 추론 성능을 희생하지 않으면서(종종 오히려 향상시키면서) 수십 배의 효율성 이득을 제공합니다.