🤖 AI

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

KV-PRM एक कुशल प्रोसेस रिवॉर्ड मॉडल है जो टेक्स्ट-आधारित री-एनकोडिंग की कम्प्यूटेशनल बाधा को समाप्त करने के लिए सीधे प्री-कंप्यूटेड KV कैश का लाभ उठाता है ताकि स्कोरिंग जटिलता को O(L²) से घटाकर O(L) किया जा सके, जिससे कई रीजनिंग बेंचमार्क में मौजूदा तरीकों के बराबर या उनसे बेहतर प्रदर्शन करते हुए गति और मेमोरी में भारी लाभ प्राप्त होता है।

मूल लेखक: Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

प्रकाशित 2026-07-13

📖 5 मिनट में पढ़ें🧠 गहराई से पढ़ें

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 से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप एक विशाल, हाई-स्टेक्स गेम शो चला रहे हैं जहाँ एआई जासूसों की एक टीम (जिसे हम "मल्टी-एजेंट स्क्वाड" कहेंगे) एक बेहद कठिन गणितीय रहस्य को सुलझाने की कोशिश कर रही है। वे केवल उत्तर का अनुमान नहीं लगाते; वे समस्या को छोटे-छोटे चरणों में तोड़ते हैं और संकेतों (clues) का आदान-प्रदान करते हैं। यह सुनिश्चित करने के लिए कि वे गलत रास्ते पर न चले जाएं, उनके पास एक "जज" (एक प्रोसेस रिवॉर्ड मॉडल, या PRM) है जो उनके द्वारा लिखे गए हर एक संकेत की जांच करता है।

पुराना तरीका: थकाऊ पुन: पाठक (The Exhaustive Re-Reader)
अतीत में, जब भी जज को किसी संकेत की जांच करनी होती थी, तो उन्हें अविश्वसनीय रूप से थकाऊ काम करना पड़ता था। कल्पना कीजिए कि जासूसों ने 5,000 शब्दों की एक कहानी लिखी है। उसे ग्रेड करने के लिए, जज को पूरी 5,000 शब्दों की कहानी को पहले शब्द से लेकर आखिरी शब्द तक फिर से पढ़ना पड़ता था, सिर्फ एक स्कोर देने के लिए। यदि कहानी लंबी होती है, तो पढ़ने का समय केवल थोड़ा ही नहीं बढ़ता; बल्कि वह विस्फोट की तरह बढ़ता है। यदि आप कहानी की लंबाई दोगुनी करते हैं, तो पढ़ने का समय चार गुना हो जाता है। यह $O(L^2)$ लागत है। यह एक घास के ढेर में सुई खोजने के लिए हर बार पूरा ढेर फिर से बनाने जैसा है। पेपर का तर्क है कि यह ऊर्जा और समय की भारी बर्बादी है, खासकर जब जासूस लंबी, जटिल कहानियाँ लिख रहे हों।

नया तरीका: KV-PRM ("मैजिक मेमोरी" रीडर)
लेखक पेंग कुआंग और उनकी टीम ने महसूस किया कि जासूस पहले से ही मुफ्त में कठिन काम कर रहे थे! जब जासूस अपनी कहानी लिखते हैं, तो उनके दिमाग (AI का आंतरिक "KV कैश") में उनके द्वारा सोचे गए हर शब्द की एक हाई-डेफिनिशन, सुपर-डिटेल्ड मेमोरी स्वाभाविक रूप से स्टोर हो जाती है। यह कहानी की केवल छपी हुई बातों के बजाय, उसकी आत्मा की एक पूर्ण, निरंतर रिकॉर्डिंग की तरह है।

यह कैसे काम करता है: जज एक एकल, छोटा "वेरिफाई टोकन" (सोचिए एक जादुई प्रश्न चिह्न की तरह, "?") लेता है और मेमोरी से पूछता है, "अब तक जो कुछ भी स्टोर किया गया है, उसके आधार पर, क्या यह रास्ता सही है?"
परिणाम: क्योंकि जज को टेक्स्ट को फिर से पढ़ने की आवश्यकता नहीं होती है, इसलिए लागत एक विशाल विस्फोट से घटकर एक सरल, लीनियर वॉक (linear walk) में बदल जाती है। पेपर गणितीय रूप से सिद्ध करता है कि इस मेमोरी में टेक्स्ट की तुलना में स्ट्रिक्टली अधिक जानकारी होती है। यह एक सपाट कागज के बजाय कहानी के 3D होलोग्राम जैसा है; होलोग्राम कम स्थान में अधिक विवरण रखता है।

