Training Flow Matching: The Role of Weighting and Parameterization

यहाँ शोध पत्र "Training Flow Matching: The Role of Weighting and Parameterization" का विस्तृत तकनीकी सारांश दिया गया है।

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

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

1. लॉस वेटिंग (Loss Weighting): विभिन्न शोर स्तरों (टाइम स्टेप्स $t$) पर प्रशिक्षण लॉस को कैसे वेट किया जाना चाहिए?
2. पैरामीट्राइजेशन (Parameterization): न्यूरल नेटवर्क को क्या प्रेडिक्ट (predict) करना चाहिए? सामान्य विकल्पों में क्लीन इमेज ($x$-prediction/denoiser), नॉइज़ ($\epsilon$-prediction), या वेलोसिटी ($v$-prediction) का प्रेडिक्शन शामिल है।

यद्यपि अनुभवजन्य ह्यूरिस्टिक्स (empirical heuristics) मौजूद हैं, लेकिन इस बात की समझ का अभाव है कि ये विकल्प मॉडल आर्किटेक्चर (जैसे, U-Net बनाम ViT), डेटा डाइमेंशनलिटी और डेटासेट आकार के साथ कैसे परस्पर क्रिया करते हैं। हालिया साहित्य विरोधाभासी निष्कर्ष प्रस्तुत करता है (जैसे, Li & He, 2025 सुझाव देते हैं कि मैनिफोल्ड धारणाओं के कारण $x$-prediction बेहतर है, जबकि शुरुआती FM कार्य $v$-prediction का पक्ष लेते हैं)। यह शोध पत्र इन कारकों को अलग करने और व्यावहारिक डिज़ाइन दिशानिर्देश प्रदान करने का लक्ष्य रखता है।

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

लेखक विभिन्न प्रशिक्षण उद्देश्यों की तुलना करने के लिए एक एकीकृत डिनोइजिंग फ्रेमवर्क (unified denoising framework) प्रस्तावित करते हैं।

• एकीकृत सूत्रीकरण (Unified Formulation): वे सभी प्रशिक्षण उद्देश्यों को एक भारित डिनोइजिंग समस्या के रूप में व्यक्त करते हैं:
  $$\min_{D \in \mathcal{C}} \mathbb{E}_{t, x_0, x_1} \left[ w_t \| D(x_t, t) - x_1 \|^2 \right]$$
  जहाँ $D$ डिनोइज़र फंक्शन है, $\mathcal{C}$ सीखने योग्य फलनों (functions) का वर्ग है जो पैरामीट्राइजेशन द्वारा निर्धारित होता है, और $w_t$ समय-निर्भर वेटिंग है।
• पैरामीट्राइजेशन क्लासेस ($\mathcal{C}$):
  • $C_{den}$: सीधे क्लीन इमेज प्रेडिक्ट करता है ($D = N_\theta$)।
  • $C_{vel}$: वेलोसिटी प्रेडिक्ट करता है, क्लीन इमेज को $D = x + (1-t)N_\theta$ के रूप में पुनर्निर्मित करता है।
  • $C_{noise}$: नॉइज़ प्रेडिक्ट करता है, क्लीन इमेज को $D = (x - (1-t)N_\theta)/t$ के रूप में पुनर्निर्मित करता है।
• वेटिंग स्कीम्स ($w_t$): वे फ्लो मैचिंग ($1/(1-t)^2$), सिग्नल-टू-नॉइज़ रेशियो (SNR, $t^2/(1-t)^2$), और क्लासिकल डिनोइजिंग वेट्स सहित मानक वेटिंग्स का विश्लेषण करते हैं।
• प्रायोगिक सेटअप (Experimental Setup):
  • डेटासेट्स: नियंत्रित ज्यामिति (इंट्रिन्सिक डाइमेंशन को नियंत्रित करने के लिए फूरियर मोड्स) वाले सिंथेटिक डेटासेट्स, CIFAR-10, CelebA-64/128, और इमेजनेट (साहित्य तुलना के माध्यम से)।
  • आर्किटेक्चर: U-Nets, विभिन्न पैच साइज़ के साथ विजन ट्रांसफॉर्मर्स (ViT), और MLPs।
  • मेट्रिक्स:
    • PSNR (Peak Signal-to-Noise Ratio): विशिष्ट शोर स्तरों पर डिनोइजिंग सटीकता को मापता है।
    • FID (Fréchet Inception Distance): जनरेटिव क्वालिटी को मापता है।
  • मुख्य नवाचार: लेखक व्यवस्थित रूप से वेटिंग्स और पैरामीट्राइजेशन को डिकपल (decouple) करते हैं, उन संयोजनों का परीक्षण करते हैं जो मानक नहीं हैं (जैसे, $v$-prediction के साथ SNR वेटिंग का उपयोग करना)।

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

