🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

यह शोध पत्र BCJR-QAT प्रस्तुत करता है, जो ट्रेलीस-कोडेड वेट क्वांटाइजेशन (trellis-coded weight quantization) का एक डिफरेंशिएबल रिलैक्सेशन है, जो गैर-डिफरेंशिएबल विटेर्बी आर्गमैक्स (Viterbi argmax) को तापमान-नियंत्रित BCJR सम-प्रोडक्ट एल्गोरिदम से बदल देता है, जिससे एंड-टू-एंड क्वांटाइजेशन-अवेयर ट्रेनिंग सक्षम होती है जो LLMs पर स्टेट-ऑफ-द-आर्ट PTQ फ्रंटियर से अनुभवजन्य रूप से बेहतर प्रदर्शन करती है।

मूल लेखक: Venugopalan Iyengar

प्रकाशित 2026-05-12

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

CC BY 4.0

मूल लेखक: Venugopalan Iyengar

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आपके पास किताबों की एक विशाल लाइब्रेरी (एक लार्ज लैंग्वेज मॉडल) है जिसे आप इतना छोटा करना चाहते हैं कि वह एक मानक लैपटॉप या फोन पर फिट हो सके। ऐसा करने के लिए, आपको लाइब्रेरी के अंदर के "शब्दों" (वजन/weights) को कंप्रेस करना होगा।

यह पेपर एक नया तरीका पेश करता है जिसे BCJR-QAT कहा जाता है, जो इन किताबों को और भी अधिक कंप्रेस करता है—सिर्फ 2 बिट्स प्रति शब्द तक (जो कि एक हाई-डेफिनिशन फोटो को एक छोटे, दानेदार थंबनेल में कंप्रेस करने जैसा है)।

यहाँ बताया गया है कि उन्होंने इसे कैसे किया, सरल उपमाओं का उपयोग करते हुए:

1. समस्या: "एकतरफा रास्ता" का जाल (The "One-Way Street" Trap)

पहले, इन मॉडल्स को सिकोड़ने का सबसे अच्छा तरीका QTIP नामक एक विधि थी। QTIP को एक बहुत ही बुद्धिमान लाइब्रेरियन के रूप में सोचें जो टेक्स्ट के एक पन्ने को देखता है और कहता है, "ठीक है, मैं इस जटिल वाक्य को अपने डिक्शनरी से सबसे करीबी सरल वाक्यांश से बदल दूँगा।"

हालाँकि, एक सीमा थी। एक बार जब लाइब्रेरियन ने वह चुनाव कर लिया, तो वह पीछे जाकर अपना मन नहीं बदल सकता था ताकि पूरी किताब बेहतर सुनाई दे सके। यह एक "एकतरफा रास्ते" की तरह था। यदि लाइब्रेरियन ने शुरुआत में थोड़ा सा गलत वाक्यांश चुना, तो बाकी की किताब प्रभावित होती, और वे इसे ठीक नहीं कर सकते थे क्योंकि निर्णय "कठोर" (hard) और अंतिम था।

इसे ठीक करने के लिए, शोधकर्ता आमतौर पर ट्रेनिंग (QAT) का उपयोग करते हैं, जहाँ आप लाइब्रेरियन को अभ्यास करने और अपने विकल्पों को समायोजित करने देते हैं। लेकिन यहाँ एक पेच है: लाइब्रेरियन की निर्णय प्रक्रिया में एक जटिल भूलभुलैया (एक "ट्रेलीस"/trellis) शामिल है। इस भूलभुलैया में सबसे अच्छा रास्ता खोजने के लिए, वे विटर्बी (Viterbi) नामक एक नियम का उपयोग करते हैं, जो "सबसे अच्छा विकल्प चुनें" बटन की तरह है।

समस्या: आप कंप्यूटर को "सबसे अच्छा विकल्प चुनें" बटन के माध्यम से सीखना नहीं सिखा सकते क्योंकि वह बटन गणितीय रूप से एक डेड एंड (dead end) है। यदि आप उस बटन को दबाते हैं, तो कंप्यूटर को यह नहीं पता चलता कि अपने विकल्पों को बेहतर बनाने के लिए उन्हें कैसे थोड़ा सा बदला जाए; वह बस निकटतम विकल्प पर झपट पड़ता है। यह एक कार चलाने की कोशिश करने जैसा है जहाँ आपको केवल अगली लेन में कूदने की अनुमति दी जाती है—आप सुचारू मोड़ नहीं ले सकते।

