COT-FM: Cluster-wise Optimal Transport Flow Matching

यहाँ शोध पत्र "COT-FM: Cluster-wise Optimal Transport Flow Matching" का विस्तृत तकनीकी सारांश दिया गया है।

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

फ्लो मैचिंग (Flow Matching - FM) एक जनरेटिव मॉडलिंग फ्रेमवर्क है जो एक सरल स्रोत वितरण (जैसे Gaussian) से एक जटिल लक्ष्य डेटा वितरण तक नमूनों (samples) को ले जाने के लिए एक समय-निर्भर वेक्टर फील्ड (vector field) सीखता है। लचीला होने के बावजूद, मानक FM मॉडल अक्सर वक्रित परिवहन प्रक्षेपवक्रों (curved transport trajectories) से ग्रस्त होते हैं, जो प्रशिक्षण के दौरान उपयोग की जाने वाली कपलिंग रणनीतियों के कारण होता है:

• रैंडम कपलिंग (Random Coupling): स्रोत और लक्ष्य नमूनों को यादृच्छिक रूप से जोड़ने से परस्पर विरोधी वेग लक्ष्य (velocity targets) उत्पन्न होते हैं, जो मॉडल को एक औसत, वक्रित वेक्टर फील्ड सीखने के लिए मजबूर करते हैं।
• बैच-वाइज ऑप्टिमल ट्रांसपोर्ट (Batch-wise Optimal Transport - OT): रैंडम कपलिंग से बेहतर होने के बावजूद, मिनी-बैच पर OT का अनुमान लगाना स्थानीयता (locality) के मामले में संघर्ष करता है, जिससे पथों का क्रॉसिंग (path crossings) और वक्रता (curvature) होती है।
• परिणाम: वक्रित पथ विवक्तीकरण त्रुटि (discretization error) को बढ़ाते हैं, जिससे उच्च-गुणवत्ता वाले उत्पादन के लिए कई एकीकरण चरणों (NFEs) की आवश्यकता होती है। यह FM मॉडलों की दक्षता को सीमित करता है, विशेष रूप से एक-चरण या कुछ-चरणों वाले उत्पादन के लिए। मौजूदा "शॉर्टकट" तरीके (जैसे MeanFlow) जो चरणों को छोड़ देते हैं, अक्सर अंतर्निहित वेक्टर फील्ड को सीधा करने में विफल रहते हैं, जिसके परिणामस्वरूप विकृत नमूने मिलते हैं।

2. कार्यप्रणाली: COT-FM (Methodology: COT-FM)

लेखक क्लस्टर-वाइज ऑप्टिमल ट्रांसपोर्ट फ्लो मैचिंग (COT-FM) का प्रस्ताव करते हैं, जो एक प्लग-एंड-प्ले फ्रेमवर्क है जो मॉडल आर्किटेक्चर को बदले बिना, अधिक सीधे और विश्वसनीय परिवहन प्रवाह प्राप्त करने के लिए प्रायिकता पथ (probability path) को नया आकार देता है।

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

यह विधि इस अवलोकन का लाभ उठाती है कि डेटा वितरण अक्सर कई मोड (clusters) से बने होते हैं। एक ही मोड के भीतर के नमूनों को समान स्रोत नमूनों से उत्पन्न किया जा सकता है। पूरे वैश्विक वितरण को एक एकल स्रोत से मैप करने के बजाय, COT-FM एक विभाजित करो और जीतो (divide-and-conquer) रणनीति अपनाता है:

1. लक्ष्य डेटा का क्लस्टरिंग (Cluster Target Data): लक्ष्य डेटासेट को $K$ क्लस्टरों में विभाजित करें (अनसुपरवाइज्ड क्लस्टरिंग जैसे K-Means का उपयोग करके या कंडीशनल जनरेशन के लिए क्लास लेबल का उपयोग करके)।
2. क्लस्टर-वाइज स्रोत वितरण (Cluster-wise Source Distributions): प्रत्येक क्लस्टर के लिए एक समर्पित स्रोत वितरण असाइन करें।
3. लोकल ऑप्टिमल ट्रांसपोर्ट (Local Optimal Transport): वैश्विक स्तर के बजाय प्रत्येक क्लस्टर के भीतर ऑप्टिमल ट्रांसपोर्ट समस्या को हल करें।

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

