ลดปัญหา KV Cache ใน LLM ด้วยกลไก Attention ใหม่

บทนำ

การใช้งาน Large Language Models (LLMs) ที่เพิ่มขึ้นและการเกิดขึ้นของกระบวนทัศน์ใหม่ในการอนุมาน ได้นำความท้าทายของการอนุมานขนาดใหญ่อย่างมีประสิทธิภาพมาสู่แนวหน้า อุปสรรคสำคัญคือ Key-Value (KV) cache ภายในกลไก attention แบบดั้งเดิม ซึ่งขยายตัวเป็นเส้นตรงตามขนาดของ batch และความยาวของลำดับ กลายเป็น "memory hog" ที่ขัดขวางการปรับขนาดและการขยายตัวของ LLMs

แม้ว่าจะมีตัวแปรต่างๆ เช่น MQA, GQA และ MLA เกิดขึ้นเพื่อแก้ไขปัญหานี้ แต่พวกเขามักจะดิ้นรนเพื่อรักษาประสิทธิภาพภายใต้ข้อจำกัดด้านหน่วยความจำที่เข้มงวด หรือแนะนำความซับซ้อนที่ก่อให้เกิดความท้าทายทางด้านวิศวกรรมและปัญหาความเข้ากันได้

Multi-matrix Factorization Attention (MFA)

งานวิจัยล่าสุดจากนักวิจัยจาก Stepes, Tsinghua University และสถาบันอื่นๆ ได้นำเสนอกลไก attention แบบใหม่: Multi-matrix Factorization Attention (MFA) และตัวแปร MFA-Key-Reuse (MFA-KR) กลไกนี้ช่วยลดต้นทุนการอนุมานของโมเดลภาษาได้อย่างมาก พร้อมทั้งปรับปรุงประสิทธิภาพไปพร้อม ๆ กัน

MFA และ MFA-KR ไม่เพียงแต่เหนือกว่า MLA ในด้านประสิทธิภาพ แต่ยังเทียบเท่าประสิทธิภาพของ MHA แบบดั้งเดิม ในขณะที่ลดการใช้ KV Cache ลงได้ถึง 93.7%

MFA ได้รับการออกแบบมาให้เรียบง่าย ทำซ้ำได้ง่าย มีความไวต่อ hyperparameters ต่ำ และเข้ากันได้กับวิธีการ Pos-embedding ต่างๆ

แนวทางและการวิเคราะห์ MFA

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

พวกเขาได้แนะนำแนวคิดของ Generalized Multi-Head Attention (GMHA) เป็นกรอบรวมสำหรับการทำความเข้าใจตัวแปร MHA ที่แตกต่างกัน

ทีมยังได้สำรวจการคำนวณและการจัดเก็บ key-values จากมุมมองของการอนุมาน และตรวจสอบความจุของโมเดลจากมุมมองของการแยกส่วนประกอบ

Fully Parameterized Bilinear Attention (FPBA) ได้รับการกำหนดให้เป็นขีดจำกัดสูงสุดทางทฤษฎีของประสิทธิภาพ พวกเขาพบว่า MHA และตัวแปรของมันเป็นการแยกส่วนประกอบแบบ low-rank ของ FPBA

การเปรียบเทียบกับ MQA และ MLA

การวิเคราะห์มุ่งเน้นไปที่สองรูปแบบการปรับปรุงที่เป็นตัวแทน: Multi-Query Attention (MQA) และ Multi-Head Latent Attention (MLA)

• MQA ใช้กลยุทธ์การแชร์พารามิเตอร์ที่ก้าวร้าวมากขึ้น โดยที่ attention heads ทั้งหมดใช้ชุดพารามิเตอร์ key-value เดียวกัน สิ่งนี้ช่วยลดการใช้หน่วยความจำ แต่อาจส่งผลต่อความสามารถในการแสดงออกของโมเดล
• MLA แนะนำพื้นที่ latent space ที่ใช้ร่วมกันสำหรับการบีบอัดพารามิเตอร์ แต่พลังการแสดงออกที่แท้จริงถูกจำกัดโดยมิติที่เล็กที่สุด ซึ่งหมายความว่าการเพิ่มมิติกลางไม่ได้ปรับปรุงประสิทธิภาพอย่างมีนัยสำคัญ

