Tensor Cache: Eviction-conditioned Associative Memory for Transformers

ملخص تقني: ذاكرة التنسور المخبئية (Tensor Cache)

بيان المشكلة
يعتمد الاستنتاج في نماذج المحولات (Transformers) ذات التوليد الذاتي (Autoregressive) على تخزين أزواج المفتاح-القيمة (KV) لتجنب إعادة حساب البوادئ (prefixes). ومع ذلك، فإن حالة الـ KV المحتفظ بها تنمو خطياً مع طول السياق، والعمق، والتزامن، مما يخلق عنق زجاجة في الذاكرة. تعمل الذاكرة المخبئية ذات النافذة المنزلقة (Sliding-window caching) على حصر استخدام الذاكرة عبر الاحتفاظ بأحدث $W$ من الرموز فقط، لكن هذا النهج يعاني من مشكلة "النسيان القاسي": فبمجرد إخراج رمز من النافذة، يتم التخلص من معلوماته تماماً، مما يجعل الأدلة ذات الصلة خارج النافذة غير قابلة للوصول. وتعتمد طرق الحالة المحدودة الحالية إما على التخلص من الرموز المُخرجة بالكامل أو محاولة ضغط كل رمز في متوسط مستمر، وهو ما قد يؤدي إلى إدخال الضجيج أو الفشل في التمييز بين السياق الحديث والبعيد بفعالية.

المنهجية: ذاكرة التنسور المخبئية (Tensor Cache - TC)
تقدم ذاكرة Tensor Cache بنية ذاكرة مخبئية ذات حالتين (two-level bounded-state architecture) تحول عملية إخراج (eviction) الـ KV من حدث حذف إلى عملية كتابة في ذاكرة ترابطية.

1. بنية المستويين:

   • المستوى الأول (L1): مخزن حلقي (ring buffer) ثابت الحجم لـ KV يحمل أحدث $W$ من الرموز. يدعم هذا المستوى آلية انتباه softmax الدقيقة للسياق المحلي.
   • المستوى الثاني (L2): مصفوفة حاصل الضرب الخارجي $A$ ثابتة الحجم (ذاكرة الأوزان السريعة) التي تخزن المعلومات حصرياً من رموز الـ KV التي تم إخراجها من نافذة L1.
   • آلية الكتابة: عندما يقوم المخزن الحلقي باستبدال مدخل، لا يتم التخلص من زوج $(k, v)$ المزاح. بدلاً من ذلك، يتم كتابته في مصفوفة L2 عبر تحديث حاصل الضرب الخارجي: $A \leftarrow \lambda A + \eta (k \otimes v)$. حيث أن $\lambda$ (التضاؤل) و $\eta$ (معدل الكتابة) هما قياسات (scalars) يتم تعلمها لكل رأس (per-head).
   • آلية القراءة: تقوم الاستعلامات المستقبلية $q_t$ بالوصول إلى الماضي المضغوط عبر حساب $r_t = q_t A$. وبالاستفادة من هوية الانتباه الخطي $q(k \otimes v) = \langle q, k \rangle v$، تقوم هذه العملية باسترجاع المعلومات من جميع الرموز المُخرجة دون الحاجة لتخزينها بشكل فردي.
   • الدمج (Fusion): يقوم بوابة (gate) قياسية متعلمة $g$ بدمج مخرجات الانتباه المحلي (L1) وقراءة الذاكرة (L2): $y_t = y_{local} + \sigma(g) m_t$.

2. تحسين التدريب (المسح المتوازي - Parallel Scan):
   حدد البحث مشكلة حرجة في اختصارات التدريب المجزأة (chunked training) التقليدية. فالنهج الساذج الذي يلخص جزءاً (chunk) عبر متوسطه ($\bar{k} \otimes \bar{v}$) يُدخل منتجات خارجية عرضية ($k_i \otimes v_j$ حيث $i \neq j$) لا توجد أثناء الاستنتاج المتدفق (streaming inference).

   • الحل: قام المؤلفون بتنفيذ مسح المجموع الموزون المتوازي (parallel weighted-sum scan) الذي يكافئ رياضياً عمليات الكتابة المتسلسلة لكل رمز. وهذا يغلق الفجوة بين التدريب والاستنتاج ضمن هامش خطأ float32 ($< 10^{-7}$)، مما يضمن تعلم النموذج للديناميكيات الترابطية الصحيحة دون ضجيج التداخل بين الرموز.

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

