ملخص تقني: تحديد الميزات عبر النواة المماسية العصبية التجريبية (eNTK)

بيان المشكلة

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

يستقصي هذا العمل ما إذا كان يمكن لاتجاهات القيم الذاتية العليا لـ النواة المماسية العصبية التجريبية (eNTK) أن تعمل كآلية للكشف عن هذه الميزات المتعلمة. تُعرف الـ eNTK بأنها النواة المتكونة من تقليص نسختين من مصفوفة جاكوبي (Jacobian matrix) للنموذج على طول اتجاه فضاء المعلمات:
$$K_{ij}(x_1, x_2) = \sum_{\mu} \frac{df_i(x_1)}{dW_\mu} \frac{df_j(x_2)}{dW_\mu}$$
حيث $f$ هو الشبكة العصبية، و $W_\mu$ تمثل الأوزان، و $i, j$ يشيران إلى فئات المخرجات. وتفترض الورقة البحثية أن الفضاءات الذاتية العليا لهذه النواة، عند تقييمها على مجموعة بيانات، تتوافق مع اتجاهات الميزات الحقيقية أو القابلة للتفسير، حتى في النماذج التي تعمل خارج نظام التدريب "الكسول" (lazy regime) الذي تنطبق عليه نظرية NTK القياسية بدقة.

المنهجية

يقترح المؤلفون خوارزمية لحساب وتحليل أعلى اتجاهات القيم الذاتية لـ eNTK في ثلاثة إعدادات مختلفة: نموذج MLP ذو طبقة واحدة، ومحول (Transformer) ذو طبقة واحدة، ونموذج لغوي كبير مدرب مسبقًا (Gemma-3-270M).

1. بناء النواة واختزالها

تكون الـ eNTK، عند تقييمها على مجموعة بيانات بحجم $N$ و $C$ من فئات المخرجات، ذات شكل $(N, N, C, C)$. ولإجراء تحليل القيم الذاتية (eigendecomposition)، يطبق المؤلفون استراتيجيتين للاختزال:

• eNTK الخاصة بالفئة (Class-specific eNTK): تحليل النواة $K_{cc}(x_1, x_2)$ لفئات محددة.
• eNTK المسطحة (Flattened eNTK): تكديس كتل الفئات المحددة في مصفوفة واحدة بحجم $NC \times NC$.
• eNTK على مستوى الطبقة (Layer-wise eNTK): جمع نواتج جاكوبي فقط فوق المعلمات التي تنتمي إلى طبقة معينة لإسناد الميزات إلى مكونات شبكية محددة.

2. القابلية للتوسع عبر تكرار لانزوس (Lanczos Iteration)

بالنسبة للنماذج الكبيرة حيث يكون تجسيد مصفوفة جاكوبي الكاملة أو الـ eNTK أمرًا غير ممكن (مثل النماذج اللغوية ذات المفردات الضخمة)، يستخدم المؤلفون تكرار لانزوس. يقومون بتقريب أعلى $k$ من اتجاهات القيم الذاتية عبر إجراء $2k$ من عمليات ضرب المصفوفة في المتجه. ومن الأهمية بمكان أنهم يحسبون $Kv = J(J^T v)$ باستخدام ضرب المتجه-جاكوبي وضرب جاكوبي-المتجه عبر التفاضل التلقائي (automatic differentiation)، مما يتجنب البناء الصريح لمصفوفة جاكوبي أو الـ eNTK.

3. إعادة البناء الفعالة للنماذج اللغوية

في تجربة Gemma-3-270M، يجعل حجم المفردات ($d_{vocab}$) الـ eNTK المسطحة عند طبقة المخرجات مستعصية حاسوبيًا. يستفيد المؤلفون من العلاقة الخطية بين مصفوفة جاكوبي لطبقة المخرجات ومصفوفة الطبقة الأخيرة (عبر مصفوفة عدم التضمين $U$). لقد اشتقوا مؤثرًا محولًا $\tilde{K} = S^{1/2} K_r S^{1/2}$ (حيث $K_r$ هي eNTK على تيار المتبقي/residual stream)، والذي يمتلك نفس القيم الذاتية لـ eNTK المخرجات الكاملة ولكنه يعمل في فضاء $d_{model}$ الأصغر. وهذا يتيح إعادة بناء أعلى اتجاهات eNTK دون تجسيد كائنات ضخمة بحجم المفردات.

4. مقياس التقييم

للتحقق من صحة الفرضية، يقيس المؤلفون التوافق بين فضاءات eNTK الذاتية ومتجهات الميزات المحددة بشكل مستقل "كحقيقة أرضية" (ground-truth).

• درجة التوافق (Alignment Score): تُحسب كمعيار فروبيني المربع (squared Frobenius norm) بين الفضاء الجزئي الذي تشغله أعلى $k$ من متجهات القيمة الذاتية لـ eNTK والفضاء الجزئي الذي تشغله ميزات الحقيقة الأرضية.
• المقارنة المرجعية (Baseline Comparison): في إعداد النموذج اللغوي، تتم مقارنة نهج eNTK مع نهج تحليل المكونات الرئيسية (PCA) المطبق على تنشيطات النموذج، باستخدام نفس الميزانية الحسابية (أعلى 25 اتجاهًا).

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

