ملخص تقني: HeatKV - ضبط مخزن KV-cache بناءً على الرؤوس لنمذجة التوليد البصري التكراري

بيان المشكلة

برزت نماذج التوليد البصري التكراري (VAR) كبديل تنافسي لنماذج الانتشار (Diffusion Models) في توليد الصور، حيث توفر سرعات استدلال عالية وخصائص توسع ملائمة من خلال توليد صور كسلسلة من خرائط الرموز (token maps) متعددة المقاييس. ومع ذلك، تواجه هذه النماذج قيودًا شديدة في الذاكرة بسبب طبيعة مخزنها من الـ KV-cache. فخلافًا لنماذج اللغات الكبيرة (LLMs) التي تولد عادةً رمزًا واحدًا في كل خطوة، تقوم نماذج VAR بتوليد خرائط رموز كاملة بمقاييس دقة متزايدة. ونتيجة لذلك، ينمو عدد الرموز في الـ KV-cache بشكل فائق الخطية مع مقياس التوليد.

يمكن لنماذج VAR الحديثة، مثل Infinity-2B وInfinity-8B، أن تتطلب ما يصل إلى 10 جيجابايت من ذاكرة الـ KV-cache لتوليد صورة واحدة بدقة 1024 × 1024. هذا الأثر المرتفع على الذاكرة يحد من الحد الأقصى لحجم الدفعة (batch size) وعدد المستخدمين المتزامنين الذين يمكن لمزود الخدمة دعمهم على مسرعات الأجهزة المتخصصة. غالبًا ما تعتمد طرق الضغط الحالية لـ VARs، مثل ScaleKV وHACK، على استراتيجيات خشنة (coarse-grained)، مثل تصنيف رؤوس الانتباه أو الطبقات إلى فئات ثنائية (على سبيل المثال، "عالية الميزانية" مقابل "منخفضة الميزانية"). ويرى المؤلفون أن هذه التصنيفات الثنائية مقيدة للغاية، حيث تختلف أنماط الانتباه داخل نفس الطبقة بشكل كبير بين الرؤوس المختلفة، مما يؤدي إلى هدر الذاटीएम أو التقليم غير الأمثل.

المنهجية: HeatKV

يقترح المؤلفون HeatKV (ضبط مخزن KV-cache بناءً على الرؤوس)، وهي استراتيجية مبتكرة تكيّف تخصيص المخزن عند مستوى تفصيل الرؤوس الفردية ومقاييس الصور المحددة. تتكون الطريقة من ثلاثة مكونات رئيسية:

1. أهمية الرأس والمقياس (المعايرة غير المتصلة - Offline Calibration)

على عكس الطرق السابقة التي تستخدم مقاييس تعتمد على التباين، يقوم HeatKV بترتيب رؤوس الانتباه بناءً على درجة الانتباه التراكمي المعتمدة على المقياس (S-CAS).

• الآلية: باستخدام مجموعة بيانات معايرة صغيرة غير متصلة، تقوم الطريقة بتجميع احتمالات الانتباه من المقاييس المستقبلية رجوعًا إلى الرموز في المقاييس السابقة.
• المقياس: لكل طبقة $\ell$، ورأس $h$، ومقياس مصدر $k$، يقيس S-CAS متوسط كتلة الانتباه من جميع المقاييس المستقبلية ($\tau > k$) إلى الرموز عند المقياس $k$.
• النتيجة: يولد هذا ترتيبًا محددًا لأزواج "الرأس-المقياس" لكل مقياس مصدر، مما يوضح مدى حيوية مخزن رأس معين لمقياس سابق معين لعملية توليد الرموز المستقبلية. وهذا يسمح للنظام بتحديد أن رأسًا معينًا قد يحتاج إلى الاحتفاظ بالمقياس 1 ولكن يمكنه التخلص من المقياس 2، وهو مستوى من التفصيل لا تستطيع الطرق الثنائية تحقيقه.

2. تخصيص الميزانية

يحسب HeatKV عدد الرؤوس ($N_k$) التي يجب تقليمها بعد توليد كل مقياس $k$ لضمان بقاء حجم المخزن الإجمالي ضمن ميزانية ذاكرة محددة من قبل المستخدم $b$.

• القيد: تلتزم الطة بمبدأ "Streaming-LLM" المتمثل في الاحتفاظ بالمقاييس $s$ الأولى كـ "بالوعة انتباه" (attention sink) (وهي ضرورية للاستقرار) وعدم تخزين المقياس الحالي الذي يتم توليده أبدًا.
• الحساب: تحدد الخوارزمية الحد الأدنى لعدد الرؤوس التي يجب تقليمها للوفاء بقيد الميزانية، مما يسمح بتقليل المخزن ديناميكيًا مع تقدم عملية التوليد.

3. تحسين ترتيب الإزالة (التقليم المبكر الجشع - Greedy Early Pruning)

التحدي الحرج في نماذج VAR هو أن حجم المخزن ينمو بشكل فائق الخطية ضمن خطوة توليد واحدة (عبر طبقات المحول/Transformer). الطرق القياسية التي تقوم بالتقليم فقط بين المقاييس قد تسبب تجاوزات مؤقتة في الذاكرة تتخطى الحدود المادية.

