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

בזמן האחרון, דגמי חלל מדינה (SSMs) הוכיחו יכולות וביצועים יוצאי דופן, והתחרו עם מודלים של ארכיטקטורת שנאים במדדי מידול בקנה מידה גדול תוך השגת מורכבות זיכרון כפונקציה של אורך הרצף והזמן ליניארי. יתרה מזאת, Mamba, מודל שטח של המדינה ששוחרר לאחרונה, הראה ביצועים יוצאי דופן במגוון של מודלים של שפות ומשימות עיבוד ברצף ארוך. במקביל, דגמי Mixture of Expert (MoE) הציגו גם ביצועים מרשימים תוך הפחתה משמעותית של השהייה ועלויות החישוביות של מסקנות, אם כי על חשבון טביעת זיכרון גדולה יותר. בהתבסס על מודלים של Mamba ו-MoE, מאמר זה ידון ב-BlackMamba, ארכיטקטורה חדשנית המשלבת את מודל ה-Mamba State Space עם מודלים של MoE כדי למנף את היתרונות שמציעות שתי המסגרות. ניסויים ב-BlackMamba הוכיחו את יכולתה לעלות על מסגרת ה-Mamba הקיימת ועל קווי הבסיס של השנאים, הן ב-FLOPs והן בהסקת האימון. הביצועים יוצאי הדופן של המסגרת של BlackMamba מראים שהיא יכולה לשלב ביעילות את היכולות של מסגרות ה-Mamba וה-MoE, ומציעה הסקה מהירה וחסכונית מ-MoE עם יצירת מורכבות ליניארית מממבה.

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

BlackMamba: מבוא ל-MoE for State Space Models

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

מאמצים משמעותיים נעשו בשנים האחרונות בניסיון להתגבר על מגבלות אלו, ותשומת הלב הוסטה לכיוון של תכנון חלופות ארכיטקטוניות לדגמי שנאי הקשב הצפופים הקנוניים, כאשר מודלים של SSM ו-MoE הם הארכיטקטורות המועמדות המבטיחות ביותר. היתרון העיקרי שנקטף על ידי העדפת מודלים של מרחב המדינה על פני מודלים של ארכיטקטורת שנאים היא המורכבות החישובית הליניארית ביחס לאורך רצף הקלט שמציעים SSMs בניגוד למורכבות הריבועית שמציעים שנאים. תיאורטית, מורכבות חישובית ליניארית ביחס לאורך רצף הקלט מאפשרת למודלים של מרחב מצבים לעבד רצפים גדולים יותר ממודלים של ארכיטקטורת שנאים עבור תקציב נתון של FLOPS או פעולות נקודה צפה לשניה, ולהפוך את היצירה האוטורגרסיבית לקבועה בחישוב ללא מטמון KV. מודלים של מרחב מדינה שפותחו לאחרונה, כולל Mamba, RetNet ועוד כמה, הוכיחו הסקה והדרכה יעילה ברצף ארוך, יחד עם ביצועי משימות דוגמנות שפות תחרותיות לשנאים בעלי מאפייני קנה מידה דומים. מצד שני, שילוב של ארכיטקטורות של דגמי מומחים צובר פופולריות כחלופה לשנאים צפופים מכיוון שהיא מאפשרת הפחתה משמעותית בהסקת מסקנות ואימון FLOPs החיוניים להשגת איכות דומה למודל צפוף. דגמי MoE (Mixture of Experts) פועלים על ידי הפעלת מבחר דליל בלבד של כל הפרמטרים במהלך מעבר אחד קדימה. הם משתמשים בפונקציית ניתוב כדי לקבוע אילו 'מומחים' נקראים לפעולה בהתבסס על ההקשר הנתון. גישה זו יוצרת הפרדה בין העלות החישובית של מסקנות לבין המספר הכולל של הפרמטרים, מה שמאפשר ביצועים משופרים במסגרת תקציב מסקנות קבוע, אם כי עם מספר מוגדל של פרמטרים ודרישת זיכרון גדולה יותר.

