שכבות זיכרון: דרך חדשה לשיפור ביצועים במודלי שפה גדולים

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

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

שכבות הזיכרון משלימות את שכבות ה-feed-forward הצפופות (Dense Feed-Forward Layers), שמהוות את הליבה של רוב מודלי השפה. שכבות feed-forward הן "צפופות" מכיוון שהן מחברות כל נוירון (יחידת עיבוד בסיסית) בשכבה אחת לכל הנוירונים בשכבה הבאה. הן מבצעות חישובים כבדים על המידע, בעוד שכבות הזיכרון מספקות קיבולת לאחסון ואחזור מידע.

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

החוקרים מראים שניתן להרחיב את השימוש בשכבות זיכרון למודלים גדולים, ולהשיג שיפורים משמעותיים בביצועים, בעיקר במשימות הדורשות ידע עובדתי (לדוגמה, מענה על שאלות על נושאים ספציפיים). הם משווים את התוצאות למודלים צפופים (Dense) – כלומר מודלים שבהם לא נעשה שימוש בשכבות זיכרון – בעלי עלויות חישוב גבוהות יותר, וכן למודלים מסוג "תערובת מומחים" (MOE - Mixture-of-Experts), שהם מודלים המשתמשים במספר מודלים קטנים יותר במקביל.

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

כדי לפתור את הבעיה הזאת, החוקרים משתמשים במפתחות מקודדים בשיטת "product quantization" (Product-Quantized Keys) – שיטה לדחיסת מידע, שמאפשרת לבצע חיפוש מהיר ויעיל בזיכרון. בנוסף, הם מבצעים הקבלה (Parallelization) של פעולות החיפוש והצבירה בין מספר יחידות עיבוד גרפיות (GPUs), שהן מעבדים מיוחדים שמיועדים לחישובים מקבילים, מה שמאפשר לעבד מידע רב במהירות.

החוקרים גם בוחנים דרכים נוספות לשיפור הביצועים, כמו הוספת "שער" תלוי קלט (Input-Dependent Gating) עם פונקציה לא ליניארית silu – פונקציה שמטפלת במידע בצורה לא ליניארית, כלומר באופן שלא ניתן לייצג על ידי קו ישר. בנוסף, הם בודקים את ההשפעה של מיקום שכבות הזיכרון במודל ואת היתרון של שימוש בזיכרון משותף (Shared Memory), שבו מספר שכבות זיכרון משתפות את אותו הזיכרון, במקום שלכל אחת מהן יהיה זיכרון נפרד.

מבחינה טכנית, החוקרים מספרים על האופטימיזציות שביצעו במימוש שכבות הזיכרון בתוכנה PyTorch, ובפרט בפעולות ה-EmbeddingBag – פעולה נפוצה בעיבוד שפה. הם פיתחו קוד מיוחד (גרעיני CUDA מותאמים אישית - Custom CUDA Kernels), שהגדיל את רוחב הפס של הזיכרון ואפשר להם להגיע לביצועים טובים יותר.

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

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

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