Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

งานวิจัย · arXiv 2312.00752

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

เผยแพร่ 2023-12-01cs.LGcs.AI arXiv PDF

อ่านพร้อมคำแปลและ AI

สรุปสามบรรทัด

🎯 Transformers สเกลเป็นกำลังสองตามความยาวของลำดับ และต้องการ KV cache ขนาดใหญ่ระหว่างการอนุมานแบบ autoregressive ในขณะที่โมเดล subquadratic ก่อนหน้านี้ เช่น LTI state space models มีปัญหาในงานที่ต้องใช้การให้เหตุผลตามเนื้อหาแบบไม่ต่อเนื่อง เช่น การสร้างโมเดลภาษา
🔧 กำหนดให้พารามิเตอร์ของ SSM ขึ้นอยู่กับอินพุต (selective) เพื่อกรองหรือส่งผ่านข้อมูลแบบไดนามิก และจับคู่สิ่งนี้เข้ากับขั้นตอนวิธี fused parallel associative scan และ recomputation ที่คำนึงถึงฮาร์ดแวร์ เพื่อหลีกเลี่ยงการสร้าง hidden state ที่ขยายขนาดแล้วลงใน GPU HBM
📈 Mamba-3B เทียบเท่าหรือเหนือกว่า Transformers ที่มีขนาดใหญ่กว่าถึงสองเท่าใน downstream zero-shot tasks (เช่น ค่าเฉลี่ย 63.3% เทียบกับ 59.1% สำหรับ Pythia-2.8B และ 61.7% สำหรับ Pythia-6.9B) ในขณะที่ให้ทรูพุตในการอนุมานสูงกว่าสูงสุดถึง 5 เท่า

บทคัดย่อ

โมเดลรากฐาน (Foundation models) ซึ่งขับเคลื่อนแอปพลิเคชันส่วนใหญ่ที่น่าตื่นเต้นในด้านดีปเลิร์นนิงปัจจุบัน แทบทั้งหมดมีพื้นฐานมาจากสถาปัตยกรรม Transformer และโมดูล attention ที่เป็นแกนหลัก สถาปัตยกรรมที่ใช้เวลาต่ำกว่ากำลังสอง (subquadratic-time) หลายตัว เช่น linear attention, gated convolution, recurrent models และ structured state space models (SSMs) ได้รับการพัฒนาขึ้นเพื่อแก้ปัญหาความไม่มีประสิทธิภาพเชิงการคำนวณของ Transformers บนลำดับข้อมูลที่มีความยาวสูง แต่โมเดลเหล่านี้ยังมีประสิทธิภาพไม่เทียบเท่ากับ attention ในรูปแบบข้อมูลที่สำคัญอย่างภาษา เราพบว่าจุดอ่อนสำคัญของโมเดลดังกล่าวคือการขาดความสามารถในการให้เหตุผลตามเนื้อหา (content-based reasoning) และได้พัฒนาข้อปรับปรุงหลายประการ ประการแรก เพียงแค่กำหนดให้พารามิเตอร์ของ SSM เป็นฟังก์ชันของอินพุต ก็สามารถแก้จุดอ่อนกับรูปแบบข้อมูลแบบไม่ต่อเนื่อง (discrete modalities) ได้ ซึ่งช่วยให้โมเดลสามารถส่งผ่านหรือละทิ้งข้อมูลตามมิติความยาวของลำดับได้อย่างจำเพาะเจาะจง (selective) โดยขึ้นอยู่กับโทเคนปัจจุบัน ประการที่สอง แม้ว่าการเปลี่ยนแปลงนี้จะทำให้ไม่สามารถใช้ convolution ที่มีประสิทธิภาพได้ แต่เราได้ออกแบบขั้นตอนวิธีแบบขนานที่คำนึงถึงฮาร์ดแวร์ (hardware-aware parallel algorithm) ในโหมด recurrence เราได้ผสานรวม selective SSMs เหล่านี้เข้ากับสถาปัตยกรรมโครงข่ายประสาทเทียมแบบ end-to-end ที่เรียบง่าย โดยไม่ต้องมี attention หรือแม้แต่บล็อก MLP (Mamba) สถาปัตยกรรม Mamba มีความสามารถในการอนุมานที่รวดเร็ว (มีทรูพุตสูงกว่า Transformers ถึง 5 $\times$ ) และมีการสเกลเป็นเส้นตรงตามความยาวลำดับข้อมูล อีกทั้งประสิทธิภาพยังดีขึ้นเมื่อใช้กับข้อมูลจริงที่มีความยาวลำดับสูงถึงระดับล้านโทเคน ในฐานะโมเดลแกนหลักสำหรับการประมวลผลลำดับทั่วไป Mamba บรรลุประสิทธิภาพระดับ state-of-the-art ในหลายรูปแบบข้อมูล เช่น ภาษา, เสียง และจีโนมิกส์ ในงานโมเดลภาษา โมเดล Mamba-3B ของเรามีประสิทธิภาพเหนือกว่า Transformers ที่มีขนาดเท่ากัน และเทียบเคียงกับ Transformers ที่มีขนาดใหญ่กว่าถึงสองเท่า ทั้งในการพรีเทรนและการประเมินผล downstream