התקדמות זו בארכיטקטורה מציעה יתרונות בולטים על פני שנאים מסורתיים ומייצגת כיוון מרגש להמשך פיתוח. אנו מניחים כי שילוב השיפורים הללו במודל Mamba-MoE משולב עשוי להאיץ משמעותית את יכולות מודל השפה והיעילות מעבר למודלים של שנאים סטנדרטיים. היתרונות הצפויים של ארכיטקטורת Mamba-MoE בהשוואה למודל שנאי צפוף מסורתי כוללים:

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

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

עם זאת, חיוני לציין שמודלים של ארכיטקטורת שנאים הם עדיין מהשורה הראשונה, והפגינו ביצועים חזקים עקביים ומרשימים במשימות דוגמנות שפה ומשימות עיבוד רצף. בבסיסה, ארכיטקטורת השנאי משתמשת בתשומת לב עצמית המבצעת השוואה ריבועית של הכל לכול של קווי הדמיון של תוצר הנקודה בין ההטמעות של אסימונים שונים ברצף, ומבצעת מפה ליניארית לוקטור פלט. מודל השנאי מורכב מגושי תשומת לב עצמית המוערמים בין בלוקים של MLP או Multi-Layer Perceptron אשר מורכבים יותר מ-MLP דו-שכבתי עם פונקציית הפעלה נתונה. 

BlackMamba: אדריכלות ומתודולוגיה

דגמי חלל מדינה

מודלים של מרחב מצבים שייכים לקבוצת מודלים של רצף בעלי מורכבות לינארית ביחס לאורך רצף הקלט. הארכיטקטורה של מודלים של מרחב המדינה מתיישרת יותר עם רשתות עצביות חוזרות ורשתות עצביות מתגלגלות ולא עם ארכיטקטורה מבוססת תשומת לב, והיא שואבת השראה ממערכת דינמית רציפה שממפה פונקציה חד-ממדית דרך מרחב סמוי סמוי. מערכת דינמית ליניארית הופכת חישובים מקבילים ליעילים באמצעות סריקה אסוציאטיבית או קונבולוציה. בתרחישים מעשיים, האופי החוזר של מודלים של שטחי מדינה הייתה הסיבה לכך שעדיין יש לאמץ אותם על חומרת AI מקבילה ביותר כמו GPUs. עם זאת, הופעתם של SSMs כמו RWKV ו ממבה השתמשו בליבת סריקה מקבילה כדי למפות פעולות חוזרות ביעילות למעבדי GPU, ובכך להקל על האימון של ארכיטקטורות חדשות ביעילות דומה לאלו שהושגו במודלים של שנאים. 

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

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

תערובת של דגמי מומחים

שילוב של מודלים של מומחה (MoE) משיג הפרדה בין עלות ההסקה לבין ספירת הפרמטרים הכוללת על ידי הפעלה סלקטיבית של פרמטרים במהלך העברה קדימה. במקום להשתמש בכל הפרמטרים, מודלים אלה מפנים אסימונים למומחים ספציפיים של Multilayer Perceptron (MLP). באופן אידיאלי, כל מומחה מותאם לעיבוד סוג מסוים של קלט, עם מנגנון ניתוב, בעצם רשת עצבית קומפקטית, הקובע את המומחה המתאים ביותר עבור כל אסימון. גישה זו שואפת לשמר את כוח הביטוי המקיף של מודל עם מספר שווה של פרמטרים בתצורה צפופה יותר, אך עם דרישות חישוב מופחתות במידה ניכרת. בדרך כלל, הנתב הוא מיפוי של השכבות הליניאריות מאסימונים ועד למדדי מומחים כאשר כל מומחה הוא פשוט שנאי סטנדרטי Multilayer Perceptron. עם זאת, המפתחים עדיין לא מצאו את שיטת האימון האופטימלית עבור הנתב מכיוון שבעיית הקצאת המומחים אינה ניתנת להבדלה, ודגמי Mixture of Expert נאבקים לעתים קרובות עם איזון עומסים ויציבות אימון בין מומחים שונים ליעילות החומרה. 

