ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

ملخص تقني: ART لعينات الانتشار (Diffusion Sampling)

بيان المشكلة
تقوم نماذج الانتشار القائمة على الدرجة (Score-based diffusion models) بتوليد العينات من خلال دمج ديناميكيات زمن عكسي متعلمة. ويعد الاختيار العددي الحرج في هذه العملية هو تقسيم الأفق الزمني إلى شبكة من الخطوات الزمنية. تعتمد النهج القياسية عادةً على شبكات موحدة أو جداول مصممة يدويًا (مثل EDM وDPM-Solver). ويرى المؤلفون أن هذه الجداول الثابتة غالبًا ما تكون غير مثالية نظرًا لعدد تقييمات الدالة (NFE) المحدد؛ فالشبكات الموحدة تخصص الجهد الحسابي بالتساوي، حتى عندما تختلف ديناميكيات الزمن العكسي بشكل كبير على طول المسار، بينما قد لا تتكيف الشبكات المصممة يدويًا مع توزيعات بيانات محددة أو تعقيدات دالة الدرجة (score function). الهدف هو إعادة توزيع الجهد الحسابي بشكل منهجي على طول مسار أخذ العينات لتقليل خطأ التجزئة الأولي (Euler discretization error) الإجمالي.

المنهجية
يقدم البحث الزمن المعاد بارامتريته التكيفي (Adaptive Reparameterized Time - ART)، وهو إطار عمل تحكم نظري لجدولة الخطوات الزمنية المثلى.

1. صياغة التحكم (ART):

• بدلًا من اختيار نقاط الشبكة مباشرة، يعامل ART "سرعة الساعة" لعملية الانتشار كمتغير تحكم.
• يقدم متغير زمن $t$ معاد بارامتريته إلى زمن الانتشار الأصلي $\tau$ عبر دالة رسم مستمر $\psi(t)$.
• يمثل التحكم $\theta(t) = \dot{\psi}(t)$ معدل التغير اللحظي لزمن الانتشار.
• الهدف هو اختيار $\theta(t)$ لتقليل خطأ التجزئة الأولي الإجمالي، مع الالتزام بالقيد بأن إجمالي الوقت محفوظ ($\int_0^T \theta(t) dt = T$).
• يتم اشتقاق تقريب الخطأ المحلي ليكون متناسبًا مع $|Q(x, \psi)| \theta(t)^2$، حيث تعتمد $Q$ على دالة الدرجة ومشتقاتها.

2. الحل عبر التعلم التعزيزي (ART-RL):

• مشكلة ART الحتمية مستعصية على الحل المباشر بسبب فضاء الحالة عالي الأبعاد وعدم وجود حل مغلق لمعادلة هاميلتون-جاكوبي-بيلمان (HJB).
• يقترح المؤلفون ART-RL، وهو نهج تعلم تعزيزي في الزمن المستمر (CTRL).
• يقدمون مشكلة مساعدة عشوائية حيث يتم أخذ عينات من التحكم $\theta$ من سياسة غاوسية $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
• يسمح هذا الصياغة باستخدام خوارزميات الممثل-الناقد (Actor-Critic) (وتحديدًا تكييف إطار عمل Jia وZhou [14]) لتعلم متوسط السياسة ودالة القيمة عبر التقريب العشوائي.

3. الجسر النظري:

• المساهمة النظرية الجوهرية هي إنشاء جسر ثنائي الاتجاه بين الحل الأمثل الحتمي لـ ART والسياسة الغاوسية المثلى في مشكلة ART-RL المساعدة.
• الكفاية (Sufficiency): التحكم الأمثل الحتمي $\mu^*$ يرتفع ليصبح سياسة غاوسية مثلى يكون متوسطها هو بالضبط $\mu^*$.
• الضرورة (Necessity): أي سياسة غاوسية مثلى للمشكلة المساعدة يجب أن يكون لها متوسط يستعيد تحكم ART الحتمي $\mu^*$.
• تبرر هذه النتيجة استخدام المتوسط الغاوسي لتمثيل الممثل (Actor) ونشره بشكل حتمي بعد التدريب، مما يوفر ضمانات صارمة بأن الجدول المتعلم هو بالفعل الأمثل للمشكلة الأصلية للتحكم.

4. التقطير (Distillation):

• بعد التدريب خارج الخط (Offline training)، يتم تقطير السياسة العشوائية إلى شبكة زمنية حتمية تعتمد على الزمن فقط. وهذا يلغي الحاجة إلى تقييمات شبكة الممثل لكل خطوة أو حسابات التباين أثناء الاستنتاج، مما يضمن عدم وجود تكلفة حسابية إضافية مقارنة بالمجدولات القياسية.

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

