Vision Mamba là gì? Kiến trúc, State Space Models và so sánh với Vision Transformer

Computer Vision · Deep Learning · Mamba

Vision Mamba (Vim) thay cơ chế self-attention toàn cục bằng các khối Mamba quét chuỗi patch ảnh theo hai hướng. Bài viết giải thích nguyên lý SSM, luồng xử lý ảnh, sự khác biệt với ViT và VMamba, cùng những giới hạn khi đánh giá hiệu năng.

Chủ đề: AI, Computer Vision, Mamba, Vision Transformer · Thời gian đọc: khoảng 9 phút

Ảnh đầu vào

H × W × C

→

Patch Embedding

N token

→

→ Forward SSM

← Backward SSM

→

Đặc trưng

Visual features

Sơ đồ khái niệm luồng xử lý trong Vim; đã giản lược các lớp chiếu tuyến tính, tích chập, cổng điều khiển và residual.

1. Vision Mamba là gì?

Vision Mamba (Vim) là kiến trúc học biểu diễn thị giác sử dụng Bidirectional State Space Models để xử lý chuỗi patch ảnh. Công trình Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model của Lianghui Zhu và cộng sự được công bố tại ICML 2024.

Hiểu ngắn gọn: ViT hỏi “mỗi patch nên chú ý đến patch nào?”, còn Vim truyền ngữ cảnh qua trạng thái ẩn được cập nhật dọc chuỗi patch. Cách truyền thông tin khác nhau dẫn đến các đánh đổi khác nhau về độ phức tạp và năng lực biểu diễn.

Tương tự Vision Transformer (ViT), Vim chia ảnh thành những patch, biến chúng thành token và bổ sung thông tin vị trí. Điểm khác biệt là Vim sử dụng các khối Mamba hai chiều thay cho Self-Attention để truyền và tổng hợp thông tin trên chuỗi token.

Ý chính: Vision Mamba không cần tạo ma trận Attention giữa mọi cặp patch. Phần Selective Scan có chi phí tăng tuyến tính theo số token khi kích thước mô hình cố định, tạo tiềm năng tốt cho ảnh có độ phân giải cao.

2. Tại sao cần một hướng đi ngoài Vision Transformer?

Trong Transformer với Self-Attention toàn cục, mỗi token có thể tương tác trực tiếp với mọi token khác. Với chuỗi N token, ma trận Attention có kích thước N × N.

Attention(Q, K, V) = softmax(QKᵀ / √d_k)V

Với patch 16 × 16 và bỏ qua token đặc biệt, ảnh 224 × 224 tạo 196 token; ảnh 1024 × 1024 tạo 4.096 token. Số phần tử của một ma trận attention tăng tương ứng từ 38.416 lên 16.777.216, tức khoảng 437 lần. Đây là minh họa cho độ lớn của ma trận, không phải phép đo trực tiếp thời gian chạy hay VRAM.

Những tối ưu như FlashAttention có thể giảm bộ nhớ trung gian đáng kể, nhưng chi phí tính toán Attention toàn cục vẫn tăng theo bậc hai của chiều dài chuỗi. Đây là động lực để khám phá các mô hình tuần tự như Mamba.

3. State Space Models và Mamba hoạt động thế nào?

3.1. State Space Models (SSM)

SSM mô tả một hệ động lực thông qua trạng thái ẩn h, đầu vào x và đầu ra y. Dạng rời rạc đơn giản của hệ tuyến tính là:

Trong công thức bên dưới, h_t là trạng thái tại bước t, x_t là đầu vào, còn y_t là đầu ra. Các ma trận có gạch ngang biểu diễn tham số đã rời rạc hóa. Đây là mô hình minh họa đơn giản, chưa thể hiện trọn vẹn cơ chế chọn lọc của Mamba.

h_t = Āh_t−1 + B̄x_t

y_t = Ch_t + Dx_t

Thay vì so sánh mọi token với nhau, SSM cập nhật một trạng thái ẩn để tổng hợp thông tin đã đi qua. Cách tiếp cận này phù hợp với việc xử lý chuỗi dài.

3.2. Mamba và Selective Scan

Mamba của Albert Gu và Tri Dao đưa vào cơ chế lựa chọn phụ thuộc đầu vào. Một số tham số như B, C và bước rời rạc hóa Δ được điều chỉnh dựa trên token hiện tại, giúp mô hình học cách giữ lại hoặc làm suy giảm thông tin.

Thuật toán hardware-aware selective scan được thiết kế để thực thi hiệu quả trên GPU. Khi các kích thước trạng thái và đặc trưng được giữ cố định, thành phần scan có chi phí tuyến tính theo độ dài chuỗi.