• الحل: يقدم HeatKV خوارزمية التقليم المبكر الجشع (الخوارزمية 1). فهي تفرض قيد الميزانية بعد كل طبقة محول، وليس فقط بين المقاييس.
• الاستراتيجية: إذا تجاوز حجم المخزن الميزانية بعد طبقة معينة، تختار الخوارزمية بجشع (greedily) رؤوسًا ومقاييس إضافية لتقليمها من تلك الطبقة أو من الطبقات السابقة. وهي تعطي الأولوية لتقليم الرؤوس ذات رتب S-CAS الأدنى، وفي حالة التساوي، تعطي الأولوية للتقليم من الطبقات المتأخرة أو المقاييس المصدرية المتأخرة لتحقيق أقصى قدر من التخفيف في الذاكرة.
• الأمثلية: يثبت المؤلفون (الافتراض 1) أن هذا النهج الجشع ينتج مجموعة ذات حد أدنى من عدد الرؤوس المقلمة مبكرًا، مما يضمن عدم انتهاك الميزانية أبدًا مع تعظيم استخدام الذاكرة المتاحة.

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

1. الملاحظة دقيقة التفصيل: يوضح المؤلفون أن أنماط الانتباه في نماذج VAR تختلف بشكل كبير بين المقاييس المختلفة داخل نفس رأس الانتباه، لكنها تظل مستقرة نسبيًا عبر مختلف النصوص المدخلة. وهذا يدفع نحو التحول من التصنيف الثنائي للرؤاء إلى تخصيص ميزانية مستمر ومعتمد على المقياس.
2. مقياس S-CAS: يسمح تقديم مقياس درجة الانتباه التراكمي المعتمد على المقياس ببناء جدول تقليم ثابت مصمم لميزانية ذاكرة محددة، مما يتيح لكل رأس التخلص من المقاييس غير الضرورية مع الاحتفاظ بالمقاييس الحرجة.
3. خوارزمية تقليم محسنة: تطوير خوارزمية إخلاء جشعة تفرض قيود الميزانية على مستوى الطبقة (بدلاً من مجرد مستوى المقياس)، مما يمنع تجاوزات الذاكرة العابرة، وهو أمر بالغ الأهمية لنمو VAR فائق الخطية.
4. أداء رائد (State-of-the-Art): تحقق الطريقة رقمًا قياسيًا جديدًا لضغط مخزن KV-cache لنماذج VAR، حيث تثبت أن التخصيص الدقيق يعطي نتائج أفضل بكثير من النهج الخشن.

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

قيم المؤلفون HeatKV على نماذج Infinity-2B وInfinity-8B باستخدام مجموعة بيانات MS-COCO 2017 للتحقق. وقارنوا HeatKV بالنماذج المرجعية بما في ذلك StreamingLLM وSnapKV وHACK وScaleKV عبر ميزانيات ذاكرة بنسبة 4% و10% و20%.

• الدقة (Fidelity): تفوق HeatKV باستمرار على جميع النماذج المرجعية. ومن الملاحظ أنه في نموذج Infinity-2B، حقق HeatKV عند ميزانية 10% درجات أفضل في مقاييس FID وLPIPS وPSNR مقارنة بـ ScaleKV عند ميزانية 20%. وعند ميزانية 4%، حافظ HeatKV على FID قدره 2.70 (Infinity-2B)، وهو أقل بكثير من أفضل طريقة تالية (ScaleKV عند 4.92).
• المواءمة مع النص (Prompt Alignment): باستخدام GenEval، حافظ HeatKV على مواءمة قوية مع النص حتى عند الميزانيات المنخفضة، مع درجات مماثلة للنموذج المرجعي ذي المخزن الكامل.
• الجودة الإدراكية: أشارت نتائج درجة التفضيل البشري (HPSv2.1) إلى أن HeatKV حافظ على جودة إدراكية عالية، متفوقًا على الطرق الأخرى عند ميزانية 4% في Infinity-2B.
• الكفاءة: حقق HeatKV نسبة ضغط أعلى بمقدار ضعفين (2×) في تخصيص الذاكرة مقارنة بالطرق الحالية مع الحفاظ على دقة صورة مماثلة أو أفضل. كما مكنت الطريقة من استخدام أحجام دفعات أكبر، مما قلل فعليًا من زمن استدلال التوليد مقارنة بالاستدلال بالمخزن الكامل.

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

يزعم البحث أن HeatKV يمثل تقدمًا كبيرًا في جعل نماذج VAR قابلة للتشغيل على الأجهزة ذات الذاكرة المحدودة، مثل وحدات معالجة الرسومات المخصصة للمستهلكين، أو لخدمة عدد أكبر من المستخدمين المتزامنين على البنية التحتية الحالية.

يؤكد المؤلفون أن طريقتهم لا تتطلب تغييرات في بنية النموذج أو عملية التدريب؛ فهي تعتمد على مجموعة معايرة غير متصلة صغيرة (كما قلي من 1-10 نصوص) لإنشاء جدول تقليم ثابت. وتشير النتائج إلى أن تخصيص المخزن الدقيق المعتمد على الرؤوس أمر ضروري للحفاظ على البنية الدلالية وجودة التوليد لنماذج VAR تحت قيود الذاكرة الصارمة، متفوقة على الطرق التي تعتمد على التصنيفات الثنائية الخشنة للرؤوس أو الطبقات.

يختتم البحث بالإشارة إلى أنه بينما يحسن HeatKV الذاكرة، فإن العمل المستقبلي يمكن أن يستكشف الجمع بينه وبين تقنيات تحسين أخرى، مثل استغلال ندرة الانتباه (attention sparsity) أو تكميم النموذج (model quantization)، لتعزيز الكفاءة بشكل أكبر.