ריבוב SGD ו-AdaHessian לצורך האצת ביצועים ללא פשרות על הדיוק

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

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

מה אם במקום לבחור צד, ניישם גישה של סינון מקדים? יש לגישה הזאת כל מיני שמות בתעשייה ובמחקר כמו Learning to Defer (L2D) או כמו Cascade Learning... וכמובן גם שמות ישירים יותר כמו Pre-Filtering.

הרעיון הוא להשתמש ב-SGD+M כברירת המחדל הגלובלית והמהירה, ולהפעיל את AdaHessian היקר והמדויק רק באזורים שבהם ה-SGD מתקשה. כך, אנחנו מקבלים האצת ביצועים משמעותית ללא פשרות על איכות הפתרון הסופי. מתי נדרשת הגישה ההיברידית הזו?

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

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

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

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

איך נגדיר את תהליך הסינון המקדים?

כדי להפוך את האינטואיציה לפרקטיקה, נגדיר מתג לוגי שמבוסס על מצב ההתכנסות. נסמן את פונקציית המטרה ב- L(w) כאשר w ∈ ℝᵈ הם משקלי המודל ונגדיר שני סוגי עדכונים:

עדכון בסיס שמשתמש ב-SGD+M:

vₜ₊₁ = β vₜ + ∇L(wₜ)

wₜ₊₁ = wₜ - ηₛ_𝓰_𝓭 vₜ₊₁

עדכון חילוץ/דיוק שמשתמש ב-AdaHessian:

wₜ₊₁ = wₜ - ηₐ_𝓭_ₐ Dₜ⁻¹ ∇L(wₜ)

כאשר D_t הוא הקירוב החלק לאלכסון של מטריצת ההסיאן.

מנגנון הסינון וההחלפה

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

התקדמות בפונקציית המטרה:

ΔLₖ = (1/k) ∑ᵢ₌₁ᵏ L(wₜ₋ᵢ) - L(wₜ)

נורמת הגרדיאנט המוחלק:

gₜ = ||∇L(wₜ)||

תנאי המעבר, פונקציית האינדיקטור Iₜ:

אם לאורך חלון הזמן ההתקדמות ב-Loss קטנה מסף מסוים (ΔLₖ < ε), אך במקביל ישנה עדיין אנרגיה במערכת (gₜ > δ), המערכת מזהה תקיעות באוסצילציה או בנקודת אוכף. במצב זה מופעל AdaHessian (כאשר Iₜ = 1) למשך M צעדים כדי לחלץ את האופטימיזציה, ולאחר מכן חוזרים ל-SGD+M.

אוטומציה של התהליך באמצעות מכונת מצבים

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

1. בזמן שהמודל רץ עם SGD, המערכת שומרת בזיכרון קצר-טווח את שני המדדים החשובים מהצעדים האחרונים, למשל, חלון של 10 צעדים: היסטוריית ה-Loss ונורמת הגרדיאנט.

2. טריגר הבדיקה האוטומטי: בכל צעד, אלגוריתם הבקרה בודק אם נתקענו, ה-Loss כמעט לא השתנה, קטן מ-ε, והאם יש עוד לאן לרדת, הגרדיאנט עדיין גדול מ-δ. אם הגרדיאנט חזק אבל ה-Loss לא יורד, אז אנחנו כנראה מזגזגים בעמק צר או תקועים באוכף.

3. החלפת ההילוכים: ברגע שהתשובה לשתי השאלות חיובית, מופעל Interrupt. האופטימיזטור מתחלף ל-AdaHessian, מחשב את קירוב ההסיאן ומבצע מספר מוגדר מראש של צעדי חילוץ ישירות אל מחוץ לאזור הבעייתי.

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

מדוע זה משתלם?

כדי להבין את הרווח, צריך להסתכל על הסיבוכיות החישובית לכל צעד.

צעד של SGD דורש העברה קדמית ואחורית אחת (עלות ≈ O(N)). צעד של AdaHessian, לעומת זאת, דורש חישוב של מכפלת הסיאן-וקטור, Hessian-Vector Product באמצעות שיטת Hutchinson, מה שמוסיף לפחות עוד העברה אחורית מלאה ולעיתים אף יותר. העלות ≈ 2.5 × O(N) בהערכה גסה.

נניח שאנחנו מפעילים את AdaHessian רק על ρ אחוזים מכלל הצעדים, למשל, רק 10% מהזמן שבו המודל נתקע. תוחלת העלות החישובית לצעד תהיה:

C_hybrid = (1 - ρ) · C_sgd + ρ · C_adahessian

המשמעות בייצור: אם ρ = 0.1, אנו מגדילים את זמן הריצה פר-צעד ב-15% בלבד לעומת SGD טהור. אולם, מכיוון ש-AdaHessian מונע את עיכובי ההיתקעות וחוסך עשרות אלפי Epochs מבוזבזים באזורים שטוחים, מספר הצעדים הכולל שדרוש להתכנסות צונח באופן דרמטי. התוצאה בפועל היא הקדמה של זמן ההתכנסות לעיתים ב-30% עד 50%, תוך הגעה למינימום לוקאלי איכותי ומוכלל יותר.

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

1. Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K., & Mahoney, M. (2021). ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning.

2. Keskar, N. S., & Socher, R. (2017). Improving Generalization Performance by Switching from Adam to SGD.

3. Jin, C., Ge, R., Netrapalli, P., Kakade, S. M., & Jordan, M. I. (2017). How to Escape Saddle Points Efficiently.

ריבוב SGD ו-AdaHessian לצורך האצת ביצועים ללא פשרות על הדיוק

יש לכם תהליכים אלגוריתמים שדורשים האצה? לא יודעים איך לעשות את זה? צריכים עזרה במחקר אלגוריתמי יישומי? אנחנו ב-Mathematic.ai הכתובת! דברו איתי:

שלמה יונה

מייסד ומדען ראשי, מתמטיקאי מחקר ופיתוח בע"מ

053-7326360

shlomo.yona@mathematic.ai 

https://mathematic.ai

פודקאסט על החברה ועליי, שלמה יונה, ואופן העבודה שלנו ואיתנו: A technical deep dive about Mathematic.ai