A. वेटिंग पर सैद्धांतिक अंतर्दृष्टि (Theoretical Insights on Weighting)

• इष्टतम वेटिंग: यह शोध पत्र इस बात का सांख्यिकीय औचित्य प्रदान करता है कि क्यों $(1-t)^{-2}$ के रूप में विचलन करने वाली वेटिंग्स (विशेष रूप से SNR और मानक फ्लो मैचिंग वेट्स) सबसे अच्छा प्रदर्शन करती हैं।
• व्युत्पत्ति (Derivation): समस्या को हेटरोसेडास्टिक रिग्रेशन (heteroscedastic regression) के रूप में फ्रेम करके और मैक्सिमम लाइकलीहुड एस्टीमेशन लागू करके, वे दिखाते हैं कि जैसे-जैसे $t \to 1$ (कम शोर) होता है, डेटा का कंडीशनल वेरिएंस कम हो जाता है। लाइकलीहुड को अधिकतम करने के लिए, लॉस को इस वेरिएंस के व्युत्क्रम (inverse) द्वारा वेट किया जाना चाहिए, जो $(1-t)^{-2}$ के रूप में स्केल करता है।
• निष्कर्ष: लो-नॉइज़ रिजीम (near $t=1$) को प्राथमिकता देने वाली वेटिंग्स सैद्धांतिक रूप से इष्टतम हैं और अनुभवजन्य रूप से सर्वश्रेष्ठ PSNR और FID प्रदान करती हैं।

B. पैरामीट्राइजेशन का पुनर्मूल्यांकन (Re-evaluating Parameterization)

• मैनिफोल्ड धारणा का विवाद: लेखक इस हालिया दावे को चुनौती देते हैं कि डेटा एक लो-डायमेंशनल मैनिफोल्ड पर स्थित होने के कारण $x$-prediction (डिनोइज़र) सार्वभौमिक रूप से श्रेष्ठ है।
• आर्किटेक्चर डिपेंडेंसी: वे प्रदर्शित करते हैं कि इष्टतम पैरामीट्राइजेशन केवल डेटा डाइमेंशनलिटी द्वारा निर्धारित नहीं होता है, बल्कि यह आर्किटेक्चरल लोकैलिटी (architectural locality) पर महत्वपूर्ण रूप से निर्भर करता है:
  • लोकल आर्किटेक्चर (U-Nets, छोटे पैच वाले ViTs): दृढ़ता से वेलोसिटी पैरामीट्राइजेशन ($C_{vel}$) का पक्ष लेते हैं।
  • ग्लोबल/कोर्स आर्किटेक्चर (बड़े पैच वाले ViTs, MLPs): डिनोइज़र पैरामीट्राइजेशन ($C_{den}$) का पक्ष लेते हैं।
• तंत्र (Mechanism): ViTs में बड़े पैच साइज टोकन की संख्या को कम करते हैं और लोकल इंडक्टिव बायस को सीमित करते हैं। इन "कोर्स" मॉडल्स में, क्लीन इमेज को सीधे प्रेडिक्ट करना अधिक स्थिर होता है। इसके विपरीत, मजबूत लोकल इंडक्टिव बायस वाले मॉडल्स (जैसे U-Nets) वेलोसिटी फॉर्मूलेशन से लाभान्वित होते हैं।

C. डेटा रिजीम और सैंपल साइज (Data Regime and Sample Size)

• लो डेटा रिजीम: सीमित प्रशिक्षण डेटा (जैसे, 10k इमेज) वाले परिदृश्यों में, डिनोइज़र पैरामीट्राइजेशन ($C_{den}$) वेलोसिटी पैरामीट्राइजेशन से काफी बेहतर प्रदर्शन करता है और बेहतर जनरलाइजेशन की ओर ले जाता है।
• हाई डेटा रिजीम: जैसे-जैसे डेटासेट का आकार बढ़ता है (जैसे, 100k+), प्रदर्शन का अंतर कम हो जाता है, और वेलोसिटी पैरामीट्राइजेशन अक्सर लोकल आर्किटेक्चर में श्रेष्ठता पुनः प्राप्त कर लेता है।

4. मुख्य परिणाम (Key Results)

