GPU เร่งความเร็วใหม่! อัลกอริทึม Gram Newton-Schulz ลดเวลาออปติไมเซอร์โมเดล MoE หลายล้านล้านพารามิเตอร์ลง 40-50%

ในสาขาการวิเคราะห์เชิงตัวเลข Newton-Schulz และวิธีการที่เกี่ยวข้องได้รับการศึกษามาหลายปี แต่ส่วนใหญ่เน้นไปที่การคำนวณความแม่นยำสูง การปรับปรุง CPU หรืออินพุตเมทริกซ์จัตุรัส

เมื่อเร็วๆ นี้ นักวิจัยสี่คนจากมหาวิทยาลัยพรินซ์ตันและมหาวิทยาลัยนิวยอร์กได้เสนออัลกอริทึม Gram Newton-Schulz การศึกษานี้ปรับโครงสร้างวิธีการ Newton-Schulz แบบคลาสสิกใหม่ เพื่อให้เข้ากับฮาร์ดแวร์ GPU และสถานการณ์การฝึกโมเดลขนาดใหญ่ได้ดีขึ้น การทดลองแสดงให้เห็นว่าอัลกอริทึมนี้สามารถลดเวลาที่ใช้ในขั้นตอนออปติไมเซอร์ลง 40% ถึง 50% เมื่อฝึกโมเดล Mixture of Experts (MoE) ขนาดหลายล้านล้านพารามิเตอร์

แนวคิดหลักของ Gram Newton-Schulz สามารถสรุปได้ดังนี้:

แทนที่จะทำการวนซ้ำบนเมทริกซ์ดั้งเดิม ( X in mathbb{R}^{n times m} ) โดยตรง ให้เปลี่ยนไปวนซ้ำบน Gram matrix ( X X^top in mathbb{R}^{n times n} ) ซึ่งมีมิติที่เล็กกว่า การดำเนินการนี้ช่วยลดปริมาณการคำนวณได้อย่างมาก และใช้ประโยชน์จากศักยภาพในการปรับปรุงการคำนวณของเมทริกซ์สมมาตรอย่างเต็มที่

ผลงานหลักของการศึกษานี้รวมถึง:
* เขียนวิธีการ Newton-Schulz มาตรฐานใหม่ในรูปแบบที่เทียบเท่าทางคณิตศาสตร์ เพื่อให้ทำงานหลักในพื้นที่ ( n times n )
* เสนอเวอร์ชันพื้นฐาน Naive Gram Newton-Schulz
* วิเคราะห์ความไม่เสถียรของอัลกอริทึมนี้ภายใต้ความแม่นยำครึ่งหนึ่ง (float16) และเสนอ กลยุทธ์รีสตาร์ท ที่สอดคล้องกัน เพื่อให้ได้เวอร์ชันที่เสถียร Stabilized Gram Newton-Schulz
* ดำเนินการ GPU kernel สำหรับการคูณเมทริกซ์สมมาตรที่มีประสิทธิภาพสูง
* สร้างออปติไมเซอร์ GramMuon ซึ่งทำให้เกิดความเร็วขึ้น 40-50% ในการฝึกโมเดลจริง โดยไม่สูญเสียความแม่นยำของโมเดล

ภาพด้านล่างแสดงการเปรียบเทียบเวลาที่ใช้ในการดำเนินขั้นตอนออปติไมเซอร์ระหว่าง AdamW และออปติไมเซอร์ Muon บนโมเดล LLaMA ขนาดต่างๆ บนแพลตฟอร์ม NVIDIA B300

Jack Zhang นักศึกษาปริญญาตรีของมหาวิทยาลัยพรินซ์ตันและหนึ่งในผู้เขียนบทความกล่าวว่า “งานของเราทำให้ความเร็วในการทำงานของออปติไมเซอร์ Muon เพิ่มขึ้นสูงสุด 2 เท่า โดยแทบไม่เพิ่มต้นทุนเพิ่มเติม Gram Newton-Schulz สามารถใช้เป็นทางเลือกแบบ plug-and-play สำหรับโมดูล Newton-Schulz ใน Muon เราได้สังเกตว่า perplexity ของโมเดลบนชุดตรวจสอบแทบไม่เปลี่ยนแปลง โดยมีข้อผิดพลาดควบคุมอยู่ภายใน 0.01”

