Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping

إليك ملخص تقني مفصل للورقة البحثية بعنوان "الضبط الدقيق لترانسفورمرات الانتشار عبر أخذ العينات الديناميكي للرقع والقفز بين الكتل بكفاءة الذاكرة" (Memory-Efficient Fine-Tuning Diffusion Transformers via Dynamic Patch Sampling and Block Skipping).

1. بيان المشكلة

أحدثت نماذج "ترانسفورمر الانتشار" (Diffusion Transformers - DiTs) ثورة في توليد الصور من النصوص (T2I)، مما مكن من إنشاء محتوى مخصص عالي الجودة. ومع ذلك، يواجه الضبط الدقيق لهذه النماذج واسعة النطاق لأغراض التخصيص (مثل DreamBooth وLoRA) عوائق كبيرة:

• استهلاك مرتفع للذاكرة: يتطلب الضبط الدقيق ذاكرة رسومية (GPU) ضخمة لتخزين معاملات النموذج، وحالات المُحسِّن (optimizer states)، وتنشيطات الانتشار الأمامي والخلفي. وهذا يحد من النشر على الأجهزة ذات الموارد المحدودة مثل الهواتف الذكية وأجهزة إنترنت الأشياء (IoT).
• محدودية الحلول الحالية:
  • الضبط الدقيق الفعال للمعاملات (PEFT): تقلل طرق مثل LoRA من عدد المعاملات القابلة للتدريب، لكنها لا تزال تتطلب انتشاراً خلفياً كاملاً عبر الشبكة بأكملها، مما يؤدي إلى استهلاك عالٍ لذاكرة التنشيط.
  • الطرق الخالية من التدرج (Gradient-Free): تقلل طرق مثل التحسين من الدرجة الصفرية (Zeroth-order optimization) من استهلاك الذاكرة، لكنها تعاني من عدم استقرار التحسين وتتطلب تكرارات تدريب مفرطة.
  • تخطي الطبقات (Layer Skipping): صُممت تقنيات تخطي الطبقات الموجودة حالياً لتسريع الاستدلال (Inference) ولا يمكن تطبيقها مباشرة على التدريب، لأنها تعطل مسار الانتشار الخلفي المطلوب للتعلم.
  • عدم توافق البنية: العديد من طرق الضبط الدقيق الفعالة (مثل HollowedNet) مصممة لبنى U-Net ولا تتعمم بشكل جيد على DiTs، حيث تكون وظائف الكتل (blocks) أقل هرمية.

2. المنهجية: DiT-BlockSkip

يقترح المؤلفون إطار عمل DiT-BlockSkip، الذي يدمج مكونين أساسيين لتقليل ذاكرة التدريب بشكل جذري مع الحفاظ على جودة التخصيص.

أ. أخذ العينات الديناميكي للرقع (Dynamic Patch Sampling)

بدلاً من معالجة صور كاملة الدقة طوال عملية الانتشار، تقوم هذه الطريقة بتعديل حجم رقعة المدخلات ديناميكياً بناءً على الخطوة الزمنية للانتشار ($t$).

• الآلية:
  • الخطوات الزمنية العالية (ضجيج عالٍ): يقوم النموذج بأخذ عينات من رقع أكبر للتركيز على تعلم البنية العالمية للموضوع.
  • الخطوات الزمنية المنخفضة (ضجيج منخفض): يقوم النموذج بأخذ عينات من رقع أصغر لالتقاط التفاصيل المحلية الدقيقة.
• التنفيذ: يتم قص الصور بناءً على دالة $f(s_{min}, s_{max}, t)$ ثم تغيير حجمها إلى دقة منخفضة ثابتة (مثلاً $256 \times 256$) قبل دخولها إلى النموذج.
• الفائدة: يقلل هذا من الأبعاد المكانية للمدخلات، مما يخفض بشكل كبير استخدام ذاكرة (التنشيطات) في العمليات الأمامية والخلفية دون تغيير بنية النموذج أو هدف إزالة الضجيج. وهو يحاكي فوائد التمثيل التي يوفرها التدريب عالي الدقة من خلال مواءمة نسب القص مع مرحلة التعلم في عملية الانتشار.

ب. تخطي الكتل مع الحساب المسبق للميزات المتبقية (Block Skipping with Residual Feature Precomputation)

يعالج هذا المكون عبء الذاثير الناتج عن كتل الترانسفورمر نفسها من خلال تخطي طبقات مختارة أثناء التدريب.

• استراتيجية اختيار الكتل:
  • على عكس U-Nets، تمتلك كتل DiT دقة مكانية ثابتة، مما يجعل من غير الواضح أي الكتل هي الأكثر حيوية.
  • يستخدم المؤلفون قناع الانتباه المتقاطع (Cross-Attention Masking) لتحديد الكتل الحيوية؛ حيث يقومون بقناع درجات الانتباه بين استعلامات الصور ومفاتيح النصوص في الكتل المتتالية.
  • النتيجة: يؤدي قناع الكتل متوسطة المستوى إلى انخفاض كبير في تشابه الموضوع، بينما يكون قناع الكتل المبكرة أو المتأخرة ذا تأثير ضئيل. لذا، تتمثل الاستراتيجية في تخطي أول $n$ وآخر $m$ من الكتل، مع الحفاظ على الكتل المتوسطة الحرجة.
