SoftWater: Class-Aware Rate Allocation for Softmax Quantization

ملخص تقني: SoftWater – تخصيص معدل الالتزام بالوعي بالفئات لتقنيات تكميم Softmax

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

عادةً ما تترك مسارات التكميم بعد التدريب (PTQ) لـ نماذج اللغات الكبيرة (LLMs) طبقة مخرجات الـ softmax النهائية (الرأس "head") بدقة عالية (مثل FP16). وبينما يتم تكميم الطبقات الخطية القياسية لتقليل متوسط مربع الخطأ الموزون (WMSE)، فإن طبقة المخرجات تنتج توزيعاً احتمالياً، مما يجعل تباعد كولباك - ليبلر (KL divergence) هو المقياس الأكثر ملاءمة للتشويه.

هذا الإغفال يزداد تكلفة في نماذج LLM الصغيرة الحديثة ذات القواميس اللغوية الكبيرة. وكما هو موضح في الجدول 1 من الورقة البحثية، يمكن أن يشكل "رأس المخرجات" ما بين 15 إلى 30% من إجمالي معلمات النموذج. وبناءً على ذلك، فإن النموذج الذي يعمل بـ "2 بت" مع رأس FP16 يخزن فعلياً عدداً أكبر بكثير من البتات لكل وزن عما يوحي به المعدل الاسمي. المحاولات السابقة لمعالجة ذلك إما استبدلت بنية الرأس بالكامل أو طبقت تكميماً يعتمد على WMSE، وهو ما يفشل في مراعاة الهندسة الخاصة لدالة softmax والتوزيع غير المنتظم للرموز (توزيع Zipfian).

تتمثل المشكلة الجوهرية في صياغة تكميم طبقة softmax كمسألة "معدل-تشويه" (rate-distortion) تحت مقياس تباعد KL، وتحديد كيفية تخصيص ميزانية بت ثابتة عبر مصفوفة الأوزان لتقليل التباعد بين التوزيعات الأصلية والمكممة.

2. المنهجية: SoftWater

يقترح المؤلفون SoftWater، وهي طريقة تشتق استراتيجية تخصيص معدل واعية بالفئات بناءً على تحليل الدرجة الثانية لتباعد KL.

اشتقاق الهدف:
يُعرَّف التشويه بأنه توقع تباعد KL بين التوزيع الأصلي $p$ والتوزيع المكمم $q$. باستخدام مفكك تايلور من الدرجة الثانية لتباعد KL حول اضطراب اللوجيت $\delta$، اشتق المؤلفون مقياس خطأ:
$$D \approx \frac{1}{2} \text{vec}(\Delta)^\top \mathbb{E}[(\nabla^2 f(0) \otimes XX^\top)] \text{vec}(\Delta)$$
حيث $\nabla^2 f(0) = \text{diag}(p) - pp^\top$ هو هسيان (Hessian) دالة softmax، و $XX^\top$ يمثل تباين الميزات المدخلة. ينتج عن هذا مصفوفة أوزان خطأ من نوع $Kn \times Kn$ تربط بين الفئات ($K$) والميزات ($n$).

التقريب القابل للتطبيق:
إن تحليل المصفوفة $Kn \times Kn$ مباشرةً أمر مستعصٍ حوسبياً. لذا قدم المؤلفون تقريبين رئيسيين:

1. حجة التذبذب (Dither Argument): بافتراض التذبذب الطرحي، تصبح الأخطاء غير مترابطة عبر الصفوف، مما يسمح بتجاهل الحدود غير القطرية للهسيان.
2. افتراض القابلية للفصل (Separability Assumption): يفترض المؤلفون أن انحناء softmax (المشار إليه بـ $\lambda$) وتباين المدخلات $XX^\top$ غير مترابطين. وهذا يسمح بتفكيك التوقع: $\mathbb{E}[\lambda \otimes XX^\top] \approx \mathbb{E}[\lambda] \otimes \mathbb{E}[XX^\top]$.

هذا التفكيك يحول المسألة إلى $K$ من عمليات تحليل تشوليسكي (Cholesky factorizations) المستقلة من نوع $n \times n$ (أو عملية تحليل واحدة من نوع $n \times n$ مع إعادة القياس)، مما يجعل المسألة قابلة للتطبيق. ويصبح تباعد الشبكة الناتج للفئة $k$ والعمود الميزي $i$ هو:
$$\alpha_{(k-1)n+i} \propto \frac{1}{\bar{\lambda}_k^{1/2} |\ell_{ii}|}$$
حيث $\bar{\lambda}_k = \mathbb{E}[p_k(1-p_k)]$ هو الانحناء المتوقع (المرتبط بتكرار الفئة والتباين) و $\ell_{ii}$ هي المدخلات القطرية لعامل تشوليسكي لتباين المدخلات.

الميزات الخوارزمية الرئيسية:

• التخصيص الواعي بالفئات: تخصص الطريقة شبكات أدق (بتات أكثر) للفئات المتكررة ومنخفضة التباين، وشبكات أخشن للفئات النادرة وعالية التباين. هذا يستغل الطبيعة الـ Zipfian لتوزيعات الرموز.
• الأساس التمهيدي للتنعيم (Smoothing Prior): لمنع عدم الاستقرار للفئات الغائبة عن بيانات المعايرة (حيث $\bar{\lambda}_k \to 0$)، تستخدم الطريقة أساساً تمهيدياً، حيث تخلط توزيع المعايرة مع توزيع منتظم. هذا يضع حداً أقصى لتباعد الشبكة، مما يضمن بقاء مفكك تايلور صالحاً.
• التنفيذ: تستخدم الخوارزمية إلغاء التداخل المتتالي (SIC) للترميز، على غرار خوارزمية WaterSIC للطبقات الخطية، ولكن مع إعادة قياس خاصة بكل فئة. وهي تتطلب تمريرة أمامية واحدة فقط على بيانات المعايرة لحساب الإحصائيات اللازمة ($\Sigma_X$ و $\bar{\lambda}$).

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

1. صياغة المشكلة: صاغت الورقة تكميم softmax كمسألة معدل-تشويه تحت تباعد KL، مشتقةً مقياس الخطأ المستحث الذي يربط إحصائيات الفئات بتباين الميزات.
2. خوارزمية SoftWater: خوارزمية عملية تستخدم تقريب القابلية للفصل لتقليل التعقيد من $O(K^3n^3)$ إلى $O(n^3)$ (بالإضافة إلى تمريرات خطية). وهي تقدم عامل قياس من جانب الفئة يكيف شبكة الشبكة (lattice grid) بناءً على تكرار الفئة والتباين عبر السياقات.
3. التحقق التجريبي: تم التحقق من صحة الطريقة عبر خمسة نماذج (من 1 مليار إلى 32 مليار معلمة) من عائلات Llama و Qwen.

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

قيم المؤلفون SoftWater مقابل مكمم WaterSIC (الذي يحسن WMSE) والأساليب المرجعية كاملة الدقة.

• تكميم الرأس فقط: عند تكميم الرأس فقط (مع إبقاء جسم النموذج في FP16)، تفوقت SoftWater على WaterSIC في 59 من أصل 60 نقطة اختبار (عبر معدلات بت ونماذج مختلفة). عند 2 بت، قللت SoftWater تباعد KL بمقدار 6.5 إلى 8.3 ضعفاً مقارنة بـ WaterSIC على مجموعة بيانات WikiText2.
• النماذج المكممة بالكامل: في نموذج Llama-3.2-1B-Instruct مع جسم مكمم (باستخدام GuidedQuant)، أدى استخدام رأس SoftWater بـ 2 بت إلى إزالة 45-60% من إجمالي البايتات المخزنة مع زيادة في الـ Perplexity بنسبة 2.9-3.7% فقط. في المقابل، تسبب خط الأساس WaterSIC في زيادة في الـ Perplexity بنسبة 9.7-17.6% لنفس مستوى الضغط.
• التكيف مع النطاق: تشير الورقة إلى أنه نظرًا لأن إحصائية جانب الفئة ($\bar{\lambda}$) مشتقة من بيانات المعايرة، فإن مطابقة نطاق المعايرة مع نطاق التشغيل يؤدي إلى أقل تباعد KL في ذلك النطاق المحدد.
• القابلية للفصل: تؤكد التجارب أن افتراض القابلية للفصل يعمل بشكل جيد، حيث يقلل النموذج البديل من التشويه الحقيقي بنسبة لا تتجاوز 10%.

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

تدعي الورقة أن SoftWater تجعل تكميم الرأس عملياً للنماذج التي يشكل فيها الرأس جزءاً كبيراً من حجم النموذج. من خلال التعامل مع الرأس كمسألة معدل-تشويه تحت KL بدلاً من WMSE، تحقق الطريقة أداءً يقترب من فقدان المعلومات عند 4 بتات وضغطاً كبيراً عند 2 بت مع حد أدنى من التدهور في الـ Perplexity.

يؤكد المؤلفون أن مكاسبهم تعود حصرياً لاستراتيجية تخصيص المعدل، حيث إنهم لم يستخدموا التحسينات الإضافية (مثل تصحيح انجراف التنشيط own، أو تصحيح تيار المتبقي) الموجودة في مسار WaterSIC، والتي تعد مستقلة عن تخصيص المعدل. تُقدم الطريقة كبديل جاهز لخطوة تكميم الرأس في مسارات PTQ الحالية، مما يوفر تقليلاً كبيراً في حجم النموذج دون الحاجة إلى تغييرات هيكلية أو ضبط دقيق (fine-tuning).