บทคัดย่อต้นฉบับ (ภาษาอังกฤษ)

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

คู่มือการอ่าน

ปัญหาและแรงจูงใจ

โมเดลรากฐาน (Foundation models) ส่วนใหญ่พึ่งพาสถาปัตยกรรม Transformer และกลไก self-attention แม้ว่า dense attention จะส่งผ่านข้อมูลข้าม context window ของลำดับได้อย่างมีประสิทธิภาพ แต่ก็มีข้อจำกัดพื้นฐานสองประการ ได้แก่ ต้นทุนการคำนวณระดับกำลังสอง $O (L^{2})$ เทียบกับความยาวของลำดับ $L$ และความจำเป็นในการจัดเก็บ key และ value ในอดีต (KV cache) ซึ่งส่งผลให้เกิดการอนุมานที่ใช้เวลาเชิงเส้นต่อสเต็ปและเกิดคอขวดด้านหน่วยความจำ

แนวทางก่อนหน้านี้ที่ใช้เวลาต่ำกว่ากำลังสอง เช่น linear attention, สถาปัตยกรรมแบบ recurrent และ structured state space models (SSMs) เช่น S4 สเกลในเชิงเส้นหรือกึ่งเชิงเส้นตามเวลา อย่างไรก็ตาม โครงสร้างแบบ SSMs ในอดีตมักไม่สามารถทำผลงานได้ทัดเทียมกับคุณภาพของ Transformer ในรูปแบบข้อมูลที่ไม่ต่อเนื่องและมีความหนาแน่นของข้อมูลสูง เช่น ข้อความ ผู้เขียนระบุสาเหตุหลักว่า: SSMs ก่อนหน้านี้เป็นระบบ linear time-invariant (LTI) เนื่องจากพลวัตการเปลี่ยนสถานะ (transition dynamics) ของระบบไม่เปลี่ยนแปลงตามเนื้อหาของอินพุต ทำให้ไม่สามารถบีบอัดบริบทแบบจำเพาะเจาะจง กรองสัญญาณรบกวนที่ไม่เกี่ยวข้อง หรือปรับการเก็บข้อมูลข้ามระยะห่างที่ไม่คงที่ได้

แนวคิดหลัก

กลไกการเลือก (Selection Mechanism) ในฐานะการบีบอัดบริบท: การสร้างโมเดลลำดับข้อมูลจำเป็นต้องบีบอัดประวัติที่ผ่านมาให้อยู่ในสถานะหนึ่งๆ SSMs แบบ LTI ก่อนหน้านี้ใช้พารามิเตอร์การเปลี่ยนสถานะที่คงที่ ทำให้ถูกบังคับให้ปฏิบัติต่อทุกช่วงเวลาเหมือนกันทั้งหมด แต่ด้วยการกำหนดพารามิเตอร์เมทริกซ์ของ SSM ตามโทเคนอินพุตปัจจุบัน โมเดลจะรับรู้บริบทของเนื้อหาเพื่อเลือกกรองหรือจดจำข้อมูลได้อย่างจำเพาะเจาะจง
การก้าวข้ามคอขวดของการไม่ใช้ Convolution: การทำให้พารามิเตอร์ขึ้นอยู่กับอินพุตทำลายความเท่ากันทางคณิตศาสตร์ระหว่าง linear recurrence และ global convolution หากไม่มี convolution การทำ recurrence แบบปกติจะต้องสร้างสถานะแฝงขนาดใหญ่ $(B, L, D, N)$ ลงใน high-bandwidth memory (HBM) ของ GPU Mamba แก้ปัญหานี้โดยการหลอมรวมการแปลงเป็นค่าไม่ต่อเนื่อง (discretization), parallel associative scan และ output projection เข้าไปทำงานบน SRAM บนชิป GPU ที่มีความเร็วสูง เพื่อลดการส่งข้อมูลผ่าน HBM ให้น้อยที่สุด
สถาปัตยกรรมที่เป็นเอกภาพ (Unified Architecture): Mamba แทนที่บล็อก Multi-Head Attention (MHA) และ MLP ที่สลับกันไปมา ด้วยบล็อกแบบเอกภาพรูปแบบเดียวที่รวมการออกแบบ gated ของ H3 เข้ากับ gated MLP สไตล์ SwiGLU

