TROISINH
FrontierNew Architectures

Hybrid Attention-SSM (Jamba, Zamba) — SSM + attention khi cần precision

Giải mã kiến trúc lai giữa Mamba SSM và Transformer Attention: tận dụng tốc độ tuyến tính của selective scan để nén context, xen kẽ attention layers để giữ độ chính xác khi cần nhớ chi tiết xa xôi.

Transformer models chết dần khi context đạt 100K tokens—KV cache nuốt chửng GPU memory. Ngược lại, mô hình SSM thuần túy như Mamba chạy nhanh qua triệu token nhưng lãng quên chi tiết như sinh viên ghi chép cẩu thả. Kiến trúc lai Hybrid Attention-SSM (Jamba, Zamba) mang đến lối thoát tinh tế: dùng SSM layers nén context tốc độ tuyến tính cho 90% công việc, xen kẽ Attention layers đúng lúc để bắt lấy những chi tiết xa xôi quan trọng. Không cần chọn giữa tốc độ và độ chính xác—bạn có thể có cả hai.

Vấn đề

Bốn năm qua, cộng đồng AI đối mặt với một nghịch lý nhị phân:

Transformer thuần túy mang lại độ chính xác tuyệt đối trong việc truy xuất thông tin (needle-in-haystack) nhưng trả giá bằng độ phức tạp O(n²). Với context 128K tokens, KV cache chiếm 32-64GB VRAM, buộc phải dùng nhiều GPU chỉ để lưu "bộ nhớ tạm"—chưa tính đến weights của mô hình. Đây là "bức tường memory" khiến inference trở nên đắt đỏ và chậm chạp.

Mamba/SSM thuần túy giải quyết vấn đề bằng selective scan—nén toàn bộ lịch sử vào một hidden state cố định (ví dụ: 4KB cho 100K tokens), đạt O(1) memory per step. Tuy nhiên, điều này tạo ra hiện tượng "amnesia có chọn lọc": mô hình giữ được "ý chính" (gist) nhưng quên mất chi tiết cụ thể như "điều khoản chấm dứt hợp đồng ở trang 847". SSM là bản tóm tắt mất mát (lossy compression), không phải bộ nhớ địa chỉ nội dung (content-addressable).

Vấn đề cốt lõi: Không có kiến trúc đơn lẻ nào tối ưu cho cả hai chế độ—xử lý thông tin dài liên tục (cần tốc độ) và truy xuất chính xác từng chi tiết xa xôi (cần độ phân giải cao).

Ý tưởng cốt lõi

Kiến trúc "Goldilocks": Nén khi có thể, nhớ khi cần.

Thay vì chọn một bên, Hybrid Attention-SSM xen kẽ hai loại layers theo tỷ lệ tối ưu—thường là 1:7 (1 Attention layer cho mỗi 8 layers tổng cộng). Đây là ý chính:

SSM layers là "context carriers": Các lớp Mamba chịu trách nhiệm truyền tải luồng thông tin (context flow) qua toàn bộ sequence với chi phí bộ nhớ cố định. Chúng nén lịch sử thành hidden state, cho phép model "đi xa" mà không bị nghẹt thở bởi KV cache.

Attention layers là "refresh stations": Mỗi lớp Attention đóng vai trò như trạm kiểm tra. Nhờ residual connections, Attention layer có thể "nhìn lại" raw text gốc, truy xuất chính xác ngày tháng, số liệu, hoặc điều khoản pháp lý từ bất kỳ đâu trong context. Sau khi lấy thông tin chính xác, SSM layers tiếp tục mang thông tin đó đi xa.

Tích hợp MoE (Mixture of Experts): Vì SSM tiết kiệm bộ nhớ KV cache (giải phóng GBs VRAM), Jamba tận dụng khoảng trống này để scale parameters. Jamba sử dụng 52B tổng tham số nhưng chỉ active 12B mỗi token, cho phép chạy context 256K trên single 80GB GPU—điều impossible với dense Transformer 52B thuần túy.

Hybrid song song (Parallel Hybrid): Thay vì xen kẽ layers, một số biến thể chia dimension bên trong cùng một layer—một nửa đi qua SSM path (tóm tắt), nửa còn lại qua Attention path (chi tiết), sau đó concatenate. Điều này duy trì cả hai dạng bộ nhớ tại mọi timestep.

That's it. Bạn không cần photographic memory cho toàn bộ cuốn sách—chỉ cần ghi chú hiệu quả cho phần lớn nội dung, và đánh dấu trang (bookmark) chính xác cho những trang quan trọng.

Tại sao nó hoạt động