2. समाधान: "सॉफ्ट" निर्णय (BCJR)

लेखकों ने एक नया तरीका आविष्कार किया जिससे लाइब्रेरियन अभ्यास कर सके। तुरंत एक कठोर "सबसे अच्छा चुनें" निर्णय लेने के बजाय, उन्होंने तापमान ( $T$ ) की एक अवधारणा पेश की।

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

यह विधि, जिसे BCJR कहा जाता है, "एकतरफा रास्ते" को एक सुचारू, चलाने योग्य सड़क में बदल देती है। कंप्यूटर अब यह सीख सकता है कि पूरे बुक के अनुभव को बेहतर बनाने के लिए वजन (weights) को कैसे समायोजित किया जाए, न कि केवल व्यक्तिगत वाक्यों को।

3. "ओवरशूट" की गलती (The "Overshoot" Mistake)

लेखकों ने एक आश्चर्यजनक विचित्रता की खोज की कि यह "तापमान" कैसे काम करता है।

पारंपरिक भौतिकी में, जब आप सबसे अच्छा समाधान खोजने की कोशिश करते हैं, तो आप आमतौर पर बहुत गर्म (बहुत अनिर्णायक) से शुरू करते हैं, और फिर धीरे-धीरे ठंडा होते हैं।

पेपर का निष्कर्ष: इस विशिष्ट प्रकार के कंप्रेशन के लिए, "बहुत गर्म" से शुरू करना एक आपदा है। यदि लाइब्रेरियन शुरुआत में बहुत अधिक अनिर्णायक है, तो वह एक बुरे इलाके (एक "खराब वोरोनोई बेसिन"/worse Voronoi basin) में भटक जाता है और वहीं फंस जाता है। जब तक वह ठंडा होता है और सबसे अच्छा रास्ता खोजने की कोशिश करता है, तब तक वह पहले से ही एक खराब क्षेत्र में खो चुका होता है और वापस नहीं आ पाता।
समाधान: उन्होंने पाया कि एक मध्यम तापमान (न बहुत गर्म, न बहुत ठंडा) से शुरू करना सबसे अच्छा काम करता है। यह एक स्पष्ट मानचित्र के साथ हाइक शुरू करने जैसा है, बजाय इसके कि आप धुंध में अंधे होकर भटकें। "सुपर हॉट" चरण को छोड़कर, उन्होंने भटकने से बचने और एक बेहतर समाधान खोजने में सफलता प्राप्त की।

4. परिणाम: छोटे कंप्यूटरों के लिए जीत

टीम ने दो अलग-अलग प्रकार के मॉडल्स पर परीक्षण किया:

"प्रॉक्सी" टेस्ट (OLMoE): उन्होंने कंप्रेशन को अनुकूलित करने के लिए केवल यह देखने की कोशिश की कि मॉडल का प्रत्येक लेयर अपने स्वयं के डेटा को कितनी अच्छी तरह से पुनर्गठित (reconstruct) करता है (जैसे यह जांचना कि क्या एक फोटोकॉपी मूल के समान दिखती है)। परिणाम: यह विफल रहा। मॉडल पुराने तरीके की तुलना में वास्तव में बदतर हो गया। इसने उन्हें सिखाया कि केवल एक "बेहतर फोटोकॉपी" बनाना यह सुनिश्चित नहीं करता कि किताब बेहतर पढ़ी जाएगी।
"असली" टेस्ट (Llama-3.2): उन्होंने मॉडल को वास्तव में टेक्स्ट को बेहतर ढंग से पढ़ने और समझने के लिए अनुकूलित किया (डिस्टिलेशन नामक एक विधि का उपयोग करके, जहाँ एक स्मार्ट शिक्षक मॉडल छात्र का मार्गदर्शन करता है)। परिणाम: यह काम कर गया!
- मॉडल के एक विशिष्ट लेयर पर, उनकी नई विधि पुराने सर्वश्रेष्ठ तरीके को एक छोटे लेकिन महत्वपूर्ण अंतर से पीछे छोड़ देती है।
- जब उन्होंने इसे कई लेयर्स पर लागू किया, तो सुधार "सुपर-एडिटिव" (super-additive) तरीके से जुड़ गए (पूरा हिस्सा उसके हिस्सों के योग से भी बड़ा हो गया)।

5. इंजन जो इसके नीचे काम करता है (The Engine Under the Hood)

