ملخص تقني: تقطير خرائط الانتقال العشوائي للاستدلال الاحتمالي السريع

بيان المشكلة

حققت نماذج الانتشار (Diffusion models) أداءً رائدًا في جودة التوليد، والتنوع، وتغطية التوزيع مقارنة بالنماذج السابقة مثل GANs وVAEs. ومع ذلك، غالبًا ما يصاحب هذا الأداء تكاليف استدلال باهظة بسبب ضرورة دمج المعادلات التفاضلية العادية (ODEs) أو المعادلات التفاضلية العشوائية (SDEs) عبر الخطوات الزمنية المستمرة.

وعلى الرغم من وجود تقنيات تسريع حالية مثل التقطير (على سبيل المثال، نماذج الاتساق/Consistency Models) والنهج القائمة على النقل الأمثل (Optimal Transport)، إلا أنها غالبًا ما تواجه مقايضة بين جودة/تنوع التوليد وتعقيد الاستدلال. ويتمثل أحد القيود الحرجة لطرق التقطير الحالية في اعتمادها على الاستدلال الحتمي (على سبيل المثال، عبر معادلة تدفق الاحتمال/probability-flow ODE). فبالنظر إلى عينة ضوضاء ثابتة، يكون المخرج فريدًا. وهذا يحد من فائدتها في المهام اللاحقة التي تتطلب بطبيعتها استدلالًا عشوائيًا، مثل أخذ العينات من التوزيع البعدي (diffusion posterior sampling) للمشكلات العكسية، أو ضبط الطاقة (energy-based fine-tuning)، أو تحسين التوافق مع استعلامات المستخدم. علاوة على ذلك، تتطلب العديد من أطر التقطير الحالية نموذج معلم (teacher model) مدرب مسبقًا، أو تحسينًا ثنائي المستوى، أو محاكاة وتخزين المسارات، مما يعقد عملية التدريب وقابلية التوسع.

المنهجية: تقطير خرائط الانتقال العشوائي (STMD)

يقترح المؤلفون تقطير خرائط الانتقال العشوائي (STMD)، وهو إطار عمل خالٍ من المعلم (teacher-free) مصمم لتسريع استدلال نماذج الانتشار مع الحفاظ على بنية الانتقال العشوائي للعملية الأساسية.

المفهوم الجوهري