आंकड़े: यह कितना तेज़ है?
टीम ने विभिन्न AI आकारों (0.6B, 4B, और 8B पैरामीटर्स) का उपयोग करके कुछ सबसे कठिन गणितीय पहेलियों (जैसे MATH, GSM8K, और AIME) पर इसका परीक्षण किया। यहाँ उन्होंने जो मापा:

गति (Speed): KV-PRM वास्तविक समय में 37 गुना तक तेज़ है। एक लंबी कहानी के लिए, एक एकल चेक जिसमें पुराने जज को 172.0 मिलीसेकंड लगते थे, उसमें नए जज को केवल 4.6 मिलीसेकंड लगे।
ऊर्जा (Energy): यह प्रति चेक 5,000 गुना कम कम्प्यूटेशनल स्टेप्स (FLOPs) का उपयोग करता है।
मेमोरी (Memory): इसे काम करने के लिए 34.2 गुना कम कंप्यूटर मेमोरी की आवश्यकता होती है।
सटीकता (Accuracy): इतनी तेज़ होने के बावजूद, इसने केवल "साथ नहीं निभाया"; इसने अक्सर पुराने, धीमे जजों की तुलना में बेहतर स्कोर प्राप्त किए।

वे स्पष्ट रूप से किसे खारिज करते हैं
पेपर बहुत स्पष्ट है कि क्या काम नहीं करता या क्या उत्तर नहीं है:

जज को छोटा बनाना: टीम ने पैसे बचाने के लिए पुराने शैली के टेक्स्ट-रीडिंग जज के रूप में एक छोटे AI (0.6B या 4B) का उपयोग करने का प्रयास किया। उन्होंने पाया कि हालांकि यह तेज़ था, लेकिन यह उतना स्मार्ट नहीं था। यहाँ तक कि एक छोटा 8B KV-PRM भी विशाल 8B टेक्स्ट-रीडिंग जज को बड़े अंतर से पीछे छोड़ गया। पेपर का तर्क है कि केवल मॉडल को छोटा करना समाधान नहीं है; बल्कि पढ़ने के तरीके को बदलना समाधान है।
अधिक टोकन पढ़ना: टीम ने सोचा, "क्या होगा यदि हम कहानी की जांच करने के लिए एक से अधिक प्रश्न चिह्नों का उपयोग करें?" उनका गणित (थ्योरम 2) और प्रयोग दिखाते हैं कि पहला प्रश्न चिह्न लगभग सारी उपयोगी जानकारी कैप्चर कर लेता है। दूसरा या तीसरा प्रश्न चिह्न जोड़ने से लगभग कोई अतिरिक्त लाभ नहीं मिलता, लेकिन इसमें अधिक लागत आती है। इसलिए, केवल एक वेरिफाई टोकन का उपयोग करना ही सबसे सटीक विकल्प (sweet spot) है।

एक बोनस ट्रिक: "KV स्टीयरिंग"
चूंकि नया जज टेक्स्ट (जो टूटे हुए और डिस्क्रीट संकेतों जैसा है) के बजाय "मैजिक मेमोरी" (जो एक स्मूथ, निरंतर सिग्नल है) को देखता है, टीम ने एक दिलचस्प साइड इफेक्ट की खोज की। वे गणितीय ग्रेडिएंट्स का उपयोग करके बातचीत को बेहतर उत्तर की ओर मोड़ सकते हैं, यानी वे जासूसों के विचारों को सोचते समय ही निर्देशित कर सकते हैं। वे इसे KV स्टीयरing कहते हैं। पेपर दिखाता है कि यह प्रूफ-ऑफ-कॉन्सेप्ट के रूप में काम कर गया, जिससे बिना किसी सर्च के AIME पहेलियों की सटीकता में 3.33 प्रतिशत अंक तक का सुधार हुआ। पेपर नोट करता है कि यह पुराने टेक्स्ट-आधारित जजों के साथ संरचनात्मक रूप से असंभव है क्योंकि आप कागज के टुकड़े को उसी तरह "स्टीयर" नहीं कर सकते जैसे आप एक मेमोरी सिग्नल को कर सकते हैं।