इस गणित को करना आमतौर पर बहुत धीमा होता है और इसके लिए महंगे सुपरकंप्यूटरों की आवश्यकता होती है। लेखकों ने एक विशेष "इंजन" (Triton kernel) भी बनाया है जो इस जटिल गणित को एक सिंगल कंज्यूमर ग्राफिक्स कार्ड (जैसे RTX 4080) पर अविश्वसनीय रूप से तेज़ चलाता है। उन्होंने इसे मानक तरीकों की तुलना में 6.5 गुना तेज़ बनाया है, जो यह साबित करता है कि इस उन्नत तकनीक को चलाने के लिए डेटा सेंटर की आवश्यकता नहीं है।

सारांश

यह पेपर AI मॉडल्स को उपभोक्ता उपकरणों (consumer devices) पर फिट होने के लिए छोटा करने का एक नया तरीका प्रस्तुत करता है।

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

लेखकों ने अपना कोड और प्रशिक्षित मॉडल्स जारी कर दिए हैं ताकि अन्य लोग भी इनका उपयोग कर सकें।

तकनीकी सारांश: BCJR-QAT

समस्या विवरण

कंज्यूमर हार्डवेयर पर लार्ज लैंग्वेज मॉडल्स (LLMs) के लिए 2 बिट्स प्रति वेट (bpw) या उससे कम प्राप्त करना एक महत्वपूर्ण बाधा है। जबकि QTIP जैसे पोस्ट-ट्रेनिंग क्वांटाइजेशन (PTQ) तरीकों ने ट्रेलिस-कोडेड क्वांटाइजेशन (TCQ) का उपयोग करके 2-बिट क्वांटाइजेशन के लिए एक मजबूत सीमा स्थापित की है, वे "PTQ सीलिंग" (PTQ ceiling) द्वारा सीमित हैं। इससे आगे बढ़ने के लिए क्वांटाइजेशन-अवेयर ट्रेनिंग (QAT) की आवश्यकता है। हालाँकि, मानक QAT तकनीकें ट्रेलिस कोड्स पर विफल हो जाती हैं क्योंकि विटर्बी डिकोडिंग एल्गोरिदम में एक्सपोनेंशियल रूप से कई पथों (paths) पर एक नॉन-डिफरेंशिएबल argmax ऑपरेशन शामिल होता है। ट्रेलिस पर स्ट्रेट-थ्रू एस्टिमेटर (STE) लागू करने से आसन्न वेट निर्णयों (adjacent weight decisions) के बीच का जुड़ाव नष्ट हो जाता है, जो कि स्केलर क्वांटाइजेशन की तुलना में ट्रेलिस कोडिंग का प्राथमिक लाभ है।

कार्यप्रणाली: BCJR-QAT

यह शोध पत्र BCJR-QAT पेश करता है, जो ट्रेलिस-कोडेड क्वांटाइजेशन का एक डिफरेंशिएबल रिलैक्सेशन है, जो हार्ड विटर्बी argmax को BCJR (Bahl-Cocke-Jelinek-Raviv) फॉरवर्ड-बैकवर्ड एल्गोरिदम से प्राप्त एक सॉफ्ट एक्सपेक्टेशन (soft expectation) से बदल देता है।

मुख्य तंत्र

सांख्यिकीय यांत्रिकी सादृश्य (Statistical Mechanics Analogy): ट्रेलिस को एक 1D इसिंग-जैसे स्पिन चेन (Ising-like spin chain) के समकक्ष एक फैक्टर ग्राफ के रूप में माना जाता है। विटर्बी argmax को पथों के ऊपर बोल्ट्ज़मैन वितरण (Boltzmann distribution) के ज़ीरो-टेम्परेचर लिमिट ( $T \to 0$ ) के रूप में देखा जाता है।
सॉफ्ट कोडवर्ड: एक परिमित तापमान $T > 0$ पर, यह विधि BCJR एल्गोरिदम का उपयोग करके प्रत्येक स्टेट के सटीक मार्जिनल पोस्टीरियर प्रोबेबिलिटीज (marginal posterior probabilities) की गणना करती है। क्वांटाइज्ड आउटपुट उत्सर्जित कोडवर्ड का बोल्ट्ज़मैन एक्सपेक्टेशन है:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
यह सॉफ्ट कोडवर्ड वेट प्री-इमेज और ट्रेलिस पैरामीटर्स के सापेक्ष सटीक रूप से डिफरेंशिएबल है।
एनीलिंग (Annealing): ट्रेनिंग के दौरान, तापमान $T$ को एक प्रारंभिक मान $T_0$ से लगभग शून्य $T_{end}$ तक एनेल किया जाता है। जैसे-जैसे $T \to 0$ होता है, सॉफ्ट कोडवर्ड हार्ड विटर्बी आउटपुट में अभिसरित (converge) हो जाता है, जो निरंतर अनुकूलन (continuous optimization) को एक असतत (discrete) समाधान में प्रभावी रूप से "क्रिस्टलाइज" कर देता है।

