จาก Dynamic Computation ไปสู่ Static Lookup: STEM ปรับโครงสร้างกลไกความจำของ Transformer อย่างไร

ในช่วงไม่กี่ปีที่ผ่านมา ด้วยการเพิ่มขึ้นอย่างต่อเนื่องของขนาดและความหนาแน่นของความรู้ในโมเดลภาษาขนาดใหญ่ นักวิจัยเริ่มกลับมาคิดทบทวนถึงปัญหาพื้นฐานที่สำคัญยิ่งขึ้น: พารามิเตอร์ในโมเดลควรถูกจัดระเบียบอย่างไร จึงจะสามารถทำหน้าที่เป็น “ความจำ” ได้อย่างมีประสิทธิภาพมากขึ้น

ในเครือข่ายฟีดฟอร์เวิร์ด (FFN) แบบมาตรฐานของ Transformer ความรู้ส่วนใหญ่ถูกเก็บไว้โดยปริยายในเมทริกซ์หนาแน่น เช่น up-projection และถูกกระตุ้นแบบไดนามิกผ่านการคูณเมทริกซ์ที่สัมพันธ์กับอินพุต วิธีนี้มีประสิทธิภาพในด้านการแสดงออก แต่มีข้อจำกัดโดยธรรมชาติในด้านความสามารถในการอ้างอิง (addressability) การแก้ไข (editability) และประสิทธิภาพของระบบของพารามิเตอร์

เกี่ยวกับปัญหานี้ ทั้งแวดวงวิชาการและอุตสาหกรรมค่อยๆ หันความสนใจไปยังวิธีการจัดระเบียบพารามิเตอร์ที่แยกส่วนและมีโครงสร้างมากขึ้น งานวิจัยหลายชิ้นล่าสุด โดยมี engram ที่เสนอโดย DeepSeek เป็นตัวแทน ได้ปลุกความสนใจต่อ “ความจำแบบค้นหาจากตาราง (lookup-based memory)” ขึ้นมาอีกครั้ง เป็นที่น่าสังเกตว่า ก่อนที่ engram จะเผยแพร่ประมาณสามเดือน งานวิจัย ICLR เรื่อง “STEM” ก็ได้สำรวจทิศทางนี้อย่างเป็นระบบเช่นกัน

• ลิงก์บทความวิจัย: https://arxiv.org/pdf/2601.10639
• หน้าแรกโปรเจกต์: https://infini-ai-lab.github.io/STEM/

แตกต่างจากงานที่เกี่ยวข้องกับ MoE STEM ไม่ได้เป็นการปรับปรุงเพียงบางส่วนบนกรอบการทำงานแบบสแปร์สรูทติ้งที่มีอยู่ แต่เริ่มต้นจากโครงสร้าง FFN โดยตรง โดยแทนที่ up-projection ด้วยตาราง embedding แบบลำดับชั้นที่จัดทำดัชนีตามโทเค็น ซึ่งเป็นการปรับโครงสร้างเส้นทางการเข้าถึง “ความจำ” ใน Transformer ใหม่ด้วยวิธีแบบสถิต (static) และไม่มีกระบวนการรูทติ้งขณะรันไทม์

จาก “คำนวณที่อยู่” สู่ “ค้นหาที่อยู่”: สัญชาตญาณหลักของ STEM

หากมอง Transformer มาตรฐานจากมุมมองของความจำแบบคีย์-วาลู (key-value memory) โครงสร้าง FFN เช่น SwiGLU มักจะแมปการแสดงแทนอินพุตไปยังพื้นที่มิติสูงผ่าน up-projection หนึ่งครั้ง เพื่อสร้าง “เวกเตอร์ที่อยู่” ที่สามารถถูกปรับด้วยเกต กระบวนการนี้โดยพื้นฐานแล้วคือการคูณเมทริกซ์หนาแน่นที่สัมพันธ์กับอินพุต ซึ่งทั้งมีค่าใช้จ่ายสูงและมีการเชื่อมโยงกันอย่างมากในระดับพารามิเตอร์