निष्कर्ष (The Bottom Line)
पेपर केवल यह सुझाव नहीं देता कि यह काम कर सकता है; उन्होंने कई डेटासेट्स और मॉडल आकारों में इसे मापा और गणितीय रूप से सिद्ध किया। उन्होंने पाया कि टेक्स्ट को फिर से पढ़ने के बजाय AI की अपनी "मैजिक मेमोरी" का पुन: उपयोग करके, हम जटिल समस्याओं को बहुत तेज़ी से, सस्ते में और अक्सर अधिक सटीकता के साथ हल कर सकते हैं। यह "पूरी किताब को फिर से पढ़ने" से "लेखक के परफेक्ट नोट्स को चेक करने" की ओर एक बदलाव है।

तकनीकी सारांश: KV-PRM

समस्या विवरण (Problem Statement)

यह शोध पत्र लार्ज लैंग्वेज मॉडल्स (LLMs) के लिए टेस्ट-टाइम स्केलिंग (TTS) में एक महत्वपूर्ण कम्प्यूटेशनल बाधा को संबोधित करता है, विशेष रूप से मल्टी-एजेंट सिस्टम्स (MAS) और रीजनिंग पाइपलाइन्स के भीतर। जबकि प्रोसेस रिवॉर्ड मॉडल्स (PRMs) मध्यवर्ती रीजनिंग स्टेप्स का मूल्यांकन करके सर्च एल्गोरिदम (जैसे बीम सर्च, MCTS) को निर्देशित करने के लिए आवश्यक हैं, मौजूदा PRMs टेक्स्ट-आधारित (text-based) हैं। उन्हें प्रत्येक स्कोरिंग कॉल के लिए पूरे ट्राजेक्टरी टेक्स्ट को फिर से एनकोड करने की आवश्यकता होती है।

चूंकि सेल्फ-अटेंशन मैकेनिज्म सीक्वेंस की लंबाई ( $L$ ) के साथ द्विघात रूप से (quadratically) स्केल करते हैं, इसलिए एक एकल PRM स्कोरिंग कॉल की कम्प्यूटेशनल लागत $O(L^2)$ होती है। लंबे-कॉन्टेक्स्ट MAS परिदृश्यों में जहाँ ट्राजेक्टरीज हजारों टोकन्स तक फैली होती हैं और सर्च में दर्जनों उम्मीदवारों को शामिल किया जाता है, यह अनावश्यक पुन: एनकोडिंग एक गंभीर बाधा उत्पन्न करती है, जो अक्सर स्वयं जेनरेशन के लिए आवश्यक FLOPs के बराबर या उससे अधिक हो जाती है। यह लंबे-कॉन्टेक्स्ट रीजनिंग में PRMs के व्यावहारिक अनुप्रयोग को सीमित करता है।

कार्यप्रणाली: KV-PRM (Methodology: KV-PRM)

लेखक KV-PRM का प्रस्ताव करते हैं, जो एक अत्यधिक कुशल प्रोसेस रिवॉर्ड मॉडल है जो बेस मॉडल के ऑटोरिग्रेसिव जेनरेशन चरण के दौरान स्वाभाविक रूप से उत्पन्न होने वाले की-वैल्यू (Key-Value - KV) कैश का सीधे उपयोग करके टेक्स्ट पुन: एनकोडिंग को समाप्त करता है।

मुख्य तंत्र (Core Mechanism)