1. نموذج MLP على الحسابات القياسية (Modular Arithmetic)

في نموذج MLP ذو طبقة واحدة تم تدريبه على الجمع القياسي ($mod\ p$) والذي أظهر ظاهرة "الاستيعاب" (grokking) - وهي مرحلة انتقالية من الحفظ إلى التعميم:

• البنية الطيفية: يُظهر طيف eNTK "منحدرين" متميزين (كتل متصلة من القيم الذاتية العالية).
• توافق الميزات: يتوافق المنحدر الأول (بحجم $4\lfloor p/2 \rfloor$) تمامًا مع ميزات فوريه (Fourier features) لمتغيرات المدخلات ($a$ و $b$). ويتوافق المنحدر الثاني مع ميزات فوريه "المجموع" و"الفرق" ($a+b$ و $a-b$)، والتي تستخدمها الطبقة الثانية للنموذج لتنفيذ الخوارزمية الحقيقية.
• ديناميكيات التدريب: يكون التوافق للمنحدر الثاني مع أنماط المجموع منخفضًا عند البداية، ولكنه يزدزد بسلاسة، حيث تصل المشتقة الأولى للتداخل إلى أقصى حد لها بالقرب من بداية مرحلة الاستيعاب (grokking).

2. المحول (Transformer) على الحسابات القياسية

في نموذج محول ذو طبقة واحدة تم تدريبه على نفس المهمة:

• الترددات المتفرقة: على عكس الـ MLP، يتعلم المحول أنماط فوريه عند مجموعة متفرقة من الترددات العشوائية المعتمدة على البذرة (seed).
• التوافق على مستوى الطبقة: تتوافق أعلى اتجاهات eNTK على مستوى الطبقة مع ميزات فوريه عند تلك الترددات الرئيسية المحددة.
  • تتوافق كتلة الانتباه (Attention block) وأوزان مدخلات الـ MLP مع مجموع ميزات فوريه المدخلات ($\cos(\omega_k a) + \cos(\omega_k b)$).
  • تتوافق مخرجات الـ MLP وأوزان عدم التضمين مع ميزات فوريه "المجموع" ($\cos(\omega_k(a+b))$).
• الديناميكيات: على غرار الـ MLP، يزداد التوافق مع أنماط المجموع أثناء التدريب ويصل إلى ذروة معدل التغير بالقرب من انتقال الاستيعاب.

3. Gemma-3-270M على اللغة الطبيعية

في نموذج Gemma-3-270M المدرب مسبقًا، عند تقييمه على مجموعة بيانات نوافذ سياق TinyStories:

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

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

تدعي هذه الورقة أن تحليل القيم الذاتية لـ eNTK يقدم نهجًا جديدًا، مدفوعًا نظريًا، ومثبتًا تجريبيًا لتحديد الميزات في النماذج المدربة.

• ما وراء النظام الكسول (Beyond the Lazy Regime): توضح الورقة أن الهياكل الطيفية لـ eNTK تظل غنية بالمعلومات وتتوافق مع الآليات الحقيقية حتى في النماذج التي لا تقع في نظام التدريب "الكسول" (حيث يكون انزياح المعلمات ضئيلاً)، وهو النظام الذي لا تنطبق عليه نظرية NTK القياسية بدقة.
• التفوق على PCA التنشيطي: في سياق النموذج اللغوي، نجح نهج eNTK في استعادة الميزات القواعدية بشكل أفضل من PCA على التنشيطات، مما يشير إلى أن بنية النواة تلتقط معلومات الميزات التي قد تخفيها التنشيطات الخام (حتى عند اختزالها عبر PCA).
• المراقبة الديناميكية: إن ملاحظة تطور توافق الفضاءات الذاتية لـ eNTK مع الميزات أثناء التدريب — وخاصة مع بلوغ ذروة معدل التغيير بالقرب من الاستيعاب — تشير إلى أن تحليل القيم الذاتية لـ eNTK يمكن أن يعمل كأداة تشخيصية لمراقبة متى يتم اكتساب ميزات معينة أثناء التدريب.

يحافظ المؤلفون على موقف متواضع، مشيرين إلى أن نتائجهم حاليًا ارتباطية. فهم لم يثبتوا بعد أن التدخلات المستوحاة من eNTK تغير سلوك النموذج بشكل سببي، كما يقرون بالقيود المتعلقة بتوسع تجربة النموذج اللغوي (Gemma-3-270M أصغر من النماذج الحديثة الرائدة) وبساطة مجموعة البيانات (TinyStories). ومع ذلك، فإن اتساق النتائج عبر المهام الخوارزمية الاصطناعية واللغة الطبيعية يشير إلى إمكانات قوية لـ eNTK في مجال التفسيرية الآلية.