COT-FM दो चरणों वाले एक वैकल्पिक अनुकूलन लूप (alternating optimization loop) के माध्यम से संचालित होता है:

• चरण 1: स्रोत अनुमान के लिए रिवर्स ODE (Reverse ODE for Source Estimation)

  • एक पूर्व-प्रशिक्षित FM मॉडल से शुरू करें (भले ही वह रैंडम कपलिंग के साथ प्रशिक्षित हो)।
  • प्रत्येक क्लस्टर $C_k$ में डेटा नमूने $x_1$ के लिए, स्रोत स्पेस तक इसे पीछे ट्रैक करने के लिए ODE इंटीग्रेशन को रिवर्स करें, जिससे एक अनुमानित स्रोत नमूना $\hat{x}_0$ प्राप्त होता है।
  • प्रत्येक क्लस्टर के लिए इन रिवर्स किए गए नमूनों का माध्य ($\mu_{0,k}$) और सहप्रसरण ($\Sigma_{0,k}$) कंप्यूट करें।
  • क्लस्टर-विशिष्ट स्रोत वितरण को एक गाऊसी वितरण (Gaussian) के रूप में परिभाषित करें: $p_{0,k} = \mathcal{N}(\mu_{0,k}, \Sigma_{0,k})$।
  • मुख्य लाभ: चूंकि पूर्व-प्रशिक्षित फ्लो गैर-प्रतिच्छेदी (non-intersecting) है, इसलिए इसे रिवर्स करने से ऐसे स्रोत नमूने मिलते हैं जो स्वाभाविक रूप से क्लस्टर संरचना के साथ संरेखित होते हैं, जिससे पथों का क्रॉसिंग कम हो जाता है।

• चरण 2: क्लस्टर-वाइज OT के साथ फाइन-ट्यूनिंग (Fine-tuning with Cluster-wise OT)

  • अनुमानित क्लस्टर-वाइज गाऊसियनों से स्रोत बिंदु लें ($x_0 \sim p_{0,k}$)।
  • स्रोत नमूनों और लक्ष्य क्लस्टर $C_k$ के बीच ऑप्टिमल ट्रांसपोर्ट मैप ($\pi_k$) की गणना करें। क्योंकि समस्या एक एकल क्लस्टर तक सीमित है, इसलिए बैच का आकार छोटा होता है, जिससे OT का अनुमान अधिक सटीक और प्रभावी हो जाता है।
  • एक अधिक सीधा वेक्टर फील्ड रिग्रेस करने के लिए इन क्लस्टर-वाइज ऑप्टिमल कपलिंग्स का उपयोग करके FM मॉडल को फाइन-ट्यून करें।

इन्फरेंस (Inference)

जनरेशन के दौरान, मॉडल पहले एक क्लस्टर इंडेक्स $k$ (क्लस्टर आकार के आनुपातिक) चुनता है, संबंधित क्लस्टर-विशिष्ट गाऊसी वितरण से एक शोर नमूना (noise sample) लेता है, और अंतिम नमूना उत्पन्न करने के लिए वेक्टर फील्ड को इंटीग्रेट करता है।

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

1. सामान्य फ्रेमवर्क (General Framework): COT-FM एक मॉडल-अग्नोस्टिक, प्लग-एंड-प्ले विधि है जो मौजूदा FM आर्किटेक्चर (जैसे Rectified Flow, MeanFlow) के साथ संगत है। यह केवल लक्ष्य प्रायिकता पथ और प्रशिक्षण कपलिंग को संशोधित करता है, नेटवर्क संरचना को नहीं।
2. वेक्टर फील्ड को सीधा करना (Straightening Vector Fields): क्लस्टरों तक OT को सीमित करके और रिवर्स ODE के माध्यम से क्लस्टर-विशिष्ट स्रोत वितरण प्राप्त करके, यह विधि प्रक्षेपवक्र की वक्रता और पथ क्रॉसिंग को काफी कम कर देती है।
3. दक्षता बनाम गुणवत्ता (Efficiency vs. Quality): यह वक्रित पथों में निहित विवक्तीकरण त्रुटियों को कम करके अत्यंत कम सैंपलिंग बजट (1-स्टेप या कुछ-स्टेप) के साथ उच्च-गुणवत्ता वाला उत्पादन सक्षम करता है।
4. सैद्धांतिक और अनुभवजन्य सत्यापन (Theoretical & Empirical Validation): शोध पत्र इस बात का कठोर विश्लेषण प्रदान करता है कि कैसे क्लस्टर-वाइज OT वैश्विक OT की तुलना में कम्प्यूटेशनल जटिलता को कम करता है जबकि परिवहन सटीकता में सुधार करता है।

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