Logic đằng sau sự kết hợp này nằm ở phân phối bất đối xứng của thông tin quan trọng trong văn bản thực tế. 90% tokens trong một tài liệu là "boilerplate"—ngữ cảnh chung, câu nối, lặp lại ý. Chỉ 10% là chi tiết mang tính định danh (tên, số, ngày tháng, điều khoản) đòi hỏi độ chính xác byte-level.

Selective Compression: SSM sử dụng cơ chế selective scan với các tham số input-dependent (B, C, Δ). Điều này cho phép mỗi token quyết định "có nên ghi đè lên trạng thái hiện tại hay không"—tương tự cơ chế ghi nhớ có cửa (gated memory). Kết quả là một vector state nhỏ gọn nhưng giàu thông tin, có thể coi là bản tóm tắt có trọng số (weighted summary) của toàn bộ quá khứ.

Content-Addressable Refresh: Attention layers cung cấp khả năng random access vào lịch sử thông qua softmax(QK^T)V. Khi SSM "nghi ngờ" về một chi tiết (tức là khi hidden state không chứa thông tin đủ cụ thể), Attention layer truy xuất trực tiếp từ KV cache của các token xa xôi—giống như tra cứu lại trang gốc thay vì tin vào bản tóm tắt.

Tính toán phân tán (Chunked Computation): Mamba-2 (trong cụm kiến trúc này) chứng minh SSM tương đương với phép nhân ma trận có cấu trúc (structured matrix multiply). Khi kết hợp với Attention, kiến trúc hybrid tận dụng được cả hai dạng song song: block-wise parallelism của SSM và sequence parallelism của Attention.

Synergy với MoE: Trong Jamba, việc giảm memory footprint nhờ SSM không chỉ để "chạy được" mà còn để "scale tham số". Bạn đánh đổi memory dành cho KV cache (giảm từ O(n) xuống O(1) per layer) lấy memory cho expert weights (tăng từ 7B lên 52B tổng tham số). Đây là sự dịch chuyển tài nguyên thông minh: từ "nhớ nhiều" sang "biết nhiều".

Ý nghĩa thực tế

Hiệu năng thực tế:

  • Jamba đạt 3× throughput so với Mixtral 8x7B trên context dài, đồng thời hỗ trợ cửa sổ context 256K tokens—mức trước đây chỉ feasible cho model 7B nhỏ hơn nhiều.
  • Zamba 7B đạt hiệu suất tương đương Llama-2 70B trên một số tác vụ reasoning với 50% ít compute training hơn, chứng minh hiệu quả data-efficiency của kiến trúc lai.

Ai đang dùng:

  • AI21 Labs (Jamba): Triển khai production với MoE scaling.
  • Zyphra (Zamba): Tối ưu cho edge devices với hiệu quả training cao.
  • Microsoft Research (Samba): Ứng dụng trong các tác vụ cần long-context retrieval.
  • Vision & Medical: MambaVision và TransMamba áp dụng hybrid cho ảnh y tế và genomics nơi cần cả pattern recognition lẫn detail preservation.

Trade-offs và giới hạn:

  • Độ phức tạp triển khai: Khó debug hơn kiến trúc thuần túy; cần careful tuning để cân bằng giữa SSM và Attention.
  • Needle-in-haystack: Dù tốt hơn pure SSM, hybrid vẫn thua Transformer thuần túy về khả năng truy xuất chính xác token ở vị trí xa (ví dụ: thông tin ở token 50K khi đang ở token 100K).
  • Tỷ lệ không universal: Tỷ lệ 1:7 phù hợp cho nhiều tác vụ nhưng có thể cần điều chỉnh (1:4 hoặc 1:10) tùy domain—RAG cần attention nhiều hơn, summarization có thể dùng ít hơn.

Tương lai: Kiến trúc lai đang trở thành template chuẩn cho "next-generation efficiency"—thay vì khai tử Attention hoàn toàn (như một số tuyên bố quá khích về Mamba), ngành đang học cách dùng Attention như "bộ nhớ cache cao cấp" trong khi SSM xử lý bulk processing.

Đào sâu hơn

Paper gốc:

Bài liên quan TroiSinh:

Cùng cụm (new-architectures):

Đọc tiếp:

External:

Mamba-2 — SSM = structured matrix multiply, kết nối với linear attention

Mamba-2 biến selective scan thành phép nhân ma trận có cấu trúc, mở khóa Tensor Cores và thống nhất SSM với Linear Attention — từ memory-bound sang compute-bound.

Linear Attention — Đổi kernel function, quadratic → linear

Thay đổi hàm kernel trong attention để giảm độ phức tạp từ O(N²) xuống O(N), mở khóa khả năng xử lý context triệu token mà không cần GPU khổng lồ.

On this page

Vấn đềÝ tưởng cốt lõiTại sao nó hoạt độngÝ nghĩa thực tếĐào sâu hơn