⚡ electrical engineering

AGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models

यह शोध पत्र JAGG (Jacobian-Aggregated Group Gradient) प्रस्तुत करता है, जो डिफ्यूजन मॉडल प्रक्षेपवक्रों (trajectories) के लगभग रैखिक परिवर्तन का लाभ उठाकर मध्यवर्ती जैकोबियन (intermediate Jacobians) को अनुमानित करने और ग्रुप ग्रेडिएंट्स को एकत्रित करने के लिए है, जिससे बिना किसी महत्वपूर्ण गुणवत्ता हानि के डिफ्यूजन ट्रांसफॉर्मर्स के लिए ग्रुप रिलेटिव पॉलिसी ऑप्टिमाइजेशन (GRPO) प्रशिक्षण की कम्प्यूटेशनल लागत को लगभग 2 $\times$ कम किया जा सकता है।

मूल लेखक: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

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

📖 4 मिनट में पढ़ें☕ कॉफ़ी ब्रेक में पढ़ें

CC BY 4.0

मूल लेखक: Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

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

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

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

यहीं पर JAGG नामक एक नया विचार आता है। इस शोध के पीछे के शोधकर्ताओं ने एक चतुर प्रश्न पूछा: "क्या हमें वास्तव में हर एक चाल के बाद रुकने और सब कुछ फिर से कैलकुलेट करने की आवश्यकता है?" उन्होंने देखा कि मूर्तिकला की प्रक्रिया के शुरुआती, अव्यवस्थित चरणों में, परिवर्तन बहुत सुचारू और अनुमानित तरीके से होते हैं—लगभग एक सीधी रेखा की तरह। यदि आप जानते हैं कि मूर्तिकला कहाँ से शुरू होती है और कुछ कदमों के बाद वह कहाँ पहुँचती है, तो आप वास्तव में भारी गणित किए बिना बीच में क्या हुआ, इसका अनुमान लगा सकते हैं।

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

शोधकर्ताओं ने गणितीय रूप से सिद्ध किया है कि यदि मूर्तिकला की प्रक्रिया पूरी तरह से सुचारू और रैखिक (linear) है, तो यह शॉर्टकट वास्तव में सटीक है—यह ठीक वही उत्तर देता है जो हर बार कठिन परिश्रम करने से मिलता है। वास्तविक दुनिया में, जहाँ चीजें पूरी तरह से सीधी नहीं होती हैं, उन्होंने पाया कि यह शॉर्टकट छवि निर्माण के "शोर वाले" शुरुआती चरणों के दौरान अविश्वसनीय रूप से अच्छा काम करता है। इस पद्धति का उपयोग करके, वे इन मॉडलों को प्रशिक्षित करने में लगने वाले समय को प्रति स्टेप लगभग 17% से 18% तक कम करने में सक्षम हुए।

सबसे अच्छी बात? चित्रों की गुणवत्ता पर कोई बुरा प्रभाव नहीं पड़ा। जब उन्होंने विभिन्न मॉडलों (जैसे Flux और QwenImage) पर JAGG का परीक्षण किया, तो उत्पन्न चित्र उतने ही अच्छे थे जितने कि धीमी, पारंपरिक विधि से बनाए गए चित्र। वास्तव में, कुछ मामलों में, इस शॉर्टकट ने मॉडल को अधिक स्थिरता से सीखने में भी मदद की। शोध पत्र दिखाता है कि भारी काम कब करना है, इस बारे में थोड़ा चतुर होकर, हम अंतिम कलाकृति की सुंदरता से समझौता किए बिना, इन शक्तिशाली एआई कलाकारों को बहुत तेज़ और सस्ता बना सकते हैं।

तकनीकी सारांश: JAGG (Jacobian-Aggregated Group Gradient)

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

जेनेरेटिव मॉडल्स को मानवीय प्राथमिकताओं के साथ संरेखित करने के लिए ग्रुप रिलेटिव पॉलिसी ऑप्टिमाइज़ेशन (GRPO) एक अग्रणी सुदृढीकरण शिक्षण (Reinforcement Learning - RL) एल्गोरिदम के रूपas उभरा है, जिसने सफलतापूर्वक लार्ज लैंग्वेज मॉडल्स (LLMs) के लिए एक अलग वैल्यू नेटवर्क की आवश्यकता को समाप्त कर दिया है। हालाँकि, डिफ्यूजन और फ्लो मैचिंग मॉडल्स (विशेष रूप से डिफ्यूजन ट्रांसफॉर्मर या DiTs) तक GRPO का विस्तार करना एक गंभीर कम्प्यूटेशनल बाधा उत्पन्न करता है।