KV-कैश ट्रांसफर: टेक्स्ट टोकन को प्रोसेस करने के बजाय, KV-PRM अनुक्रम के अंत में एक एकल "वेरिफाई टोकन" (जैसे, "?") जोड़ता है। इसके बाद यह इस टोकन को बेस मॉडल के पूर्व-मौजूद KV कैश पर आधारित एक हल्के वजन वाले LoRA एडॉप्टर के माध्यम से प्रोसेस करता है।
जटिलता में कमी: प्री-कंप्यूटेड कैश को पढ़ने से, स्कोरिंग लागत $O(L^2)$ (टेक्स्ट पुन: एनकोडिंग) से घटकर $O(L)$ ( $L$ कैश किए गए पदों के विरुद्ध एकल टोकन प्रोसेसिंग) हो जाती है।
आर्किटेक्चर: बेस मॉडल ( $f_\theta$ ) ट्राजेक्टरी और KV कैश जनरेट करता है। रिवॉर्ड मॉडल ( $f_{\theta+\Delta\theta}$ ) उन्हीं बेस वेट्स का उपयोग करता है जिसमें एक सक्रिय LoRA एडॉप्टर होता है जो कैश की व्याख्या करता है और वेरिफाई टोकन के लॉजिट्स के माध्यम से स्कोर आउटपुट करता है। इंफरेंस के दौरान बेस मॉडल के माध्यम से कोई ग्रेडिएंट प्रवाहित नहीं होता है; केवल LoRA पैरामीटर्स को ट्रेनिंग के दौरान अपडेट किया जाता है।

सैद्धांतिक आधार (Theoretical Foundation)

पेपर इस दृष्टिकोण को न्यायसंगत ठहराने के लिए एक सैद्धांतिक ढांचा प्रदान करता है:

सूचना क्षमता (Proposition 1): KV कैश में डिकोड किए गए टेक्स्ट की तुलना में स्पष्ट रूप से अधिक सूचना क्षमता होती है। लंबाई- $L$ के KV कैश की सूचना को टेक्स्ट के माध्यम से लॉसलेस रूप से व्यक्त करने के लिए $\Omega(d \cdot L / \log |V|)$ टोकन की आवश्यकता होगी, जिसका अर्थ है कि कैश काफी अधिक सूचना-सघन (information-dense) है।
त्रुटि अपघटन (Theorem 1): वेरिफिकेशन एरर को एक बेयस-ऑप्टिमल फ्लोर और एक एप्रोक्सिमेशन गैप में विभाजित किया गया है। KV कैश पर बेयस-ऑप्टिमल प्रेडिक्टर, समृद्ध प्रतिनिधित्व के कारण टेक्स्ट की तुलना में समान या कम त्रुटि प्राप्त करता है।
घटते प्रतिफल (Theorem 2): अधिक वेरिफाई टोकन ( $k$ ) जोड़ने से मिलने वाला मार्जिनल इंफॉर्मेशन गेन तेजी से घटता है। एक एकल वेरिफाई टोकन ( $k=1$ ) अधिकांश निष्कर्षण योग्य रिवॉर्ड जानकारी को कैप्चर कर लेता है, जिससे $O(L)$ डिज़ाइन लगभग इष्टतम (near-optimal) हो जाता है।

मुख्य योगदान (Key Contributions)

सैद्धांतिक ढांचा: वेरिफिकेशन के लिए टेक्स्ट की तुलना में KV कैश के प्रतिनिधित्व संबंधी लाभ को औपचारिक रूप देता है और सिद्ध करता है कि एप्रोक्सिमेशन गैप रीडआउट डेप्थ के साथ तेजी से घटता है, जो सिंगल-टोकन रीडआउट की दक्षता को न्यायसंगत ठहराता है।
एल्गोरिदम (KV-PRM): एक नवीन PRM आर्किटेक्चर पेश करता है जो जेनरेशन KV कैश का पुन: उपयोग करता है, जिससे स्कोरिंग जटिलता $O(L^2)$ से घटकर $O(L)$ हो जाती है और टेक्स्ट-आधारित PRMs की आर्किटेक्चरल बाधा को समाप्त करती है।
अनुभवजन्य सत्यापन (Empirical Validation): MATH, GSM8K, और AIME बेंचमार्क पर कई मॉडल स्केल्स (Qwen3-0.6B, 4B, 8B) और सर्च एल्गोरिदम (बीम सर्च, MCTS, वेटेड वोटिंग) पर KV-PRM की सटीकता को टेक्स्ट-आधारित PRMs के बराबर या उससे बेहतर प्रदर्शित करता है।
KV स्टीयरिंग (KV Steering): लेटेंट इंटर-एजेंट संदेशों के ग्रेडिएंट-आधारित अनुकूलन को सक्षम करने के लिए KV-कैश स्कोरिंग की डिफरेंशिएबिलिटी (differentiability) का अन्वेषण करता है, जो कि डिस्क्रीट टेक्स्ट-आधारित PRMs के साथ संरचनात्मक रूप से असंभव है।