प्रमुख तकनीकी नवाचार

डिफरेंशिएबिलिटी: BCJR फॉरवर्ड-बैकवर्ड रिकर्सन (जटिलता $O(L \cdot S^2)$ ) ट्रेलिस के माध्यम से सटीक ग्रेडिएंट प्रदान करता है, जिससे इंटर-स्टेट कपलिंग सुरक्षित रहती है जिसे स्केलर STE विधियाँ खो देती हैं।
Triton कर्नेल: लेखकों ने एक फ्यूज्ड ट्राइटन (Triton) कर्नेल लागू किया है जो गेदर (gather), लॉग-सम-एक्सप (log-sum-exp), और एडिशन ऑपरेशंस को एक सिंगल लॉन्च में जोड़ता है। यह फॉरवर्ड/बैकवर्ड पास को सैकड़ों कर्नेल लॉन्च से घटाकर एक में बदल देता है, जिससे FP32 न्यूमेरिकल पैरिटी बनाए रखते हुए एक रेफरेंस पायटॉर्च ऑटोग्रेड कार्यान्वयन की तुलना में 6.57× स्पीडअप प्राप्त होता है।
शेड्यूल करेक्शन (Skip-High-T): शोध पत्र मानक सिम्युलेटेड एनीलिंग शेड्यूल्स (शुरुआत $T_0=1.0$ से) में एक महत्वपूर्ण विफलता मोड का निदान करता है। उच्च तापमान पर, सॉफ्ट कोडवर्ड ग्रेडिएंट "स्मियर" (smear) और अनइंफॉर्मेटिव हो जाता है, जिससे ऑप्टिमाइज़र कूलिंग चरण के ठीक होने से पहले ही एक खराब वोरोनोई बेसिन (Voronoi basin) में चला जाता है। प्रस्तावित समाधान हाई- $T$ चरण को स्किप करना है, यानी एक निचले तापमान (जैसे, $T_0=0.3$ ) पर शुरुआत करना जहाँ ग्रेडिएंट पहले से ही इंफॉर्मेटिव हों।

प्रयोगात्मक परिणाम

शोध पत्र दो आर्किटेक्चर पर BCJR-QAT का मूल्यांकन करता है: OLMoE-1B-7B (MoE) और Llama-3.2-1B (Dense)।

1. OLMoE-1B-7B (नकारात्मक परिणाम: प्रॉक्सी गैप)

सेटअप: प्रति-लेयर ग्रीडी QAT जो रिकंस्ट्रक्शन MSE (एक प्रॉक्सी ऑब्जेक्टिव) के लिए अनुकूलित है।
परिणाम: BCJR-QAT ने QTIP-PTQ की तुलना में प्रति-लेयर रिकंस्ट्रक्शन MSE को सफलतापूर्वक 17-20% कम किया। हालाँकि, यह सुधार एंड-टास्क परफॉरमेंस में परिवर्तित नहीं हुआ। मॉडल WikiText-2 पर QTIP-PTQ बेसलाइन की तुलना में $\sim 1.3$ पर्कप्लेक्सिटी (PPL) से पीछे रह गया।
निष्कर्ष: यह एक "प्रॉक्सी गैप" की पुष्टि करता है जहाँ स्थानीय रिकंस्ट्रक्शन एरर को कम करने का अर्थ बेहतर एंड-टास्क परफॉरमेंस की गारंटी नहीं है।

2. Llama-3.2-1B (सकारात्मक परिणाम: एंड-टू-एंड KL)