Tri Dao ผู้ช่วยศาสตราจารย์แห่งมหาวิทยาลัยพรินซ์ตัน ผู้ร่วมก่อตั้งและหัวหน้านักวิทยาศาสตร์ของ Together AI และผู้เขียนอีกคนหนึ่งให้ความเห็นว่า “นี่คืองานประเภทที่ฉันชอบมากที่สุด – การผสมผสานระหว่างความเข้าใจเชิงลึกเกี่ยวกับพีชคณิตเชิงเส้นและการนำโอเปอเรเตอร์ประสิทธิภาพสูงไปใช้ เราใช้เวลาหลายเดือนในการวิเคราะห์ลักษณะเฉพาะและเวกเตอร์ลักษณะเฉพาะของเมทริกซ์ที่เกี่ยวข้องในกระบวนการกลางของ Muon อย่างลึกซึ้ง และในที่สุดก็เสนออัลกอริทึมที่เรียบง่ายและสง่างาม ซึ่งทำให้ความคิดนี้เป็นจริงได้”

ข้อได้เปรียบของอัลกอริทึม: จาก Newton-Schulz มาตรฐานสู่ Gram Newton-Schulz

เพื่อให้เข้าใจข้อดีของ Gram Newton-Schulz มากขึ้น ก่อนอื่นต้องทบทวนภูมิหลังของ ออปติไมเซอร์ Muon Muon สามารถเข้าใจได้ว่าเป็นวิธีการลงที่เร็วที่สุดภายใต้สเปกตรัมนอร์ม:

โดยที่ ( mu ) แสดงถึงสัมประสิทธิ์โมเมนตัม ( eta ) แสดงถึงอัตราการเรียนรู้ และ ( text{polar} ) แสดงถึงการแยกขั้ว

การแยกขั้วถูกกำหนดเป็น: ถ้า ( X = U Sigma V^top ) ดังนั้น:
[
text{polar}(X) = U V^top
]
เนื่องจากการคำนวณการแยกขั้วที่แม่นยำมีค่าใช้จ่ายสูง Muon จึงใช้ การวนซ้ำ Newton-Schulz เพื่อประมาณค่า

รูปแบบการวนซ้ำของ Newton-Schulz มาตรฐานมีดังนี้:

การวนซ้ำนี้โดยพื้นฐานแล้วเป็นการทำให้ค่าเอกพจน์ของเมทริกซ์เป็นมาตรฐานทีละน้อย เพื่อให้เข้าใกล้ 1 อย่างไรก็ตาม Newton-Schulz มาตรฐานมีคอขวดในการคำนวณ: แต่ละรอบการวนซ้ำประกอบด้วยการคูณเมทริกซ์สามครั้ง ปริมาณการคำนวณทั้งหมด (FLOPs) คือ:

สิ่งนี้นำไปสู่ปัญหาหลายประการ: การคูณเมทริกซ์สี่เหลี่ยมจำนวนมากที่ไม่มีประสิทธิภาพ การใช้ประโยชน์จากความสมมาตรของเมทริกซ์ไม่เต็มที่ และในที่สุดทำให้ขั้นตอนออปติไมเซอร์กลายเป็นคอขวดในการฝึก

ดังนั้น แนวคิดหลักของ Gram Newton-Schulz คือการย้ายการวนซ้ำจากเมทริกซ์ดั้งเดิม ( X ) ไปยัง Gram matrix ( R = X X^top ):

การเปลี่ยนแปลงทางทฤษฎีที่สำคัญคือ: พหุนามคี่ใดๆ ( p(t) ) สามารถเขียนใหม่ในรูปแบบ ( t cdot q(t^2) ) ได้ ผ่านการเปลี่ยนแปลงนี้ กระบวนการวนซ้ำจะถูกแปลงเป็นการวนซ้ำพหุนามเมทริกซ์บน ( R ) ซึ่งโดยพื้นฐานแล้วเป็นการประมาณ ( Y^{-1/2} )