• الحساب المسبق للميزات المتبقية:
  • مجرد تخطي الكتل يؤدي إلى عدم تطابق بين مسارات التدريب والاستدلال.
  • الحل: قبل الضبط الدقيق، يقوم المؤلفون بحساب الميزات المتبقية ($\Delta f = f_{out} - f_{in}$) مسبقاً لتسلسل الكتل التي سيتم تخطيها.
  • التدريب: أثناء الضبط الدقيق، يتم إخراج الكتل التي تم تخطيها من وحدة معالجة الرسومات (GPU) (لتوفير ذاكرة المعاملات والمُحسِّن). تُضاف الميزات المتبقية المحسوبة مسبقاً إلى مدخل الكتلة النشطة التالية، مما يسمح للشبكة بتجاوز الطبقات المخطاة مع الحفاظ على المسار الأمامي الصحيح للانتشار الخلفي.

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

1. إطار عمل DiT-BlockSkip: مخطط ضبط دقيق مبتكر مصمم خصيصاً لترانسفورمرات الانتشار يجمع بين أخذ العينات الدينمايكي للرقع وتخطي الكتل.
2. أخذ العينات الديناميكي للرقع: استراتيجية تكيف دقة المدخلات بناءً على الخطوات الزمنية للانتشار، مما يمكّن النموذج من تعلم كل من البنى العالمية والتفاصيل الدقيقة من مدخلات منخفضة الدقة.
3. الحساب المسبق للميزات المتبقية: تقنية تسمح بتخطي كتل الترانسفورمر أثناء التدريب عن طريق تخزين وإعادة استخدام الميزات المتبقية، مما يؤدي فعلياً إلى إخراج الأوزان من ذاكرة الـ GPU.
4. اختيار الكتل القائم على الانتباه المتقاطع: استراتيجية مؤتمتة لتحديد والحفاظ على أكثر الكتل متوسطة المستوى حيوية للتخصيص، مما يحل مشكلة الطبيعة "الصندوق الأسود" لأهمية كتل DiT.
5. قابلية التشغيل على الأجهزة: يقلل النهج بشكل كبير من متطلبات الذاكرة، مما يجعله قابلاً للتطبيق لتنفيذ نماذج DiT المخصصة على أجهزة الحافة (Edge devices).

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

تم تقييم الطريقة على نموذجي FLUX.1-dev و SANA باستخدام مجموعات بيانات DreamBooth و CustomConcept101.

• تقليل الذاكرة:
  • FLUX: قلل ذروة ذاكرة التدريب بنسبة 65.8% (من حوالي 22.84 GiB إلى حوالي 11.82 GiB) وقلل من عمليات TFLOPs بشكل كبير.
  • SANA: قلل ذاكرة التدريب بنسبة ~58% مقارنة بـ LoRA القياسي.
  • المقارنة: تفوق على النماذج المرجعية مثل HollowedNet (الذي فشل في التعامل مع DiTs دون تكييف)، وLISA، وLoRA-FA من حيث كفاءة الذاكرة.
• الأداء (الدقة/Fidelity):
  • دقة الموضوع (DINO, CLIP-I): حقق أداءً مقارباً للضبط الدقيق الكامل لـ LoRA. على سبيل المثال، في نموذج FLUX مع تخطي 50% من الكتل، ظلت درجات DINO تنافسية (~0.696 مقابل 0.732 لـ LoRA الكامل).
  • دقة النص (CLIP-T): حافظ على توافق عالٍ مع المطالبات النصية.
  • دراسة المستخدم: فضل المستخدمون الطريقة المقترحة على LoRA-FA وHollowedNet، ووجدوا أنها تضاهي LoRA القياسي في دقة الموضوع والنصوص.
• دراسات الاستئصال (Ablation Studies):
  • موضع الكتلة: أدى تخطي أول أو آخر 50% من الكتل فقط إلى أداء ضعيف، مما أكد ضرورة الحفاظ على الكتل متوسطة المستوى.
  • أخذ عينات الرقع: تفوق أخذ العينات الديناميكي على تغيير الحجم الثابت، حيث حافظ على الأنسجة والتفاصيل بشكل أفضل من مجرد خفض العينة (downsampling).
  • الميزات المتبقية: بدون الميزات المتبقية المحسوبة مسبقاً، تسبب تخطي الكتل في انحراف الميزات وفشل تام في التخصيص.

5. الأهمية

يقلص هذا العمل الفجوة بين الأداء العالي لترانسفورمرات الانتثار واسعة النطاق والقيود الصارمة للذاكرة في أجهزة الحافة. من خلال إثبات أن تخطي الكتل (وهي تقنية تستخدم تقليدياً في الاستدلال) يمكن تكييفها لعملية التدريب عبر الحساب المسبق للميزات المتبقية، ومن خلال تقديم أخذ عينات الرقع المدرك للخطوة الزمنية، قدم المؤلفون مساراً لـ:

• التخصيص على الأجهزة: تمكين المستخدمين من إجراء ضبط دقيق لمولدات الصور الحديثة مباشرة على الهواتف الذكية أو أجهزة إنترنت الأشياء دون الاعتماد على السحابة.
• التدريب الفعال: تقليل البصمة الكربونية وتكاليف الأجهزة المرتبطة بتدريب النماذج التوليدية الكبيرة.
• الأبحاء المستقبلية: فتح آفاق جديدة لاستراتيجيات التدريب الفعالة للذاكرة المصممة خصيصاً لنماذج الترانسفورمر التوليدية، وتجاوز بنى U-Net.