सेटअप: एक FP16 टीचर के विरुद्ध सिंगल-लेयर QAT जो एंड-टू-एंड फॉरवर्ड-KL डिस्टिलेशन के लिए अनुकूलित है।
परिणाम:
- शेड्यूल ओवरशूट: एक नेटिव हाई- $T$ शेड्यूल ( $T_0=1.0$ ) का उपयोग करने से परफॉरमेंस PTQ से थोड़ा खराब (+0.005 से +0.022 PPL) रहा।
- Skip-High-T फिक्स: सुधारा गया शेड्यूल ( $T_0=0.3$ ) उपयोग करने से ओवरशूट समाप्त हो गया। लेयर 4 पर, मॉडल ने 10.1347 PPL प्राप्त किया, जो QTIP-PTQ बेसलाइन (10.2189) की तुलना में -0.084 PPL का सुधार है।
- सुपर-एडिटिविटी (Super-Additivity): जब प्रशिक्षित लेयर्स (लेयर 4 और लेयर 8) को जोड़ा गया, तो संयुक्त लाभ (-0.077 PPL) व्यक्तिगत लाभों के योग (-0.062 PPL) से अधिक था, जो एक "कोऑपरेशन सरप्लस" को दर्शाता है जहाँ BCJR-प्रशिक्षित कोडवर्ड एक-दूसरे को सुदृढ़ करते हैं।

3. सैद्धांतिक विश्लेषण

ड्रिफ्ट-बजट थ्योरी: लेखक PTQ वोरोनोई बेसिन से बाहर निकलने की स्थिति को व्युत्पन्न करते हैं: संचयी वेट ड्रिफ्ट ( $\eta \cdot N \cdot g_{max}$ ) को बेसिन रेडियस ( $r_{Voronoi} \approx 10^{-3}$ ) से अधिक होना चाहिए। प्रयोग पुष्टि करते हैं कि जो शेड्यूल्स इस शर्त को पूरा नहीं करते वे PTQ समाधान से बाहर नहीं निकल पाते, जबकि जो इसे पार करते हैं वे बेहतर समाधान पा सकते हैं यदि ऑब्जेक्टिव और शेड्यूल सही हों।

महत्व और दावे

शोध पत्र तीन प्राथमिक योगदानों का दावा करता है:

एल्गोरिद्मिक: BCJR-QAT पहला तरीका है जो ट्रेलिस-कोडेड क्वांटाइजेशन को स्वाभाविक रूप से डिफरेंशिएबल बनाता है, जिससे पाथ कपलिंग (path coupling) का लाभ खोए बिना TCQ पर QAT संभव होता है।
इंजीनियरिंग: उनका फ्यूज्ड ट्राइटन कर्नेल इसे कंज्यूमर हार्डवेयर (सिंगल 12GB GPU) पर कम्प्यूटेशनल रूप से व्यवहार्य बनाता है, जिससे रिसर्च के लिए इंफरेंस-टाइम डिप्लॉयमेंट बाधा दूर हो जाती है।
अनुभवजन्य अंतर्दृष्टि (Empirical Insight):
- यह पहचानता है कि प्रति-लेयर रिकंस्ट्रक्शन MSE, 2-बिट QAT के लिए एक खराब प्रॉक्सी है, जबकि एंड-टू-एंड KL डिस्टिलेशन प्रभावी है।
- यह सिम्युलेटेड एनीलिंग के पारंपरिक ज्ञान को चुनौती देता है, यह दिखाते हुए कि डिस्क्रीट रिलैक्सेशन के लिए "हॉट" ( $T=1.0$ ) से शुरुआत करना BCJR-QAT के लिए हानिकारक है क्योंकि हाई-टेम्परेचर ग्रेडिएंट शोर-प्रधान (noise-dominated) होता है। इष्टतम रणनीति उस तापमान पर शुरू करना है जहाँ ग्रेडिएंट पहले से ही इंफॉर्मेटिव हो।

शोध पत्र निष्कर्ष निकालता है कि हालांकि सिंगल-लेयर परिणाम आशाजनक हैं, BCJR-QAT की पूर्ण क्षमता संभवतः सभी लेयर्स के माध्यम से संयुक्त एंड-टू-एंड ट्रेनिंग में निहित है, जो कि जारी किए गए कर्नेल द्वारा सुलभ बनाई गई है लेकिन 1B से बड़े मॉडल्स के लिए क्लाउड-स्केल हार्डवेयर (H100/H200) की आवश्यकता है। लेखक स्पष्ट रूप से कहते हैं कि उन्होंने इस कार्य में इस पूर्ण-स्तरीय संयुक्त प्रशिक्षण को नहीं किया है।