ระเบียบวิธี

โมเดล state space แบบต่อเนื่องจะแมปอินพุต $x (t) \in R$ ไปยัง $y (t) \in R$ ผ่านสถานะแฝง $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

การใช้การแปลงเป็นค่าไม่ต่อเนื่องแบบ zero-order hold (ZOH) ด้วยขนาดสเต็ป $Δ$ จะแปลงระบบต่อเนื่องให้เป็นพารามิเตอร์ recurrence แบบไม่ต่อเนื่อง: $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

กลไกการเลือก (S6)

แทนที่จะกำหนดให้ $(Δ, B, C)$ คงที่ตามเวลา Mamba ทำให้พารามิเตอร์เหล่านี้ขึ้นอยู่กับอินพุต:

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ โดยที่ $τ_{Δ} = softplus$ และ $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

สิ่งนี้เป็นการขยายแนวคิด gating ของ RNN ดั้งเดิม: เมื่อ $N = 1, A = - 1, B = 1$ การทำ recurrence จะลดรูปกลายเป็นการอัปเดตแบบ gated โดยตรง $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ โดยที่ $g_{t} = σ (Linear (x_{t}))$

การขยายสถานะที่คำนึงถึงฮาร์ดแวร์ (Hardware-Aware State Expansion)

เพื่อรักษาทรูพุตระดับสูงโดยไม่ต้องสร้างสถานะขนาด $(B, L, D, N)$ ใน HBM:

พารามิเตอร์ $(Δ, A, B, C)$ จะถูกโหลดจาก HBM เข้าสู่ SRAM บนชิปโดยตรง
การแปลงเป็นค่าไม่ต่อเนื่องให้เป็น $(\overline{A}, \overline{B})$ และ parallel associative scan ที่มีประสิทธิภาพการคำนวณสูงจะถูกคำนวณภายใน SRAM ทั้งหมด
สถานะแฝงจะถูกย่อขนาดด้วย $C$ ให้เหลือขนาด $(B, L, D)$ ก่อนจะเขียนกลับไปยัง HBM
สถานะขั้นกลางจะถูกคำนวณซ้ำ (recomputed) ใน SRAM ระหว่าง backward pass แทนที่จะบันทึกไว้ใน HBM

ผลลัพธ์

การคาดการณ์จากข้อมูลสังเคราะห์ (Synthetic Extrapolation): ในภารกิจ Induction Heads สถาปัตยกรรม Mamba ที่เทรนบนความยาวลำดับ 256 สามารถคาดการณ์ได้อย่างสมบูรณ์แบบ (ความแม่นยำ 100%) จนถึงความยาว $2^{20} = 1, 048, 576$ (ความยาวลำดับเพิ่มขึ้น $4000 \times$ ) ในขณะที่ baseline ของ attention ล้มเหลวหรือหน่วยความจำเต็ม (OOM) เกินกว่า $2^{14}$ สำหรับ Selective Copying โมเดล Mamba มีความแม่นยำถึง 99.8% เมื่อเทียบกับ 57.0% ของ H3 และ 18.3% ของ S4
การสร้างโมเดลภาษา: เมื่อประเมินภายใต้กฎ Chinchilla scaling สูงสุดถึง 1.3B พารามิเตอร์บนชุดข้อมูล Pile พบว่า Mamba เทียบเท่าหรือเหนือกว่า Transformer++ (rotary embeddings, SwiGLU, RMSNorm) สำหรับ downstream zero-shot tasks (เช่น LAMBADA, PIQA, HellaSwag, WinoGrande, ARC) Mamba-2.8B บรรลุความแม่นยำเฉลี่ย 63.3% ซึ่งสูงกว่า Pythia-2.8B (59.1%), Pythia-6.9B (61.7%) และเทียบเท่าหรือเหนือกว่าโมเดลขนาด 6B–7B พารามิเตอร์
จีโนมิกส์และเสียง: ในการจำแนก DNA ของ Great Apes ที่ความยาวลำดับ 1M ( $2^{20}$ ) Mamba บรรลุความแม่นยำ 71.67% (1.4M พารามิเตอร์) และ 81.31% (7M พารามิเตอร์) เทียบกับ 54.87% ของ HyenaDNA สำหรับการสร้างเสียงพูด SC09 โมเดล Mamba ลดค่า FID จาก 1.99 (SaShiMi) เหลือ 0.94 (6.1M พารามิเตอร์) และ 0.67 (24.3M พารามิเตอร์)
ประสิทธิภาพ: fused scan kernel ทำงานได้เร็วกว่าการทำ scan ใน PyTorch แบบมาตรฐาน 20–40 เท่า และเร็วกว่า FlashAttention-2 สูงสุด 7 เท่าที่บริบทขนาด 32K ทรูพุตในการอนุมานสูงกว่า Transformers ที่เทียบเคียงกันสูงสุดถึง 5 เท่า