4. Kiến trúc Vision Mamba: Từ ảnh đến biểu diễn hai chiều

4.1. Patch Embedding

Với ảnh có chiều cao H, chiều rộng W và patch kích thước P × P (giả sử H và W chia hết cho P), số patch là:

N = HW / P²

Mỗi patch được ánh xạ thành một vector có D chiều, tạo chuỗi X ∈ ℝ^N×D.

4.2. Positional Embedding

Vì việc biến ảnh 2D thành chuỗi 1D có thể làm yếu thông tin về vị trí, Vim cộng positional embedding vào chuỗi token để mô hình nhận biết thứ tự các patch.

4.3. Bidirectional Mamba Block

Mamba ban đầu được thiết kế để xử lý chuỗi theo một hướng. Nhưng khi nhận diện một vật thể, thông tin từ các vùng ảnh ở cả trước và sau vị trí đang xét đều quan trọng. Do đó, Vim dùng hai hướng quét:

• Forward scan: xử lý chuỗi patch từ đầu đến cuối.
• Backward scan: xử lý chuỗi patch theo chiều ngược lại.

Hai luồng thông tin được kết hợp trong khối mạng. Ngoài SSM, khối Vim thực tế còn có các phép chiếu tuyến tính, convolution cục bộ, normalization, gating và kết nối residual. Nhờ đó, mô hình vừa xử lý đặc trưng cục bộ vừa tổng hợp ngữ cảnh từ hai hướng.

Điều quan trọng là quét hai chiều trên chuỗi không đồng nghĩa mô hình giữ nguyên hoàn toàn cấu trúc hai chiều của ảnh. Thứ tự sắp xếp patch vẫn ảnh hưởng tới những quan hệ không gian mà mô hình cần học.

4.4. Luồng xử lý tổng quát

1. Chia ảnh thành patch: chuyển ảnh 2D thành chuỗi token.
2. Gắn thông tin vị trí: bổ sung positional embedding trước khi xử lý tuần tự.
3. Quét hai chiều: mô hình thu nhận thông tin từ đầu đến cuối chuỗi và ngược lại.
4. Kết hợp biểu diễn: dùng các lớp biến đổi và nhánh residual để tạo đặc trưng cho tác vụ đầu ra.

Ví dụ trực quan: khi một patch chứa một phần bánh xe, thông tin từ các patch ở trước và sau nó trong chuỗi có thể giúp mô hình nhận ra bối cảnh chiếc xe. Tuy nhiên, hai patch gần nhau trong chuỗi chưa chắc gần nhau trong không gian 2D.

5. So sánh Vision Transformer, Vision Mamba và VMamba

Vision Mamba (Vim) và VMamba không phải cùng một kiến trúc. VMamba là công trình riêng, sử dụng 2D Selective Scan (SS2D) để xây dựng các tuyến quét trên biểu diễn ảnh.

Trong VMamba, SS2D khai thác bốn tuyến quét để xử lý đặc trưng không gian. Bảng dưới đây đối chiếu ý tưởng cốt lõi, không nhằm xếp hạng mô hình nào luôn chính xác hoặc nhanh hơn.

Tiêu chí	ViT	Vision Mamba (Vim)	VMamba
Cơ chế	Self-Attention	Bidirectional SSM	2D Selective Scan
Xử lý không gian	Tương tác token trực tiếp	Quét chuỗi hai hướng	Nhiều tuyến quét trên đặc trưng 2D
Chi phí theo số token	Bậc hai với Attention toàn cục	Tuyến tính cho Selective Scan	Tuyến tính cho SS2D
Điểm mạnh	Học tương tác trực tiếp	Tiềm năng tốt cho chuỗi dài	Thiết kế sát cấu trúc không gian ảnh hơn
Công bố	ViT: ICLR 2021	ICML 2024	NeurIPS 2024

Độ phức tạp ở đây mô tả cơ chế xử lý token cốt lõi, không phải toàn bộ chi phí tính toán của cả mạng.

6. Kết quả thực nghiệm, ưu điểm và hạn chế

Nhóm tác giả Vim đánh giá mô hình trên ImageNet (phân loại ảnh), COCO (phát hiện đối tượng) và ADE20K (phân đoạn ngữ nghĩa). Theo bài báo, Vim cho kết quả cạnh tranh so với các baseline DeiT cùng thiết lập.

Theo bài báo gốc, trong phép thử batch inference để trích xuất đặc trưng ở độ phân giải 1248 × 1248, Vim đạt tốc độ cao hơn 2,8 lần và sử dụng ít hơn 86,8% bộ nhớ GPU so với DeiT theo cấu hình được nhóm tác giả báo cáo. Không nên áp dụng trực tiếp các tỷ lệ này cho mọi GPU, batch size, độ phân giải hoặc tác vụ. Xem công bố chính thức tại ICML 2024.