ข้อสังเกตสำคัญของ STEM คือ: หากบทบาทของ FFN มีความใกล้เคียงกับการ “เข้าถึงความจำตามโทเค็น” มากขึ้น แล้วเวกเตอร์ที่อยู่เหล่านี้จำเป็นต้องถูกสร้างขึ้นแบบไดนามิกจากอินพุตจริงๆ หรือ? จากคำถามนี้ STEM เลือกเส้นทางที่ตรงไปตรงมาอย่างยิ่ง: ไม่คำนวณเวกเตอร์ที่อยู่ผ่าน up-projection อีกต่อไป แต่แทนที่จะรักษาตาราง embedding ที่จัดทำดัชนีตามโทเค็นสำหรับแต่ละเลเยอร์ และในระหว่างการแพร่ไปข้างหน้า (forward propagation) จะค้นหาจากตารางโดยตรงตาม id ของโทเค็นเพื่อรับเวกเตอร์ที่สอดคล้องกัน

ในทางรูปแบบ STEM ใช้เวกเตอร์แบบสถิตที่จัดทำดัชนีตามโทเค็น แทนที่ ใน FFN แบบดั้งเดิม ส่วนเกตและ down-projection ยังคงถูกเก็บไว้อย่างครบถ้วน เพื่อใช้สำหรับการปรับบริบทและการบีบอัดเวกเตอร์ที่ได้จากการค้นหาตาราง การแทนที่นี้ทำให้โครงสร้างของ FFN เกิดการเปลี่ยนแปลงพื้นฐาน: “ความจุความจำ” และ “ปริมาณการคำนวณต่อโทเค็น” ถูกแยกออกจากกันอย่างชัดเจนเป็นครั้งแรก

การเปลี่ยนแปลงโครงสร้างเล็กน้อย แต่ก่อให้เกิดผลกระทบต่อเนื่องหลายประการ

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

1. ความสามารถในการแก้ไข: ปรับเปลี่ยน “เวกเตอร์ความรู้” โดยตรง

คุณลักษณะหนึ่งของ STEM ที่ดึงดูดความสนใจเป็นพิเศษคือ ความสัมพันธ์ที่ชัดเจนระหว่างโทเค็นและพารามิเตอร์

เนื่องจาก embedding ในแต่ละเลเยอร์ถูกผูกไว้โดยตรงกับ id ของโทเค็น นักวิจัยสามารถแทนที่เวกเตอร์ STEM ของโทเค็นใดโทเค็นหนึ่งได้โดยตรง โดยไม่ต้องฝึกฝนใหม่ จึงสามารถเปลี่ยนพฤติกรรมการส่งออกของโมเดลในข้อเท็จจริงที่เกี่ยวข้องได้ ตัวอย่างเช่น ตัวอย่างในภาพแสดงให้เห็นว่าเพียงแค่แทนที่เวกเตอร์ของ “Spain” และ “Germany” ก็สามารถทำให้คำตอบของโมเดลต่อคำถามเกี่ยวกับเมืองหลวงเปลี่ยนแปลงไปตามนั้น ความสามารถนี้ให้แนวทางใหม่สำหรับการแก้ไขความรู้ของโมเดลและการทำความเข้าใจโมเดล

2. ความเสถียรในการฝึก: ความเบาบางแบบสถิตดีกว่าการรูทติ้งแบบไดนามิก

แตกต่างจากโซลูชันเช่น MoE ที่พึ่งพาการรูทติ้งขณะรันไทม์ STEM เป็นโครงสร้างแบบสแปร์สแบบสถิต: embedding ที่แต่ละโทเค็นเข้าถึงในแต่ละเลเยอร์เป็นสิ่งที่กำหนดไว้ ไม่เกี่ยวข้องกับการปรับสมดุลโหลด (load balancing) หรือการสื่อสารแบบ all-to-all การออกแบบนี้หลีกเลี่ยงปัญหาการเอียงของโหลด (load skew) และการพุ่งสูงขึ้นของค่าความสูญเสีย (loss spike) ที่พบบ่อยใน MoE

