אז מה הקטע עם מודלי שפה גדולים? (חלק ד' - מטריצות, הרבה מטריצות)

מבוא

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

1. המרה של כל טוקן לוקטור שמייצג אותו.

2. הפעלה של שלב שנקרא Attention שבו כל וקטור אוסף מידע מוקטורים שבאו לפניו בסדרה.

3. הפעלה של שלב שנקרא MLP שבו מתוך כל וקטור מחושב וקטור חדש שאמור לתאר את "מה שצריך לבוא אחריו".

4. תרגום של כל וקטור אל דירוג של כל הטוקנים שנקרא logits.

את שלבים 1 ו-4 ראינו בפוסט הקודם. בשני המקרים אפשר לחשוב עליהם פשוט בתור כפל של מטריצה (גדולה) בוקטור. בשלב 1, יש לנו וקטור שהאורך שלו הוא כמספר הטוקנים שה-LLM מכיר (גודל המילון ) שלו ויש בו 1 בכניסה שמתאימה לטוקן שעובדים עליו כרגע, ו-0 ביתר הכניסות, וכופלים אותו במטריצה שמספר השורות בה הוא גודל המילון ומספר העמודות בה הוא הגודל של הוקטור שמייצג את הטוקן. התוצאה היא אחת מהשורות של המטריצה הזו (בפועל השלב הזה לא ממומש בתור כפל מטריצות כי יש דרכים יעילות יותר לעשות את זה). את הוקטור שמתקבל אני מסמן ב-\(h\) וקורא לו hidden state. עבור המודל הספציפי שאנחנו מסתכלים עליו, גודל המילון הוא \(49,152\) וגודל ה-hidden state הוא \(2,048\), אז המטריצה של שלב 1 היא מסדר \(49152\times2048\) .

בשלב 4 עושים משהו מאוד דומה: יש לנו hidden state, \(h\), ואנחנו כופלים אותו במטריצה מסדר \(2048\times49512\) כדי לקבל בחזרה וקטור שהאורך שלו הוא כמספר הטוקנים; כל כניסה שלו מכונה logit והיא מין דירוג מספרי של כמה הטוקן הספציפי הזה רלוונטי אל ה-hidden state. במודל הספציפי שלנו המטריצה של שלב 4 היא אותה מטריצה כמו בשלב 1 (עד כדי ביצוע של שחלוף שהופך כל שורה לעמודה וההפך). בפוסט הקודם שיחקתי את המשחק של להפעיל את שלב 1 ואז מייד להפעיל את שלב 4 ולראות מה מקבלים; מילה כמו cat החזירה הרבה וריאציות על cat(למשל עם אותיות גדולות או רווח בהתחלה) וגם למשל את feline שהיא מילה שונה ואפילו לא דומה בכלל בתור מחרוזת (אין לה אותיות משותפות עם cat) אבל מבחינת הסמנטיקה היא בוודאי מילה רלוונטית. זה משכנע אותנו ששלבים 1 ו-4 הם באמת תרגום טוב מ"עולם הטוקנים" אל "עולם הוקטורים", אבל מה שלא קורה כאן הוא מה ש-LLM אמור לעשות - להציע את הטוקן הבא אחרי הטוקן הנוכחי, בהתבסס על הטוקן הנוכחי אבל גם על ההיסטוריה של יתר הטוקנים עד כה. את האפקט הקסום הזה משיגים על ידי חזרה שוב ושוב על שלבים 2 ו-3 (במודל שלנו - חוזרים 24 פעמים על השלבים הללו, כל פעם עם פרמטרים אחרים), אז בואו נבין בפוסט הזה מה קורה בהם.

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

מה קורה בשלב ה-MLP?

בגדול, מאוד בגדול, שלב ה-MLP הוא

• מכפילים את \(h\) במטריצה ומקבלים וקטור ארוך יותר.

• מפעילים על הוקטור הארוך יותר פונקציה לא לינארית.

• מכפילים אותו במטריצה שמחזירה אותנו לאורך המקורי של \(h\) .

כלומר, יש לנו שתי מטריצות \(A,B\) ופונקציה לא לינארית כלשהי \(f\), ואנחנו מחשבים את \(B\left(f\left(Ah\right)\right)\) . הרעיון הוא ש-\(f\) היא פונקציה ספציפית שבה משתמשים תמיד, אבל \(A,B\) הן מטריצות שמורכבות מפרמטרים שמחושבים כחלק מתהליך האימון של המודל, ושונים בין חזרות שונות של שלב ה-MLP.

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

אז הנה אינטואיציה שכזו. כשכופלים את \(h\) במטריצה, אפשר לחשוב על זה כאילו אנחנו מחשבים הרבה פונקציות לינאריות על \(h\) שכל אחת מהן מחזירה ערך מספרי בודד (כל שורה של המטריצה מקודדת את החישוב של אחת הפונקציות הללו). על כל פלט כזה אפשר לחשוב בתור "פיצ'ר" של \(h\) - ערך שמתאר אותו עבור תכונה ספציפית: האם \(h\) הוא שם עצם? האם הוא טורף? האם הוא מילה גסה? וכו' וכו'. כלומר אנחנו מחלצים הרבה סוגי מידע שונים ומשונים על \(h\) לוקטור גדול יותר. עכשיו אנחנו מפעילים את \(f\) על סוגי המידע הללו; הרעיון פה הוא לחזק את מה שכבר חזק, בצורה לא לינארית. למה לא לינארית? כי בשלב הבא אנחנו שוב כופלים במטריצה; אם כל מה שהיינו עושים היה לכפול במטריצה וא לכפול בעוד מטריצה, היה אפשר להחליף את שני השלבים הללו בשלב יחיד של כפל במטריצה (המכפלה של שתי המטריצות הללו), כלומר הייתה לנו פחות אינפורמציה שיכולה לבוא לידי ביטוי. המבנה הזה של "מטריצה-לא לינארי-מטריצה" מאפשר לנו ליצור טרנספורמציות עשירות יותר (זו אבחנה שהייתה מוכרת כבר זמן רב לפני שהיו טרנספורמרים).

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

במודל הספציפי שלנו השלב הזה הוא טיפה יותר מורכב. יש מטריצה שלישית מעורבת, שהתפקיד שלה היא לשמש בתור מעין "שער" - להחליט על פיצ'רים מסויימים שאנחנו הולכים להתעלם מהם יותר ופחות, בהתבסס על \(h\) המקורית. בואו נקרא למטריצה הזו \(C\) . מה שעושים הוא:

• מחשבים את \(Ah\)

• מחשבים את \(Ch\) ומקבלים את וקטור ה"שער"

• מפעילים את \(f\) על השער, \(f\left(Ch\right)\)

• ההפעלה של השער עצמו היא על ידי כפל נקודתי: \(\left(Ah\right)\odot\left(f\left(Ch\right)\right)\), כלומר כופלים כל זוג כניסות באותו מקום בשני הוקטורים.

• את כל זה כופלים ב-\(B\) כדי לחזור לוקטור בגודל ה-hidden state

עכשיו אפשר לתת לדברים את השמות שיש להם במודל עצמו. המטריצה \(A\) נקראת \(W_{up}\) (כי היא "מרימה" את ה-hidden state אל רמת ה"פיצ'רים"). המטריצה \(C\) נקראת \(W_{gate}\) ואילו \(B\) נקראת \(W_{down}\) . הפונקציה \(f\) נקראת \(\text{SiLU}\), ולכן הפעולה הכוללת ששלב MLP מבצע הוא

\(\text{MLP}\left(h\right)=W_{down}\left[\text{SiLU}\left(W_{gate}h\right)\odot W_{up}h\right]\)

מה זו SiLU? ה-Si מגיע מ-Sigmoid, שזה שם כללי לפונקציות עם גרף דמוי האות \(S\), וה-LU אומר Linear Unit. הרעיון הוא שלוקחים את הפונקציה הלינארית הפשוטה \(f\left(x\right)=x\) ומוסיפים לה כפל בפקטור של פונקציה לא לינארית, במקרה הזה הפונקציה \(\frac{1}{1+e^{-x}}\) (זו פונקציה שכאשר \(x=0\) היא בדיוק \(\frac{1}{2}\), כאשר \(x\) שואף לאינסוף היא בדיוק 1 וכשהוא שואף למינוס אינסוף היא בדיוק 0). כלומר בסך הכל

\(\text{SiLU}\left(x\right)=\frac{x}{1+e^{-x}}\)

זהו, זה כל התיאור של השלב הזה! ה"מוח" האמיתי שלו הוא שלוש המטריצות \(W_{up},W_{down},W_{gate}\) . מה הגודל שלהן? מימד אחד שלהן צריך להיות \(2048\), כגודל ה-hidden state, אבל יש למודל חופש לבחור את המימד השני; במודל הספציפי שלנו זה \(8192\) (כלומר, גדול פי 4 מה-hidden state). זה אומר ש-\(W_{up},W_{gate}\) שתיהן ממימד \(8192\times2048\) ואילו \(W_{down}\) ממימד \(2048\times8192\) . בכל שלושת המקרים, מספר הפרמטרים שלהן הוא 16,777,216, כך שיוצא שכל שלב MLP משתמש בערך בחמישים מיליון פרמטרים. זה לא מעט! שלב התרגום של טוקן לוקטור דרש בערך מאה מיליון פרמטרים, אבל אצלנו שלב ה-MLP לא קורה פעם אחת אלא 24 פעמים, ובכל פעם עם פרמטרים שונים, כך שבסך הכל יש לנו בפעולה 1.2 מיליארד פרמטרים! כזכור, המודל שלנו נקרא SmolLM2-1.7B-Instruct, כשה-1.7B שם אומר "בערך 1.7 מיליארד פרמטרים". זה עתה ראינו שרובם המוחץ נמצא בשלב ה-MLP. איפה היתר שטרם דיברנו עליהם נמצאים? בשלב ה-Attention, אז בואו נדבר על השלב הזה... רגע רגע רגע, עוד לא סיימנו. עדיין צריך לדבר על משהו שלא נחשב בדיוק חלק משלב ה-MLP אבל מגיע מייד לפניו ואחריו, אחרת לא באמת נבין עד הסוף מה קורה פה.

ראשית, התוצר של שלב ה-MLP לא הופך להיות ה-hidden state החדש שלנו - הוא מתווסף אליו. כלומר, אנחנו מבצעים

\(h\leftarrow h+\text{MLP}\left(h\right)\)

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

שנית, בואו נעשה ניסוי קטן. אני מקודד את "cat" בתור וקטור ומדפיס אותו. ההתחלה היא \(\left[0.0845,0.0258,0.0012,\ldots\right]\) . אוקיי, נחמד. עכשיו אני מפעיל את שלב ה-MLP פעם אחת, כולל החיבור בסוף; אני מקבל \(\left[-0.2773,1.5938,-0.8242,\ldots\right]\) שזה שונה אבל אוקיי, סבבה, אני אמור לקבל משהו שונה. עכשיו אני מפעיל את שלב ה-MLP שוב ומקבל \(\left[158.,306.,-220,\ldots\right]\) וזה... קצת מוזר. קודם היו לי מספרים קטנים, עכשיו פתאום המספרים הם מסדר גודל של מאה. מה הולך פה? אחרי 6 שלבים אני מקבל NaN בכל הכניסות - מה שקרה הוא שהמספרים התנפחו יותר מדי, מעבר לגודל הייצוג של מספרי נקודה צפה שפייתון יודע להתמודד איתם.

זה לא דבר מפתיע בשום צורה. זה קורה באופן כללי באלגברה לינארית - טרנספורמציות לינאריות שרירותיות יכולות "להגדיל" את הוקטור. יש לנו מושג קונקרטי של נורמה כדי לתאר את זה - גם של וקטור וגם של הטרנספורמציה עצמה. יש מחלקה שלמה של טרנספורמציות שמשמרות נורמה - טרנספורמציות אוניטריות, שמשחקות תפקיד חשוב בחישוב קוונטי, ובגלל שאני מתעסק עם חישוב קוונטי ביומיום אני כבר רגיל לכך שאפשר לפעול על וקטור שוב ושוב ושוב והוא יישאר נחמד - אבל אצלנו זה לא המצב. אז מה עושים? מנרמלים. עושים משהו על הוקטור כדי להקטין את הכניסות שלו מבלי לפגום ב"משמעות" שאנחנו מייחסים לו. יש כל מני דרכים לעשות את זה, ובדרך שבה משתמשים כאן לא סתם מפעילים על הוקטור פונקצייה קונקרטית כדי לנרמל אותו - משתמשים בפונקציה שתלויה בפרמטרים שהמודל לומד. זה משהו שעדיין לא סיפרתי עליו - אמרתי שהמודל לומד פרמטרים בשביל התרגום הראשוני מ-token והחזרה אל מרחב ה-token-ים בסיום, והוא לומד פרמטרים לשלב ה-MLP והוא לומד פרמטרים גם לשלב ה-Attention; אבל בנוסף לכך, וזה כבר ממצה סופית את רשימת הפרמטרים, הוא לומד (יחסית מעט מאוד) פרמטרים עבור שלבי הנרמול שמופעלים גם לפני MLP וגם לפני Attention. כפי שאנחנו רואים עכשיו, אלו אולי שלבים "קטנים" אבל הם קריטיים; בלעדיהם המודל לא יעבוד.

הפונקציה של הנירמול שבה משתמש המודל שלנו נקראת RMSNorm, קיצור של Root Mean Square Normalization. זה הולך ככה: יש לנו וקטור \(x=\left(x_{1},x_{2},\ldots,x_{N}\right)\) של מספרים ממשיים, אז קודם כל נחשב את הממוצע החשבוני של סדרת הריבועים של אברי הוקטור, כלומר נחשב

\(m=\frac{1}{N}\sum^{N}_{i=1}x^{2}_{i}\)

ועכשיו נוציא שורש לדבר הזה ונחלק בו את הוקטור המקורי - אבל בעצם, לא, זהירות. ה-\(m\) הזה עשוי להיות קרוב לאפס (אם כי אני מודה שזה לא נראה לי הכי סביר שזה יקרה) ואז חלוקה בו תגרום לתוצאה שלנו "להתפוצץ". כדי למנוע את זה מרחיקים אותו מאפס על ידי הוספה של \(\varepsilon\) קטן כלשהו - במודל שלנו זה \(\varepsilon=10^{-5}\) . כלומר, החישוב מבצע

\(x\leftarrow\frac{x}{\sqrt{m+\varepsilon}}\)

אבל לא סיימנו - בשלב הזה אנחנו לוקחים וקטור של \(N\) פרמטרים וכופלים אותו רכיב-רכיב ב-\(x\) ; זה שלב הנירמול הסופי שתלוי בפרמטרים שהזכרתי. פורמלית, לכל \(1\le i\le N\), מה שעשינו הוא לחשב את

\(x_{i}\leftarrow\frac{\gamma_{i}x_{i}}{\sqrt{\frac{1}{N}\sum^{N}_{j=1}x^{2}_{j}+\varepsilon}}\)

ואצלנו \(N=2048\) .

עכשיו אפשר לעשות את חשבון הפרמטרים: כל נרמול משתמש ב-2048 פרמטרים. יש שני שלבי נירמול בכל סיבוב ויש 24 סיבובים. בנוסף לכך, ממש לפני התרגום ל-logits בסוף עושים נרמול אחד אחרון. בסך הכל 49 נרמולים, כל אחד עם סט הפרמטרים שלו - בסך הכל בערך מאה אלף פרמטרים שזה... כמעט כלום. ראינו שמטריצה אחת של שלב ה-MLP היא כבר עם סביבות ה-16 מיליון פרמטרים, כך שבספירת הפרמטרים הכוללת הנירמול כמעט ולא משפיע.

עכשיו באמת אפשר לעבור לדבר על שלב ה-Attention, קיוויתי שאפשר יהיה אולי לעשות ניסוי נחמד ולהריץ את המודל בלי השלב הזה, אבל התוצאות שאני מקבל נראות כמו ג'יבריש וזה לא מפתיע כי גם המודל עצמו לא אמור לתת תוצאה טובה על "cat" לבדו (הרי יש פורמט שלם של פנייה אליו, כבר דיברתי על זה) אבל גם כי שלבי ה-Attention משנים את ה-hidden state בצורה שאי אפשר סתם להחליק. אז נוותר על הניסוי הזה כאן.

מה קורה בשלב ה-Attention?

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

ראשית, מה המטרה של השלב הזה? אם נשים לב, כל מה שעשינו עד עכשיו התמקד בכולו ברמת ה-token הבודד. הפכנו אותו ל-hidden state ובשלב ה-MLP פעלנו על hidden state בודד. לא הייתה שום התייחסות לכך שהמודל לא עובד על טוקן בודד אלא על רצף של טוקנים כאלו, שנותן קונטקסט קריטי. התשובה לשאלה "האם החיה החמודה בעולם היא חתול?" תהיה שונה מהתשובה לשאלה "האם החיה הגדולה בעולם היא חתול?" למרות שבשני המקרים בסופו של דבר המודל יפעל על ה-hidden state שנוצר מתוך "חתול". אנחנו חייבים שה-hidden state שאנחנו פועלים עליו בשלב ה-MLP יכיל עוד מידע שאפשר לעבוד איתו על "מה היה קודם" - בלי זה ברור ששום דבר לא יעבוד.

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

הרעיון הבסיסי היה ונשאר כזה: מכיוון שאנחנו מתחילים עם סדרה של token-ים, בשלב שאליו אנחנו מגיעים נוצרה מהם סדרה של hidden states שאני אסמן ב-\(h_{1},h_{2},h_{3},\ldots,h_{N}\) . שימו לב שה-\(N\) הזה הוא לא מספר שקשור למודל עצמו, כמו ה-2048 שהוא הגודל של hidden state או ה-\(49512\) שהוא גודל המילון. זה מספר שמתאר את כמות הטוקנים שהמודל מתייחס אליהם - מה שנמצא ב"קונטקסט" שלו. מודלים צריכים להגביל את גודל הקונטקסט שלהם ועוד מעט גם נבין למה, אבל \(N\) יכול להיות גם נמוך בהרבה מההגבלה הזו. אני כנראה בהמשך אתייחס בטעות הרבה אל ה-hidden states הללו בתור "טוקנים" - תסלחו לי, אתם מבינים את הכוונה.

עכשיו, בשלב ה-attention שבו אנחנו נמצאים, יש שלוש מטריצות, כל אחת עם 2048 עמודות ומספר שורות שנתון לבחירת המודל שלנו ובינתיים בשביל הכלליות אני אסמן אותו ב-\(d\) . המטריצות הן

• מטריצת השאילתא (query), \(W_{Q}\) .

• מטריצת המפתח (key), \(W_{K}\) .

• מטריצת הערך (value), \(W_{V}\) .

ניקח hidden state כלשהו, \(h_{i}\), נכפול אותו בשלוש המטריצות ונקבל שלושה וקטורים מאורך \(d\) :

• \(q_{i}=W_{Q}h_{i}\)

• \(k_{i}=W_{K}h_{i}\)

• \(v_{i}=W_{V}h_{i}\)

השאילתא \(q_{i}\) היא הדרך של \(h_{i}\) לשאול \(h_{j}\) -ים אחרים - "עד כמה המידע שיש בכם רלוונטי למה שאני צריך?". המפתח \(k_{i}\) הוא הדרך של \(h_{i}\) לענות לשאילתא הזו כשאחרים שואלים אותה. הערך \(v_{i}\) הוא המידע ש-\(h_{i}\) נותן לאחרים כשהתשובה שלו לשאילתא שלהם מוצאת חן בעיניהם. כל זו כמובן אינטואיציה שאין לה משמעות אמיתית; יש משמעות אמיתית לחישוב שעושים בפועל, שהוא לחשב את המכפלה הסקלרית\(q_{i}\cdot k_{j}\) ולהשתמש בסקלר שמתקבל ממנו בתור המקדם של \(v_{j}\) . פורמלית, עבור המצב \(h_{i}\), אנחנו בונים את הערך

\(\sum^{N}_{j=1}\left(q_{i}\cdot k_{j}\right)v_{j}\)

והערך הזה הוא הפלט של שלב ה-attention עבור \(h_{i}\) . את הפלט הזה מחברים אל \(h_{i}\), בדיוק כמו שעשינו בסוף שלב ה-MLP. ולפני שהשלב מתחיל, מנרמלים עם RMSNorm, כמו בשלב ה-MLP. אז האם סיימנו? וואו, ממש לא, כמה שחיפפתי בפרטים. בואו נתחיל לסבך.

ראשית, משהו שכבר הזכרתי בפוסט המקורי: אנחנו לא באמת משתמשים ב-\(q_{i}\cdot k_{j}\) בתור המקדמים של ה-\(v_{j}\) -ים. אנחנו מבצעים לערכים הללו נירמול שהופך אותם להתפלגות של הסתברות, כלומר לסדרת ערכים \(p_{1},\ldots,p_{N}\) שכולם אי שליליים והסכום שלהם הוא 1, ואת זה עושים בעזרת פעולה שנקראת softmax שהולכת ככה:

\(\text{softmax}\left(z_{1},\ldots,z_{N},i\right)=p_{i}=\frac{e^{z_{i}}}{\sum^{N}_{j=1}e^{z_{j}}}\)

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

אם כן, דה פקטו מה שעושים, כשבאים לחשב attention ל-\(h_{i}\), זה פשוט להתעלם מכל ה-\(j\gt i\) -ים במסגרת החישובים.

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

ראשית, אם נסתכל על מה שקורה כרגע בחישוב Attention, נראה שהדבר היחיד שמעניין טוקן בנוגע למיקום של טוקנים אחרים זה אם הם אחריו (ואז הוא מתעלם מהם) או לפניו (ואז הוא מתייחס לכולם באותו האופן). אין חשיבות נוספת למיקום שלהם, אבל זה מן הסתם רלוונטי. במשפט "הבוקר היה יפה למרות הכיעור של החתול" יש בדיוק אותן מילים כמו "למרות הכיעור של הבוקר היה יפה החתול", אבל מבחינת "החתול" המשפטים כמובן שונים מאוד - אחד משמיץ את החתול ואחד מהלל אותו. כש"יפה" יותר קרוב לחתול, המשפט מהלל אותו; כש"יפה" רחוק יותר ודווקא "הכיעור" קרוב, המשפט מבזה אותו. אבל מה שהצגנו עד כה לא יודע להבדיל בין שני המשפטים בכלל, מנקודת מבטו של החתול. אז מה עושים? משנים קצת את וקטורי המפתח והשאילתא של כל טוקן בהתאם למיקום שלו במשפט. הטריק הזה לא היה במאמר המקורי של Attention is All You Need; הוא שיפור שהגיע אחר כך.

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

\(\left(\begin{array}{cc} \cos\alpha & -\sin\alpha\\ \sin\alpha & \cos\alpha \end{array}\right)\)

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

הסיבוך הזה הוא שאנחנו לא מבצעים חישוב Attention בודד, אלא קונספטואלית מבצעים 32 חישובי Attention נפרדים, קטנים יותר, ואז משלבים איכשהו את התוצאות שלהם - זה מה שנקרא Multi-head attention. תחשבו על זה ככה: במקום שתהיה לנו שאילתא \(q_{i}=W_{Q}h_{i}\) בודדת, יש לנו סדרה \(q^{1}_{i},q^{2}_{i},\ldots,q^{32}_{i}\) של שאילתות; אם במקור \(q_{i}\) היה וקטור מגודל \(d\), אז עכשיו כל \(q^{r}_{i}\) יהיה מגודל \(\frac{d}{32}\) . מעשית כל מה שזה אומר הוא שמבצעים את החישוב \(q_{i}=W_{Q}h_{i}\) שתיארתי קודם, אבל מחלקים את \(q_{i}\) "לחלקים". במודל שלנו \(T=2048\) ולכן הגודל של כל שאילתא הוא \(64\) .

את אותו הדבר בדיוק עושים גם עבור המפתח: מחשבים את \(k_{i}=W_{K}h_{i}\) ומחלקים אותו ל-32 חלקים. וכך גם עבור הערך \(v_{i}=W_{V}h_{i}\) .

אז מה ההבדל? ההבדל הוא בשלב שבו מרכיבים את הערך החדש.

קודם אמרנו: נניח שאנחנו טוקן \(i\) ועכשיו רוצים לראות מה טוקן \(j\) תורם לנו. אז נחשב את \(q_{i}\cdot k_{j}\) לכל \(j\) כזה, ואז נפעיל softmax על כולם ונקבל את המקדמים של ה-\(v_{j}\) -ים.

עכשיו אנחנו אומרים: עבור הראש ה-\(r\), נחשב את כל ה-\(q^{r}_{i}\cdot k^{r}_{j}\) -ים לכל \(j\) ונפעיל softmax על כולם ונקבל את המקדמים של ה-\(v^{r}_{j}\) -ים.

כלומר: אנחנו כבר לא מחשבים מכפלה סקלרית בודדת אחת לכל אחד מ-\(j\) ה-tokens; אנחנו מחשבים 32 מכפלות כאלו (על וקטורים קטנים יותר) ומנרמלים אותן בנפרד. בסיטואציה המקורית, כל \(v_{j}\) היה משתתף בפלט הסופי באופן שמוכתב על ידי מקדם יחיד ואילו עכשיו \(v_{j}\) מחולק לחלקים וכל חלק בנפרד משתתף בצורה שונה. די ברור שהגישה הזו נותנת אקסטרה גמישות שהמודל יכול לנצל.

עוד דבר שלא אמרתי במפורש הפעם הוא שלא סתם מבצעים מכפלה פנימית; כדי למנוע מהערכים להתנפח מדי, מנרמלים אותה על ידי חלוקה בשורש של אורך הוקטור, מה שבפוסט המקורי סימנתי ב-\(\sqrt{d}\) . כאן זה טיפה יותר טריקי כי כאמור, מטפלים בכל ראש בנפרד והגודל של כל ראש הוא קטן יותר; הוא מאורך 64 ולכן מנרמלים על ידי חלוקה ב-\(\sqrt{64}=8\) .

ועכשיו אפשר גם לחזור אל איך RoPE עובד: הוא עובד על כל וקטור מאורך 64 בנפרד. הוא מחלק את הוקטור לשני חצאים ואת הכניסות מזווג על פי החצאים (הראשונים בשני החצאים; השניים בשני החצאים וכו'). עכשיו כל זוג מאופיין על ידי שני פרמטרים: האינדקס \(j\) של הטוקן ברשימת הטוקנים (זה מספר שיכול להיות גדול למדי), והאינדקס \(t\) של הזוג ברשימת הזוגות של הוקטור (מספר שרץ מ-0 עד 31). במודל שלנו, הזווית שמתאימה לזוג הזה היא

\(\alpha_{j,t}=\frac{j}{130000^{2t/64}}\)

אינטואיטיבי בסך הכל!

נשאר רק דבר אחד שטרם תיארתי: הסיום. אחרי שמקבלים מכל ראש attention בנפרד וקטור מאורך 64, משרשרים את כולם ומקבלים וקטור מאורך 2048, ואז, פורמלית, כופלים אותו בעוד מטריצה עם פרמטרים נלמדים - מטריצה שטרם הזכרתי עד כה שנקראת ה-output projection. למה? ובכן, זה מעביר אותנו שוב מרמת הפורמליזם לרמת האינטואיציות ונפנופי הידיים: מה שקרה הוא שלפני הכפל במטריצה הזו היו לנו 32 וקטורים נפרדים ש"לא קשורים אחד לשני" במובן מסוים, אבל אנחנו רוצים לאפשר למידע הרלוונטי מכולם להתערבל יחד כדי לקבל וקטור יחיד מאורך 2048 שמייצג את השילוב של כל המידע. אנחנו לא יודעים איך נכון לעשות את זה, אז אנחנו נותנים למודל ללמוד את זה בעצמו, במחיר של עוד מטריצה עתירת פרמטרים, במקרה הזה \(2048^{2}\), בדיוק כמו שלוש המטריצות האחרות שמעורבות בשלב הזה.

זה מעביר אותנו לשלב "ספירת הפרמטרים", שהוא למעשה די פשוט: יש לנו כאמור 4 מטריצות של \(2048^{2}=4,194,304\) פרמטרים ולכן בסך הכל \(16,777,216\) פרמטרים לכל אחד מ-24 שלבי ה-Attention ולכן בסך הכל \(402,653,184\) פרמטרים: הרבה, אבל עדיין בערך שליש משלב ה-MLP.

רגע, זהו? סיימנו... את הכל? ובכן, כן; חוץ מלהבין את התמונה העוד יותר גדולה

סוף דבר - התמונה העוד יותר גדולה

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

• בשלב ה-embedding (הפיכה של טוקן לוקטור) משתמשים ב-\(100,663,296\) פרמטרים, ובשלב ה-LM head (הפיכה של וקטור ל-logits) במודל הספציפי שלנו משתמשים בדיוק באותה מטריצה ולכן אין עוד פרמטרים.

• בשלבי ה-attention משתמשים בסך הכל ב-\(402,653,184\) פרמטרים.

• בשלבי ה-MLP משתמשים בסך הכל ב-\(1,207,959,552\) פרמטרים.

• בשלבי הנירמול משתמשים ב-\(100,352\) פרמטרים.

הסך הכל הסופי הוא \(1,711,376,384\) פרמטרים; זה ה-1.7B בשם המודל.

עכשיו אפשר לקחת את כל הידע שכבר יש לנו ולכתוב קוד שמבצע לבד את כל החישובים, רק בעזרת המטריצות של המודל, בלי שום ספרייה מסביב. מה שיפה זה שהקוד אפילו לא יהיה ארוך במיוחד - לא צריך אפילו מאתיים שורות (טוב, תלוי עד כמה רוצים לממש הכל מאפס - למשל, האם גם את softmax נממש?). כמובן שלי לא היה כוח לעשות את זה וביקשתי מ-LLM שיעשה את זה בשבילי; המטרה הייתה לקבל קוד שמייצר בדיוק את אותו פלט כמו המודל ה"רגיל" שמריצים בפייתון. זה התגלה בתור אתגר לא טריוויאלי, כי אפילו שגיאות עיגול קטנות ברמת ה-floating point יכלו ליצור הבדלים בפלט הסופי ולכן גם כשכותבים קוד "עצמאי", אם רוצים להשיג את אותה תוצאה כמו הקוד ה"רגיל" חייבים לנקוט בכל מני תעלולים שהוא מבצע (למשל, הוא לא עובד טוקן-טוקן על הפרומפט המקורי אלא על כולם בבת אחת, בתור מטריצה גדולה). אז ויתרתי על זה.

עדיין, אם ננסה לכתוב את הכל מאפס ולהריץ בצורה שבה ה-LLM יוצר תשובה ארוכה, קופצים לנו לעיניים העניינים האחרונים שטרם דיברתי עליהם. ספציפית, משהו שנקרא ה-KV cache והוא דווקא לא משעמם כמו שזה נשמע!

העניין הוא כזה: נניח שנתנו לי פרומפט של 1,000 טוקנים. אני רוצה לייצר לו תשובה, ואני יודע שמה שה-LLM שלי יודע לעשות, בהינתן טוקן, הוא להגיד מה יהיה הטוקן הבא. אז אני מתעניין בעצם רק במה שיקרה כשאריץ את המודל על הטוקן ה-1,000; אני אקבל פלט שהוא טוקן חדש, אצרף אותו לתגובה שלי, ואז אריץ את המודל שוב, רק על הטוקן החדש, וכן הלאה וכן הלאה.

אבל איפה פה באים לידי ביטוי 999 הטוקנים הראשונים?

כבר אמרנו קודם שהם באים ליד ביטוי בשלב ה-Attention. מה שקורה הוא שאנחנו מייצרים \(q\) עבור הטוקן החדש, ואז מחשבים את \(q\cdot k_{i}\) עבור כל הטוקנים הקודמים, ואז מקבלים מקדמים עבור \(v_{i}\) - כל הערכים של הטוקנים הקודמים. כלומר, אנחנו לא צריכים את כל המידע על הטוקנים הקודמים: אנחנו רק צריכים, לכל טוקן \(i\), את שני הוקטורים הללו: \(k_{i},v_{i}\) (ליתר דיוק, את כל 24 הזוגות של \(k_{i},v_{i}\) כי יש 24 שלבי Attention וצריך לזכור את המידע הזה מכל אחד מהם). כל המידע על ה-hidden state שבהם הטוקן היה במהלך החישוב, וגם כל המידע על ה-\(q_{i}\) -ים שבהם הוא השתמש? זה כבר לא חשוב. אבל הערכים הללו של \(k_{i},v_{i}\) ? נצטרך אותם בהמשך. נצטרך אותם עבור כל טוקן חדש שאנחנו מנסים ליצור. גם אם אנחנו בטוקן ה-10,000, עדיין נצטרך להשתמש בהם. אז אנחנו שומרים אותם במקום מיוחד שנקרא ה-KV cache.

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

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

באופן הולם למדי, בזמן שחלף מאז שהתחלתי לכתוב את סדרת הפוסטים הזו LLM-ים גרמו לזעזועים בעולם המתמטי כשהתחילו למצוא עצמאית פתרונות למאות בעיות מורכבות, ולכו תדעו איפה זה ייעצר (האם יפתרו את בעיית \(\text{P=NP}\) סוף כל סוף?!). יש שלל דברים שאפשר לומר כנגד הוכחות מתמטיות מבוססות LLM אבל יש לי גם דבר חיובי אחד להגיד - כשלוקחים הוכחות כאלו, ומפשטים עוד יותר את מה שאפשר, ומחדדים עוד יותר את מה שאפשר, והופכים את התוצר לנגיש ככל הניתן, אז אנחנו יכולים לקרוא אותו ולהבין גם ולחוש את הפליאה וההתלהבות שמרגישים כשנתקלים במשהו יפה במתמטיקה. ואת אותן פליאה והתלהבות אני מרגיש גם כשאני מסתכל על האופן שבו LLM-ים עובדים שתיארתי כאן, ויכול להגיד לעצמי שאני מבין איך זה עובד אבל זה עדיין נראה לי כאילו יושב בפנים איזה חד קרן קסום שמשליך קשת בענן לכל עבר.