डिफ्यूजन-आधारित GRPO में, जनरेशन प्रक्रिया को $W$ टाइमस्टेप्स के सैंपलिंग प्रक्षेपवक्र (trajectory) पर अनुकूलित किए जाने वाले एक मल्टी-स्टेप निर्णय समस्या के रूप में माना जाता है। मानक प्रशिक्षण के लिए पॉलिसी ग्रेडिएंट की गणना करने हेतु प्रक्षेपवक्र के प्रत्येक सिंगल टाइमस्टेप के लिए हाई-कैपेसिटी DiT बैकबोन के माध्यम से एक पूर्ण फॉरवर्ड और बैकवर्ड पास की आवश्यकता होती है। हाई-रेज़ोल्यूशन टेक्स्ट-टू-इमेज (T2I) जनरेशन के लिए, जहाँ प्रक्षेपवक्र अक्सर दर्जनों स्टेप्स तक फैले होते हैं, यह अत्यधिक खर्चीला प्रशिक्षण लागत का परिणाम देता है। लैंग्वेज मॉडल्स में टोकन स्टेप्स हल्के होते हैं, जबकि आधुनिक DiT के माध्यम से प्रत्येक बैकवर्ड पास एक मॉडर्न ट्रांसफॉर्मर फॉरवर्ड पास की पूरी लागत वहन करता है, जिससे प्रति-स्टेप बैकवर्ड की आवश्यकता प्रमुख प्रशिक्षण व्यय बन जाती है।

जबकि ट्रेनिंग-फ्री इन्फरेंस एक्सेलेरेशन विधियों (जैसे $\Delta$ -DiT, ScalingCache) ने इस अवलोकन का सफलतापूर्वक लाभ उठाया है कि DiT हिडन स्टेट्स और वेलोसिटी प्रेडिक्शन्स डिनोइजिंग प्रक्षेपवपथ के दौरान सुचारू रूप से और लगभग रैखिक रूप से (linearly) बदलते हैं, यह एक खुला प्रश्न बना हुआ था कि क्या इसी रैखिकता (linearity) का उपयोग RL प्रशिक्षण में बैकवर्ड-पास लागत को कम करने के लिए किया जा सकता है।

कार्यप्रणाली: JAGG

लेखक JAGG (Jacobian-Aggregated Group Gradient) प्रस्तावित करते हैं, जो एक गणितीय रूप से सिद्धांतबद्ध त्वरण विधि है जो $W$ निरंतर डिनोइजिंग स्टेप्स के समूह के लिए बैकवर्ड पास की संख्या को $W$ से घटाकर 2 कर देती है।

मुख्य अंतर्दृष्टि (Core Insight)

यह विधि इस अंतर्दृष्टि पर आधारित है कि मॉडल पैरामीटर्स $\theta$ के सापेक्ष डिनोइजिंग वेलोसिटी फील्ड $v_\theta(z, t)$ का जैकोबियन (Jacobian), इनपुट्स $(z, t)$ के प्रति अनुमानित रैखिकता प्रदर्शित करता है, विशेष रूप से डिनोइजिंग प्रक्रिया के हाई-नॉइज़ रिजीम के दौरान। यह एंडपॉइंट जैकोबियन्स ( $J_0$ और $J_{W-1}$ ) का उपयोग करके लीनियर इंटरपोलेशन के माध्यम से मध्यवर्ती-स्टेप जैकोबियन्स का सटीक अनुमान लगाने की अनुमति देता है।

एल्गोरिदम वर्कफ़्लो (Algorithmic Workflow)

JAGG चार चरणों के माध्यम से $W$ टाइमस्टेप्स को प्रोसेस करता है:

फॉरवर्ड पासेस (Forward Passes): मॉडल सभी $W$ स्टेप्स के लिए फॉरवर्ड पास करता है। हालाँकि, ग्रेडिएंट ट्रैकिंग केवल दो एंडपॉइंट्स (स्टेप 0 और स्टेप $W-1$ ) के लिए सक्षम की जाती है। मेमोरी और कंप्यूटेशन बचाने के लिए मध्यवर्ती $W-2$ स्टेप्स को torch.no_grad() का उपयोग करके निष्पादित किया जाता है।
अपस्ट्रीम ग्रेडिएंट कंप्यूटेशन (Upstream Gradient Computation): प्रत्येक स्टेप $j$ के लिए, मॉडल आउटपुट का एक डिटैच्ड प्रॉक्सी बनाया जाता है। प्रति-स्टेप लॉस $L_j$ (जो PPO-स्टाइल ऑब्जेक्टिव से प्राप्त होता है) की गणना की जाती है, और ऑटोमैटिक डिफरेंशिएशन के माध्यम से अपस्ट्रीम ग्रेडिएंट $s_j = \partial L_j / \partial \hat{v}_j$ की गणना की जाती है। महत्वपूर्ण रूप से, इस चरण में कोई ट्रांसफॉर्मर बैकवर्ड पास नहीं होता है; ग्रेडिएंट केवल हल्के SDE स्टेप और PPO ऑब्जेक्टिव के माध्यम से प्रवाहित होता है।
टाइमस्टेप-वेटेड एग्रीगेशन (Timestep-Weighted Aggregation): अपस्ट्रीम ग्रेडिएंट्स $s_j$ को टाइमस्टेप $t_j$ से प्राप्त इंटरपोलेशन वेट्स $\alpha_j$ का उपयोग करके दो कंपोजिट सिग्नल्स, $s_{base}$ और $s_{corr}$ में एकत्रित किया जाता है:
$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$
$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$
यह एकत्रीकरण लीनियर इंटरपोलेशन धारणा का उपयोग करके पूर्ण ग्रेडिएंट सिग्नल को प्रभावी रूप से पुनर्गठित करता है।
जॉइंट बैकवर्ड (Joint Backward): दो एंडपॉइंट आउटपुट्स $[\hat{v}_0, \hat{v}_{W-1}]$ पर एकत्रित ग्रेडिएंट्स $[s_{base}, s_{corr}]$ के साथ एक एकल torch.autograd.backward कॉल जारी की जाती है। यह ठीक एक FSDP ऑल-गैदर को ट्रिगर करता है, जो $W$ अलग-अलग बैकवर्ड पासेस को बदल देता है जिनकी मानक प्रशिक्षण में आवश्यकता होती है।

सैद्धांतिक गारंटी (Theoretical Guarantee)

लेखक एक औपचारिक प्रमाण (थ्योरम 1) प्रदान करते हैं जिसमें कहा गया है कि यदि वेलोसिटी फील्ड $v_\theta(z, t)$ $(z, t)$ में पूरी तरह से रैखिक है, तो $t$ -वेटेड जैकोबियन इंटरपोलेशन सटीक (exact) है, जिसमें शून्य एप्रोक्सिमेशन एरर है। इन शर्तों के तहत, दो-एंडपॉइंट बैकवर्ड पास सटीक ग्रुप ग्रेडिएंट को रिकवर करता है।

एडैप्टिव रूटिंग ( $jagg\_frac$ )

यह पहचानते हुए कि डिनोइजिंग के निम्न-नॉइज़ (विलंबित) चरणों में रैखिकता की धारणा कम हो जाती है, लेखक एक रूटिंग पैरामीटर $jagg\_frac$ पेश करते हैं। यह पैरामीटर JAGG को केवल डिनोइजिंग समूहों के पहले अंश (जहाँ हाई-नॉइज़ रैखिकता धारणा बनी रहती है, जिसे ग्रेडिएंट्स की कोसाइन समानता के माध्यम से अनुभवजन्य रूप से सत्यापित किया गया है) पर लागू करता है और शेष समूहों के लिए सटीक प्रति-स्टेप बैकवर्ड पास का उपयोग करता है। परीक्षण किए गए T2I मॉडल्स के लिए, $jagg\_frac$ को 0.6 पर सेट किया गया था।

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