लेखकों ने 2D सिंथेटिक डेटा, इमेज जनरेशन और रोबोटिक मैनिपुलेशन पर COT-FM का मूल्यांकन किया।

• 2D पॉइंट क्लाउड्स (Mixture of Gaussians, Two Moons, Checkerboard):

  • COT-FM ने सभी बेसलाइन्स (Rectified Flow, OT-CFM, MeanFlow) में सबसे कम Wasserstein दूरी और सबसे कम वक्रता प्राप्त की।
  • उदाहरण: 5-Gaussian मिश्रण पर, Wasserstein दूरी Rectified Flow के 0.5421 से घटकर COT-FM के साथ 0.1995 हो गई।

• अनकंडीशनल इमेज जनरेशन (CIFAR-10):

  • 1-स्टेप FID: COT-FM के साथ 378.0 (Rectified Flow) से सुधरकर 205.0 हो गया।
  • 50-स्टेप FID: 4.45 से सुधरकर 3.97 हो गया।
  • इसने MeanFlow को महत्वपूर्ण रूप से पछाड़ दिया (1-स्टेप FID: 2.92 $\to$ 2.60)।

• कंडीशनल इमेज जनरेशन (ImageNet 256x256):

  • SiT बैकबोन का उपयोग करते हुए, COT-FM ने सभी NFE सेटिंग्स में लगातार FID को कम किया, जिसमें लो-NFE व्यवस्था में सबसे महत्वपूर्ण लाभ देखा गया (जैसे, SiT-B/2 के लिए 1-स्टेप FID 264.36 से घटकर 231.99 हो गया)।

• रोबोटिक मैनिपुलेशन (LIBERO):

  • केवल 1 NFE के साथ LIBERO-Spatial पर 96.1% सफलता दर और LIBERO-Long पर 94.5% प्राप्त की।
  • इसने FLOWER पॉलिसी (जिसे 97.1% और 93.5% तक पहुँचने के लिए 4 NFE की आवश्यकता थी) और अन्य सिंगल-स्टेप बेसलाइन्स को पछाड़ दिया।

5. महत्व और प्रभाव (Significance and Impact)

• अंतराल को पाटना (Bridging the Gap): COT-FM फ्लो-आधारित मॉडलों में सैंपलिंग गति और जनरेशन गुणवत्ता के बीच मौलिक ट्रेड-ऑफ को संबोधित करता है। यह सिद्ध करता है कि वेक्टर फील्ड को सीधा करना ही उच्च-गुणवत्ता वाले एक-चरण जनरेशन को सक्षम करने की कुंजी है।
• स्केलेबिलिटी (Scalability): वैश्विक OT समस्या को छोटे क्लस्टर-वाइज समस्याओं में विघटित करके, यह बड़े डेटासेट के लिए ऑप्टिमल ट्रांसपोर्ट को संभव बनाता है, जो वैश्विक OT सॉल्वर की क्यूबिक जटिलता से बचता है।
• बहुमुखी प्रतिभा (Versatility): यह विधि अनकंडीशनल और कंडीशनल जनरेशन (क्लास लेबल या सीखे गए कंडीशन एम्बेडिंग का उपयोग करके) दोनों के लिए लागू है, जो इसे इमेज सिंथेसिस से लेकर रोबोट नियंत्रण नीतियों तक, जनरेटिव कार्यों की एक विस्तृत श्रृंखला के लिए प्रासंगिक बनाती है।
• भविष्य की दिशा (Future Direction): यह कार्य सुझाव देता है कि कपलिंग रणनीति (कैसे स्रोत और लक्ष्य को जोड़ा जाता है) को समझना और अनुकूलित करना, जनरेटिव AI को आगे बढ़ाने के लिए मॉडल आर्किटेक्चर जितना ही महत्वपूर्ण है।