Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Bài báo · arXiv 2312.00752

Mamba: Mô hình hóa chuỗi thời gian tuyến tính với không gian trạng thái chọn lọc

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Xuất bản 2023-12-01cs.LGcs.AI arXiv PDF

Đọc với bản dịch và AI

Tóm tắt trong ba dòng

🎯 Transformer mở rộng với độ phức tạp bậc hai theo độ dài chuỗi và yêu cầu bộ nhớ đệm KV cache lớn khi suy luận tự hồi quy, trong khi các mô hình dưới bậc hai trước đây như SSM bất biến theo thời gian tuyến tính (LTI) lại gặp khó khăn trong các tác vụ suy luận nội dung rời rạc như mô hình hóa ngôn ngữ.
🔧 Làm cho các tham số SSM phụ thuộc vào đầu vào (chọn lọc) để lọc hoặc truyền thông tin một cách linh hoạt, kết hợp với thuật toán parallel associative scan hợp nhất tối ưu phần cứng và tính toán lại (recomputation) nhằm tránh việc hiện thực hóa trạng thái ẩn mở rộng trong GPU HBM.
📈 Mamba-3B đạt hiệu năng tương đương hoặc vượt trội so với các Transformer có kích thước gấp đôi trong các tác vụ zero-shot downstream (ví dụ: trung bình 63,3% so với 59,1% của Pythia-2.8B và 61,7% của Pythia-6.9B) trong khi mang lại thông lượng suy luận cao hơn tới 5 lần.

Tóm tắt

Các mô hình nền tảng, hiện đang vận hành hầu hết các ứng dụng đột phá trong học sâu, hầu như đều dựa trên kiến trúc Transformer và mô-đun chú ý cốt lõi của nó. Nhiều kiến trúc có thời gian dưới bậc hai như linear attention, gated convolution, các mô hình hồi quy và mô hình không gian trạng thái có cấu trúc (SSM) đã được phát triển nhằm giải quyết sự kém hiệu quả về mặt tính toán của Transformer trên các chuỗi dài, nhưng chúng chưa đạt hiệu năng tốt như cơ chế chú ý trên các phương thức quan trọng như ngôn ngữ. Chúng tôi nhận thấy điểm yếu cốt lõi của các mô hình này là không có khả năng suy luận dựa trên nội dung, và từ đó đưa ra một số cải tiến. Đầu tiên, việc cho phép các tham số của SSM phụ thuộc vào đầu vào sẽ giải quyết điểm yếu của chúng đối với dữ liệu rời rạc, cho phép mô hình truyền tải hoặc quên thông tin một cách chọn lọc dọc theo chiều dài chuỗi tùy thuộc vào token hiện tại. Thứ hai, mặc dù thay đổi này ngăn cản việc sử dụng các phép tích chập hiệu quả, chúng tôi thiết kế một thuật toán song song nhận biết phần cứng ở chế độ hồi quy. Chúng tôi tích hợp các SSM chọn lọc này vào một kiến trúc mạng nơ-ron đầu-cuối đơn giản hóa mà không cần khối chú ý hay thậm chí là khối MLP (Mamba). Mamba sở hữu tốc độ suy luận nhanh (thông lượng cao gấp 5 $\times$ so với Transformer) và khả năng mở rộng tuyến tính theo độ dài chuỗi, đồng thời hiệu năng được cải thiện trên dữ liệu thực tế với độ dài chuỗi lên tới hàng triệu. Với vai trò là một kiến trúc nền tảng xử lý chuỗi tổng quát, Mamba đạt hiệu năng vượt trội trên nhiều phương thức như ngôn ngữ, âm thanh và hệ gen. Trong mô hình hóa ngôn ngữ, mô hình Mamba-3B của chúng tôi vượt trội hơn các Transformer cùng kích thước và tương đương với các Transformer có kích thước gấp đôi, ở cả giai đoạn tiền huấn luyện lẫn đánh giá downstream.

Tóm tắt gốc (tiếng Anh)

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.

Hướng dẫn đọc

Vấn đề và động lực

