Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning

إليك ملخص تقني مفصل لورقة البحث بعنوان "Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning" (ترانسفورمرز ذات عنق الزجاجة: توحيد ذاكرة التخزين المؤقت KV بشكل دوري للاستدلال المعمم).

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

أظهرت النماذج اللغوية الكبيرة (LLMs) القائمة على الـ Transformers قدرات استدلال قوية، لا سيوت عند زيادة الحوسبة أثناء الاستدلال عبر سلاسل "سلاسل التفكير" (توليد في فضاء الرموز/Tokens). ومع ذلك، يستكشف جسم متنامٍ من الأبحاث الحوسبة المساعدة في الفضاء الكامن (ALSC)، حيث تقوم النماذج بعمليات حوسبة إضافية في حالاتها المستمرة الداخلية (الفضاء الكامن) بدلاً من إصدار رموز وسيطة.

تندرج طرق ALSC الحالية عموماً تحت ثلاث فئات:

1. التدحرجات الكامنة بوساطة الرموز (Token-mediated latent rollouts): توليد "رموز كامنة" داخلية تمد التسلسل.
2. توجيه البقايا/التنشيط (Residual/activation steering): تعديل الحالات المخفية للتأثير على المخرجات.
3. ضغط الذاكرة (KV compression): تقليم أو دمج مدخلات ذاكرة المفتاح-القيمة (KV cache) لتوفير الذاكرة.

الفجوة: تركز الطرق الحالية بشكل كبير على الضغط (تقليل بصمة الذاكرة) أو التمديد (إضافة المزيد من الرموز). هناك اتجاه غير مستكشف بشكل كافٍ وهو: توحيد وإعادة توحيد الذاكرة (Memory Consolidation and Reconsolidation). في علم الأعصاب، يعمل "التوحيد" (Consolidation) على تثبيت الذكريات الجديدة، بينما تسمح "إعادة التوحيد" (Reconsolidation) باستدعاء الذكريات وتحديثها بسياق جديد قبل إعادة تثبيتها. يجادل المؤلفون بأن نماذج الـ Transformers ذاتية الانحدار (Autoregressive) القياسية تفتقر إلى آلية لإعادة معالجة و"إعادة كتابة" ذاكرتها العاملة (KV cache) بشكل دوري للتخلص من التفاصيل غير ذات الصلة ودمج السياق الجديد، مما يعيق التعميم في مهام الاستدلال المعقدة.

2. الدافع النظري: نظرية عنق الزجاجة المعلوماتي (IB Theory)

يقدم المؤلفون تبريراً نظرياً لإعادة كتابة الذاكرة باستخدام نظرية عنق الزجاجة المعلوماتي (IB).

• الهدف: تحسين التمثيل الكامن $Z$ ليكون غنياً بالمعلومات بحد أقصى حول المخرج $Y$ (خطوة الاستدلال التالية) مع أن يكون بأقل قدر من المعلومات حول المدخل $X$ (التاريخ الخام). هذا التوازن ($I(X;Z)$ مقابل $I(Z;Y)$) هو ما يدفع التعميم.
• المشكلة في الـ Transformers التقليدية: تعمل ذاكرة الـ KV cache ($C_{0:n}$) في عملية التنبؤ بالرمز التالي كـ "عنق زجاجة نهائي". إن هدف التدريب (تعظيم احتمالية الرمز التالي) يحفز النموذج على الحفاظ على كل المعلومات التاريخية للتنبؤ بالرمز التالي بدقة. نتيجة لذلك، تراكم الذاكرة تفاصيل زائدة وغير تنبؤية ($I(X;Z)$ تظل مرتفعة)، مما قد يعيق قدرة النموذج على التعميم لأنماط استدلال جديدة.
• الحل المقترح: يمكن أن تعمل إعادة كتابة الـ KV cache بشكل دوري كخطوة لإعادة المعالجة. من خلال ضغط $I(X;Z)$ بشكل انتقائي (إزالة الضجيج/التفاصيل غير ذات الصلة) مع الحفاظ على أو تعزيز $I(Z;Y)$ (المنفعة التنبؤية)، يمكن للنموذ de تحقيق عنق زجاجة معلوماتي أفضل، مما يؤدي إلى تحسين التعميم.

