Mamba: การปฏิวัติแบบจำลองลำดับและเอาชนะสถาปัตยกรรม Transformer

สำรวจ Mamba แบบจำลองปริภูมิสถานะ (SSM) ที่ Albert Gu และ Tri Dao พัฒนาขึ้นเพื่อปรับปรุงการสร้างแบบจำลองลำดับ Mamba ประมวลผลลำดับซับซ้อนในงานภาษา จีโนมิกส์ และเสียงอย่างมีประสิทธิภาพ โดยใช้ปริภูมิสถานะแบบเลือกสรรเพื่อให้ต้นทุนการคำนวณเพิ่มเป็นเชิงเส้นตามความยาวลำดับ

Mamba รับมือกับปัญหาการคำนวณของ Transformer แบบดั้งเดิมเมื่อลำดับยาวขึ้น แนวทางแบบเลือกสรรช่วยให้อนุมานได้เร็วและเพิ่มปริมาณงานที่ประมวลผลได้ จุดเด่นคือการประมวลผลรวดเร็ว ชั้น SSM แบบเลือกสรร และการออกแบบที่เหมาะกับฮาร์ดแวร์ซึ่งได้รับแรงบันดาลใจจาก FlashAttention ผลที่บทความต้นฉบับนำเสนอแสดงว่า Mamba ทำได้ดีกว่าแบบจำลองเดิมหลายแบบ รวมถึงแบบที่ใช้ Transformer

Transformer กับ Mamba

Transformer เช่น GPT-4 ยกระดับมาตรฐานการประมวลผลภาษาธรรมชาติ แต่ต้นทุนเพิ่มขึ้นเมื่อลำดับยาว Mamba เสนอสถาปัตยกรรมที่ต่างออกไปเพื่อประมวลผลลำดับยาวอย่างมีประสิทธิภาพ

Transformer จัดการข้อมูลลำดับ เช่น ข้อความ ต่างจากแบบจำลองก่อนหน้าที่ประมวลผลทีละส่วน โดยสามารถประมวลผลตำแหน่งต่าง ๆ แบบขนานระหว่างการฝึกเพื่อจับความสัมพันธ์ซับซ้อน กลไกความสนใจ (attention) ช่วยเลือกส่วนของข้อมูลที่สำคัญต่อการทำนาย โดยสร้างตัวแทน query, key และ value จากข้อมูลเข้า แล้วเปรียบเทียบองค์ประกอบเพื่อคำนวณน้ำหนักความสำคัญ

ในโครงสร้างตัวเข้ารหัส–ตัวถอดรหัสที่อธิบายที่นี่ ตัวเข้ารหัสประมวลผลข้อมูลเข้าและตัวถอดรหัสสร้างผลลัพธ์ คำอธิบายนี้เป็นโครงสร้าง Transformer รูปแบบหนึ่ง ไม่ใช่โครงสร้างของทุกรุ่น แต่ละชั้นตัวเข้ารหัสมี self-attention หลายหัวและเครือข่ายป้อนไปข้างหน้าแบบเชื่อมต่อเต็มที่ซึ่งใช้แยกกับแต่ละตำแหน่ง การปรับมาตรฐานและการเชื่อมต่อส่วนตกค้างช่วยฝึกเครือข่ายลึก ตัวถอดรหัสเพิ่ม attention ต่อผลลัพธ์ตัวเข้ารหัส และใช้หน้ากากเชิงเหตุเพื่อให้การทำนายในแต่ละตำแหน่งอ้างอิงเฉพาะตำแหน่งก่อนหน้า รักษาคุณสมบัติการสร้างแบบอัตถดถอย

Mamba ใช้ปริภูมิสถานะแบบเลือกสรรแทนแนวทาง attention นี้ หมายเหตุเกี่ยวกับต้นฉบับ: ประโยคเปรียบเทียบภาษาอังกฤษถูกตัดกลางคำ จึงไม่ได้เติมข้อสรุปที่หายไปเอง

ขั้นตอนสำคัญของ Transformer ได้แก่

• ประมวลผลข้อมูลเข้า: แปลงข้อมูลเป็นตัวแทนแบบ embedding พร้อมข้อมูลตำแหน่งในลำดับ
• กลไก attention: คำนวณคะแนนว่าควรให้น้ำหนักแก่ส่วนอื่นมากเพียงใดเพื่อทำความเข้าใจองค์ประกอบปัจจุบัน
• ตัวเข้ารหัสและตัวถอดรหัส: ในรูปแบบนี้ หลายชั้นช่วยปรับตัวแทนข้อมูลเข้าและสร้างผลลัพธ์
• Attention หลายหัว: พิจารณาความสัมพันธ์หลายแบบพร้อมกันจากปริภูมิตัวแทนต่างกัน เพื่อเรียนรู้บริบทหลากหลาย
• เครือข่ายป้อนไปข้างหน้าแยกตามตำแหน่ง: ใช้การแปลงเดียวกันกับแต่ละตำแหน่งอย่างอิสระ แล้วรวมผ่านการเชื่อมต่อส่วนตกค้างและปรับมาตรฐานระดับชั้น
• สร้างผลลัพธ์: ตัวถอดรหัสทำนายโดยอาศัยบริบทจากตัวเข้ารหัสและสิ่งที่สร้างไปแล้ว