אדריכלות

בליבה, BlackMamba משתמשת במודל שנאי סטנדרטי המורכב מגושי MLP משולבים ובלוקי קשב שנוספו ברצף לאורך זרם שיורי. כעת, רוב הדגמים של Mixture of Expert פשוט מחליפים את בלוקי הפרצפטרון הרב-שכבתיים בשכבת מומחה מנותבת. מצד שני, המסגרת של BlackMamba לא רק מחליפה את בלוק הפרצפטרון הרב-שכבתי בשנאי בשכבת מומחה מנותבת, אלא גם מחליפה את שכבת הקשב בשכבת Mamba State Space Model. הארכיטקטורה של מסגרת BlackMamba מודגמת באיור הבא. 

הדרכה ומערך נתונים

מודל BlackMamba מאומן על למעלה מ-300 מיליארד אסימונים על מערך נתונים מותאם אישית, ומשתמש בפונקציית ההפעלה של SwiGLU עבור התפיסות הרב-שכבתיות המומחים. המסגרת מתאמנת עם 8 מומחים, מספר שהמפתחים מצאו שהוא האיזון הנכון ומתפשר בין טביעת הרגל של הזיכרון לבין עלות ההסקה של המודל. מערך הנתונים המותאם אישית המשמש לאימון המסגרת של BlackMamba מורכב מתערובת של מערכי נתונים קיימים כבר בקוד פתוח כולל Starcoder, SlimPajama, Pile ועוד. הטבלה הבאה מדגימה את המשקלים של כל אחד ממערכי הנתונים המשמשים לאימון מסגרת BlackMamba. בסך הכל, ישנם 1.8 טריליון אסימונים במערך הנתונים. 

BlackMamba: תוצאות

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

יתר על כן, יתרונות ההשהיה של שני דגמי ה-Mixture of Expert ו-Mamba משולבים במסגרת BlackMamba וכתוצאה מכך זמני הסקה מהירים יותר באופן משמעותי בהשוואה לדגמי שנאים, דגמי Mamba טהורים ומודלים של MoE. יתר על כן, יתרון ההסקה של המסגרת של BlackMamba עומד ביחס ישר לאורכי הרצף, מה שהופך את BlackMamba ליעילה ביותר ביצירת רצף ארוך. בהמשך, האיור הבא ממחיש את מספר האסימונים שהוקצו לדגמי BlackMamba עם 340 מיליון ו-640 מיליון פרמטרים בהתאמה. כפי שניתן לראות, רוב השכבות מדגימות רמה גבוהה של איזון מומחים כתוצאה מהאלגוריתם המשופר של Sinkhorn המיושם על ידי דגמי BlackMamba. 

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

מחשבות סופיות

במאמר זה דיברנו על BlackMamba, ארכיטקטורה חדשנית המשלבת את מודל החלל של Mamba State עם שילוב של מודלים של מומחים כדי לקצור את היתרונות שמציעות שתי המסגרות הללו. ניסויים ב-BlackMamba הוכיחו שהיא מתעלה על מסגרת ה-Mamba הקיימת ועל קווי הבסיס של השנאים, הן ב-FLOPs והן בהסקת האימון. הביצועים יוצאי הדופן של מסגרת BlackMamba מוכיחים שהיא מסוגלת לרשת ולשלב את היכולות של ה-Mamba ו-MoE בצורה יוצאת דופן, שכן היא משלבת את ההסקה הזולה והמהירה מ-MoE עם יצירת מורכבות ליניארית מממבה. דיברנו על האופן שבו הארכיטקטורה של מסגרת BlackMamba מסוגלת להתעלות על מודלים גדולים מאומנים חזקים, מסגרת Mamba קיימת ומודלים של Mixture of Expert במונחים של FLOPs של אימון ועלות ההסקה. יתרה מזאת, המסגרת של BlackMamba גם יורשת את ה-FLOPs של הדור והכשרה מופחתת משני דגמי ה-Mixture of Expert ו-Mamba בו-זמנית.