على عكس نماذج الانتشار القائمة على الدرجة (score-based) التي تعمل عادةً على نمذجة متوسط التوزيع البعدي، يقوم STMD بتقطير خريطة الانتقال الكاملة المرتبطة بـ SDE لأخذ العينات. تتعلم الطريقة مولدًا شرطيًا يتنبأ بخريطة الانتقال $p(x_{t'} | x_t)$ لـ $t' < t$ في خطوة واحدة أو خطوات قليلة، متجنبة التكامل العددي المكلف لـ SDE العكسي.

الصيغة التقنية

1. تدفق المتوسط الشرطي (Conditional Mean Flow - MF): يقوم الإطار بترميز انتقالات SDE باستخدام نموذج تدفق متوسط شرطي. فبدلاً من تعلم خريطة تدفق حتمية، يتعلم النموذج متوسط السرعة على طول مسار عملية الانتشار.
2. ترميز $x_0$ ($x_0$-Parametrization): لنمذجة الانتقال $p(x_{t'} | x_t)$، يستخدم المؤلفون ترميز $x_0$. حيث يعبرون عن الانتقال كتحليل تكاملي فوق التوزيع البعدي $p(x_0 | x_t)$ ومقياس الجسر المعروف $p(x_{t'} | x_0, x_t)$. وبما أن مقياس الجسر لنماذج SDE الحافظة للتباين (VP SDEs) متاح في صيغة مغلقة، فإن المشكلة تتقلص إلى تعلم نموذج شرطي $p_\theta(x_0 | x_t)$.
3. هدف التدريب: يتم تدريب النموذج $u_\theta$ لتقريب متوسط السرعة الشرطي. يعرّف المؤلفون دالة خسارة التدفق المتوسط الشرطي (CMF) كالتالي:
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   حيث $z_s$ هو استكمال خطي مستقيم بين البيانات $x_0$ وعينة سابقة $z_1$، مشروطًا بحالة الانتشار $x_t$ والزمن $t$. ويمثل العامل $\text{sg}$ مؤثر إيقاف التدرج (stop-gradient).
4. تدريب خالٍ من المعلم: لا تتطلب عملية التدريب نموذج معلم مدرب مسبقًا. فهي تعتمد على عينات من التوزيع المشترك $p(x_0, x_t)$، والتي يمكن توليدها عن طريق أخذ عينات من $x_0$ من توزيع البيانات وتطبيق نواة SDE الأمامي (المعروفة بصيغتها المغلقة).

الاستدلال

يتم إجراء الاستدلال عن طريق دمج تدفق المتوسط الشرطي المتعلم عكس اتجاه الزمن. وتتضمن العملية:

1. أخذ عينة من الضوضاء الأولية $x_1$ (أو $z_1$).
2. تطبيق حقل السرعة المتعلم $u_\theta$ بشكل متكرر للانتقال من الزمن $s$ إلى $s-\Delta s$.
3. عند كل خطوة زمنية للانتشار $t$، يتم أخذ عينة من توزيع الجسر $p(x_{t-\Delta t} | x_0, x_t)$ لتحديث الحالة، مما يحاكي فعليًا الانتقال العشوائي في خطوات كبيرة.

المساهمات الرئيسية

1. إطار عمل مبتكر: اقتراح STMD، وهو إطار تقطير قائم على تدفقات المتوسط الشرطية يتيح تعلم خرائط الانتقال لعمليات SDE عالية الأبعاد باستخدام مولد خطوة واحدة أو خطوات قليلة عبر هدف الانحدار، دون الحاجة إلى معلم مدرب مسبق.
2. التقارب النظري: اشتقاق حدود التقارب لخوارزمية التدفق المتوسط في مسافة 2-Wasserstein. وبناءً على علم المؤلفين، يعد هذا أول تحليل تقارب لخوارزمية التدفق المتوسط في هذا المقياس. كما قاموا بتوسيع ذلك لإثبات أن خريطة الانتقال الشرطية المتعلمة تتقارب مع نواة SDE العكسي.
3. التحقق التجريبي: إثبات الفعالية على مهام توليد الصور وإكمال الصور (inpainting) عبر مجموعات بيانات MNIST وCIFAR-10 وCelebA، مما أظهر أداءً تنافسيًا مع النماذج المرجعية مع الحفاظ على العشوائية.

النتائج التجريبية

قيم المؤلفون STMD على ثلاث مجموعات بيانات:

• MNIST: باستخدام بنية U-Net، تمت مقارنة النموذج بخوارزمية التدفق المتوسط الأصلية ونموذج DDPM. استخدم المؤلفون مسافة Fréchet (FD) في الفضاء الكامن للمصنف (لأن FID غير موثوق لـ MNIST). أظهرت النتائج أن STMD يعمل بشكل مشابه للنماذج المرجعية مع دعم أخذ العينات العشوائي.
• CIFAR-10: أُجريت تجارب مماثلة باستخدام U-Net. حقق النموذج درجات Fréchet Inception Distance (FID) تنافسية مقارنة بـ DDPM وMean Flow لعدد ثابت من دورات التدريب، مع الحفاظ على العشوائية الكاملة لعينات الانتشار.
• CelebA: استخدمت التجارب بنية DiT كامنة. ولد النموذج عينات غير مشروطة بـ FID قدره 8.28 (باستخدام 4 خطوات استدلال وخطوتين من التدفق المتوسط) ونجح في مهام إكمال الصور (inpainting).

في جميع الحالات، تم تدريب النماذج من الصفر باستخدام نفس البنى الشبكية وعدد التكرارات مثل النماذج المرجعية لضمان مقارنة عادلة.

الأهمية والادعاءات

يدعي البحث أن STMD يعالج الفجوة بين التقطير الحتمي السريع والحاجة إلى الاستدلال العشوائي في المهام اللاحقة. ومن خلال تقطير خريطة الانتقال بدلاً من مجرد خريطة التدفق، يحافظ STMD على البنية العشوائية لعملية الانتشار.

يؤكد المؤلفون أن نهجهم خالٍ من المعلم، مما يلغي الحاجة إلى التحسين ثنائي المستوى أو تخزين المسارات، مما يجعل عملية التدريب أكثر كفاءة وقابلية للتوسع. توفر المساهمة النظرية المتمثلة في وضع حدود تقارب 2-Wasserstein أساسًا صارمًا لمنهج التدفق المتوسط، مما يسد الفجوة بين الضمانات النظرية والتطبيق العملي.

يضع هذا العمل الحل للمهام التي تتطلب استدلالًا احتماليًا، مثل المشكلات العكسية وأخذ العينات البعدي، حيث تفشل طرق التقطير الحتمية. ويظل المؤلفون متواضعين بشأن الأداء الحالي، مشيرين إلى أن العدد المحدود من دورات التدريب يمنع الوصول إلى نتائج رائدة عالميًا (state-of-the-art)، لكن الدراسة المقارنة تثبت فعالية الطريقة بالنسبة للنماذج المرجعية.