Mamba: Linear-Time Sequence Modeling with Selective State Spaces

זה קורה עכשיו, אחרי 9 סקירות שחלקם היו די לא פשוטות הגענו למטרתנו הקדושה שזה Mamba. מכיוון שאני מפרסם סקירות בשלשות (באתר MDLI) אני אוציא עוד 2 סקירות נוספות של שכלולי ממבה (אחד מהם Mamba MoE והשני עוד הוחלט).

האמת שאחרי שאנו הבנו מה- (SSM (space-state models ואיך ניתן לבנות ארכיטקטורה מבוססת עליהם לעיבוד דאטה סדרתי, השכלול המוצע על ידי mamba הוא די אינטואיטיבי ומתבקש. כמו שאתם זוכרים SSM ממומשת בתור מערכת דינמית(DLS) לינארית כאשר הקלט למערכת זו היא ייצוג וקטורי (embeddings) של איברי הסדרה (= טוקנים). 

בשלב הראשון המערכת הדינמית מחשבת וקטור s הוא הוא ייצוג דחוס של זיכרון כלומר וקטור ״הזוכר״ את המידע הרלוונטי עבור כל הטוקנים הקודמים לטוקן הנוכחי. בשלב השני מחשבים את הפלט עבור טוקן זה המוזן לשכבה הבאה (שיכולה להיות גם שכבת שמייצרת פלט סופי). כל חישובים אלו מתבצעים באמצעות מיפויים לינאריים כלומר מכפלות במטריצות. חשוב להבין שכל המעברים בין ייצוגי הזכרון בין הטוקנים הם לינאריים ונשלטים על ידי אותה מטריצה A ווקטורי B, C וסקלר $\delta$. ד״א פרמטר $\delta$ מגדיר (באופן גס) את קצב דעיכה של הזכרון (כלומר ככל ש $\delta$ גבוה יותר אנו נוטים ״לזכור״ פחות מהטוקנים הקודמים).

מה היתרונות של הארכיטקטורה הזו? היא בעלת תכונה הדואליות המיוחלת המשלבת 2 התכונות הבאות:

• ניתן לחזות באופן מקבילי (בו זמנית) כמה טוקנים במהלך אימון (כמו בטרנספורמרים)
• חיזוי מהיר של טוקן במהלך היסק (ללא התחשבות בכל הטוקנים בחלון ההקשר כמו בטרנספורמרים שמביא לנו את הסיבוכיות הריבועית).

כלומר הארכיטקטורות מסוג זה הם יעילות בזמן האימון ומהירות בזמן ההיסק. אבל כמו שאתם יכולים לנחש יש לנו מחיר לשלם על כל התכונות הנחמדות האלו. ומחיר הוא כמובן יכולת של  המודל למדל תלויות מורכבות של הדאטה. עקב כך מאמרים כמו Hyena, H3, S4 ניסו ניסו לבנות את הפרמטרים של DLS (המגדירה מעברים בין ייצוגי הזכרון ויצירת הפלט) בצורה חכמה (ודי מורכבת).

אבל מתברר שזה לא מספיק. מעברים לינאריים עם פרמטרי DLS קבועים לא מסוגלת למדל דאטה מורכב (כמו שפה טבעית). אחד המשימות שמודל כזה נכשל עליה הוא העתקת טוקנים הבאים אחרי טוקן ספציפי (וזה די הגיוני לאור הפרמטרים הקבועים של DLS). כמו שאתם יכולים כבר לנחש אולי מחברי ממבה מציעים לעשות חלק מהפרמטרים (B, C ו-$\delta$) תלוים בייצוג הטוקן הנוכחי. התלות הזו היא לינארית עם מטריצות נלמדות. וזה עוזר לנו להתחשב בפיסת הקלט הנוכחית בצורה יותר טובה. כאמור B, C מגדירות את האופן בו ייצוג הזכרון והפלט עבור הטוקן הנוכחי בהתאמה ואז יש לנו סיכוי יותר טוב להצליח במשימות מהסוג שתיארתי לפני. בנוסף תלות של $\delta$ בייצוג הטוקן הנוכחי מקנה לנו אפשרות לשחק עם קצב דעיכה בצורה יותר גרנולרית שמקנה לנו יכולת ״לשכוח״ ו״לזכור״ איפה שצריך.

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

מה שכן קורה הוא זה החישובים הופכים לקצת יותר מורכבים (שימו לב שהחלק הכי בעייתי בחישוב שהוא העלאה של מטריצה בחזקה לא השתנה כי A נותרה קבועה). בסוף המאמר מציע כמה שכלולים לאופן חישוב המייעלים ומזרזים אותו (המשחקים בין זכרון מהיר ואיטי של GPU).

זה וזה עכשיו אתם יודעים מה זה ממבה.

Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

המאמר שנסקור היום בודק את האם מודלים שמבוססים על ארכיטקטורת ממבה על למידת in-context (או ICL). למעשה ICL היא יכולת של מודל לבצע למידה על בסיס כמה דוגמאות בודדות (גם נקרא למידת few-shot) ללא שינוי של משקלי המודל. בגדול יכולת זו של הטרנספורמרים לא מאוד מפתיעה  כי ה״חיזויים״ שלהם תלויים ביחסים בין חלקי הדאטה השונים(טוקנים) באופן מפורש באמצעות מנגנון ה-attention שלהם. כמובן יש מחקרים לא מעטים ומעניינים שחוקרים את התופעה המרתקת הזו ואני ממליץ לכם בחום להעיף מבט.

לעומת זאת הארכיטקטורה של ממבה לא לוקחת את היחסים בין הטוקנים השונים של הדאטה באופן מפורש ודוחסת את ה״עבר״ בוקטור אחד אז היכולת שלה לבצע ICL היא פחות אינטואיטיבי.  זה אכן פחות קורה. המאמר בדק כמה ארכיטקטורות מבוססות (SSM (state-space models כמו ממבה, S4 וגם S4-mamba ללא מנגנון attention של הטרנספורמרים והשוו את יכולות ICL שלהם עם ארכיטקטורות הברידיות: כלומר שילוב של ממבה עם מנגנון ה-attention של הטרנספורמרים.

איך משלבים ממבה עם הטרנספורמר? המאמר בדק שתי גישות (די דומות). בגישה הראשונה הוא החליף את MLP שיש בבלוקי טרנספורמר אחרי ה-attention במנגנון של ממבה. הגישה השנייה (הכי מוצלחת) הנקראת MambaFormer בנוסף מחליפה את הקידוד המיקומי (positional encoding) של עוברים הטוקנים בממבה נוספת. 

כאמור MambaFormer הגיע לביצועים הטובים ביותר מכל הארכיטקטורות הלא הברדיות (הטרנספורמר הטהור וכמה וריאנטים של SSM) באופן לא מפתיע בכלל. הרי MLP (רק 2 שכבות) ממדלים הפעולה די פשוטה ו-mamba היא למעשה מנגנון של זכרון הדוחס את המידע המהותי (בתקווה) של העבר (בטוקנים הקודמים). לא פלא שזה ניצח את כולם.

נשאר לנו רק לציין איזה משימות ניתנו למודלים אלו כדי לבחון את יכולות ICL שלהם. אחת המשימות היא לתת למודל כמה זוגות של  $(x, f(x))$  עבור פונקציה f לינארית ולבקש ממנו לחשב  $f(x)$ עבור $x$-ים נוספים. משימה אחרת היתה לתת לה נקודות שנדגמו מ Gaussian Mixture מסוים ולבקש ממנו לדגום עוד נקודות. טבלה עם כל המשימות מצורפת לפוסט.

נתראה בסקירה ממבה הבאה והאחרונה (לא בחרתי עדיין).

VMamba: Visual State Space Model

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

הפעם לא אספק לכם סקירה על ארכיטקטורות מבוססת SSM (עשיתי זאת לפחות ב 3 הסקירות הקודמות). כמו שאתם יודעים הטרנספורמרים השתלטו היום גם על תחום הויז'ן והחלק הארי של מודלי SOTA בתחום מבוסס על הטרנספורמרים. הטרנספורמרים החליפו את רשתות הקונבולוציה(CNN) ששלטו בתחום הויז'ן עד 2020 בערך. למרות שיש טוענים שבכל מודל ויז'ן עובד יש או איזה backbone מבוסס CNNs או שמכיל רכיב כמו attention לוקאלי (דרך להתחשב ב-inductive bias שיש בדאטה ויזואלי שאותו מנצלים CNNs), עדיין השליטה של הטרנספורמרים בויז'ן נראית די מוחלטת.

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

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

מה הבעיה העיקרית עם הכנסה של קצת -SSM -ים למודלים ויזואליים? הכיוון!! הרי עבור שפה טבעית וגם עבור אודיו די ברור שעבור  טוקן נתון צריך ״לזכור״ את הטוקנים מתחילת הטקסט/אודיו. בתמונה לטוקן נתון (פאץ') ניתן לבנות זיכרון מכיוונים שונים (הסדר חשוב בממבה וב-SSM-ים אחרים). אפשר להתחיל מלמעלה, או מלמטה שלה תמונה, ללכת ימינה או שמאלה. לא ברור מה הכי טוב מבחינת ביצועים ואז VMamba משלב אותם. לפאץ' נתון מתחילים מפאץ' השמאלי והעליון והפאצ'ים נכנסים ל-SSM בכיוון ימין-מטה. משלבים את ה-SSM הזה עם ה-SSM שנבנה החל מהפאץ' התחתון מימין כאשר כאשר הפאצ'ים נכנסים ל-SSM מכיוון שמאל-למעלה. ככה בונים את החלק ה-SSM של VMamba. שאר הדברים בארכיטקטורה המוצעת הם די סטנדרטיים.

שאר הדברים הדי סטנדרטיים: חלוקת תמונה לפאצ'ים, הפעלה כמה סבבי downsampling המבוצע עם בלוקים המכילים קונבולוציות 3×3, SSM, חיבור ResNet וכמה שכבות לינאריות.

זהו זה, סוף הסקירה מקווה שנהניתם לקרוא את סדרת ממבה…