Các mô hình nền tảng hiện nay phụ thuộc áp đảo vào kiến trúc Transformer và cơ chế tự chú ý (self-attention). Mặc dù cơ chế chú ý dày đặc định tuyến thông tin hiệu quả qua cửa sổ ngữ cảnh của chuỗi, nó lại vấp phải hai hạn chế cơ bản: chi phí tính toán bậc hai $O (L^{2})$ theo độ dài chuỗi $L$ , và yêu cầu lưu trữ các khóa (key) cùng giá trị (value) trong quá khứ (KV cache), dẫn đến thời gian suy luận tăng tuyến tính theo từng bước và gây nghẽn bộ nhớ.

Các phương pháp dưới bậc hai trước đây—chẳng hạn như linear attention, các kiến trúc hồi quy và mô hình không gian trạng thái có cấu trúc (SSM) như S4—mở rộng tuyến tính hoặc giả tuyến tính theo thời gian. Tuy nhiên, các SSM có cấu trúc trong lịch sử đã không thể theo kịp chất lượng của Transformer trên các phương thức rời rạc, mật độ thông tin cao như văn bản. Các tác giả chỉ ra nguyên nhân gốc rễ: các SSM trước đây là hệ tuyến tính bất biến theo thời gian (LTI). Vì động lực chuyển tiếp của chúng không thay đổi theo nội dung đầu vào, chúng không thể nén ngữ cảnh một cách chọn lọc, lọc bỏ nhiễu không liên quan, hoặc điều chỉnh khả năng ghi nhớ qua các khoảng cách biến thiên.

Ý tưởng cốt lõi

Cơ chế chọn lọc dưới dạng nén ngữ cảnh: Mô hình hóa chuỗi đòi hỏi phải nén lịch sử vào một trạng thái. Các SSM LTI trước đây sử dụng các tham số chuyển tiếp cố định, buộc chúng phải xử lý mọi bước thời gian như nhau. Bằng cách tham số hóa các ma trận SSM dựa trên token đầu vào hiện tại, mô hình có được nhận thức về nội dung cần thiết để lọc hoặc ghi nhớ thông tin một cách chọn lọc.
Vượt qua nút thắt cổ chai không thể dùng tích chập: Việc làm cho các tham số phụ thuộc vào đầu vào sẽ phá vỡ sự tương đương toán học giữa hồi quy tuyến tính và tích chập toàn cục. Nếu không có tích chập, phép hồi quy ngây thơ sẽ đòi hỏi phải hiện thực hóa trạng thái ẩn mở rộng có kích thước $(B, L, D, N)$ trong bộ nhớ băng thông cao (HBM) của GPU. Mamba giải quyết vấn đề này bằng cách hợp nhất việc rời rạc hóa, parallel associative scan và phép chiếu đầu ra vào bộ nhớ SRAM nhanh trên chip GPU, giảm thiểu tối đa lưu lượng truy cập HBM.
Kiến trúc hợp nhất: Mamba thay thế các khối Multi-Head Attention (MHA) và MLP xen kẽ bằng một khối đồng nhất duy nhất, kết hợp thiết kế cổng của H3 với MLP có cổng theo kiểu SwiGLU.

Phương pháp

Các mô hình không gian trạng thái liên tục ánh xạ đầu vào $x (t) \in R$ sang $y (t) \in R$ thông qua một trạng thái ẩn $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

Sử dụng phép rời rạc hóa zero-order hold (ZOH) với kích thước bước $Δ$ , hệ thống liên tục chuyển đổi thành các tham số hồi quy rời rạc: $\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}$

Cơ chế chọn lọc (S6)

Thay vì giữ cố định $(Δ, B, C)$ theo thời gian, Mamba biến chúng thành các hàm phụ thuộc vào đầu vào:

$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}))$ , với $τ_{Δ} = softplus$ và $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Điều này khái quát hóa cơ chế gating của RNN cổ điển: khi $N = 1, A = - 1, B = 1$ , phép hồi quy rút gọn trực tiếp thành phép cập nhật có cổng $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , trong đó $g_{t} = σ (Linear (x_{t}))$ .

Mở rộng trạng thái tối ưu phần cứng (Hardware-Aware State Expansion)

Để duy trì thông lượng cao mà không cần hiện thực hóa trạng thái $(B, L, D, N)$ trong HBM:

Các tham số $(Δ, A, B, C)$ được nạp từ HBM trực tiếp vào SRAM trên chip.
Quá trình rời rạc hóa thành $(\overline{A}, \overline{B})$ và thuật toán parallel associative scan hiệu quả được tính toán hoàn toàn bên trong SRAM.
Trạng thái ẩn được co rút với $C$ về kích thước $(B, L, D)$ trước khi ghi ngược lại HBM.
Các trạng thái trung gian được tính toán lại trong SRAM trong quá trình lan truyền ngược thay vì lưu trong HBM.

Kết quả

Ngoại suy tổng hợp: Trên tác vụ Induction Heads, Mamba được huấn luyện trên chuỗi độ dài 256 có thể ngoại suy hoàn hảo (độ chính xác 100%) lên tới độ dài $2^{20} = 1, 048, 576$ (gấp $4000 \times$ độ dài chuỗi), trong khi các mô hình cơ sở chú ý thất bại hoặc tràn bộ nhớ (OOM) sau $2^{14}$ . Trên tác vụ Selective Copying, Mamba đạt độ chính xác 99,8% so với 57,0% của H3 và 18,3% của S4.
Mô hình hóa ngôn ngữ: Được đánh giá theo quy luật mở rộng Chinchilla lên đến 1,3 tỷ tham số trên tập dữ liệu Pile, Mamba đạt hiệu năng tương đương hoặc vượt trội so với Transformer++ (rotary embeddings, SwiGLU, RMSNorm). Trên các tác vụ zero-shot downstream (ví dụ: LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B đạt độ chính xác trung bình 63,3%, vượt qua Pythia-2.8B (59,1%), Pythia-6.9B (61,7%), và tương đương hoặc vượt các mô hình 6B–7B tham số.
Hệ gen và Âm thanh: Trong phân loại DNA của vượn người lớn ở độ dài chuỗi 1M ( $2^{20}$ ), Mamba đạt độ chính xác 71,67% (1,4M tham số) và 81,31% (7M tham số) so với 54,87% của HyenaDNA. Trong tác vụ tạo tiếng nói SC09, Mamba giảm FID từ 1,99 (SaShiMi) xuống 0,94 (6,1M tham số) và 0,67 (24,3M tham số).
Hiệu quả tính toán: Kernel scan hợp nhất thực thi nhanh hơn 20–40 lần so với các cài đặt scan chuẩn trong PyTorch và chạy nhanh hơn tới 7 lần so với FlashAttention-2 ở ngữ cảnh 32K. Thông lượng suy luận cao hơn tới 5 lần so với các mô hình Transformer tương đương.

Hạn chế và các câu hỏi mở

Hiệu năng trên tín hiệu liên tục: Mặc dù cơ chế chọn lọc phụ thuộc đầu vào giúp cải thiện khả năng suy luận rời rạc (văn bản, DNA), các phân tích loại trừ trên việc tiền huấn luyện dạng sóng âm thanh thô cho thấy các mô hình tuyến tính bất biến theo thời gian thuần túy (LTI) (như S4 chuẩn) hoạt động tốt hơn trên dữ liệu liên tục, được lấy mẫu đều đặn.
Quy mô mô hình: Các mô hình được đánh giá chỉ giới hạn ở mức 3B tham số và 300B token tiền huấn luyện; hiệu năng thực nghiệm ở quy mô 7B–70B+ trở lên vẫn cần nghiên cứu thêm.
Hệ sinh thái downstream: Transformer sở hữu các công cụ phong phú cho kỹ thuật tạo prompt, can thiệp in-context, lượng tử hóa KV-cache và tinh chỉnh bằng adapter. Việc tìm hiểu cách SSM thích ứng trên các chế độ tương tác này vẫn là hướng nghiên cứu mở.

Thuật ngữ chính

Structured State Space Models (SSMs): Một lớp mô hình chuỗi sâu bắt nguồn từ các hệ thống động lực tuyến tính liên tục, ánh xạ các đầu vào 1D thông qua một trạng thái ẩn tiềm ẩn.
Linear Time Invariance (LTI): Một đặc tính toán học trong đó các tham số chuyển tiếp của hệ thống không đổi qua mọi bước thời gian, cho phép biểu diễn tích chập toàn cục nhanh nhưng hạn chế khả năng suy luận phụ thuộc vào nội dung.
Selection Mechanism: Việc tham số hóa động lực chuyển tiếp SSM (Δ, B, C) dưới dạng hàm của token đầu vào, cho phép mô hình truyền tải, đặt lại hoặc bỏ qua ngữ cảnh một cách linh hoạt dọc theo chiều dài chuỗi.
S6: Tên viết tắt của các mô hình không gian trạng thái có cấu trúc chọn lọc được tính toán thông qua parallel associative scan thay vì tích chập.
Parallel Associative Scan: Thuật toán tính toán tiền tố hiệu quả giúp đánh giá phép hồi quy tuyến tính trên toàn bộ chiều dài chuỗi một cách song song trên phần cứng tăng tốc.
Hardware-Aware State Expansion: Chiến lược hợp nhất kernel GPU và tính toán lại (recomputation) nhằm thực hiện rời rạc hóa và quét bên trong bộ nhớ SRAM nhanh trên chip để tránh việc luân chuyển các trạng thái trung gian lớn sang bộ nhớ băng thông cao (HBM).
Induction Heads: Một cơ chế gợi nhớ liên kết tổng hợp nhằm đánh giá xem mô hình có thể tái tạo lại một token đã quan sát trước đó theo sau token ngữ cảnh liền trước hay không.
Mamba Block: Một khối mô hình chuỗi đồng nhất, đơn giản hóa, tích hợp các phép chiếu đầu vào, tích chập 1D, SSM chọn lọc và gating nhân tử mà không cần các lớp attention hoặc MLP riêng biệt.

Câu hỏi thường gặp

Tại sao các SSM chọn lọc không thể được huấn luyện bằng tích chập FFT nhanh?

Tích chập FFT toàn cục nhanh đòi hỏi động lực hệ thống (các ma trận A, B và C) phải cố định theo thời gian (tính chất tuyến tính bất biến theo thời gian - LTI). Vì Mamba biến B, C và Δ thành các hàm phụ thuộc vào token đầu vào tại mỗi bước thời gian, kernel tích chập sẽ biến thiên tại mỗi vị trí, làm phá vỡ sự tương đương tích chập và buộc phải sử dụng parallel scan hồi quy thay thế.

Mamba tránh việc bùng nổ bộ nhớ do mở rộng trạng thái hồi quy trong quá trình huấn luyện bằng cách nào?

Phép hồi quy thông thường với chiều mở rộng trạng thái N đòi hỏi phải hiện thực hóa một trạng thái trung gian có kích thước (B, L, D, N) trong HBM của GPU. Mamba sử dụng kỹ thuật hợp nhất kernel để chỉ nạp các đầu vào số chiều thấp vào SRAM, tính toán quá trình rời rạc hóa và parallel scan bên trong SRAM tốc độ cao, chỉ ghi đầu ra đã thu gọn (B, L, D) ngược lại HBM, và tính toán lại các kích hoạt trung gian trong quá trình lan truyền ngược.

Mối quan hệ giữa tham số Δ của Mamba và cơ chế cổng (gating) hồi quy là gì?

Tham số kích thước bước Δ hoạt động như một cơ chế cổng khái quát hóa. Khi được rời rạc hóa theo zero-order hold (ZOH) với N=1 và A=-1, phép hồi quy chọn lọc rút gọn về mặt toán học thành một cổng RNN tiêu chuẩn: Δ lớn làm cho mô hình tập trung vào đầu vào hiện tại và đặt lại trạng thái, trong khi Δ nhỏ duy trì trạng thái lịch sử và bỏ qua đầu vào.

Mamba có sử dụng KV cache trong quá trình sinh văn bản không?

Không. Vì Mamba hoạt động như một hệ hồi quy không gian trạng thái tự hồi quy trong quá trình tạo văn bản, nó chỉ duy trì một trạng thái ẩn có kích thước cố định (chiều D × N trên mỗi lớp) thay vì lưu bộ nhớ đệm khóa và giá trị cho tất cả các token trong quá khứ, giúp suy luận đạt thời gian không đổi O(1) ở mỗi bước.

Phần giải thích trên trang này do AI viết dựa trên bài báo và có thể có sai sót. Hãy đọc bài báo gốc để biết chi tiết.

Đọc với bản dịch và AI