SR-TTT: Surprisal-Aware Residual Test-Time Training

إليك ملخص تقني مفصل لورقة البحث بعنوان "SR-TTT: التدريب في وقت الاختبار المتبقي الواعي بالمفاجأة (Surprisal-Aware Residual Test-Time Training)."

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

السياق: تعاني النماذج اللغوية الكبيرة (LLMs) تقليديًا من قيود تتعلق بذاكرة $O(N)$ والتعقيد الحسابي $O(N^2)$ لذاكرة المفاتيح والقيم (KV cache) في بنية المحول (Transformer). يقدم التدريب في وقت الاختبار (TTT) حلاً عبر استبدال ذاكرة KV الصريحة بـ "أوزان سريعة" ($W_{fast}$) يتم تحديثها عبر التعلم الخاضع للإشراف الذاتي أثناء الاستدلال، مما يحقق بصمة ذاكرة قدرها $O(1)$ ونافذة سياق لا نهائية نظريًا.

المشكلة الجوهرية: تعاني بنيات TTT الصرفة من فشل كارثي في الاستدعاء (Recall Failure) في مهام الاستدعاء الدقيق (مثل مهمة "الإبرة في كومة القش" - Needle-in-a-Haystack).

• آلية الفشل: تقوم الأوزان السريعة بضغط السياق بأكمله بقوة داخل عنق زجاجة معلوماتي. وتتم عملية استبدال وتجاهل الرموز (tokens) الفريدة أو المفاجئة (مثل الإبر، كمعرفات محددة أو أسماء) بسرعة بواسطة التحديثات المتدرجة للرموز الخلفية عالية التكرار.
• القيود الحالية: تعتمد النهج الهجينة الحالية غالبًا على نوافذ منزلقة ثابتة أو استدلالات تعتمد على درجات الانتباه (attention scores)، والتي تفتقر إلى إشارة ذاتية الإشراف مبدئية للتمييز بين الضجيج الخلفي القابل للضغط والمعلومات الحرجة غير القابلة للضغط.

2. المنهجية: بنية SR-TTT

يقترح المؤلفون بنية SR-TTT (التدريب في وقت الاختبار المتبقي الواعي بالمفاجأة)، وهي بنية هجينة تحافظ على فوائد الذاكرة $O(1)$ لـ TTT مع دمج آلية ذاكرة متفرقة محكومة بالخسارة (loss-gated sparse memory mechanism) للحفاظ على الاستدعاء الدقيق.

المكونات الرئيسية:

1. عمود فقري لـ TTT (TTT Backbone): نموذج TTT القياسي الذي يحدث الأوزان السريعة لتلخيص تاريخ التسلسل.
2. مرشح المفاجأة (آلية التوجيه - Surrounding Filter):
   • بدلاً من استخدام درجات الانتباه، يستخدم SR-TTT خسارة إعادة البناء في الحلقة الداخلية لـ TTT ($L_t = \|z_t - v_t\|^2$) كإشارة لعدم القابلية للضغط.
   • يتم وسم الرمز على أنه "مفاجئ" (وبالتالي يتم توجيهه إلى الذاكرة) إذا استوفى شرط المسار المزدوج:
     • تجاوزت خسارة الرمز لكل وحدة تجاوز العتبة المئوية الملساء باستخدام المتوسط المتحرك الأسي ($\tau_{EMA}$).
     • تجاوز متوسط الخسارة للكتلة المحلية التي تحتوي على الرمز نسبة معينة ($0.8 \cdot \tau_{EMA}$).
3. الذاكرة المتبقية (Residual Cache):
   • يتم تخزين الرموز الموسومة (تحديدًا قيم ومفاتيح ما بعد RoPE الخاصة بها) في ذاكرة متبقية (Residual Cache) ذات سعة ثابتة.
   • تعمل هذه الذاكرة كمسار ذاكرة موازٍ، متجاوزة الاختناق المتكرر للرموز الحرجة.
   • وتستخدم سياسة طرد قائمة على الأولوية.
4. بوابة دمج ألفا (Alpha Fusion Gate):
   • يتم دمج مخرجات العمود الفقري لـ TTT وانتباه الذاكرة المتبقية عبر متجه بوابة متعلم $\alpha$:
     $$\text{Output} = \text{TTT}(x) + \alpha \cdot \text{CacheAttention}(x)$$
   • ابتكار الاستقرار: لمنع "تلاشي التدرجات" الشائع في بوابات Sigmoid، استخدم المؤلفون توصيف التثبيت المباشر ($\alpha = \text{clamp}(\theta_{gate}, 0, \alpha_{max}$)) لضمان تدفق مستقر للتدرجات أثناء دمج الذاكرة.

استراتيجية التدريب: منهج تعليمي على مرحلتين

