גלאי "Attention Sink" חדש מזהה טוקנים שצוברים את רוב מסת ה-attention במהלך שלב ה-prefill של הסקת (inference) מודלי שפה גדולים (LLM). הכלי מראה שהסרת טוקני ה-sink הללו מתוך ה-KV cache משביתה את ביצועי המודל, מה שהופך אותם לעוגנים חיוניים ליצירה (generation) יציבה.

למה שלב ה-prefill חשוב

רוב המחקר על LLM מתמקד בלולאת היצירה טוקן-אחר-טוקן, אך העבודה שמתרחשת לפני הטוקן הראשון — ה-prefill — קובעת את הטון לכל מה שיבוא אחריה. במהלך ה-prefill, המודל מעבד את כל ה-prompt, בונה את ה-KV cache ומפיץ את ה-attention על פני כל מיקום. מכיוון שה-softmax שמפיק את משקלי ה-attention חייב להסתכם ב-1, המודל אינו יכול לומר "שום דבר אינו רלוונטי". הוא "זורק" את מסת ההסתברות הנותרת על טוקן שקל להתייחס אליו, בדרך כלל הטוקן הראשון ברצף. טוקן זה הופך ל-attention sink.

מהו attention sink, במילים פשוטות

ב-Transformer, כל head מחשב משקל עבור כל זוג טוקנים. כאשר ההתפלגות מוטה מאוד, טוקן בודד עשוי לקבל נתח לא פרופורציונלי ממשקל ה-attention. התופעה אינה באג; היא נובעת מהמגבלה של ה-softmax. הטוקן הראשון (לעתים קרובות סימן תחילת משפט) משמש כ"שסתום שחרור" להסתברות היתרה שלא ניתן להקצות לשום מקום אחר.

ההשלכות על ניהול ה-KV-cache

KV caches שומרים את וקטורי ה-key וה-value עבור כל טוקן שעובד, מה שמאפשר חיפושים מהירים במהלך היצירה. בתרחישים של הקשר ארוך (long-context), אנשי מקצוע מבצעים pruning (גיזום) ל-cache כדי להישאר במסגרת מגבלות זיכרון ה-GPU. הגלאי החדש מראה שמחיקה חסרת הבחנה של טוקנים שנראים "לא חשובים" מוחקת גם את ה-sinks עצמם שהמודל מסתמך עליהם, מה שגורם לקריסת הביצועים. קיבוע (pinning) של טוקני ה-sink ב-cache שומר על שיווי המשקל הפנימי של המודל ושומר על יצירה יציבה.

איך הגלאי עובד

  • Eager attention: המימוש עוקף kernels מהירים כגון FlashAttention, שמכווצים את מטריצת ה-attention המלאה, ובמקום זאת רושם ציוני attention גולמיים עבור כל head.
  • Layer-wide aggregation: הציונים ממוצעים על פני כל השכבות וה-heads כדי להפיק פרופיל מסת-attention יחיד לכל טוקן.
  • Log-median absolute deviation (MAD): מכיוון שמשקלי ה-attention מוטים מאוד ימינה, מבחן mean-standard-deviation פשוט יסווג שונות רגילה בטעות כחריגים (outliers). טרנספורמציה לוגריתמית של המשקלים מנרמלת את ההתפלגות; שימוש ב-MAD לאחר מכן מסמן טוקנים שמסת ה-attention שלהם חורגת מהפיזור הטיפוסי.

שני סוגי sinks שנחשפו

  1. True sinks – טוקן תחילת המשפט (BOS) סופג בעקביות חלק עצום מה-attention, ללא קשר לתוכן ה-prompt.
  2. Structural sinks – טוקנים השייכים ל-prompts ברמת המערכת, כגון הסימנים המשמשים ב-ChatML (<im_start>, <im_end>), יוצרים אשכולות קטנים שגם הם מושכים attention. הם פועלים כגבולות לוגיים, המסייעים למודל להפריד בין הודעות משתמש להוראות מערכת.

נקודת מבט נגדית: האם אנחנו באמת זקוקים למטריצה הגולמית?

המחבר מציין כי ללא המספרים הגולמיים, תופעת ה-sink נותרת חבויה, וכל מדיניות pruning של ה-cache המבוססת על נתונים חסרים מסתכנת בערעור יציבות המודל.

מה כדאי לעקוב אחריו בהמשך

הגלאי הוא החלק השני בסדרה בת ארבעה חלקים שהחלה עם עוקב אנטרופיה לאיכות היצירה. החלק הבא ישלב זיהוי sinks עם speculative decoding מונחה אנטרופיה. המחקר האחרון יהיה מחקר אמפירי.

בשורה התחתונה: Attention sinks אינם מוזרויות נדירות; הם עמודים מבניים השומרים על התפלגות ה-attention של ה-Transformer מתנהגת היטב במהלך ה-prefill. כל אסטרטגיית pruning של KV-cache שמתעלמת מהם תסכן את יציבות המודל. זיהוי ושימור הטוקנים הללו הוא אמצעי הגנה בעלות נמוכה שיכול להפוך הסקת הקשר ארוך לאמינה ויעילה כאחד.