💻 computer science

SAFE-DiT: Semantics-Aware Fast-path Execution for High-Resolution Diffusion Transformers

SAFE-DiT एक प्रशिक्षण-मुक्त (training-free) फ्रेमवर्क है जो अनावश्यक अटेंशन मास्क को हटाने और प्रॉम्प्ट-कंडीशन्ड टोकन विभाजन के उपयोग के माध्यम से "मास्क-इंड्यूस्ड डिस्पैच टैक्स" को समाप्त करके हाई-रेज़ोल्यूशन डिफ्यूजन ट्रांसफॉर्मर इन्फरेंस को तेज़ करता है, जिससे विजुअल क्वालिटी से समझौता किए बिना 5.09× तक की गति वृद्धि और काफी कम मेमोरी उपयोग प्राप्त होता है।

मूल लेखक: Xuanhua Yin, Yuxuan Jia, Chuanzhi Xu, Weidong Cai

प्रकाशित 2026-06-30

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

CC BY 4.0

मूल लेखक: Xuanhua Yin, Yuxuan Jia, Chuanzhi Xu, Weidong Cai

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

कल्पना कीजिए कि आप एक विशाल ऑर्केस्ट्रा के कंडक्टर हैं जो एक विशाल, हाई-रिज़ॉल्यूशन मास्टरपीस (उत्कृष्ट कृति) बनाने की कोशिश कर रहे हैं। संगीतकार "टोकन" (छवि के छोटे टुकड़े) हैं, और कंडक्टर एक AI मॉडल (डिफ्यूजन ट्रांसफॉर्मर) है।

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

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

SAFE-DiT एक नया सिस्टम है जो इसे ठीक करने के लिए एक स्मार्ट, कुशल कंडक्टर के रूप में कार्य करता है जो बिल्कुल जानता है कि कब कागजी कार्रवाई को छोड़ना है और ऑर्केस्ट्रा की ऊर्जा पर कैसे ध्यान केंद्रित करना है।

यह कैसे काम करता है, इसे सरल अवधारणाओं में यहाँ दिया गया है:

1. "लाल फीताशाही" की समस्या (मास्क-इंड्यूस्ड डिस्पैच टैक्स - MIDT)

यह पेपर एक छिपे हुए बॉटलनेक (रुकावट) की पहचान करता है जिसे MIDT कहा जाता है।

उपमा: एक कॉन्सर्ट में सुरक्षा गार्ड की कल्पना करें जो हर किसी का टिकट चेक करता है। भले ही टिकट पर "ओपन एडमिशन" (यानी सभी प्रवेश कर सकते हैं) लिखा हो, फिर भी गार्ड भीड़ को रोकने के लिए उसे स्कैन करता है। यह पूरी भीड़ को धीमा कर देता है।
वास्तविकता: AI में, "गार्ड" वह कंप्यूटर कोड है जो "मास्क" की जांच करता है। यदि मास्क कहता है कि "सभी को अनुमति है," तो कोड अभी भी इसे जांचने के लिए एक धीमा, जटिल रास्ता अपनाता है। SAFE-DiT महसूस करता है कि यदि नियम यह है कि "सभी को अनुमति है," तो यह नियम पुस्तिका को फेंक सकता है और संगीतकारों को तुरंत खेलने दे सकता है। यह "लाल फीताशाही" को हटा देता है और चीजों को काफी तेज कर देता है।

2. "स्मार्ट फोकस" रणनीति (SAFE-Core)

हर एक क्षण में हर संगीतकार से उनके गाने का हिस्सा अपडेट करने के लिए कहने के बजाय, SAFE-DiT एक रणनीति का उपयोग करता है जिसे प्रॉम्प्ट-कंडीशन्ड सेंसिटिविटी पार्टीशनिंग कहा जाता है।

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

3. "रिफ्रेश ब्रेक" (कॉन्टेक्स्ट एंकर रिफ्रेश)