बॉटलनेक की पहचान: पेपर ने सैंपलिंग प्रक्षेपवक्र के माध्यम से प्रति-टाइमस्टेप बैकवर्ड पास के कारण डिफ्यूजन GRPO प्रशिक्षण में कम्प्यूटेशनल बॉटलनेक को औपचारिक रूप से पहचाना और लक्षणबद्ध किया है।
JAGG फ्रेमवर्क: यह एक जैकोबियन-एग्रीगेशन फ्रेमवर्क प्रस्तावित करता है जो वेलोसिटी जैकोबियन रैखिकता के कठोर सैद्धांतिक विश्लेषण पर आधारित, बैकवर्ड पास को $W$ से घटाकर 2 कर देता है।
सैद्धांतिक प्रमाण: लेखक सिद्ध करते हैं कि रैखिक वेलोसिटी स्थितियों के तहत, प्रस्तावित इंटरपोलेशन सटीक है, जो एक ठोस सैद्धांतिक आधार प्रदान करता है।
अनुभवजन्य सत्यापन (Empirical Validation): कई डिफ्यूजन बैकबोन (Flux, QwenImage) और RL एल्गोरिदम (DanceGRPO, Flow-GRPO-Fast) पर व्यापक प्रयोग JAGG की सामान्यता और प्रभावशीलता को प्रदर्शित करते हैं।

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

लेखकों ने HPSv2.1, PickScore और CLIPScore जैसे मेट्रिक्स का उपयोग करके T2I बेंचमार्क पर JAGG का मूल्यांकन किया।

दक्षता (Efficiency): JAGG ने मानक DanceGRPO सेटअप पर प्रति प्रशिक्षण स्टेप 17.3% से 18.2% की निरंतर वॉल-क्लॉक स्पीडअप हासिल की (QwenImage पर समय ~835s से घटाकर ~691s और Flux पर ~279s से घटाकर ~228s)। जब इसे Flow-GRPO-Fast पर लागू किया गया, तो अपडेट स्टेज में ही 29.9% से 30.0% की स्पीडअप देखी गई।
गुणवत्ता (Quality): इस पद्धति ने जनरेशन की गुणवत्ता में नगण्य गिरावट के साथ ये स्पीडअप दिए। कई मामलों में, JAGG ने बेसलाइन रिवॉर्ड कर्व्स को करीब से ट्रैक किया और कन्वर्जेंस पर बेसलाइन के बराबर या उससे थोड़ा बेहतर प्रदर्शन किया।
एब्लेशन (Ablation): एक एब्लेशन वेरिएंट जिसने एंडपॉइंट्स पर बैकवर्ड पास किया लेकिन बिना जैकोबियन इंटरपोलेशन के (एंडपॉइंट्स को स्वतंत्र रूप से पूर्ण लॉस असाइन किया), लगातार खराब प्रदर्शन करता रहा। यह पुष्टि करता है कि जैकोबियन इंटरपोलेशन के माध्यम से मध्यवर्ती ग्रेडिएंट सूचना का एकत्रीकरण स्थिर प्रशिक्षण के लिए महत्वपूर्ण है।
ग्रेडिएंट गुणवत्ता: अनुभवजन्य विश्लेषण ने दिखाया कि हाई-नॉइज़ समूहों (0–2) में, JAGG-एकत्रित ग्रेडिएंट और सटीक ग्रेडिएंट के बीच कोसाइन समानता उच्च (0.74–0.85) थी, जो रैखिकता धारणा को मान्य करती है। बाद के समूहों में समानता कम हो गई, जो $jagg\_frac$ रूटिंग रणनीति को उचित ठहराती है।

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

पेपर का दावा है कि JAGG बड़े पैमाने पर डिफ्यूजन GRPO प्रशिक्षण को व्यावहारिक बनाने की दिशा में एक महत्वपूर्ण कदम है। वेलोसिटी जैकोबियन्स की हाई-नॉइज़ रिजीम में निकट-रैखिकता का लाभ उठाकर, यह विधि जनरेशन की गुणवत्ता से समझौता किए बिना RL प्रशिक्षण के कम्प्यूटेशनल बोझ को कम करती है।

लेखक इस बात पर जोर देते हैं कि JAGG केवल एक ह्यूरिस्टिक शॉर्टकट नहीं है बल्कि डिनोइजिंग प्रक्षेपवक्र के गणितीय गुणों पर आधारित है। इस पद्धति को मौजूदा इन्फरेंस एक्सेलेरेशन तकनीकों के ऑर्थोगोनल (orthogonal) के रूप में प्रस्तुत किया गया है, जो विशेष रूप से डिफ्यूजन RL में अद्वितीय बैकवर्ड-पास बॉटलनेक को लक्षित करती है। परिणाम बताते हैं कि JAGG प्रशिक्षण लागत की बाधा को काफी कम करके विजुअल जेनेरेटिव मॉडल्स के लिए RL-आधारित अलाइनमेंट के व्यापक प्रसार को सक्षम बनाता है।