ข้อจำกัดและคำถามปลายเปิด

ประสิทธิภาพกับสัญญาณต่อเนื่อง: แม้ว่าการเลือกที่ขึ้นอยู่กับอินพุตจะช่วยปรับปรุงการให้เหตุผลแบบไม่ต่อเนื่อง (ข้อความ, DNA) แต่การทดลอง ablation บนการพรีเทรนรูปคลื่นเสียงดิบแสดงให้เห็นว่า โมเดลที่เป็น linear time-invariant (LTI) ล้วน (เช่น S4 มาตรฐาน) ทำงานได้ดีกว่ากับข้อมูลต่อเนื่องที่ถูกสุ่มอย่างสม่ำเสมอ
สเกลของโมเดล: โมเดลที่ได้รับการประเมินจำกัดอยู่ที่ประมาณ 3B พารามิเตอร์ และโทเคนการพรีเทรน 300B โทเคน ประสิทธิภาพเชิงประจักษ์ที่ระดับ 7B–70B+ พารามิเตอร์ขึ้นไปยังต้องมีการศึกษาเพิ่มเติม
ระบบนิเวศ Downstream: Transformers มีเครื่องมือมากมายสำหรับการทำ prompt engineering, การปรับแต่ง in-context, การทำ quantize บน KV-cache และการปรับจูน adapter การทำความเข้าใจว่า SSMs จะปรับตัวเข้ากับรูปแบบการทำงานร่วมกันเหล่านี้อย่างไรยังคงเป็นหัวข้อวิจัยที่มีการสำรวจอย่างต่อเนื่อง

คำศัพท์สำคัญ

Structured State Space Models (SSMs): กลุ่มของโมเดลลำดับข้อมูลแบบลึกที่พัฒนามาจากระบบ dynamic เชิงเส้นแบบต่อเนื่อง ซึ่งแปลงอินพุต 1 มิติผ่านสถานะแฝงโดยนัย
Linear Time Invariance (LTI): คุณสมบัติทางคณิตศาสตร์ที่พารามิเตอร์การเปลี่ยนสถานะของระบบยังคงคงที่ในทุกช่วงเวลา ช่วยให้สามารถแทนค่าด้วย global convolution ได้อย่างรวดเร็ว แต่จำกัดความสามารถในการให้เหตุผลตามเนื้อหา
Selection Mechanism: การกำหนดให้พลวัตการเปลี่ยนสถานะของ SSM (Δ, B, C) เป็นฟังก์ชันของโทเคนอินพุต ช่วยให้โมเดลสามารถส่งผ่าน รีเซ็ต หรือละทิ้งบริบทตามความยาวลำดับได้อย่างไดนามิก
S6: คำย่อของ selective structured state space models ที่คำนวณผ่าน associative parallel scan แทนที่จะเป็น convolution
Parallel Associative Scan: ขั้นตอนวิธี prefix computation ที่มีประสิทธิภาพการคำนวณสูง ซึ่งประมวลผล linear recurrence ข้ามความยาวลำดับแบบขนานบนฮาร์ดแวร์ตัวเร่งความเร็ว
Hardware-Aware State Expansion: กลยุทธ์การทำ kernel fusion และ recomputation บน GPU ซึ่งทำการแปลงเป็นค่าไม่ต่อเนื่องและ scan ใน SRAM บนชิปที่มีความเร็วสูง เพื่อหลีกเลี่ยงการส่งสถานะขั้นกลางขนาดใหญ่ไปกลับยัง high-bandwidth memory (HBM)
Induction Heads: กลไก associative recall สังเคราะห์ที่ใช้ประเมินว่าโมเดลสามารถสร้างโทเคนที่เคยสังเกตมาก่อนซ้ำได้หรือไม่เมื่อพบกับโทเคนบริบทก่อนหน้า
Mamba Block: บล็อกโมเดลลำดับแบบเอกภาพที่เรียบง่าย ซึ่งรวม input projections, 1D convolution, selective SSMs และ multiplicative gating เข้าด้วยกันโดยไม่มีเลเยอร์ attention หรือ MLP แยกต่างหาก