การประมวลผลขนานและ attention ช่วยงานแปลภาษาและสร้างข้อความ ส่วน Mamba ใช้ปริภูมิสถานะแบบเลือกสรรเพื่อแก้ต้นทุนของลำดับยาว อนุมานเร็วขึ้นและขยายแบบเชิงเส้นตามความยาวลำดับ

Mamba

Mamba ต่างจากโครงสร้างที่ใช้บล็อก attention และเพอร์เซปตรอนหลายชั้น (MLP) แบบดั้งเดิม โดยออกแบบให้เบา เร็ว และเป็นเนื้อเดียวกันมากขึ้น

• SSM แบบเลือกสรร: กรองข้อมูลที่ไม่เกี่ยวข้องและเน้นข้อมูลที่มีประโยชน์ รองรับการให้เหตุผลตามเนื้อหา
• อัลกอริทึมที่คำนึงถึงฮาร์ดแวร์: ใช้อัลกอริทึมขนานที่เหมาะกับฮาร์ดแวร์สมัยใหม่ โดยเฉพาะ GPU เพื่อคำนวณเร็วและใช้หน่วยความจำน้อยลง
• สถาปัตยกรรมเรียบง่าย: รวม SSM แบบเลือกสรรและตัดบล็อก attention กับ MLP แบบเดิม ช่วยการขยายขนาดและประสิทธิภาพ

ผลที่รายงานครอบคลุมภาษา เสียง และจีโนมิกส์ ทั้งการฝึกล่วงหน้าและงานเฉพาะโดเมน ในการเปรียบเทียบด้านภาษา Mamba ทำได้เทียบเท่าหรือดีกว่า Transformer ที่ใหญ่กว่า โค้ดและแบบจำลองที่ฝึกล่วงหน้าเปิดให้ใช้งานบน GitHub

งานคัดลอกทั่วไปง่ายสำหรับแบบจำลองเชิงเส้น แต่งานคัดลอกแบบเลือกสรรและหัวแบบอุปนัยต้องใช้หน่วยความจำแบบพลวัตที่คำนึงถึงเนื้อหา

แบบจำลองปริภูมิสถานะแบบมีโครงสร้าง S4 ผสานคุณสมบัติของเครือข่ายวนซ้ำ (RNN) เครือข่ายคอนโวลูชัน (CNN) และแบบจำลองปริภูมิสถานะดั้งเดิม ได้แรงบันดาลใจจากระบบต่อเนื่องที่แมปฟังก์ชันหนึ่งมิติหรือลำดับผ่านสถานะแฝง เป็นแนวทางสร้างแบบจำลองลำดับที่มีประสิทธิภาพและปรับตัวได้ในงานเรียนรู้เชิงลึก

พลวัตของ S4

SSM แบบมีโครงสร้างรับลำดับ x และสร้างผลลัพธ์ y ด้วยพารามิเตอร์ที่เรียนรู้ A, B, C และช่วงเวลา Δ ขั้นแรกแปลงพารามิเตอร์ต่อเนื่องเป็นแบบไม่ต่อเนื่อง จากนั้นใช้การดำเนินการ SSM ในรูปแบบพื้นฐานที่ไม่แปรตามเวลา พารามิเตอร์เหมือนกันทุกขั้นเวลา

ความสำคัญของการทำให้เป็นแบบไม่ต่อเนื่อง

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

ความเป็นเชิงเส้นและความไม่แปรตามเวลา (LTI)

คุณสมบัติหลักของ S4 คือพลวัตคงเดิมตามเวลา และพารามิเตอร์คงที่ทุกขั้นเวลา LTI เป็นพื้นฐานของการคำนวณแบบวนซ้ำและคอนโวลูชัน เป็นกรอบที่เรียบง่ายแต่มีประสิทธิภาพสำหรับลำดับ

เอาชนะข้อจำกัดพื้นฐาน

LTI จำกัดการสร้างแบบจำลองข้อมูลที่ต้องปรับพลวัตตามบริบท งานวิจัยนำพารามิเตอร์ที่เปลี่ยนตามเวลามาใช้เพื่อยกเลิกข้อจำกัดนี้ ทำให้รองรับลำดับและงานได้หลากหลายขึ้น