يفشل التدريب المباشر من البداية للنهاية بسبب "ضجيج البداية الباردة" (Cold Start Noise). في بداية التدريب، ينتج عمود TTT الفقري تمثيلات غير معايرة، مما يجعل الشبكة تقلل الخسارة عبر إجبار بوابات $\alpha$ على أن تكون 0.0 (مما يعطل الذاكرة فعليًا).

• المرحلة 1 (الخطوات 1–7,000): تدريب عمود TTT الفقري مع تعطيل الذاكرة المتبقية.
• المرحلة 2 (الخطوات 7,001–10,000): تجميد معاملات العمود الفقري وتفعيل الذاكرة المتبقية. هذا يجبر الشبكة على توجيه التدرجات حصريًا عبر وحدة $\alpha$ لتقليل الخسارة المتبقية، مما ينجح في "فتح" البوابات.

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

• التوجيه الواعي بالمفاجأة: تقديم طريقة مبتكرة لتحديد الرموز غير القابلة للضغط باستخدام خسارة إعادة البناء لحلقة TTT نفسها، بدلاً من درجات الانتباه الاستدلالية.
• بنية هجينة $O(1)$: إثبات إمكانية تحقيق الاستدعاء الدقيق في نماذج ذات سياق لا نهائي دون التضحية بمزايا الذاكرة $O(1)$ في الحالة القصوى، وذلك عبر استخدام ذاكرة متبقية متفرقة ومتوازية.
• حل التعلم المنهجي: اقتراح منهج تعليمي من مرحلتين لحل مشكلة "البداية الباردة"، مما يسمح بالدمج الديناميكي لوحدات الذاكرة التي قد يتم تجاهلها لولا ذلك أثناء التدريب.
• تنفيذ مفتوح المصدر: توفير الكود الكامل، وسكريبتات التدريب، والأوزان المدربة مسبقًا لتسهيل إعادة الإنتاج.

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

تم تقييم النموذج في مهمة "الإبرة في كومة القش" باستخدام مجموعة بيانات TinyStories مع إبرة مكونة من 8 أحرف أبجدية رقمية مدرجة في خلفية منخفضة الإنتروبيا.

• مواصفات النموذج: 15.8 مليون معلمة، 4 طبقات، $d_{model}=256$، تم التدريب على طول تسلسل 2048.
• مكاسب الأداء (عند العمق 0.50 و 0.75):
  • العمق 0.50: تحسن التطابق التام (Exact Match) من 10% (Pure TTT) إلى 33% (SR-TTT) (+23%).
  • العمق 0.75: تحسن التطابق التام من 17% (Pure TTT) إلى 37% (SR-TTT) (+20%).
• التحقق من الآلية: أكد التحليل أن المنهج التعليمي ذو المرحلتين نجح في دفع بوابات $\alpha$ لتفتح لتصل إلى حوالي 10% في الطبقات الدلالية العميقة، مما يؤكد فرضية التوجيه الانتقائي.
• حالة الفشل (RoPE): انهارت كلتا البنيتين (تطابق تام 0%) عند طول سياق 4096. ويُعزى ذلك إلى حدود استقراء تضمين الموضع الدوار (Rope)، وليس إلى آلية SR-TTT نفسها، حيث تم تدريب النماذج فقط على طول 2048.

5. القيود

1. النطاق: أُجريت التجارب على نطاق صغير (15.8 مليون معلمة). ولم يثبت بعد ما إذا كانت عتبات المفاجأة والمنهج التعليمي سينتقلان إلى نماذج بمليارات المعلمات.
2. استقراء الموضع: لا يمكن للنموذج إجراء استدلال صفري (zero-shot) يتجاوز طول التدريب (2048) بسبب قيود RoPE.
3. سياسة الطرد: قد تتدهور سياسة الطرد القائمة على الأولوية الحالية في الذاكرة المتبقية لتصبح سلوك FIFO قياسي تحت أحجام سياق قصوى، مما قد يؤدي إلى التخلص من "الإبر" المهمة المبكرة. هناك حاجة لسياسة طرد متعلمة لمزيد من المتانة.

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

يوفر SR-TTT إثبات مفهوم قوي على أن الاستدعاء الدقيق وضغط السياق اللانهائي ليسا متنافيين. ومن خلال الاستفادة من خطأ إعادة البناء الخاص بالنموذج كإشارة للاحتفاظ بالذاكرة، فإنه يقدم بديلاً مبدئيًا للأنظمة الهجينة القائمة على الاستدلال.

الاتجاهات المستقبلية:

• تنفيذ YaRN أو Dynamic NTK interpolation للتغلب على جدران استقراء RoPE.
• استبدال الطرد الاستدلالي بـ آليات تسجيل متعلمة (مستوحاة من TRIM-KV).
• توسيع التجارب لتشمل نماذج أكبر وسياقات أطول للتحقق من القدرة على التعميم.