1. वेटिंग प्रदर्शन:

   • SNR वेटिंग ($w_t = t^2/(1-t)^2$) और फ्लो मैचिंग वेटिंग ($w_t = 1/(1-t)^2$) लगातार उच्चतम PSNR और निम्नतम FID प्राप्त करते हैं।
   • क्लासिकल डिनोइजिंग वेट्स (हाई नॉइज़ को प्राथमिकता देने वाले) जनरेटिव कार्यों के लिए उप-इष्टतम (suboptimal) हैं।
   • उच्च PSNR (डिनोइजिंग सटीकता) और निम्न FID (जनरेटिव क्वालिटी) के बीच एक मजबूत सहसंबंध है।

2. पैरामीट्राइजेशन बनाम आर्किटेक्चर:

   • U-Nets: डेटा रेजोल्यूशन की परवाह किए बिना (128x128 तक टेस्ट किया गया) $C_{vel}$ श्रेष्ठ है।
   • ViTs: प्रदर्शन पैच साइज के आधार पर बदलता है।
     • छोटे पैच (4x4): $C_{vel}$ जीतता है।
     • बड़े पैच (16x16): $C_{den}$ जीतता है; $C_{vel}$ हाई नॉइज़ स्तरों पर विनाशकारी रूप से विफल हो जाता है।
   • यह Li & He (2025) (जिन्होंने बड़े-पैच वाले ViTs का उपयोग किया था) और शुरुआती FM कार्यों (जिन्होंने U-Nets का उपयोग किया था) के बीच के अंतर को स्पष्ट करता है।

3. मैनिफोल्ड डाइमेंशन:

   • नियंत्रित इंट्रिन्सिक डाइमेंशन ($m$) वाले सिंथेटिक डेटा पर, "मैनिफोल्ड धारणा" (कि $C_{den}$ को लो-डायमेंशनल डेटा पर जीतना चाहिए) केवल कोर्स मॉडल्स (MLP, ViT-16) के लिए ही मान्य है।
   • U-Nets $C_{den}$ और $C_{vel}$ के बीच चुनाव के संबंध में डेटा की इंट्रिन्सिक डाइमेंशन के प्रति काफी हद तक असंवेदनशील हैं।

4. डिकपलिंग स्ट्रैटेजी (Decoupling Strategy):

   • सर्वोत्तम परिणाम डिकपलिंग से प्राप्त होते हैं। उदाहरण के लिए, SNR वेटिंग का उपयोग वेलोसिटी पैरामीट्राइजेशन के साथ करने पर अक्सर मानक "नॉइज़ प्रेडिक्शन + SNR" या "वेलोसिटी + FM वेटिंग" जोड़ों की तुलना में बेहतर परिणाम मिलते हैं।

5. महत्व (Significance)

यह शोध पत्र फ्लो मैचिंग और डिफ्यूजन मॉडल को प्रशिक्षित करने के लिए एक महत्वपूर्ण "डिज़ाइन मैनुअल" प्रदान करता है:

• संघर्षों का समाधान: यह यह दिखाकर विरोधाभासी साहित्य को सुलझाता है कि आर्किटेक्चरल इंडक्टिव बायस (लोकैलिटी), मैनिफोल्ड धारणा की तुलना में पैरामीट्राइजेशन सफलता का एक अधिक मजबूत निर्धारक है।
• सैद्धांतिक आधार: यह वेटिंग को अनुभवजन्य ह्यूरिस्टिक्स से हटाकर हेटरोसेडास्टिक रिग्रेशन में इनवर्स-वेरिएंस वेटिंग पर आधारित एक सांख्यिकीय रूप से ठोस व्युत्पत्ति की ओर ले जाता है।
• व्यावहारिक मार्गदर्शन:
  • लगभग सभी सेटिंग्स के लिए SNR या FM वेटिंग ($\propto (1-t)^{-2}$) का उपयोग करें।
  • यदि आप U-Nets या छोटे पैच वाले ViTs का उपयोग कर रहे हैं, तो वेलोसिटी पैरामीट्राइजेशन का उपयोग करें।
  • यदि आप बड़े पैच वाले ViTs, MLPs का उपयोग कर रहे हैं, या बहुत छोटे डेटासेट्स पर प्रशिक्षण दे रहे हैं, तो डिनोइज़र पैरामीट्राइजेशन का उपयोग करें।
  • प्रदर्शन को अनुकूलित करने के लिए वेटिंग और पैरामीट्राइजेशन को डिकपल (decouple) करें (जैसे, SNR वेट + वेलोसिटी पैरामि)।

यह कार्य इस बात पर जोर देता है कि कोई "एक-आकार-सभी-के-लिए" (one-size-fits-all) समाधान नहीं है; इष्टतम प्रशिक्षण उद्देश्य को मॉडल आर्किटेक्चर और डेटा रिजीम के बीच विशिष्ट परस्पर क्रिया के अनुसार तैयार किया जाना चाहिए।