מנגנון קשב חדש מפחית מטמון KV

מבוא

השימוש הגובר במודלי שפה גדולים (LLMs) והופעת פרדיגמות חדשות בהסקה העלו את האתגר של הסקה יעילה בקנה מידה גדול. צוואר בקבוק משמעותי הוא מטמון המפתח-ערך (KV) במנגנוני הקשב המסורתיים, שמתרחב באופן ליניארי עם גודל האצווה ואורך הרצף, והופך ל'זולל זיכרון' המעכב את קנה המידה וההרחבה של LLMs.

בעוד שגרסאות כמו MQA, GQA ו-MLA הופיעו כדי לטפל בכך, הן לרוב מתקשות לשמור על ביצועים תחת מגבלות זיכרון קפדניות או מציגות מורכבויות המציבות אתגרים הנדסיים ובעיות תאימות.

Multi-matrix Factorization Attention (MFA)

מאמר שפורסם לאחרונה על ידי חוקרים מ-Stepes, אוניברסיטת Tsinghua ומוסדות אחרים מציג ארכיטקטורת מנגנון קשב חדשה: Multi-matrix Factorization Attention (MFA) וגרסתה MFA-Key-Reuse (MFA-KR). מנגנון זה מפחית משמעותית את עלות ההסקה של מודל השפה תוך שיפור הביצועים בו זמנית.

MFA ו-MFA-KR לא רק עולים על MLA בביצועים אלא גם מתאימים לביצועים של MHA מסורתי תוך הפחתת השימוש במטמון KV בעד 93.7%.

MFA מתוכנן לפשטות, שכפול קל, רגישות נמוכה להיפרפרמטרים ותאימות לשיטות Pos-embedding שונות.

גישת MFA וניתוח

צוות המחקר ניתח את העיצוב הכללי והקיבולת של מנגנוני הקשב, וזיהה שני ממדים קריטיים הקשורים לקיבולת. ניתוח זה הוביל לפיתוח שיטות ניתוח ועקרונות עיצוב חדשים.

הם הציגו את המושג של Generalized Multi-Head Attention (GMHA) כמסגרת מאחדת להבנת גרסאות MHA שונות.

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

Fully Parameterized Bilinear Attention (FPBA) הוקם כגבול העליון התיאורטי של הביצועים. הם גילו ש-MHA וגרסאותיו הם פירוקים בדרגה נמוכה של FPBA.

השוואה עם MQA ו-MLA

הניתוח התמקד בשתי תוכניות שיפור מייצגות: Multi-Query Attention (MQA) ו-Multi-Head Latent Attention (MLA).

MQA משתמשת באסטרטגיית שיתוף פרמטרים אגרסיבית יותר, כאשר כל ראשי הקשב חולקים את אותו סט של פרמטרי מפתח-ערך. זה מפחית את השימוש בזיכרון אך עשוי להשפיע על יכולת הביטוי של המודל.

MLA מציגה מרחב סמוי משותף לדחיסת פרמטרים, אך כוח הביטוי בפועל מוגבל על ידי הממד הקטן ביותר, מה שאומר שהגדלת ממדים ביניים לא משפרת משמעותית את הביצועים.

חידושי מפתח של MFA

הפיתוח של MFA נבע מהמטרה ליצור מנגנון קשב שממזער את צריכת המשאבים תוך התקרבות לגבולות הביצועים התיאורטיים. העיצוב של MFA משלב שלושה חידושים עיקריים:

1. הגדלה משמעותית של מספר וממד ראשי הקשב כדי למקסם את קיבולת המודל.
2. שימוש באסטרטגיית פירוק דרגה נמוכה אגרסיבית כדי לשמור על יעילות הפרמטרים תוך הרחבת מספר ראשי הקשב והממדים.
3. שימוש בעיצוב ראש מפתח-ערך יחיד כדי לשמור על צריכת זיכרון מינימלית, גם עם מורכבות מודל מוגברת.

מדידת קיבולת והשוואה

כדי לנתח עוד יותר את MFA ומנגנוני קשב אחרים, הצוות הציג שני מדדים מרכזיים:

• Total Effective Rank (TER): מכפלת מספר ראשי הקשב ודרגת הפירוק לכל ראש (FRH).
• Shared Latent Subspace Dimension (SLSD): ממד המרחב הנסתר המשותף לכל ראשי הקשב.

MFA משיגה SLSD ו-TER גבוהים יותר בהשוואה ל-MQA.

בהשוואה ל-MLA, MFA משיגה גודל מטמון KV קטן יותר ו-TER גבוה יותר עם תקציבי פרמטרים דומים, תוך שמירה על SLSD דומה.

בהשוואה ל-MHA מסורתי, ל-MFA יש TER גבוה יותר, למרות שה-SLSD שלו קטן יותר.

תוצאות ניסיוניות

ניסויים נרחבים נערכו כדי להעריך את ביצועי הארכיטקטורה החדשה בסדרי גודל גדולים יותר, ובדקו מודלים שנעים בין 1B ל-7B פרמטרים ונתוני אימון מ-10B ל-1T.

MFA הדגימה יכולות קנה מידה דומות ל-MHA מסורתי, ושמרה על ביצועים מצוינים גם בסדרי גודל גדולים יותר.

בעוד של-MFA-KR היו ביצועים נמוכים מעט יותר, מגמת קנה המידה שלה התיישבה עם MHA. יתרונות החיסכון בזיכרון של MFA ו-MFA-KR המשיכו להתרחב עם גודל המודל, כאשר MFA השיגה 87.5% חיסכון בזיכרון ו-MFA-KR הפחיתה את השימוש בזיכרון ל-6.25% בסדר הגודל הגדול ביותר.

מחקרי אבלציה

מחקרי אבלציה אישרו את האפקטיביות של MFA ו-MFA-KR. יתרונות הביצועים שלהם אושרו גם בשיטות קידוד מיקום מרכזיות שונות.

תחזית

MFA מציעה שיפורים משמעותיים עם עיצוב פשוט, ומטפלת ביעילות בצוואר הבקבוק של הזיכרון בהסקה של LLM מבלי להוסיף מורכבות הנדסית נוספת. היא משתלבת בצורה חלקה במערכת האקולוגית הקיימת של Transformer, ומאיצה את היישום של LLMs בתרחישים שונים.