• ذاكرة ترابطية مشروطة بالإخراج: على عكس ذاكرات حاصل الضرب الخارجي السابقة (مثل Infini-attention أو mLSTM) التي تستوعب كل رمز، تستخدم TC ذاكرة حاصل الضرب الخارجي حصرياً كذاكرة مخبئية من المستوى الثاني (L2) تتغذى فقط من عمليات إخراج النافذة المنزلقة. هذا يفصل الرموز بوضوح إلى "موجودة في L1" و"مضغوطة في L2".
• تصحيح الآثار الاصطناعية للتدريب: حدد البحث مشكلة (chunked-mean) التي تُدخل منتجات خارجية عرضية، وقام بحلها. يزيل المسح المتوازي المقترح هذا التباين بنفس التكلفة الحسابية.
• كفاءة الحالة المحدودة: توفر TC آلية تجعل السياق القديم يؤثر على التنبؤات دون الاحتفاظ بالمدخلات الفردية المُخرجة، مما يوفر بديلاً عملياً للاحتفاظ الكامل بـ KV.

النتائج التجريبية
قيم المؤلفون TC عبر مقاييس توسع الأنظمة، والاستدعاء الترابطي الاصطناعي، ونمذجة اللغة ذات السياق الطويل (OpenWebText و Shakespeare).

• الاستدعاء الترابطي الاصطناعي: في المهام التي تتطلب استرجاع معلومات تتجاوز النافذة المنزلقة، حققت TC دقة بنسبة 100%، مطابقة لأداء الـ Full KV. كما قللت من حالة الـ KV المحتفظ بها بنسبة 72-84% مقارنة بـ Full KV، مع تفوق ملحوظ على Infini-attention (التي تدهورت مع زيادة الفجوات) وعلى نماذج النافذة المنزلقة (التي كان أداؤها عند مستوى الصدفة).
• نمذجة اللغة ذات السياق الطويل: في OpenWebText، حققت TC أدنى متوسط لـ Negative Log-Likelihood (NLL) عند كل طول سياق تقييمي من 1,024 إلى 32,768 رمزاً، باستثناء الحالة $L=2,048$ حيث كانت غير متمايزة إحصائياً عن Infini-attention. عند $L=32,768$، وصلت TC إلى NLL قدره 5.14، متفوقة على Full KV (6.00) و Infini-attention (5.42)، مع استخدام 2.4 ضعف أقل من ذروة ذاكرة GPU مقارنة بـ Full KV.
• الإنتاجية (Throughput): تسببت TC في عبء إضافي بنسبة ~30% في الإنتاجية مقارنة بطرق النافذة المنزلقة النقية (بسبب قراءة وتحديث L2)، لكنها ظلت أسرع باستمرار من Infini-attention عبر جميع أطوال السياق المختبرة.
• تشخيص السعة: أظهرت التجارب على سعة الذاكرة الترابطية الخام أن أداء الاسترجاع يتدهور مع ضغط المزيد من الارتباطات غير المرتبطة في الحالة الثابتة، مما يؤكد الطبيعة المحدودة للذاكرة.

الأهمية والادعاءات
يزعم البحث أن Tensor Cache تنجح في سد الفجوة بين كفاءة الذاكرة في النافذة المنزلقة وقدرات الاستدعاء في السياق الكامل. ومن خلال معاملة إخراج (eviction) الـ KV كعملية ذاكرة مهيكلة بدلاً من كونها عملية حذف مدمرة، تمكن TC المحولات من الوصول إلى الأدلة ذات الصلة خارج النافذة المباشرة دون التكلفة الخطية لذاكرة الـ KV الكاملة. ويؤكد المؤلفون أن مساهمتهم الأساسية ليست في ابتكار ذاكرة حاصل الضرب الخارجي نفسها، بل في تطبيقها المحدد كذاكرة مخبئية L2 مشروطة بالإخراج، وفي التصحيح الصارم للآثار الاصطناعية للتدريب التي أعاقت مثل هذه الطرق سابقاً. تُقدم TC كبديل عملي محدود الحالة يحسن حدود جودة الذاكرة للاستنتاج طويل السياق.