Ưu điểm

• Tránh ma trận Self-Attention toàn cục kích thước N × N.
• Khả năng mở rộng thuận lợi về mặt lý thuyết với chuỗi token dài.
• Biểu diễn ngữ cảnh từ cả hai chiều quét.
• Tiềm năng ứng dụng trong các bài toán hình ảnh độ phân giải cao.

Hạn chế

• Việc làm phẳng ảnh thành chuỗi vẫn gây ra thách thức bảo toàn cấu trúc không gian.
• Thông tin cần được tổng hợp qua trạng thái ẩn hữu hạn thay vì tương tác token trực tiếp như Attention.
• Hiệu năng thực tế còn phụ thuộc GPU, batch size, kernel tối ưu và kiến trúc đầy đủ.
• Không thể kết luận Mamba luôn vượt Transformer chỉ từ độ phức tạp lý thuyết.

7. Vision Mamba phù hợp với những ứng dụng nào?

Vim và các backbone dựa trên SSM là hướng đáng khảo sát khi chuỗi token dài: ảnh y tế, ảnh viễn thám, phân đoạn ảnh và phân tích video. Đây là các kịch bản tiềm năng, không phải cam kết rằng Vim đã tốt nhất cho từng lĩnh vực. Cần đánh giá trên bộ dữ liệu, độ phân giải, cấu hình phần cứng và tiêu chí đo cụ thể.

Một hướng nghiên cứu hứa hẹn là kết hợp Attention và State Space Models: tận dụng tương tác trực tiếp của Attention ở những vị trí cần thiết, đồng thời sử dụng scan để xử lý chuỗi dài hiệu quả.

8. Câu hỏi thường gặp về Vision Mamba

Vision Mamba có phải là Vision Transformer không?

Không. Cả hai đều có thể dùng patch embedding để biểu diễn ảnh thành token, nhưng Vision Mamba sử dụng các khối SSM hai chiều thay cho Self-Attention toàn cục trong backbone được đề xuất.

Vision Mamba có nhanh hơn ViT không?

Vim cho thấy lợi thế trong một số benchmark, nhất là ở độ phân giải cao. Tuy nhiên, tốc độ thực tế phụ thuộc kích thước mô hình, phần cứng, kernel và tác vụ nên không thể khẳng định luôn nhanh hơn.

Vision Mamba và VMamba khác nhau thế nào?

Vim xử lý chuỗi patch bằng Mamba hai chiều; VMamba sử dụng 2D Selective Scan để mô hình hóa đặc trưng không gian qua nhiều tuyến quét. Đây là hai công trình độc lập.

Vision Mamba có độ phức tạp O(N) không?

Thành phần Selective Scan có chi phí tuyến tính theo số token N khi giữ cố định kích thước trạng thái và đặc trưng. Tổng chi phí của cả mô hình còn chứa các phép biến đổi khác.

Khi nào nên thử Vision Mamba thay vì ViT?

Khi dữ liệu có độ phân giải cao hoặc cần khảo sát một backbone không dùng attention toàn cục. Nên so sánh công bằng về số tham số, độ chính xác, thông lượng, độ trễ và mức sử dụng bộ nhớ trên cùng hệ thống.

9. Kết luận

Vision Mamba chứng minh rằng Self-Attention không phải con đường duy nhất để học biểu diễn thị giác giàu ngữ cảnh. Thông qua Selective State Space Models và quét hai chiều, Vim mang lại một hướng thiết kế vision backbone đáng chú ý, đặc biệt ở các tình huống phải xử lý chuỗi patch dài.

Thay vì xem Mamba là đối thủ thay thế tuyệt đối Transformer, có thể nhìn nhận đây là một họ công cụ bổ sung với những đánh đổi khác nhau về khả năng biểu diễn, chi phí tính toán và triển khai thực tế.

10. Tài liệu tham khảo và mã nguồn

1. Zhu et al. (2024), Bản công bố chính thức Vision Mamba tại PMLR / ICML 2024.
2. Liu et al. (2024), Bản công bố chính thức VMamba tại NeurIPS 2024.
3. Zhu et al. (2024), Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model, ICML 2024.
4. Gu & Dao (2023), Mamba: Linear-Time Sequence Modeling with Selective State Spaces.
5. Liu et al. (2024), VMamba: Visual State Space Model, NeurIPS 2024.
6. Mã nguồn Vision Mamba (Vim).
7. Mã nguồn VMamba.

Bài viết tổng hợp kiến thức phục vụ học tập và nghiên cứu. Khi trích dẫn kết quả thực nghiệm, hãy tham khảo điều kiện đánh giá trong các bài báo gốc.