यदि आप बहुत लंबे समय तक केवल महत्वपूर्ण हिस्सों को अपडेट करते हैं, तो बैकग्राउंड अजीब दिख सकता है या मूल योजना से भटक सकता है।

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

4. "वॉल्यूम कंट्रोल" (SW-CFG)

अंत में, सिस्टम छवि के विभिन्न हिस्सों के लिए निर्देशों का कितनी मजबूती से पालन करना है, इसे एडजस्ट करता है।

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

परिणाम: उन्होंने क्या हासिल किया?

पेपर ने इसे एक बहुत शक्तिशाली AI Lumina-Next पर टेस्ट किया और पाया:

गति: यह 1024x1024 रिज़ॉल्यूशन पर 2.7 गुना तेज़ है और 2560x2560 रिज़ॉल्यूशन पर मानक विधि की तुलना में 5 गुना तेज़ है।
मेमोरी: यह बहुत कम कंप्यूटर मेमोरी का उपयोग करता है। वास्तव में, मानक विधि 3072x3072 इमेज बनाने की कोशिश करते समय क्रैश (मेमोरी खत्म) हो जाती है, लेकिन SAFE-DiT इसे आसानी से कर सकता है।
गुणवत्ता: छवियां मानक धीमी विधि जितनी ही अच्छी दिखती हैं। ब्लाइंड टेस्ट में, इंसान अंतर नहीं बता सके, और AI के अपने स्कोरिंग सिस्टम ने भी दिखाया कि छवियां उतनी ही अच्छी थीं।

सारांश

SAFE-DiT एक "ट्रेनिंग-फ्री" अपग्रेड है। इसे AI को कुछ भी नया सिखाने की आवश्यकता नहीं है। इसके बजाय, यह केवल इस बात को बदलता है कि AI शो कैसे चलाता है:

यह अनावश्यक "नियम पुस्तिका जांच" (मास्क) को हटा देता है जो चीजों को धीमा करती है।
यह अपनी ऊर्जा केवल छवि के उन हिस्सों पर केंद्रित करता है जिन्हें ध्यान देने की आवश्यकता है।
यह सब कुछ सटीक रखने के लिए समय-समय पर "रियलिटी चेक" लेता है।

परिणाम यह है कि आप बहुत बड़ी, उच्च-गुणवत्ता वाली छवियां बहुत तेज़ी से और उन कंप्यूटरों पर बना सकते हैं जो पहले इसे संभालने में सक्षम नहीं थे, और वह भी बिना किसी दृश्य गुणवत्ता को खोए।

तकनीकी सारांश: SAFE-DiT

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

डिफ्यूजन ट्रांसफॉर्मर्स (DiTs) में उच्च-रिज़ॉल्यूशन इन्फरेंस (inference) एक दोहरी चुनौती का सामना करता है: इमेज जनरेशन में पर्याप्त स्थानिक अतिरेक (spatial redundancy) और इस अतिरेक को संभालने के तरीके के कारण उत्पन्न होने वाला एक महत्वपूर्ण सिस्टम बॉटलनेक। जबकि स्थानिक रूप से अनुकूलित (spatially adaptive) विधियाँ विशिष्ट क्षेत्रों पर ध्यान केंद्रित करके गणना को कम करने का प्रयास करती हैं, मौजूदा कार्यान्वयन अक्सर इन निर्णयों को अटेंशन मास्क (attention masks) के रूप में एनकोड करते हैं।