นวัตกรรมหลักของ MFA

การพัฒนา MFA เกิดจากเป้าหมายในการสร้างกลไก attention ที่ลดการใช้ทรัพยากรให้เหลือน้อยที่สุด ในขณะที่เข้าใกล้ขีดจำกัดประสิทธิภาพทางทฤษฎี การออกแบบของ MFA ประกอบด้วยสามนวัตกรรมหลัก:

• เพิ่มจำนวนและมิติของ attention heads อย่างมีนัยสำคัญ เพื่อเพิ่มความจุของโมเดลให้สูงสุด
• ใช้กลยุทธ์การแยกส่วนประกอบแบบ low-rank ที่ก้าวร้าว เพื่อรักษาประสิทธิภาพของพารามิเตอร์ในขณะที่ขยายจำนวนและมิติของ attention heads
• ใช้การออกแบบ key-value head เดียว เพื่อให้การใช้หน่วยความจำน้อยที่สุด แม้ว่าโมเดลจะมีความซับซ้อนเพิ่มขึ้น

การวัดและความจุการเปรียบเทียบ

เพื่อวิเคราะห์ MFA และกลไก attention อื่น ๆ ทีมงานได้แนะนำตัวชี้วัดที่สำคัญสองตัว:

• Total Effective Rank (TER): ผลคูณของจำนวน attention heads และ Factorization rank per head (FRH)
• Shared Latent Subspace Dimension (SLSD): มิติของ hidden space ที่ใช้ร่วมกันโดย attention heads ทั้งหมด

MFA มี SLSD และ TER ที่สูงกว่าเมื่อเทียบกับ MQA

เมื่อเทียบกับ MLA, MFA มีขนาด KV cache ที่เล็กกว่าและมี TER ที่สูงกว่าด้วยงบประมาณพารามิเตอร์ที่คล้ายกัน ในขณะที่ยังคงรักษา SLSD ที่เทียบเคียงกันได้

เมื่อเทียบกับ MHA แบบดั้งเดิม MFA มี TER ที่สูงกว่า แม้ว่า SLSD จะเล็กกว่า

ผลการทดลอง

มีการทดลองอย่างกว้างขวางเพื่อประเมินประสิทธิภาพของสถาปัตยกรรมใหม่ในขนาดที่ใหญ่ขึ้น โดยทดสอบโมเดลตั้งแต่ 1B ถึง 7B พารามิเตอร์ และข้อมูลการฝึกอบรมตั้งแต่ 10B ถึง 1T

MFA แสดงให้เห็นถึงความสามารถในการปรับขนาดได้เทียบเท่ากับ MHA แบบดั้งเดิม โดยยังคงรักษาประสิทธิภาพที่ดีเยี่ยมแม้ในขนาดที่ใหญ่ขึ้น

ในขณะที่ MFA-KR แสดงประสิทธิภาพที่ต่ำกว่าเล็กน้อย แต่แนวโน้มการปรับขนาดของมันสอดคล้องกับ MHA ข้อดีของการประหยัดหน่วยความจำของ MFA และ MFA-KR ยังคงขยายตัวตามขนาดของโมเดล โดย MFA สามารถประหยัดหน่วยความจำได้ 87.5% และ MFA-KR ลดการใช้หน่วยความจำลงเหลือ 6.25% ที่ขนาดใหญ่ที่สุด

การศึกษา Ablation

การศึกษา Ablation ได้ตรวจสอบความถูกต้องของ MFA และ MFA-KR ข้อได้เปรียบด้านประสิทธิภาพของพวกเขายังได้รับการยืนยันในวิธีการเข้ารหัสตำแหน่ง mainstream ต่างๆ

แนวโน้ม

MFA นำเสนอการปรับปรุงที่สำคัญด้วยการออกแบบที่เรียบง่าย แก้ปัญหาคอขวดด้านหน่วยความจำในการอนุมาน LLM ได้อย่างมีประสิทธิภาพ โดยไม่ต้องเพิ่มความซับซ้อนทางวิศวกรรมเพิ่มเติม มันผสานรวมเข้ากับระบบนิเวศ Transformer ที่มีอยู่ได้อย่างราบรื่น ช่วยเร่งการประยุกต์ใช้ LLMs ในสถานการณ์ต่างๆ