לולאות באמצע השכבות חוסכות עד 18% בחישוב אימון

חוקרים הציגו מתכון בשם SMELT שמוכיח כי שימוש חוזר בשכבות אמצעיות של MoE Transformer משפר ביצועים גם כשמתאימים בדיוק את תקציב ה-FLOPs, הפרמטרים וה-KV cache.

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

מודלי Looped Transformers מגדילים עומק אפקטיבי על ידי הרצה חוזרת של בלוק שכבות משותף. עד כה רוב הבדיקות השוו באותו גודל מודל, כך שיתרון ארכיטקטוני התערבב עם תוספת חישוב טהורה. המאמר החדש בודק לולאות על Mixture-of-Experts תוך נעילת שלושת התקציבים, FLOPs לטוקן, פרמטרים שאינם הטמעות, ומטמון KV, כדי לבודד את תרומת הארכיטקטורה עצמה.

המתכון: SMELT

לאחר סדרת אבלציות הגיעו החוקרים ל-SMELT (Sparse MoE Transformer, middle layers Loop Twice): הלולאה מקיפה את מחצית השכבות האמצעיות ומריצה אותן פעמיים. המודל הלא-לולאתי (Baseline) וה-SMELT חולקים בדיוק אותו מספר פרמטרים לא-הטמעתיים, אותה עלות חישוב לטוקן ואותו גודל KV cache. הניסויים נערכו בארבעה גדלים עד 54 מיליארד פרמטרים לא-הטמעתיים.

חוקי סקיילינג נפרדים

לכל ארכיטקטורה הותאם חוק סקיילינג בסגנון Chinchilla. העקומה של SMELT יורדת מהר יותר: על החזית האופטימלית-חישובית היא חוסכת 6.8%-18.0% מ-FLOPs האימון הנדרשים להגיע לאותו הפסד אימות. הפער גדל ככל שהמודל גדל, מה שמצביע על יתרון שמתעצם בסקייל.

העברה לבנצ'מרקים וניתוח מכניסטי

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

שורה תחתונה

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