คำว่าแบบจำลองปริภูมิสถานะใช้กว้าง ๆ กับกระบวนการวนซ้ำที่มีสถานะแฝงในหลายสาขา ในงานเรียนรู้เชิงลึก S4 หรือ SSM แบบมีโครงสร้างหมายถึงกลุ่มเฉพาะที่ปรับให้คำนวณมีประสิทธิภาพแต่ยังแทนลำดับซับซ้อนได้ สามารถรวมในเครือข่ายที่ฝึกตั้งแต่ต้นจนจบเป็นชั้นแปลงลำดับ โดยมีบทบาทคล้ายชั้นคอนโวลูชันใน CNN

SSM เทียบกับ SSM ที่เพิ่มกลไกเลือกสรร

เหตุใดการสร้างแบบจำลองลำดับจึงต้องเลือกสรร

SSM แบบมีโครงสร้าง

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

SSM แบบเลือกสรรทำให้พารามิเตอร์ขึ้นกับข้อมูลเข้า เกิดระบบที่แปรตามเวลาและใช้วิธีคอนโวลูชันแบบ SSM ที่ไม่แปรตามเวลาไม่ได้ จึงใช้กลไกวนซ้ำที่มีต้นทุนเชิงเส้นแทน

ใน SSM ที่เพิ่มการเลือกสรร (S6) พารามิเตอร์ B, C และ Δ ขึ้นกับข้อมูลเข้า x การทำให้เป็นแบบไม่ต่อเนื่องคำนึงถึงการเลือกนี้ แล้วใช้การสแกนประมวลผลลำดับและปรับจุดเน้นในแต่ละขั้นเวลา

ผลการประเมินสำคัญของ Mamba

แผนภูมิของต้นฉบับระบุว่า Mamba ดีที่สุดในกลุ่มสำหรับผลประเมินทุกรายการที่แสดง

ในการทดสอบที่รายงาน Mamba มีความเร็วอนุมานและความแม่นยำสูง ใช้บริบทยาวได้ดีในงาน DNA และเสียง และทำได้ดีกว่าแบบจำลองก่อนหน้าในงานซับซ้อนที่ต้องพึ่งความสัมพันธ์ระยะไกล การประเมินแบบ zero-shot หลายงานยังแสดงความยืดหยุ่น ประสิทธิภาพ และความสามารถในการขยายขนาด

เริ่มใช้งาน Mamba

ข้อกำหนดที่บทความปี 2023 ระบุคือ Linux, GPU NVIDIA, PyTorch 1.12+ และ CUDA 11.6+ ใช้ pip ติดตั้งแพ็กเกจจาก คลังโค้ด Mamba หากมีปัญหาความเข้ากันได้ของ PyTorch ต้นฉบับเสนอแฟล็ก --no-build-isolation แบบจำลองฝึกบนข้อมูลขนาดใหญ่ เช่น Pile และ SlimPajama เพื่อตอบโจทย์ทรัพยากรคำนวณหลายระดับ

อินเทอร์เฟซมีตั้งแต่ชั้น SSM แบบเลือกสรร บล็อก Mamba ไปจนถึงแบบจำลองภาษาเต็มรูปแบบ บล็อกหลักใช้ Conv1d แบบเชิงเหตุและรวมเข้ากับเครือข่ายประสาทได้ง่าย หมายเหตุเกี่ยวกับต้นฉบับ: บทความอังกฤษกล่าวถึงตัวอย่าง Python สำหรับสร้างแบบจำลองและประมวลผลข้อมูล แต่ไม่มีโค้ดตัวอย่างนั้นอยู่จริง

แบบจำลองที่ฝึกล่วงหน้าบน Hugging Face มีตั้งแต่ 130 ล้านถึง 2.8 พันล้านพารามิเตอร์ ฝึกบน Pile และ SlimPajama และอิงมาตรฐานขนาดมิติของ GPT-3 ออกแบบเพื่อความต้องการด้านการคำนวณและประสิทธิภาพที่หลากหลาย ต้นฉบับรายงานปริมาณงานและความแม่นยำสูงสำหรับงานภาษาและการใช้งานอื่น

ผลกระทบของ Mamba

Mamba เป็นทางเลือกที่มีศักยภาพต่อ Transformer สำหรับข้อมูลที่มีสารสนเทศหนาแน่น ออกแบบให้ใช้หน่วยความจำและการประมวลผลขนานบนฮาร์ดแวร์สมัยใหม่ได้ดี โค้ดโอเพนซอร์สและแบบจำลองที่ฝึกล่วงหน้าช่วยให้นักวิจัยและนักพัฒนา AI กับการเรียนรู้เชิงลึกเข้าถึงได้