3. المنهجية: الـ Transformer ذو عنق الزجاجة (Bottlenecked Transformer)

يقدم المؤلفون الـ Bottlenecked Transformer، وهو بنية تعزز نموذج (backbone) LLM مدرب مسبقاً بإضافة معالج ذاكرة التخزين المؤقت (Cache Processor).

البنية

• النموذج الأساسي (Backbone): نموذج Transformer قياسي من نوع decoder-only (يكون مجمدًا خلال المرحلة الثانية من التدريب).
• معالج الذاكرة (Cache Processor): نموذج Transformer مساعد أصغر يعمل بالتوازي مع النموذج الأساسي. يتم استدعاؤه في فترات محددة (على سبيل المثال، عند رموز السطر الجديد التي تمثل نهاية خطوة استدلال).
• الآلية:
  1. الاختيار (Selection): عند الاستدعاء، يختار المعالج مجموعتين من مدخلات KV:
     • الرموز الحديثة (Recent Tokens): مدخلات KV من خطوة الاستدلال الأخيرة (التوحيد - Consolidation).
     2. الرموز المستدعاة (Recalled Tokens): مجموعة من أعلى $k$ من مدخلات KV السابقة المختارة بناءً على كتلة الانتباه (attention mass) بالنسبة للخطوة الأخيرة (إعادة التوحيد - Reconsolidation).
  2. المعالجة (Processing): يتم تحويل هذه المدخلات المختارة إلى "رموز KV" وتغذيتها إلى معالج الذاكرة. يقوم المعالج بإجراء إعادة كتابة غير سببية وفي مكانها (non-causal, in-place rewrites) على هذه المدخلات.
  3. إعادة الكتابة (Rewrite): يخرج المعالج تحديثات ($\Delta K, \Delta V$) تُضاف إلى مدخلات KV الأصلية عبر اتصال متبقي بوابي (gated residual connection). ومن المهم ملاحظة أنه لا يتم تطبيق أي تقليل للأبعاد؛ الهدف هو إعادة هيكلة المحتوى وليس الضغط.

استراتيجية التدريب

يتم التدريب في مرحلتين:

1. المرحلة الأولى (SFT): يتم ضبط النموذج الأساسي (fine-tuning) على بيانات الاستدلال باستخدام التنبؤ بالرمز التالي القياسي.
2. المرحلة الثانية (تدريب المعالج): يتم تجميد النموذج الأساسي. يتم تدريب معالج الذاكرة لتقليل خسارة التقاطع (cross-entropy loss) لخطوة الاستدلال التالية، مشروطة بالذاكرة المعاد كتابتها.
   • يتعلم المعالج إعادة تنظيم الذاكرة لتحسين التنبؤات المستقبلية دون دالة خسارة صريحة لـ IB؛ حيث يظهر الضغط بشكل ضمني عبر تحسين كفاءة التنبؤ.

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

1. إطار نظري: تطبيق مبتكر لنظرية عنق الزجاجة المعلوماتي على نماذج الـ Transformers من نوع decoder-only، مع الحجة بأن إعادة كتابة الـ KV cache بشكل دوري ضروري لكسر طبيعة "الدقة العالية ولكن غير الانضغاطية" لذاكرات الـ autoregressive القياسية، مما يحسن التعميم.
2. بنية مبتكرة: الـ Bottlenecked Transformer، الذي ينفذ مفاهيم علم الأعصاب (التوحيد/إعادة التوحيد) عبر معالج ذاكرة مخصص يقوم بإعادة كتابة مدخلات KV بشكل غير سببي وفي مكانها.
3. التحقق التجريبي: تقييم واسع النطاق يظهر تفوق هذا الهيكل على نماذج الـ Transformers التقليدية وعلى خطوط الأساس الأخرى لـ ALSC (مثل pause tokens و latent rollouts) في مهام الاستدلال الرياضي.

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