3. “พื้นที่ความจำ” ที่มีประสิทธิภาพมากขึ้น: การเปลี่ยนแปลงโครงสร้างเรขาคณิตของ embedding

จากมุมมองของพื้นที่การแสดงแทน การใช้ embedding ที่จัดทำดัชนีตามโทเค็นแทนที่ up-projection จะเปลี่ยนการกระจายตัวทางเรขาคณิตของเวกเตอร์อินพุตของ FFN อย่างมีนัยสำคัญ ตาราง embedding ใน STEM แสดงการกระจายเชิงมุมที่กว้างขึ้น (large angular spread) เวกเตอร์ของโทเค็นที่แตกต่างกันมีแนวโน้มตั้งฉากกันมากขึ้น จึงลดการรบกวนระหว่างพารามิเตอร์ (cross-talk) ได้อย่างมีประสิทธิภาพ คุณสมบัตินี้หมายความว่า ด้วยต้นทุนการคำนวณที่เท่ากันหรือต่ำกว่า โมเดลสามารถรักษา “ช่องความจำที่สามารถอ้างอิงได้” มากขึ้น ซึ่งเป็นประโยชน์อย่างยิ่งสำหรับการจัดเก็บและค้นคืนความรู้

4. ประสิทธิภาพการคำนวณและระบบ: การลดลงสองเท่าของ FLOPs และ IO

จากมุมมองความซับซ้อนในการคำนวณ หลังจากลบ up-projection ออกแล้ว แต่ละเลเยอร์สามารถประหยัดค่าใช้จ่ายในการคูณเมทริกซ์ระดับประมาณ d·d_ff ได้ ที่สำคัญกว่านั้น ตาราง embedding ขนาดใหญ่สามารถถูกโอนไปยัง CPU และเข้าถึงได้ผ่านกลยุทธ์การดึงข้อมูลล่วงหน้าแบบอะซิงโครนัส (asynchronous prefetch) และการแคช

ผลการทดลอง: การตรวจสอบอย่างเป็นระบบในด้านขนาด งาน และบริบทยาว

ในสองขนาดโมเดลคือ 350M และ 1B STEM ได้รับการเปรียบเทียบอย่างเป็นระบบกับฐานข้อมูลแบบหนาแน่น (dense baseline) ผลลัพธ์แสดงว่า:

• ประสิทธิภาพโดยเฉลี่ยโดยรวมเพิ่มขึ้นประมาณ 3–4% ในบางงานที่เกี่ยวกับความรู้สามารถเพิ่มขึ้นได้ถึง 9–10%;
• ในการประเมินบริบทยาว เช่น Needle-in-a-Haystack และ LongBench ข้อได้เปรียบของ STEM จะขยายใหญ่ขึ้นเมื่อความยาวของบริบทเพิ่มขึ้น

มุมมองทางวิศวกรรม: จะนำไปใช้งานจริงได้อย่างไร

บทความวิจัยยังได้ให้การอภิปรายที่ค่อนข้างครบถ้วนเกี่ยวกับการนำระบบ STEM ไปปฏิบัติ:

• ตำแหน่งการแทนที่: การแทนที่เฉพาะ up-projection เป็นสิ่งสำคัญ การแทนที่ gate-projection อย่างมืดบอดจะทำลายความสามารถในการปรับบริบท;
• กลยุทธ์การจัดเก็บ: ตาราง embedding สามารถโอนไปยัง CPU ได้ ในระหว่างการฝึกจำเป็นต้องเขียนเกรเดียนต์กลับไปยังสถานะของออปติไมเซอร์ที่สอดคล้องกัน;
• โซลูชันประนีประนอม: ระหว่างสัดส่วนพารามิเตอร์และแรงกดดันต่อหน่วยความจำ สามารถใช้การแทนที่บางเลเยอร์หรือตัวแปรแบบผสมเพื่อสร้างสมดุลได้

สรุป