• الصياغة: يقترح ART، وهو أول إطار عمل تحكم نظري لجدولة خطوات زمن الانتشار عبر إعادة بارامترة الزمن التكيفية.
• النظرية: يثبت التكافؤ ثنائي الاتجاه بين مشكلة التحكم الحتمي الإجرائي وإجراء تعلم (Actor-Critic) في الزمن المستمر مع سياسات غاوسية. وهذا يحول تطبيقات التعلم التعزيزي التجريبية إلى مسار مبدئي لإيجاد الحلول المثلى الحتمية.
• الخوارزمية: يطور ART-RL، وهي خوارزمية (Actor-Critic) تتعلم جداول زمنية تكيفية بناءً على البيانات.
• التعميم: يوضح أن الجدول الزمني الذي تم تقطيره يمكن نقله وتعميمه عبر مجموعات بيانات مختلفة وأعداد خطوات زمنية مختلفة دون الحاجة لإعادة التدريب.

النتائج التجريبية
قيم المؤلفون ART-RL ضمن مسار EDM الرسمي عبر CIFAR-10، وAFHQv2، وFFHQ، وImageNet.

• CIFAR-10: يتفوق ART-RL باستمرار على جداول Uniform وDPM-Solver log-SNR وجداول EDM القياسية عبر مجموعة واسعة من ميزانيات NFE.
  • عند الميزانيات المنخفضة (مثل NFE=3)، ينتج ART-RL أجسامًا يمكن التعرف عليها بينما تظل النماذج المرجعية قريبة من الضجيج.
  • عند NFE=35، يحقق ART-RL قيمة FID قدرها 1.82، متفوقًا على EDM (1.85) وDPM-logSNR (1.89).
  • تم تأكيد الأهمية الإحصائية عبر بذور (seeds) متعددة، حيث تفوقت أسوأ نتيجة لـ ART-RL على أفضل نتيجة لـ EDM.
• التعميم:
  • الاستكمال الداخلي والخارجي (Interpolation/Extrapolation): الجدول الزمني المتعلم عند $K=18$ خطوة على CIFAR-10 يحتفظ بميزته في الأداء عند استكماله أو استكماله خارجيًا لأعداد خطوات أخرى ($K \in \{4, 9, 12, 15, 20\}$).
  • النقل عبر مجموعات البيانات: تم تطبيق جدول CIFAR-10 المقطر على AFHQv2 وFFHQ وImageNet دون إعادة تدريب. وقد حسن باستمرار قيمة FID مقارنة بـ EDM وDPM-logSNR، خاصة عند قيم NFE المنخفضة والمتوسطة.
• التحليل أحادي البعد: في حالة وجود دالة درجة تحليلية، قلل ART-RL بشكل كبير من مسافة Wasserstein ($W_2$) مقارنة بالجداول المصممة يدويًا، مما يؤكد أن التحسن ناتج عن تخصيص أفضل للزمن وليس عن أخطاء في تقدير الدرجة.

الأهمية والادعاءات
يدعي البحث أنه الأول في صياغة جدولة خطوات زمن الانتشار كمسألة تحكم عبر إعادة بارامترة الزمن التكيفية، والأول في تقديم جسر ثنائي الاتجاه صارم بين الحل الأمثل الحتمي وإجراء تعلم (Actor-Critic) مع السياسات الغاوسية.

يؤكد المؤلفون أن طريقتهم تعتمد على البيانات ومستردة التكلفة (Amortized):

• تعتمد على البيانات: يتم تعلم الجدول الزمني من دالة الهدف وديناميكيات البيانات، وليس تصميمًا يدويًا.
• التكلفة المستردة (Amortized Transfer): تكلفة التدريب خارج الخط (حوالي 1-2 ساعة على Colab T4 لـ CIFAR-10) هي تكلفة تُدفع لمرة واحدة. والشبكة الحتمية المقطرة الناتجة قابلة لإعادة الاستخدام عبر مجموعات بيانات وأعداد خطوات مختلفة مع لا تكلفة إضافية عند الاستنتاج، مما يطابق وقت تشغيل المجدولات القياسية مثل Uniform أو EDM.

يضع هذا العمل إعادة بارامترة الزمن التكيفية كاتجاه ناشئ وواعد لتصميم الجداول الزمنية المنهجية في نماذج الانتشار التوليدية، مما يوفر نقطة انطلاق من منظور نظرية التحكم للأبحاث المستقبلية.