วิธีนี้ให้ข้อได้เปรียบที่สำคัญ:
* ลดมิติ: จาก ( n times m ) เป็น ( n times n )
* การปรับปรุงการคำนวณ: สามารถใช้ GPU kernel ที่ออกแบบมาสำหรับเมทริกซ์สมมาตรโดยเฉพาะ
* ลดการดำเนินการที่ไม่มีประสิทธิภาพ: ลดจำนวนการคูณเมทริกซ์ทั่วไปสี่เหลี่ยม (GEMM)

ความท้าทายด้านความเสถียรและแนวทางแก้ไข

อย่างไรก็ตาม Naive Gram Newton-Schulz พื้นฐานไม่เสถียรภายใต้ความแม่นยำ float16 Gram matrix อาจสร้างค่าลักษณะเฉพาะที่เป็นลบ ทำให้เกิดการพุ่งสูงขึ้นของ loss (loss spike) หรือค่าอนันต์ (Inf) ในการฝึก

(ภาพแสดงปรากฏการณ์ไม่เสถียรเมื่อใช้ Naive Gram Newton-Schulz บนโมเดล Llama-430M)

เพื่อแก้ไขปัญหานี้ นักวิจัยได้เสนอเวอร์ชัน Stabilized Gram Newton-Schulz การปรับปรุงหลักรวมถึง:
1. ใช้ float16 ในขั้นตอนสำคัญเพื่อเร่งความเร็ว
2. แนะนำกลไกการรีสตาร์ท
3. เริ่มต้น Gram matrix ใหม่

ขั้นตอนของอัลกอริทึมเวอร์ชันเสถียรมีดังนี้:

ผลลัพธ์แสดงให้เห็นว่าอัลกอริทึมเวอร์ชันเสถียรยังคงเร็วกว่าอัลกอริทึมเดิมอย่างมีนัยสำคัญ ในขณะที่รักษาเสถียรภาพเชิงตัวเลข

การเปรียบเทียบประสิทธิภาพและการประยุกต์ใช้จริง

ความซับซ้อนในการคำนวณของสามวิธี ได้แก่ Newton-Schulz มาตรฐาน, Naive Gram และ Stabilized Gram เปรียบเทียบได้ดังนี้:

ผลลัพธ์แสดงให้เห็นว่าเมื่ออัตราส่วนรูปร่างเมทริกซ์ ( alpha = m/n > 1 ) (ซึ่งพบได้บ่อยมากในโมเดลการเรียนรู้เชิงลึก) วิธี Gram จะดีกว่าในแง่ของปริมาณการคำนวณอย่างชัดเจน โดยสามารถลด FLOPs ได้ประมาณ 42% ถึง 58%

สุดท้าย การศึกษาได้ประเมินอัลกอริทึมนี้บนโมเดลจริงขนาดหลายล้านล้านพารามิเตอร์ Kimi K2 Kimi K2 เป็นโมเดล Mixture of Experts (MoE) ที่เบาบางและละเอียด ประกอบด้วยผู้เชี่ยวชาญ 384 คนต่อเลเยร์ มิติของเลเยร์ที่ซ่อนอยู่คือ 7168 และมิติของเลเยร์กลางของผู้เชี่ยวชาญคือ 2048 เนื่องจากใช้สถาปัตยกรรม MoE ที่ละเอียดกว่าและใช้ Muon ในการฝึก จึงเป็นสถานการณ์ที่เหมาะสำหรับการประเมิน Gram Newton-Schulz

ในการทดสอบบนฮาร์ดแวร์ NVIDIA H100 และ B300 ล่าสุด ในการกำหนดค่าขนานแบบไปป์ไลน์ของ Kimi K2 ความเร็วของ Gram Newton-Schulz เป็น 2 เท่า ของ Newton-Schulz มาตรฐาน

(การเปรียบเทียบประสิทธิภาพบนแพลตฟอร์ม NVIDIA H100 Hopper)

(การเปรียบเทียบประสิทธิภาพบนแพลตฟอร์ม NVIDIA B300 Blackwell)

สำหรับรายละเอียดทางเทคนิคเพิ่มเติม โปรดดูบล็อกต้นฉบับที่เผยแพร่โดยทีมวิจัย

ที่อยู่บล็อก: https://dao-lab.ai/blog/2026/gram-newton-schulz/

⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง

☕ สนับสนุนค่ากาแฟทีมงาน

หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay

SCAN TO PAY WITH ANY BANK