STEM แสดงให้เห็นถึงแนวโน้มที่น่าสนใจ: ด้วยการเปลี่ยน “วิธีการจัดระเบียบ” ของพารามิเตอร์ แทนที่จะเพิ่มขนาดหรือการคำนวณเพียงอย่างเดียว โมเดลก็สามารถได้รับความสามารถที่เพิ่มขึ้นอย่างมีนัยสำคัญได้เช่นกัน ในบริบทที่สถาปัตยกรรมโมเดลขนาดใหญ่ค่อยๆ ก้าวไปสู่ความซับซ้อนมากขึ้น แนวทางการออกแบบที่เรียบง่าย เสถียร และเป็นมิตรต่อวิศวกรรมเช่นนี้ อาจเป็นเส้นทางที่ไม่สามารถละเลยได้ในวิวัฒนาการของโมเดลฐานรากในขั้นตอนต่อไป

แนะนำผู้เขียน

ผู้เขียนหลักของบทความนี้คือ Ranajoy Sadhukhan นักศึกษาปริญญาเอกที่ห้องปฏิบัติการ InfiniAI Lab มหาวิทยาลัยคาร์เนกีเมลลอน (CMU) ภายใต้การดูแลของศาสตราจารย์ Bei Chen งานนี้เสร็จสิ้นในช่วงที่เขาเป็นนักศึกษาฝึกงานที่ Meta AI โดยมีผู้ให้คำปรึกษาระหว่างฝึกงานรวมถึงนักวิจัยเช่น Zechun Liu, Sheng Cao (Rick Cao) และ Yuandong Tian

InfiniAI Lab ก่อตั้งโดยศาสตราจารย์ Bei Chen มุ่งมั่นในการออกแบบร่วมกันของโมเดล ระบบ และฮาร์ดแวร์ วิจัยอัลกอริธึมและระบบ AI ที่มีประสิทธิภาพและขยายขนาดได้ โดยเน้นทิศทางหลักได้แก่ การสร้างแบบจำลองหลายรูปแบบ (multimodal) บริบทยาว สถาปัตยกรรมโมเดล新一代ที่ก้าวข้ามกฎการสเกลแบบดั้งเดิม และการเพิ่มความสามารถในการเข้าใจและให้เหตุผลของโมเดลพื้นฐาน พร้อมทั้งส่งเสริมการเพิ่มประสิทธิภาพในระดับอัลกอริธึมและระบบ เพื่อส่งเสริมการทำให้เทคโนโลยี AI เป็นที่แพร่หลาย

Zechun Liu เป็นนักวิทยาศาสตร์วิจัยของ Meta AI ทิศทางการวิจัยครอบคลุมการฝึกโมเดลฐาน การบีบอัดโมเดลขนาดใหญ่ การทำให้เบาบาง และการปรับใช้บนอุปกรณ์ปลายทาง โดยมุ่งเน้นการอนุมานที่มีประสิทธิภาพของโมเดลและการออกแบบร่วมกันของระบบ

Sheng Cao (Rick Cao) เป็นนักวิจัยของ Meta AI วิจัยหลักเกี่ยวกับการเพิ่มประสิทธิภาพระบบโมเดลขนาดใหญ่และการออกแบบสถาปัตยกรรมการอนุมานที่มีประสิทธิภาพ โดยให้ความสนใจกับปัญหาการปรับใช้และเร่งความเร็วของโมเดลขนาดใหญ่ในสภาพแวดล้อมระบบจริง

Yuandong Tian เป็นนักวิทยาศาสตร์วิจัยอาวุโสของ Meta AI ทำงานวิจัยด้านการเรียนรู้แบบเสริมกำลัง (reinforcement learning) และโมเดลขนาดใหญ่มาเป็นเวลานาน เคยมีส่วนร่วมในการพัฒนาระบบการเรียนรู้แบบเสริมกำลังเช่น AlphaZero และให้ความสนใจกับความสามารถในการให้เหตุผลและการตัดสินใจของโมเดลพื้นฐาน

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