प्रयोगात्मक परिणाम (Experimental Results)

MATH, GSM8K, AIME 2024, और AIME 2025 पर मूल्यांकित:

सटीकता (Accuracy): KV-PRM, टेक्स्ट-आधारित PRMs के बराबर या उनसे बेहतर प्रदर्शन करता है। उदाहरण के लिए, MATH पर MCTS और Qwen3-8B के साथ, KV-PRM ने 69.57% सटीकता प्राप्त की, जबकि Text-PRM के लिए यह 68.92% थी।
दक्षता (FLOPs): प्रति कॉल स्कोरिंग FLOPs में 5,000× तक की कमी हासिल की (उदाहरण के लिए, MATH पर 8B मॉडल के लिए Text-PRM की लागत का $1/2782$ भाग)।
विलंबता (Latency): वॉल-क्लॉक स्कोरिंग लेटेंसी को 15× से 37× तक कम किया गया। 4,096 की सीक्वेंस लेंथ के साथ 8B पैरामीटर्स पर, KV-PRM ने 4.6 ms में स्कोर किया, जबकि Text-PRM के लिए यह 172.0 ms था।
मेमोरी: $O(L^2)$ अटेंशन मैप्स के निर्माण (materialization) से बचकर प्रति-अनुक्रम मेमोरी फुटप्रिंट को 34.2× कम किया।
रीडआउट डेप्थ (Readout Depth): प्रयोगों ने पुष्टि की कि अधिक वेरिफाई टोकन ( $k > 1$ ) जोड़ने से घटते प्रतिफल मिलते हैं, जो इस सैद्धांतिक दावे को मान्य करता है कि $k=1$ पर्याप्त है।

महत्व और दावे (Significance and Claims)

पेपर का दावा है कि KV-PRM वेरिफायर डिज़ाइन में एक मौलिक बदलाव का प्रतिनिधित्व करता है, जो रेडंडेंट टेक्स्ट पुन: एनकोडिंग से हटकर रिप्रेजेंटेशन रियूज (representation reuse) की ओर बढ़ता है।

स्केलेबिलिटी: स्कोरिंग लागत को जेनरेशन लागत के नगण्य अंश तक कम करके, KV-PRM उन लंबे-कॉन्टेक्स्ट, मल्टी-एजेंट रीजनिंग परिदृश्यों में व्यावहारिक अनुप्रयोग को सक्षम बनाता है जिन्हें पहले कम्प्यूटेशनल रूप से अत्यधिक कठिन माना जाता था।
नई क्षमताएं: KV-कैश स्कोरिंग सिग्नल की डिफरेंशिएबिलिटी, इंफरेंस-टाइम ऑप्टिमाइजेशन (KV स्टीयरिंग) के लिए एक नया मार्ग खोलती है, जो सर्च के बिना ग्रेडिएंट-आधारित एजेंट इंटरैक्शन के रिफाइनमेंट की अनुमति देती है, जो डिस्क्रीट टेक्स्ट आउटपुट्स के साथ असंभव है।
मामूली सीमाएं: लेखक नोट करते हैं कि इस पद्धति के लिए जेनेरेटर और वेरिफायर का एक ही आर्किटेक्चर साझा करना आवश्यक है, यह थ्योरिटिकल बाउंड्स के लिए लीनियर रिप्रेजेंटेशन हाइपोथेसिस पर निर्भर करता है, और KV स्टीयरिंग वर्तमान में एक प्रूफ-ऑफ-कॉन्सेप्ट है।

संक्षेप में, KV-PRM जेनरेशन के दौरान पहले से गणना किए गए उच्च-फिडेलिटी (high-fidelity) प्रतिनिधित्व का लाभ उठाकर टेस्ट-टाइम स्केलिंग की कम्प्यूटेशनल बाधा को हल करता है, जो प्रदर्शन में समझौता किए बिना, और अक्सर सुधार करते हुए, ऑर्डर्स-ऑफ-मैग्निट्यूड दक्षता प्रदान करता है।