คำถามที่พบบ่อย

ทำไม selective SSMs ถึงไม่สามารถเทรนโดยใช้ FFT convolutions ที่รวดเร็วได้?

การทำ global FFT convolutions ที่รวดเร็วจำเป็นต้องให้พลวัตของระบบ (เมทริกซ์ A, B และ C) มีค่าคงที่ตามเวลา (linear time invariance) แต่เนื่องจาก Mamba กำหนดให้ B, C และ Δ เป็นฟังก์ชันของโทเคนอินพุตในแต่ละช่วงเวลา convolution kernel จึงเปลี่ยนไปในทุกตำแหน่ง ทำให้สูญเสียความเท่ากันทางคณิตศาสตร์แบบ convolutional และจำเป็นต้องใช้ parallel scan แบบ recurrent แทน

Mamba หลีกเลี่ยงปัญหาหน่วยความจำระเบิดจากการขยายสถานะแบบ recurrent ระหว่างการเทรนได้อย่างไร?

การทำ recurrence แบบปกติที่มีมิติการขยายสถานะ N จะต้องสร้างสถานะขั้นกลางขนาด (B, L, D, N) ใน GPU HBM แต่ Mamba ใช้ kernel fusion เพื่อโหลดเฉพาะอินพุตมิติต่ำเข้าสู่ SRAM จากนั้นทำการแปลงเป็นค่าไม่ต่อเนื่องและ parallel scan ภายใน SRAM ที่มีความเร็วสูง แล้วเขียนเฉพาะผลลัพธ์ที่ลดขนาดแล้ว (B, L, D) กลับไปยัง HBM รวมถึงคำนวณ activation ขั้นกลางใหม่ระหว่างการทำ backpropagation

พารามิเตอร์ Δ ของ Mamba มีความสัมพันธ์กับ recurrent gating อย่างไร?

พารามิเตอร์ขนาดสเต็ป Δ ทำหน้าที่เสมือนกลไก gating ทั่วไป เมื่อแปลงเป็นค่าไม่ต่อเนื่องภายใต้ zero-order hold (ZOH) โดยที่ N=1 และ A=-1 การทำ selective recurrence จะลดรูปทางคณิตศาสตร์เป็นเกตของ RNN มาตรฐาน: โดยที่ค่า Δ ขนาดใหญ่จะทำให้โมเดลเน้นไปที่อินพุตปัจจุบันและรีเซ็ตสถานะเดิม ในขณะที่ค่า Δ ขนาดเล็กจะรักษาประวัติสถานะในอดีตไว้และละทิ้งอินพุตปัจจุบัน

Mamba ใช้ KV cache ระหว่างการสร้างข้อความหรือไม่?

ไม่ใช้ เนื่องจาก Mamba ทำงานในรูปแบบ state space recurrence แบบ autoregressive ระหว่างการสร้างผลลัพธ์ จึงรักษาเพียง hidden state ที่มีขนาดคงที่ (มิติ D × N ต่อเลเยอร์) แทนที่จะต้องแคช key และ value สำหรับโทเคนในอดีตทั้งหมด ส่งผลให้การอนุมานใช้เวลาคงที่ O(1) ต่อสเต็ป

คำอธิบายในหน้านี้เขียนโดย AI จากงานวิจัยและอาจมีข้อผิดพลาด โปรดอ่านต้นฉบับเพื่อดูรายละเอียด

อ่านพร้อมคำแปลและ AI