تم تقييم النموذج على سبعة اختبارات مرجعية (GSM8K, MATH, SVAMP, TheoremQA, LogiQA, Gaokao-MathQA, GSM-Hard) عبر أربعة نماذج أساسية (Llama 3.2 1B/3B, Llama 3.1 8B, Qwen 3 0.6B).

• مكاسب الأداء: تفوق الـ Bottlenecked Transformer باستمرار على النماذج المرجعية.
  • Llama 3.2 1B: حسن دقة SVAMP بنسبة +6.6% (من 38.0% إلى 44.6%).
  • Llama 3.2 3B: حسن GSM8K بنسبة +4.6%.
  • Llama 3.1 8B: حسن LogiQA بنسبة +3.1%.
• المقارنة مع الخطوط المرجعية (Baselines):
  • الرموز المتوقفة (Pause Tokens): أظهرت أداءً متذبذبًا، وغالبًا ما كان أداؤها أقل من SFT القياسي ما لم تقترن بالتدريب المسبق المستمر.
  • التدحرجات الكامنة (Latent Rollouts): كانت أداؤها ضعيفًا بشكل عام، وأدت أحيانًا إلى عدم استقرار النموذج (خاصة في النماذج الكبيرة مثل Llama 3.1 8B).
• دراسات الاستئصال (Ablation Studies):
  • إعادة التوحيد بـ Top-$k$: كانت قيم $k \approx 32-64$ هي الأمثل لمعظم المهام، بينما استفادت المهام الأصعب (MATH) من قيم $k$ أكبر.
  • نافذة الخطوة الأخيرة ($R$): كان الأداء قويًا عبر نطاق من أحجام النوافذ، مما يشير إلى أن الآلية تستفيد من إعادة تشكيل السياق متوسط المدى بدلاً من التحديثات لكل رمز على حدة.
  • مقادولات إعادة الكتابة: أظهر التحليل أن المعالج يعدل بشكل أساسي متجهات القيمة (Value vectors) (محتوى الذاكرة) بدلاً من متجهات المفتاح (Key vectors) (عناوين الذاكرة)، مع حدوث التغييرات الأكثر أهمية في الطبقات الأولى من الشبكة.

6. الأهمية والعمل المستقبلي

• الأهمية: يجسّر هذا العمل الفجوة بين علم الأعصاب (توحيد الذاكرة) وبنية الـ LLM. إنه يثبت أن إعادة معالجة الحالات الداخلية (بدلاً من مجرد توليد المزيد من الرموز أو ضغط الذاكرة) هي مسار قابل للتطبيق لتحسين تعميم الاستدلال. كما يتحدى الافتراض بأن الذاكرات الأكبر هي دائمًا الأفضل، مقترحًا أن جودة الذاكرة (عبر إعادة التنظيم) أهم من كميتها.
• القيود والتوجهات المستقبلية:
  • تخصيص الائتمان (Credit Assignment): يمكن أن يؤدي التدريب حصريًا على خسارة الخطوة التالية إلى تباين عالٍ في تخصيص الائتمان لإعادة كتابة الذاكرة.
  • الدقة البيولوجية: يدمج النموذج الحالي التوحيد وإعادة التوحيد في خطوة واحدة عبر الإنترنت (online). العمل المستقبلي يمكن أن يفصل بينهما إلى عمليات "أوفلاين" (قائمة على إعادة التشغيل) و"أونلاين" (بوابة خطأ التنبؤ).
  • خسارة IB الصريحة: يمكن للإصدارات المستقبلية استكشاف حقن الضجيج الصريح أو أهداف إزالة الضجيج لتحسين مقايضة عنق الزجاجة المعلوماتي بشكل مباشر.

باختصار، يقدم الـ Bottlenecked Transformer بديلاً مقنعًا للاستدلال المعتمد على كثافة الرموز، حيث يثبت أن إعادة الكتابة الذكية والدورية للذاكرة الداخلية للنموذج يمكن أن تعزز قدرات الاستدلال بشكل كبير دون زيادة زمن انتقال الاستدلال عبر توليد الرموز.