लेखक एक विशिष्ट सिस्टम अक्षमता की पहचान करते हैं जिसे मास्क-इंड्यूस्ड डिस्पैच टैक्स (Mask-Induced Dispatch Tax - MIDT) कहा जाता है। आधुनिक इन्फरेंस स्टैक (जैसे PyTorch SDPA) इनपुट बाधाओं के आधार पर निष्पादन बैकएंड (execution backends) का चयन करते हैं। जब एक गणितीय रूप से रेडंडेंट मास्क (जैसे, एक 'ऑल-वैलिड' इमेज सेल्फ-अटेंशन मास्क) को अटेंशन ऑपरेटर में पास किया जाता है, तो यह अनजाने में सिस्टम को एक धीमे, नॉन-फ्लैश-अटेंशन (non-FlashAttention) पाथ पर डिस्पैच करने के लिए मजबूर कर सकता है। यह टैक्स लेटेंट सीक्वेंस लंबाई के साथ द्विघातीय (quadratically) रूप से बढ़ता है, जिससे स्थानिक अनुकूलन के सैद्धांतिक लाभ कम हो जाते हैं। इसके अलावा, सिमेंटिक बाधाओं (जैसे टेक्स्ट पैडिंग) को रूटिंग मेटाडेटा के साथ एक ही मास्क इंटरफेस में मिलाने से मॉडल के अर्थ (semantics) को तोड़े बिना रेडंडेंट ऑपरेशन्स को हटाना जटिल हो जाता है।

कार्यप्रणाली (Methodology)

SAFE-DiT (सेमेंटिक्स-अवेयर फास्ट-पाथ एक्जीक्यूशन फॉर डिफ्यूजन ट्रांसफॉर्मर्स) इसे सटीक मास्क एलिजन (exact mask elision) को अनुमान-आधारित स्थानिक शेड्यूलिंग (approximation-based spatial scheduling) से अलग करके संबोधित करता है। यह फ्रेमवर्क चार मुख्य घटकों के माध्यम से कार्य करता है:

1. सटीक सेमेंटिक्स-अवेयर फास्ट-पाथ रीराइट (Exact Semantics-Aware Fast-path Rewrite)

सिस्टम मास्क हटाने के लिए एक प्रमाणित मानदंड प्रस्तुत करता है। एक एडिटिव मास्क $m$ तब हटाया जा सकता है योग्य है यदि वह अटेंशन लॉजिट्स ( $L_{i,:} + c_i$ ) में केवल एक रो-वाइज कांस्टेंट शिफ्ट (row-wise constant shift) उत्पन्न करता है। सॉफ्टमैक्स फंक्शन के गुणों के कारण, ऐसा शिफ्ट नॉर्मलाइजेशन के दौरान रद्द हो जाता है, जिससे अटेंशन डिस्ट्रीब्यूशन अपरिवर्तित रहता है।

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

2. प्रॉम्प्ट-कंडीशनड सेंसिटिविटी पार्टीशनिंग (Prompt-Conditioned Sensitivity Partitioning - PCSP)

सेल्फ-अटेंशन मास्क पर निर्भर रहने के बजाय, स्थानिक अनुकूलन को सक्षम करने के लिए, SAFE-DiT शुरुआती इमेज-टू-टेक्स्ट क्रॉस-अटेंशन रिस्पॉन्स का उपयोग करके उन क्षेत्रों की पहचान करता है जिन्हें बार-बार अपडेट की आवश्यकता होती है।

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

3. सेंसिटिव-रीजन स्टेट अपडेट (SRSU) और कॉन्टेक्स्ट एंकर रिफ्रेश (CAR)

स्थानिक अनुकूलन को स्टेट शेड्यूलिंग के माध्यम से निष्पादित किया जाता है, न कि मास्किंग के माध्यम से:

SRSU: नॉन-एंकर स्टेप्स पर, मॉडल केवल सेंसिटिव सेट ( $S$ ) के क्वेरी रोज़ को अपडेट करता है। महत्वपूर्ण रूप से, ये क्वेरीज़ अभी भी ग्लोबल सेट के कीज़ (keys) और वैल्यूज़ (values) (वर्तमान सेंसिटिव स्टेट्स और कैश किए गए कॉन्टेक्स्ट स्टेट्स दोनों) पर ध्यान केंद्रित करती हैं, जिससे ग्लोबल कॉन्टेक्स्ट बना रहता है।
CAR: पुराने कॉन्टेक्स से होने वाली त्रुटि के संचय को रोकने के लिए, मॉडल आवधिक एंकर स्टेप्स पर एक "डेंस" अपडेट करता है, जो सभी टोकन्स के लिए कैश किए गए कॉन्टेक्स्ट स्टेट्स को रिफ्रेश करता है।

