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

טבלת קישורים

תקציר ומבוא 1

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

3 מודלים של מרחב מצב סלקטיבי ו-3.1 מוטיבציה: בחירה כאמצעי דחיסה

3.2 שיפור SSMs עם בחירה

3.3 יישום יעיל של SSMs סלקטיבי

3.4 ארכיטקטורת SSM פשוטה

3.5 מאפיינים של מנגנוני בחירה

3.6 פרטי דגם נוספים

4 הערכה אמפירית ו-4.1 משימות סינתטיות

4.2 מודל שפה

4.3 דוגמנות DNA

4.4 דוגמנות ויצירת אודיו

4.5 מדדי מהירות וזיכרון

4.6 אבלציות דגם

5 דיון

6 מסקנה והפניות

דיון: מנגנון בחירה

B עבודה קשורה

C מכניקה של SSMs סלקטיבי

D אלגוריתם מודע לחומרה עבור SSMs סלקטיבי

E פרטים ניסויים ותוצאות נוספות

תַקצִיר

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

1 הקדמה

מודלים של יסוד (FMs), או מודלים גדולים שהוכשרו מראש על נתונים מסיביים שהותאמו לאחר מכן למשימות במורד הזרם, הופיעו כפרדיגמה יעילה בלמידת מכונה מודרנית. עמוד השדרה של FMs אלה הם לרוב מודלים של רצף, הפועלים על רצפים שרירותיים של קלט ממגוון רחב של תחומים כגון שפה, תמונות, דיבור, אודיו, סדרות זמן וגנומיקה (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2016; Sutskever, Vinyals, and Quoc V Le 2014). בעוד שמושג זה אגנוסטי לבחירה מסוימת של ארכיטקטורת מודל, מכשירי FM מודרניים מבוססים בעיקר על סוג יחיד של מודל רצף: הרונאי (Vaswani et al. 2017) ושכבת תשומת הלב המרכזית שלו (Bahdanau, Cho, and Bengio 2015). יעילותה של תשומת לב עצמית מיוחסת ליכולתה לנתב מידע בצפיפות בתוך חלון הקשר, מה שמאפשר לה לדגמן נתונים מורכבים. עם זאת, תכונה זו מביאה עם חסרונות מהותיים: חוסר יכולת לדגמן שום דבר מחוץ לחלון סופי, וקנה מידה ריבועי ביחס לאורך החלון. גוף מחקר עצום הופיע על וריאנטים יעילים יותר של תשומת לב כדי להתגבר על החסרונות הללו (Tay, Dehghani, Bahri, et al. 2022), אך לעתים קרובות על חשבון עצם התכונות שהופכות אותו לאפקטיבי. נכון לעכשיו, אף אחת מהווריאציות הללו לא הוכחה כיעילה אמפירית בקנה מידה בין תחומים.

לאחרונה, מודלים של רצפי מרחב מצבים מובנים (SSMs) (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021) הופיעו כמעמד מבטיח של ארכיטקטורות למידול רצף. ניתן לפרש מודלים אלו כשילוב של רשתות עצביות חוזרות (RNNs) ורשתות עצביות קונבולוציוניות (CNNs), עם השראה ממודלים של חלל מצבים קלאסיים (Kalman 1960). ניתן לחשב מחלקה זו של מודלים ביעילות רבה כחזרה או קונבולולוציה, עם קנה מידה ליניארי או כמעט ליניארי באורך הרצף. בנוסף, יש להם מנגנונים עקרוניים למידול תלות ארוכת טווח (Gu, Dao, et al. 2020) באופני נתונים מסוימים, והם שלטו במדדים כגון Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). טעמים רבים של SSMs (Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington, and Linderman 2023) הצליחו בתחומים הכוללים נתוני אות רציפים כגון אודיו וחזון (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta, and Cui 2023). עם זאת, הם היו פחות יעילים במודלים של נתונים דיסקרטיים וצפופים במידע כגון טקסט.

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

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

אלגוריתם מודע לחומרה. השינוי הפשוט הזה מציב אתגר טכני לחישוב המודל; למעשה, כל המודלים הקודמים של SSMs חייבים להיות בלתי משתנים בזמן ובקלט כדי להיות יעילים מבחינה חישובית. אנו מתגברים על כך באמצעות אלגוריתם מודע לחומרה שמחשב את המודל באופן חוזר ונשנה עם סריקה במקום קונבולציה, אך אינו מממש את המצב המורחב על מנת למנוע גישה ל-IO בין רמות שונות של היררכיית זיכרון ה-GPU. היישום המתקבל מהיר יותר משיטות קודמות הן בתיאוריה (קנה מידה ליניארי באורך הרצף, בהשוואה לפסאודו-לינארית עבור כל SSMs מבוססי קונבולוציה) והן בחומרה מודרנית (עד פי 3 מהיר יותר ב-A100 GPUs).

ארכיטקטורה . אנו מפשטים ארכיטקטורות קודמות של מודל רצף עמוק על ידי שילוב עיצוב של ארכיטקטורות SSM קודמות (Dao, Fu, Saab, et al. 2023) עם בלוק MLP של רובוטריקים לבלוק אחד, מה שמוביל לעיצוב ארכיטקטורה פשוט והומוגני (Mamba) המשלב מרחבי מדינה סלקטיביים.

SSMs סלקטיבי, ובהרחבה ארכיטקטורת Mamba, הם מודלים חוזרים לחלוטין עם מאפייני מפתח שהופכים אותם למתאימים כעמוד השדרה של מודלים בסיסיים כלליים הפועלים על רצפים. (i) איכות גבוהה: סלקטיביות מביאה לביצועים חזקים בשיטות צפופות כמו שפה וגנומיקה. (ii) אימון מהיר והסקת מסקנות: חישוב וזיכרון משנים באופן ליניארי באורך רצף במהלך האימון, ופתיחת המודל באופן אוטורגרסיבי במהלך הסקת מסקנות דורשת רק זמן קבוע לכל שלב, שכן הוא אינו דורש מטמון של אלמנטים קודמים. (iii) הקשר ארוך: האיכות והיעילות יחד מניבות שיפורי ביצועים בנתונים אמיתיים עד אורך רצף של 1M.

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

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

• אודיו וגנומיקה. Mamba מעלה ביצועים קודמים של מודלים חדישים כגון SaShiMi, Hyena, ו-Transformers במודלים של צורות גל אודיו ורצפי DNA, הן במדדי איכות טרום אימון והן במדדים במורד הזרם (למשל הפחתת FID במערך נתונים מאתגר ליצירת דיבור ביותר ממחצית ). בשתי ההגדרות, הביצועים שלו משתפרים עם הקשר ארוך יותר של עד מיליון רצפים.

• מודל שפה. Mamba הוא מודל רצף הזמן הליניארי הראשון שבאמת משיג ביצועים באיכות שנאי, הן בתמיהה לפני אימון והן בהערכות במורד הזרם. עם חוקי קנה מידה של עד 1B פרמטרים, אנו מראים שממבה עולה על הביצועים של מגוון רחב של קווי בסיס, כולל מתכוני אימון טרנספורמטור מודרניים חזקים מאוד המבוססים על LLaMa (Touvron et al. 2023). למודל שפת ה-Mamba שלנו יש תפוקה של 5× דור בהשוואה לרובוטריקים בגודל דומה, והאיכות של Mamba-3B תואמת את האיכות של רובוטריקים פי שניים מגודלו (למשל, ממוצע גבוה יותר ב-4 נקודות בהיגיון בריא בהשוואה ל-Pythia-3B ואף עולה על Pythia-7B ).