State Space Models (SSM)

Mô hình không gian trạng thái (SSM) ánh xạ một chuỗi đầu vào thành chuỗi đầu ra thông qua vector trạng thái ẩn được cập nhật ở mỗi bước, vì vậy bộ nhớ của mô hình luôn giữ nguyên kích thước dù chuỗi dài thêm bao nhiêu. Ý tưởng này bắt nguồn từ lý thuyết điều khiển, trong đó một hệ tuyến tính được biểu diễn bằng phương trình cập nhật trạng thái h' = Ah + Bx và đầu ra y = Ch + Dx; bộ lọc Kalman đã ước lượng các trạng thái như vậy từ phép đo nhiễu từ năm 1960. SSM trong deep learning học A, B và C từ dữ liệu rồi xếp chồng nhiều lớp như vậy; từ S4 năm 2021 và Mamba năm 2023, chúng đã cạnh tranh với transformer trên các chuỗi dài gồm văn bản, âm thanh và dữ liệu chuỗi thời gian, trong khi compute tăng tuyến tính theo độ dài chuỗi.

Cách SSM hoạt động#

Một lớp SSM neural chuyển các phương trình liên tục thành truy hồi rời rạc với bước Δ được học, sử dụng quy tắc giữ bậc không Ā = exp(ΔA). Sau đó, mỗi bước tính h_t = Ā h_(t-1) + B̄ x_t và y_t = C h_t. Khi A, B và C giữ cố định qua các bước thời gian, khai triển truy hồi cho thấy đầu ra là phép tích chập giữa đầu vào và kernel (CB̄, CĀB̄, CĀ²B̄, ...). Vì vậy, cùng một lớp có thể chạy như phép tích chập trong quá trình training, xử lý toàn bộ chuỗi song song trên GPU, và như phép truy hồi trong quá trình sinh, với một lần cập nhật trạng thái có kích thước cố định cho mỗi token. Trong PyTorch, hai cách biểu diễn cho cùng đầu ra:

import torch

torch.manual_seed(0)
L, N = 64, 16  # độ dài chuỗi, kích thước trạng thái
x = torch.randn(L)

# SSM thời gian liên tục dạng đường chéo, rời rạc hóa với bước delta (giữ bậc không)
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# Cách biểu diễn truy hồi: một lần cập nhật trạng thái mỗi bước, như khi sinh
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# Cách biểu diễn tích chập: tính trước kernel K_k = C . (A_bar^k * B_bar), như khi training
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Kết quả in ra là Same outputs: True. Phép truy hồi chỉ giữ 16 giá trị trạng thái ở mỗi bước, trong khi kernel được tính trước cho phép training tính song song toàn bộ 64 đầu ra.

Từ S4 đến Mamba#

Các SSM deep learning đầu tiên quá chậm và ngốn bộ nhớ để training ở quy mô lớn. S4, do Albert Gu, Karan Goel và Christopher Ré phát triển, tạo cấu trúc cho ma trận trạng thái để tính kernel tích chập với chi phí thấp, đồng thời dùng cách khởi tạo được thiết kế để ghi nhớ lịch sử dài. S4 đạt kết quả tiên tiến nhất trên mọi tác vụ của benchmark Long Range Arena, trở thành model đầu tiên giải được tác vụ Path-X gồm chuỗi 16K bước và sinh nhanh hơn transformer 60 lần.

Kernel cố định xử lý mọi đầu vào theo cùng một cách, nên không thể quyết định bỏ qua phần thừa và giữ lại thông tin quan trọng. Mamba, do Albert Gu và Tri Dao phát triển, biến Δ, B và C thành các hàm của token hiện tại, cho phép model chọn thông tin nào cần ghi vào trạng thái và thông tin nào cần quên. Các tham số có tính chọn lọc này phá vỡ cách biểu diễn tích chập, nên Mamba tính truy hồi bằng phép quét song song có nhận biết phần cứng. Mamba đạt thông lượng inference gấp năm lần transformer; Mamba-3B vượt transformer cùng kích thước và ngang bằng transformer lớn gấp đôi. Mỗi block Mamba bao quanh SSM chọn lọc bằng một phép tích chập ngắn và một gate SiLU:

Mamba-2 cho thấy SSM chọn lọc và một dạng attention tuyến tính là hai cách biểu diễn đối ngẫu của cùng một phép tính. Các tác giả gọi khuôn khổ này là đối ngẫu không gian trạng thái và dùng nó để làm lớp lõi nhanh hơn từ 2 đến 8 lần.

SSM so với các kiến trúc liên quan#

• Transformer: Self-attention so sánh từng token với mọi token khác, nên có thể sao chép chính xác bất kỳ token nào trước đó; đổi lại, chi phí tăng theo bình phương và KV cache tăng theo cửa sổ ngữ cảnh. SSM nén lịch sử vào trạng thái cố định, tiết kiệm hơn nhưng mất mát thông tin; bài báo năm 2024 Repeat After Me cho thấy transformer sao chép và truy xuất từ đầu vào dài tốt hơn SSM có trạng thái cố định.
• Mạng nơ-ron hồi quy (RNN): RNN hoặc LSTM áp dụng cập nhật phi tuyến, nên quá trình training phải đi qua từng token một. SSM cập nhật tuyến tính theo trạng thái, cho phép dùng phép tích chập hoặc quét song song.
• Mô hình Markov ẩn (HMM): HMM có một tập hữu hạn các trạng thái rời rạc cùng xác suất chuyển tiếp. Mô hình không gian trạng thái tuyến tính-Gaussian đứng sau bộ lọc Kalman là phiên bản liên tục tương ứng; SSM deep learning duy trì vector trạng thái liên tục đó.

Các ứng dụng của SSM#

Phần lớn ứng dụng production nằm trong mô hình ngôn ngữ lớn (LLM) lai, giữ lại một vài lớp attention để truy hồi chính xác. Jamba của AI21 xen kẽ các lớp transformer và Mamba với lớp mixture of experts (MoE), chạy được trên một GPU 80 GB và đạt kết quả tốt với ngữ cảnh lên đến 256K token; các model Granite 4.0-H lai của IBM kết hợp Mamba-2 và các lớp transformer theo tỷ lệ 9:1. Vision Mamba áp dụng phép quét chọn lọc cho các patch ảnh; SSM cũng mô hình hóa âm thanh thô và chuỗi DNA quá dài để xử lý bằng attention toàn phần. Dạng cổ điển được dùng trong tracking của Ultralytics YOLO26: tracker TrackTrack mặc định, cũng như ByteTrack và BoT-SORT, dự đoán bounding box tiếp theo của mỗi đối tượng bằng bộ lọc Kalman có trạng thái 8 chiều chứa tâm, kích thước bounding box và vận tốc tương ứng. Trạng thái được cập nhật bằng ma trận chuyển tiếp vận tốc không đổi và ma trận quan sát tuyến tính đọc bốn tọa độ bounding box.