4. सेंसिटिविटी-वेटेड क्लासिफायर-फ्री गाइडेंस (SW-CFG)

यह फ्रेमवर्क एक वैकल्पिक ऑपरेटिंग पॉइंट प्रदान करता है जो स्थानिक रूप से गाइडेंस स्ट्रेंथ को एडजस्ट करता है।

तंत्र: PCSP से प्राप्त सेंसिटिविटी स्कोर का उपयोग करके, गाइडेंस स्केल $\gamma$ को प्रति टोकन मॉड्यूलेट किया जाता है ( $\gamma_i = \gamma_{ctx} + (\gamma_{sens} - \gamma_{ctx})s_i$ )।
अंतर: यह मॉड्यूल फॉरवर्ड पास के बाद प्रेडिक्शन आउटपुट पर काम करता है, जिसका अर्थ है कि यह स्थानिक सेल्फ-अटेंशन मास्क को फिर से पेश नहीं करता है।

लेखक SAFE-Core को सटीक मास्क एलिजन और ग्लोबल गाइडेंस के साथ SRSU/CAR का उपयोग करने वाले कॉन्फ़िगरेशन के रूप में परिभाषित करते हैं, और SAFE-DiT+SW को सेंसिटिविटी-वेटेड गाइडेंस मॉड्यूल सहित कॉन्फ़िगरेशन के रूप में परिभाषित करते हैं।

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

MIDT की पहचान: पेपर 'मास्क-इंड्यूस्ड डिस्पैच टैक्स' को कैरेक्टराइज करता है, यह प्रदर्शित करते हुए कि रेडंडेंट मास्क लंबे सीक्वेंस पर प्रदर्शन को 4.1–5.8× तक कम कर सकते हैं क्योंकि वे सब-ऑप्टिमल कर्नेल डिस्पैच को मजबूर करते हैं।
सटीक मास्क एलिजन मानदंड: एक औपचारिक प्रमाण कि एडिटिव मास्क जो रो-वाइज कांस्टेंट शिफ्ट उत्पन्न करते हैं, उन्हें नॉर्मलाइज्ड अटेंशन को बदले बिना हटाया जा सकता है, जो इसे अनुमानित विधियों से अलग करता है।
SAFE-DiT फ्रेमवर्क: एक ट्रेनिंग-फ्री आर्किटेक्चर जो सटीक निष्पादन रीराइट को मास्क-मुक्त स्थानिक शेड्यूलिंग (PCSP, SRSU, CAR) और वैकल्पिक प्रॉम्प्ट-अलाइनमेंट गाइडेंस (SW-CFG) के साथ जोड़ता है।
अलग मूल्यांकन (Decoupled Evaluation): SAFE-Core और SAFE-DiT+SW को अलग-अलग रिपोर्ट करके, यह कार्य त्वरण तंत्र की शुद्धता को गाइडेंस मॉड्यूल के प्रभावों से अलग करता है।

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

मुख्य रूप से Lumina-Next और अन्य DiT बैकबोन (SD3, FLUX.1, PixArt-Σ) पर PyTorch SDPA का उपयोग करके मूल्यांकन किया गया:

गति और मेमोरी:
- 1024² रिज़ॉल्यूशन पर, SAFE-DiT 2.69× एंड-टू-एंड त्वरण प्राप्त करता है।
- 2560² पर, त्वरण 5.09× तक पहुँच जाता है।
- 2560² पर पीक मेमोरी (Peak memory) 'डेंस' (Dense) के 94.1 GB से घटकर 27.9 GB हो जाती है।
- SAFE-DiT 3072² पर जनरेशन सक्षम बनाता है, एक ऐसा रिज़ॉल्यूशन जहाँ डेंस इन्फरेंस मेमोरी की कमी (OOM) के कारण विफल हो जाता है।
गुणवत्ता और फिडेलिटी (Quality and Fidelity):
- SAFE-Core डेंस फास्ट-पाथ बेसलाइन के साथ विजुअल नॉन-इन्फेरियोरिटी (visual non-inferiority) दिखाता है, जिसमें पेयर्ड मेट्रिक्स (ImageReward, HPSv2, LPIPS) सांख्यिकीय रूप से समान या थोड़े बेहतर हैं।
- SAFE-DiT+SW एक अलग ऑपरेटिंग पॉइंट प्रदान करता है जो बेहतर प्रॉम्प्ट अलाइनमेंट (Dense के ऊपर ImageReward +0.043) प्रदान करता है, बिना स्थानिक मास्क को फिर से पेश किए।
- एब्लेशन स्टडीज पुष्टि करती हैं कि मास्क एलिजन उच्च रिज़ॉल्यूशन पर गति लाभ का प्राथमिक चालक है, जबकि SRSU/CAR वास्तविक शेड्यूलिंग लाभ प्रदान करता है।
क्रॉस-बैकबोन जनरलाइजेशन:
- स्पीडअप MIDT बॉटलनेक की उपस्थिति के साथ सह-संबंधित हैं। मास्क वाले बैकबोन (Lumina-Next) कंपाउंडिंग लाभ देखते हैं (2.69× से 5.1×), जबकि मास्क-मुक्त बैकबोन (SD3, FLUX) मुख्य रूप से शेड्यूलिंग से लाभ देखते हैं (मॉडल की स्टेप-हैवी प्रकृति के आधार पर 1.3× से 5.14× तक)।

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

पेपर दावा करता है कि SAFE-DiT स्थानिक अनुकूलन को विशुद्ध रूप से एल्गोरिथमिक के बजाय एक सेमेंटिक्स-अवेयर एक्जीक्यूशन समस्या के रूप में पुनर्गठित करता है। इसका महत्व यहाँ है:

सिस्टम-लेवल ऑप्टिमाइजेशन: यह दर्शाता है कि "फंक्शनली रेडंडेंट" मास्क को हटाना स्थानिक अनुकूलन की दक्षता को अनलॉक करने के लिए एक आवश्यक कदम है, जो सब-ऑप्टिमल कर्नेल डिस्पैच के "टैक्स" को रोकता है।
सटीकता बनाम अनुमान (Exactness vs. Approximation): यह सटीक रीराइट (मास्क एलिजन) को अनुमानित शेड्यूलिंग (SRSU/CAR) से अलग करता है, यह सुनिश्चित करता है कि त्वरण की फिडेलिटी टोकन को छोड़ने की अनुमानित त्रुटि के साथ उलझी हुई नहीं है।
स्केलेबिलिटी: मेमोरी प्रेशर और लेटेंसी को कम करके, यह उच्च-रिज़ॉल्यूशन जनरेशन (3072² तक) को सक्षम बनाता है जो वर्तमान में डेंस इन्फरेंस के साथ असंभव है।

लेखक अपने गाइडेंस मॉड्यूल पर एक मध्यम रुख बनाए रखते हैं, यह नोट करते हुए कि हालांकि SW-CFG प्रॉम्प्ट अलाइनमेंट में सुधार करता है, मूल त्वरण (SAFE-Core) जनरेशन ट्राजेक्टरी को बदले बिना बैकबोन के मूल व्यवहार को सुरक्षित रखता है। परिणाम इस सिद्धांत का समर्थन करते हैं कि स्थानिक अनुकूलन को अटेंशन मास्क पर निर्भर रहने के बजाय सेमेंटिक्स-प्रिजर्विंग शेड्यूलिंग का उपयोग करना चाहिए, जो सिस्टम-लेवल पेनल्टी लागू करते हैं।