Các Kiến trúc Thay thế: Mixture of Experts, Attention Tuyến tính và Mamba (SSM)

Khi đọc một câu hỏi về lập trình Python, liệu một mô hình ngôn ngữ có thực sự cần kích hoạt các nơ-ron lưu trữ kiến thức về ẩm thực Pháp hay lịch sử La Mã? Và khi đọc một cuốn sách dài $500.000$ token, liệu mô hình có bắt buộc phải giữ lại nguyên vẹn từng cặp vector khóa – giá trị $(k_t, v_t)$ của mọi từ từ trang đầu tiên?

Hai trăn trở đó dẫn dắt toàn bộ cuộc cách mạng kiến trúc hậu-Transformer xoay quanh hai câu hỏi cốt lõi:

• Làm thế nào để tách rời tổng dung lượng tri thức (số lượng tham số $N$) khỏi chi phí tính toán trên mỗi token (FLOPs), giúp mô hình có hàng trăm tỷ tham số nhưng suy luận nhanh như một mô hình nhỏ?
• Làm thế nào để vượt qua nút cổ chai bậc hai $\mathcal{O}(L^2)$ của Softmax Attention, kết hợp khả năng huấn luyện song song của Transformer với chi phí suy luận trạng thái hằng số $\mathcal{O}(1)$ của RNN?

1. Hỗn hợp Chuyên gia Thưa (Sparse Mixture of Experts — MoE)

Trong Chương 2 (Mục 5), khi phân rã số lượng tham số của một khối Transformer chuẩn, ta đã thấy rằng các lớp mạng truyền thẳng (Feed-Forward Network — FFN / SwiGLU) chiếm tới 2/3 tổng số tham số của mô hình. Trong một mô hình đặc (Dense model) như Llama-3 70B, mọi token — dù là dấu chấm câu , hay một thuật ngữ chuyên ngành phức tạp — đều phải nhân ma trận qua toàn bộ các trọng số FFN này.

1.1 Từ mạng FFN đặc đến tính toán có điều kiện (Conditional Computation)

Ý tưởng của Hỗn hợp Chuyên gia Thưa (Sparse Mixture of Experts — MoE) là thay thế một lớp FFN khổng lồ duy nhất bằng $E$ mạng FFN độc lập nhỏ hơn gọi là các chuyên gia (experts) $\{E_1, E_2, \dots, E_E\}$, đi kèm một mạng cổng nhẹ gọi là bộ định tuyến (router hoặc gating network).

Với vector biểu diễn $x \in \mathbb{R}^{d_{\text{model}}}$ của một token ở đầu vào lớp FFN, bộ định tuyến dùng ma trận trọng số $W_g \in \mathbb{R}^{d_{\text{model}} \times E}$ để tính điểm tương thích cho từng chuyên gia, sau đó chỉ kích hoạt Top-$k$ chuyên gia có điểm cao nhất (thường $k = 1$ hoặc $k = 2$, với $k \ll E$):

trong đó $\mathcal{T}_k(x)$ là tập hợp chỉ số của $k$ chuyên gia có xác suất $p_i(x)$ lớn nhất. Đầu ra của lớp MoE là tổng có trọng số chỉ trên $k$ chuyên gia được chọn:

Vì $G_i(x) = 0$ với $E - k$ chuyên gia còn lại, ta không cần thực hiện phép nhân ma trận của các chuyên gia đó! Nhờ vậy, ta có thể tăng số lượng chuyên gia $E$ lên 8, 64, hay 256 để mở rộng dung lượng ghi nhớ tri thức của mô hình mà chi phí FLOPs trên mỗi token chỉ tương đương với việc chạy $k$ chuyên gia.

1.2 Sụp đổ định tuyến (Routing Collapse) và Hàm mất mát cân bằng tải

Tuy nhiên, khi huấn luyện lớp Sparse MoE bằng hạ gradient thông thường, một hiện tượng bệnh lý gần như chắc chắn sẽ xảy ra: sụp đổ định tuyến (routing collapse).

Nếu ở những bước khởi tạo đầu tiên, chuyên gia $E_1$ tình cờ có trọng số nhỉnh hơn một chút và được bộ định tuyến chọn nhiều hơn, thì $E_1$ sẽ nhận được nhiều cập nhật gradient hơn các chuyên gia khác. Kết quả là $E_1$ học nhanh hơn, dự đoán tốt hơn, và bộ định tuyến lại càng tăng xác suất $p_1(x)$ gửi token vào $E_1$! Chỉ sau vài trăm bước, toàn bộ token đều dồn vào 1–2 chuyên gia (“hiệu ứng người giàu càng giàu”), trong khi $E - 2$ chuyên gia còn lại bị bỏ hoang hoàn toàn (“chuyên gia chết” — dead experts), biến mô hình MoE đắt đỏ thành một mô hình đặc cỡ nhỏ.

Để ngăn chặn hiệu ứng tự củng cố này, Switch Transformer (Fedus & Shazeer, 2022) chuẩn hóa một hàm mất mát cân bằng tải phụ trợ (load-balancing auxiliary loss) cực kỳ thanh lịch. Xét một batch gồm $N$ token:

1. Gọi $f_i$ là tỷ lệ token thực sự được định tuyến tới chuyên gia $i$: $$f_i = \frac{1}{k N} \sum_{x \in \text{Batch}} \mathbb{I}\big(i \in \mathcal{T}_k(x)\big), \qquad \text{với } \sum_{i=1}^{E} f_i = 1.$$
2. Gọi $P_i$ là xác suất định tuyến trung bình mà cổng Softmax phân bổ cho chuyên gia $i$ trên toàn batch: $$P_i = \frac{1}{N} \sum_{x \in \text{Batch}} p_i(x), \qquad \text{với } \sum_{i=1}^{E} P_i = 1.$$

Vì hàm chỉ thị $\mathbb{I}(\cdot)$ trong $f_i$ không khả vi, ta nhân $f_i$ với $P_i$ (khả vi theo trọng số cổng $W_g$) để tạo ra hàm mất mát phụ trợ:

Tại sao công thức này hoạt động? Khi cả 8 chuyên gia được chia tải đồng đều hoàn hảo, ta có $f_i = \frac{1}{E}$ và $P_i = \frac{1}{E}$ với mọi $i$, suy ra:

Ngược lại, nếu toàn bộ token bị dồn vào một chuyên gia duy nhất ($f_1 = 1, P_1 = 1$), hàm mất mát sẽ vọt lên $\mathcal{L}_{\text{aux}} = E \cdot 1 = E$. Do tích vô hướng $\sum f_i P_i$ phạt nặng những chuyên gia vừa có tần suất được chọn $f_i$ cao vừa có xác suất $P_i$ lớn, gradient của $\mathcal{L}_{\text{aux}}$ sẽ lập tức kéo giảm logit của các chuyên gia quá tải và đẩy token sang các chuyên gia đang rảnh rỗi.

1.3 Từ Mixtral 8x7B đến Chuyên gia dùng chung (Shared Experts) trong DeepSeek-V3

Trong thực tế hiện đại, kiến trúc Sparse MoE đã phát triển qua hai thế hệ tiêu biểu:

• Mixtral 8x7B (Jiang & others, 2024) sử dụng $E = 8$ chuyên gia SwiGLU lớn và kích hoạt $k = 2$ chuyên gia cho mỗi token. Tổng số tham số của mô hình là 46.7B, nhưng mỗi token chỉ đi qua 12.9B tham số hoạt động (active parameters), giúp Mixtral đạt chất lượng vượt Llama-2 70B với tốc độ suy luận nhanh gấp 6 lần.
• DeepSeek-V3 (DeepSeek-AI, 2024) tiến thêm hai bước đột phá: (1) Chuyên gia tinh mịn (Fine-grained Experts) — chia nhỏ các chuyên gia thành $E = 256$ chuyên gia nhỏ và kích hoạt $k = 8$ chuyên gia mỗi token để tạo ra hàng tỷ tổ hợp chuyên môn hóa đa dạng; và (2) Chuyên gia dùng chung (Shared Experts) — tách riêng $E_s = 1$ chuyên gia luôn luôn được kích hoạt cho mọi token mà không cần qua bộ định tuyến:

Nhờ có $E_{\text{shared}}(x)$ gánh vác các tri thức ngữ pháp và từ vựng phổ quát, các chuyên gia định tuyến $E_i$ không cần phải học lặp lại kiến thức cơ bản mà có thể tập trung 100% dung lượng vào các lĩnh vực chuyên biệt. Ngoài ra, DeepSeek-V3 còn thay thế hàm phạt $\mathcal{L}_{\text{aux}}$ bằng cơ chế tự động điều chỉnh độ lệch định tuyến (auxiliary-loss-free load balancing), tránh việc hàm phạt phụ trợ làm suy giảm chất lượng ngôn ngữ.

Hãy tương tác với sơ đồ dưới đây để quan sát cách bộ định tuyến Top-2 phân phối các token của một đoạn mã Python vào 8 chuyên gia, và điều gì xảy ra khi bạn tắt hàm mất mát cân bằng tải:

Bây giờ, hãy tự tay cài đặt lớp SparseMoELayer bằng PyTorch và kiểm chứng bằng thực nghiệm sự khác biệt giữa huấn luyện có và không có $\mathcal{L}_{\text{aux}}$.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(42)
B, L, d_model, d_ff = 16, 32, 64, 128
num_experts, top_k = 8, 2

# Tạo batch token đầu vào và nhãn mục tiêu giả lập
x = torch.randn(B, L, d_model)
target = torch.randn(B, L, d_model)
print(f"Input shape: {tuple(x.shape)} ({B * L} tokens)")

Input shape: (16, 32, 64) (512 tokens)

=== Không có Load Balancing Loss (alpha = 0.0) ===
Tỷ lệ chọn f_i (%): [56.2, 35.4, 3.1, 2.0, 1.4, 0.8, 0.6, 0.6]
L_aux cuối: 3.685 | Độ lệch chuẩn tải: 20.1%

=== Có Load Balancing Loss (alpha = 0.05) ===
Tỷ lệ chọn f_i (%): [13.1, 12.5, 12.7, 12.1, 12.3, 12.9, 12.1, 12.3]
L_aux cuối: 1.002 | Độ lệch chuẩn tải: 0.4%

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(42)
B, L, d_model, d_ff = 16, 32, 64, 128
num_experts, top_k = 8, 2

# Tạo batch token đầu vào và nhãn mục tiêu giả lập
x = torch.randn(B, L, d_model)
target = torch.randn(B, L, d_model)
print(f"Input shape: {tuple(x.shape)} ({B * L} tokens)")

Input shape: (16, 32, 64) (512 tokens)

class SparseMoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts=8, top_k=2):
        super().__init__()
        self.E = num_experts
        self.k = top_k
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.SiLU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        N, D = x.numel() // x.size(-1), x.size(-1)
        x_flat = x.view(N, D)
        probs = F.softmax(self.router(x_flat), dim=-1)  # (N, E)
        top_vals, top_idx = torch.topk(probs, self.k, dim=-1)
        weights = top_vals / top_vals.sum(dim=-1, keepdim=True)

        # (1) Hàm mất mát cân bằng tải L_aux = E * sum(f_i * P_i)
        one_hot = F.one_hot(top_idx, self.E).float()    # (N, k, E)
        f_i = one_hot.sum(dim=1).mean(dim=0) / self.k   # (E,)
        P_i = probs.mean(dim=0)                         # (E,)
        aux_loss = self.E * torch.sum(f_i * P_i)

        # (2) Tính đầu ra tổ hợp thưa từ k chuyên gia được chọn
        out = torch.zeros_like(x_flat)
        for slot in range(self.k):
            idx_s = top_idx[:, slot]
            w_s = weights[:, slot].unsqueeze(-1)
            for e in range(self.E):
                mask = idx_s == e
                if mask.any():
                    out[mask] += w_s[mask] * self.experts[e](
                        x_flat[mask]
                    )
        return out.view_as(x), aux_loss, f_i.detach()


def run_experiment(alpha, title):
    torch.manual_seed(42)
    moe = SparseMoELayer(d_model, d_ff, num_experts, top_k)
    # Tạo độ lệch khởi tạo nhỏ khiến chuyên gia 0, 1 được ưu ái ban đầu
    with torch.no_grad():
        moe.router.weight[0] *= 1.8
        moe.router.weight[1] *= 1.5
    opt = torch.optim.Adam(moe.parameters(), lr=1e-2)
    for _ in range(80):
        out, aux_loss, f_i = moe(x)
        loss = F.mse_loss(out, target) + alpha * aux_loss
        opt.zero_grad()
        loss.backward()
        opt.step()
    pct = [round(v * 100, 1) for v in f_i.tolist()]
    std = (f_i * 100).std().item()
    print(f"=== {title} ===")
    print(f"Tỷ lệ chọn f_i (%): {pct}")
    print(f"L_aux cuối: {aux_loss.item():.3f} | Độ lệch chuẩn tải: {std:.1f}%")


run_experiment(0.0, "Không có Load Balancing Loss (alpha = 0.0)")
print()
run_experiment(0.05, "Có Load Balancing Loss (alpha = 0.05)")

=== Không có Load Balancing Loss (alpha = 0.0) ===
Tỷ lệ chọn f_i (%): [56.2, 35.4, 3.1, 2.0, 1.4, 0.8, 0.6, 0.6]
L_aux cuối: 3.685 | Độ lệch chuẩn tải: 20.1%

=== Có Load Balancing Loss (alpha = 0.05) ===
Tỷ lệ chọn f_i (%): [13.1, 12.5, 12.7, 12.1, 12.3, 12.9, 12.1, 12.3]
L_aux cuối: 1.002 | Độ lệch chuẩn tải: 0.4%

2. Attention Tuyến tính và Sự đối ngẫu RNN–Attention (Linear Attention)

Sparse MoE giúp giảm chi phí tại các lớp FFN, nhưng còn bản thân lớp Self-Attention thì sao? Khi chiều dài chuỗi $L$ tăng lên $128\text{K}$ hay $1\text{M}$ token, ma trận điểm chú ý $Q K^\top \in \mathbb{R}^{L \times L}$ đòi hỏi số phép tính tăng theo bậc hai $\mathcal{O}(L^2 d)$, và trong lúc sinh văn bản tự hồi quy, bộ nhớ KV Cache phải lưu toàn bộ các cặp $(k_1, v_1), \dots, (k_t, v_t)$, tốn thời gian đọc bộ nhớ $\mathcal{O}(t)$ ở bước thứ $t$.

Tại sao Self-Attention lại bị khóa chặt vào độ phức tạp $\mathcal{O}(L^2)$? Thủ phạm nằm ở hàm mũ phi tuyến Softmax bao quanh tích vô hướng $q_t^\top k_j$.

2.1 Phép ma thuật của tính kết hợp phép nhân ma trận

Trong cơ chế Causal Self-Attention chuẩn, đầu ra $y_t \in \mathbb{R}^d$ tại vị trí $t$ là trung bình có trọng số của các vector giá trị $v_j$ ($j \le t$):

Vì hàm $\exp(q_t^\top k_j)$ không thể tách thành tích độc lập của một hàm theo $q_t$ và một hàm theo $k_j$ (trừ khi khai triển chuỗi Taylor vô hạn), ta bắt buộc phải tính tích vô hướng $q_t^\top k_j$ cho mọi cặp $(t, j)$ với $j \le t$, dẫn đến phép nhân ma trận theo thứ tự:

Năm 2020, Katharopoulos et al. (Katharopoulos & others, 2020) đặt một câu hỏi táo bạo: điều gì sẽ xảy ra nếu ta thay hàm tương đồng $\exp(q_t^\top k_j / \sqrt{d})$ bằng một hàm nhân tách được (kernel feature map) $\text{sim}(q_t, k_j) = \phi(q_t)^\top \phi(k_j)$, trong đó $\phi: \mathbb{R}^d \to \mathbb{R}_{+}^d$ là một hàm kích hoạt dương áp dụng theo từng phần tử, chẳng hạn $\phi(x) = \text{elu}(x) + 1$?

Khi đó, nhờ tính phân phối và kết hợp của phép nhân ma trận, ta có thể đưa $\phi(q_t)^\top$ ra ngoài dấu tổng $\sum_{j=1}^{t}$:

Ở dạng ma trận toàn chuỗi (không mặt nạ), chỉ cần đổi vị trí dấu ngoặc đơn từ $\big(\phi(Q)\phi(K)^\top\big)V$ sang $\phi(Q)\big(\phi(K)^\top V\big)$, ma trận trung gian kích thước $L \times L$ biến mất và được thay bằng ma trận trung gian kích thước cố định $d \times d$:

Khi $L \gg d$ (ví dụ chiều dài ngữ cảnh $L = 65.536$ trong khi chiều một đầu chú ý chỉ là $d = 64$), thay $L^2 d$ bằng $L d^2$ giúp giảm số phép tính đi hơn 1.000 lần!

2.2 Biểu diễn Truy hồi Trạng thái $S_t \in \mathbb{R}^{d \times d}$ và Sự đối ngẫu RNN–Attention

Điều kỳ diệu nhất của công thức trên xuất hiện khi ta đặt hai biến trạng thái tích lũy theo thời gian $t$:

• Ma trận trạng thái bộ nhớ kết hợp (associative memory state matrix) $S_t \in \mathbb{R}^{d \times d}$: tổng các tích ngoài (outer product) giữa khóa $\phi(k_j)$ và giá trị $v_j$.
• Vector chuẩn hóa $z_t \in \mathbb{R}^d$: tổng các vector khóa $\phi(k_j)$.

Tại mỗi bước thời gian $t$, trạng thái mới chỉ phụ thuộc vào trạng thái ở bước $t - 1$ và token hiện tại $(q_t, k_t, v_t)$:

Hãy dừng lại một nhịp để thưởng thức hệ phương trình này! Đây chính là Sự đối ngẫu giữa Attention và RNN (Attention–RNN Duality):

1. Khi huấn luyện (Chế độ Song song): Toàn bộ chuỗi đầu vào đã biết trước, ta tính song song trên GPU bằng các phép nhân ma trận khối (chunkwise / parallel scan) mà không cần vòng lặp tuần tự theo $t$.
2. Khi suy luận tự hồi quy (Chế độ Truy hồi — Recurrent mode): Thay vì phải lưu trữ KV Cache phình to vô hạn theo $L$, ta chỉ cần giữ duy nhất một ma trận trạng thái cố định $S_t \in \mathbb{R}^{d \times d}$ và một vector $z_t \in \mathbb{R}^d$. Mỗi khi sinh thêm một token mới, chi phí cập nhật $S_t$ và truy xuất $y_t$ là hằng số $\mathcal{O}(d^2)$ — tức $\mathcal{O}(1)$ theo chiều dài chuỗi $L$!

Hướng tiếp cận biến RNN thành kiến trúc huấn luyện song song như Transformer cũng là nền tảng của dòng mô hình RWKV (Reinventing RNNs for the Transformer Era) (Peng & others, 2023), nơi tác giả bổ sung hệ số suy giảm theo thời gian $e^{-w}$ vào trạng thái tích lũy để các token gần hiện tại có trọng số cao hơn các token ở xa.

Hãy khám phá trực quan hai cách đặt dấu ngoặc của Linear Attention trong biểu đồ tương tác dưới đây:

Trong bài tập dưới đây, bạn sẽ tự tay cài đặt Causal Linear Attention ở cả hai chế độ và chứng minh rằng chúng cho ra cùng một tensor đầu ra tới từng chữ số thập phân.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
B, L, d = 4, 256, 32

# Khởi tạo Q, K, V ngẫu nhiên và hàm ánh xạ đặc trưng dương phi(x)
Q = torch.randn(B, L, d)
K = torch.randn(B, L, d)
V = torch.randn(B, L, d)

def phi(x):
    return F.elu(x) + 1.0

print(f"Q, K, V shape: {tuple(Q.shape)}")

Q, K, V shape: (4, 256, 32)

Output shape: (4, 256, 32)
Trạng thái truy hồi cuối S_L shape: (4, 32, 32)
Sai số tuyệt đối cực đại |Y_parallel - Y_recurrent|: 4.77e-07
Kiểm chứng trùng khớp (< 1e-6): True

import torch
import torch.nn.functional as F

torch.manual_seed(0)
B, L, d = 4, 256, 32

# Khởi tạo Q, K, V ngẫu nhiên và hàm ánh xạ đặc trưng dương phi(x)
Q = torch.randn(B, L, d)
K = torch.randn(B, L, d)
V = torch.randn(B, L, d)

def phi(x):
    return F.elu(x) + 1.0

print(f"Q, K, V shape: {tuple(Q.shape)}")

Q, K, V shape: (4, 256, 32)

# (1) Dạng ma trận song song (Parallel form) khi huấn luyện: O(L^2 d)
def causal_linear_attention_parallel(Q, K, V, eps=1e-6):
    Qp, Kp = phi(Q), phi(K)                        # (B, L, d)
    scores = torch.matmul(Qp, Kp.transpose(-1, -2))  # (B, L, L)
    causal_mask = torch.tril(torch.ones(L, L, dtype=torch.bool))
    scores = scores.masked_fill(~causal_mask, 0.0)
    num = torch.matmul(scores, V)                  # (B, L, d)
    den = scores.sum(dim=-1, keepdim=True) + eps   # (B, L, 1)
    return num / den


# (2) Dạng truy hồi trạng thái (Recurrent form) khi suy luận: O(L d^2)
def causal_linear_attention_recurrent(Q, K, V, eps=1e-6):
    Qp, Kp = phi(Q), phi(K)
    S = torch.zeros(B, d, d)   # Ma trận trạng thái S_t in R^{d x d}
    z = torch.zeros(B, d, 1)   # Vector tổng khóa chuẩn hóa z_t in R^d
    outputs = []

    for t in range(L):
        q_t = Qp[:, t, :].unsqueeze(1)             # (B, 1, d)
        k_t = Kp[:, t, :].unsqueeze(-1)            # (B, d, 1)
        v_t = V[:, t, :].unsqueeze(1)              # (B, 1, d)

        S = S + torch.matmul(k_t, v_t)             # k_t * v_t^T
        z = z + k_t
        num_t = torch.matmul(q_t, S).squeeze(1)    # (B, d)
        den_t = torch.matmul(q_t, z).squeeze(1) + eps
        outputs.append(num_t / den_t)

    return torch.stack(outputs, dim=1), S


y_par = causal_linear_attention_parallel(Q, K, V)
y_rec, S_final = causal_linear_attention_recurrent(Q, K, V)
max_err = (y_par - y_rec).abs().max().item()

print(f"Output shape: {tuple(y_rec.shape)}")
print(f"Trạng thái truy hồi cuối S_L shape: {tuple(S_final.shape)}")
print(f"Sai số tuyệt đối cực đại |Y_parallel - Y_recurrent|: {max_err:.2e}")
print(f"Kiểm chứng trùng khớp (< 1e-6): {max_err < 1e-6}")

Output shape: (4, 256, 32)
Trạng thái truy hồi cuối S_L shape: (4, 32, 32)
Sai số tuyệt đối cực đại |Y_parallel - Y_recurrent|: 4.77e-07
Kiểm chứng trùng khớp (< 1e-6): True

3. Mô hình Không gian Trạng thái Chọn lọc (SSM & Mamba)

Trong lý thuyết điều khiển học và xử lý tín hiệu, Mô hình Không gian Trạng thái (State Space Model — SSM) mô tả cách một tín hiệu đầu vào một chiều $x(t) \in \mathbb{R}$ điều khiển một vector trạng thái ẩn $h(t) \in \mathbb{R}^N$ và sinh ra tín hiệu đầu ra $y(t) \in \mathbb{R}$ thông qua hệ phương trình vi phân tuyến tính bậc nhất:

trong đó $A \in \mathbb{R}^{N \times N}$ là ma trận chuyển trạng thái (quyết định tốc độ suy giảm hoặc dao động của bộ nhớ), $B \in \mathbb{R}^{N \times 1}$ là ma trận chiếu đầu vào, và $C \in \mathbb{R}^{1 \times N}$ là ma trận đọc trạng thái ra đầu ra. (Khi áp dụng cho vector embedding $D$ chiều của LLM, ta chạy độc lập một hệ SSM như vậy trên từng kênh trong $D$ kênh).

3.1 Rời rạc hóa Zero-Order Hold (ZOH) và Mô hình S4

Vì văn bản là một chuỗi các token rời rạc $x_1, x_2, \dots, x_L$ thay vì hàm liên tục $x(t)$, ta cần rời rạc hóa (discretize) hệ phương trình vi phân trên bằng một tham số bước nhảy thời gian (step size) $\Delta > 0$. Theo quy tắc giữ bậc không (Zero-Order Hold — ZOH), các ma trận rời rạc $(\bar{A}, \bar{B})$ được tính từ $(A, B, \Delta)$ như sau:

trong đó xấp xỉ $\bar{B} \approx \Delta B$ là khai triển Euler bậc nhất cực kỳ trực quan. Khi đó, hệ vi phân trở thành một mạng RNN tuyến tính theo thời gian rời rạc $t$:

Năm 2022, Albert Gu, Karan Goel và Christopher Ré giới thiệu mô hình S4 (Structured State Spaces for Sequence Modeling) (Gu & Ré, 2022). Họ nhận ra rằng nếu các ma trận $(A, B, C, \Delta)$ là bất biến theo thời gian (Linear Time-Invariant — LTI, tức không phụ thuộc vào bước $t$ hay nội dung $x_t$), thì việc mở rộng đệ quy của $h_t$ sẽ biến toàn bộ chuỗi đầu ra $y$ thành một phép tích chập toàn cục (global convolution):

Nhờ định lý tích chập và cấu trúc ma trận đường chéo HiPPO của $A$, S4 có thể tính toàn bộ chuỗi dài $L$ song song bằng biến đổi Fourier nhanh (FFT) trong thời gian $\mathcal{O}(L \log L)$ khi huấn luyện, và chạy như RNN trong thời gian $\mathcal{O}(1)$ khi suy luận!

3.2 Đột phá của Mamba: Cơ chế Quét Chọn lọc (Selective Scan)

S4 đạt kết quả xuất sắc trên các chuỗi tín hiệu vật lý (âm thanh, sóng điện tâm đồ), nhưng lại thua kém Transformer trên mô hình hóa ngôn ngữ tự nhiên. Tại sao?

Trong bài báo Mamba (Gu & Dao, 2023), Albert Gu và Tri Dao đã chỉ ra căn nguyên cốt lõi: ngôn ngữ tự nhiên đòi hỏi sự nén thông tin phụ thuộc vào ngữ cảnh (content-aware reasoning), điều mà mọi hệ thống bất biến thời gian (LTI) đều bất lực.

Hãy tưởng tượng mô hình đọc câu: “Mã xác thực của bạn là A7 … uhm … à … K9 … ờm … M2” và phải nhắc lại mã A7-K9-M2 ở cuối câu (bài toán Selective Copying).

• Trong một mô hình LTI như S4, vì bước nhảy $\Delta$ và các ma trận $\bar{A}, \bar{B}, C$ là hằng số cố định cho mọi bước thời gian, mô hình đối xử với từ rác "uhm" giống hệt như từ khóa "A7"! Nó không thể nhìn vào nội dung của $x_t$ để quyết định: “Token này là nhiễu, đừng ghi vào trạng thái $h_t$!”.
• Giải pháp của Mamba (Selective State Space Model — S6) vô cùng gọn gàng: biến bước nhảy rời rạc hóa $\Delta_t$ cũng như các ma trận chiếu $B_t, C_t$ thành hàm phụ thuộc trực tiếp vào đầu vào $x_t$:

trong đó $\text{softplus}(z) = \ln(1 + e^z) > 0$ đảm bảo bước nhảy thời gian luôn dương, còn ma trận $A$ được chọn là ma trận đường chéo âm $A = -\exp(\text{log\_A}) < 0$ để hệ thống luôn ổn định. Khi đó, ở mỗi bước $t$:

3.3 Đối ngẫu Không gian Trạng thái có Cấu trúc (SSD) trong Mamba-2

Năm 2024, Tri Dao và Albert Gu tiếp tục công bố Mamba-2 (Dao & Gu, 2024) với lý thuyết Đối ngẫu Không gian Trạng thái có Cấu trúc (Structured State Space Duality — SSD). Họ chứng minh rằng khi ma trận chuyển trạng thái $\bar{A}_t = a_t I$ là tích của một vô hướng $a_t \in (0, 1)$ với ma trận đơn vị, phương trình SSM của Mamba và phương trình Linear Attention ở Mục 2 trở thành một thể thống nhất:

trong đó $C, B, X$ tương ứng trực tiếp với $Q, K, V$ của Attention, còn $L_{\text{decay}}$ là ma trận mặt nạ nhân quả bán tách được (1-semiseparable matrix) chứa tích các hệ số suy giảm chọn lọc $a_s$. Nhờ chia chuỗi thành các khối nhỏ (chunkwise decomposition), Mamba-2 tận dụng trực tiếp các nhân nhân ma trận Tensor Cores của GPU, giúp huấn luyện nhanh hơn Mamba-1 từ 2 đến 8 lần trong khi cho phép tăng kích thước trạng thái $N$ lên gấp 8 lần ($N = 128$).

Hãy tương tác với biểu đồ dưới đây để thấy rõ cách cổng chọn lọc $\Delta_t(x_t)$ của Mamba giải quyết bài toán Sao chép Chọn lọc mà SSM bất biến thời gian (LTI) bó tay:

Bây giờ, hãy tự tay lập trình cả hai mô hình (LTI SSM và Selective SSM) trên bài toán Sao chép Chọn lọc trong PyTorch!

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(7)

# Bài toán Nhớ có Chọn lọc (Selective Copying):
# Chuỗi gồm L=24 bước; chỉ có 1 token tín hiệu (marker=1) ở vị trí ngẫu nhiên,
# các bước còn lại là nhiễu (marker=0). Mô hình phải xuất lại giá trị tín hiệu ở bước cuối.
def sample_selective_copy_batch(batch_size=64, seq_len=24):
    values = torch.randn(batch_size, seq_len, 1)
    markers = torch.zeros(batch_size, seq_len, 1)
    pos = torch.randint(0, seq_len - 4, (batch_size,))
    for b in range(batch_size):
        markers[b, pos[b], 0] = 1.0
    x = torch.cat([values, markers], dim=-1)      # (B, L, 2)
    y = values[torch.arange(batch_size), pos, 0]  # (B,)
    return x, y, pos

x_demo, y_demo, pos_demo = sample_selective_copy_batch(4, 12)
print(f"Batch x shape: {tuple(x_demo.shape)}, target y shape: {tuple(y_demo.shape)}")

Batch x shape: (4, 12, 2), target y shape: (4,)

[SSM Bất biến thời gian (LTI)] MSE cuối: 0.8421
[Selective SSM (Mamba)]        MSE cuối: 0.0014

Kiểm tra cổng chọn lọc Delta_t của Mamba:
  Trung bình Delta_t tại token TÍN HIỆU (marker=1): 2.841
  Trung bình Delta_t tại token NHIỄU    (marker=0): 0.018

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(7)

# Bài toán Nhớ có Chọn lọc (Selective Copying):
# Chuỗi gồm L=24 bước; chỉ có 1 token tín hiệu (marker=1) ở vị trí ngẫu nhiên,
# các bước còn lại là nhiễu (marker=0). Mô hình phải xuất lại giá trị tín hiệu ở bước cuối.
def sample_selective_copy_batch(batch_size=64, seq_len=24):
    values = torch.randn(batch_size, seq_len, 1)
    markers = torch.zeros(batch_size, seq_len, 1)
    pos = torch.randint(0, seq_len - 4, (batch_size,))
    for b in range(batch_size):
        markers[b, pos[b], 0] = 1.0
    x = torch.cat([values, markers], dim=-1)      # (B, L, 2)
    y = values[torch.arange(batch_size), pos, 0]  # (B,)
    return x, y, pos

x_demo, y_demo, pos_demo = sample_selective_copy_batch(4, 12)
print(f"Batch x shape: {tuple(x_demo.shape)}, target y shape: {tuple(y_demo.shape)}")

Batch x shape: (4, 12, 2), target y shape: (4,)

class SimpleSSM(nn.Module):
    def __init__(self, d_in=2, d_state=16, selective=True):
        super().__init__()
        self.selective = selective
        self.d_state = d_state
        self.in_proj = nn.Linear(d_in, d_state)
        self.log_A = nn.Parameter(torch.zeros(d_state))

        if selective:
            # Mamba: Delta_t, B_t, C_t là hàm tuyến tính của đầu vào x_t
            self.proj_delta = nn.Linear(d_in, d_state)
            self.proj_B = nn.Linear(d_in, d_state)
            self.proj_C = nn.Linear(d_in, d_state)
        else:
            # LTI SSM: Delta, B, C cố định, bất biến theo thời gian
            self.raw_delta = nn.Parameter(torch.full((d_state,), -1.0))
            self.B = nn.Parameter(torch.randn(d_state) * 0.2)
            self.C = nn.Parameter(torch.randn(d_state) * 0.2)

        self.out_head = nn.Linear(d_state, 1)

    def forward(self, x):
        B_sz, L, _ = x.shape
        u = self.in_proj(x)                        # (B, L, N)
        A = -torch.exp(self.log_A)                 # (N,) luôn âm để ổn định

        if self.selective:
            delta = F.softplus(self.proj_delta(x)) # (B, L, N)
            B_t = self.proj_B(x)
            C_t = self.proj_C(x)
        else:
            delta = F.softplus(self.raw_delta).view(1, 1, -1).expand(B_sz, L, -1)
            B_t = self.B.view(1, 1, -1).expand(B_sz, L, -1)
            C_t = self.C.view(1, 1, -1).expand(B_sz, L, -1)

        # Rời rạc hóa và quét trạng thái h_t = A_bar * h_{t-1} + B_bar * u_t
        h = torch.zeros(B_sz, self.d_state)
        for t in range(L):
            a_bar = torch.exp(delta[:, t, :] * A)
            b_bar = delta[:, t, :] * B_t[:, t, :]
            h = a_bar * h + b_bar * u[:, t, :]

        y_hat = self.out_head(h * C_t[:, -1, :]).squeeze(-1)
        return y_hat, delta


def train_ssm(selective):
    torch.manual_seed(7)
    model = SimpleSSM(selective=selective)
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
    for _ in range(250):
        bx, by, _ = sample_selective_copy_batch(64, 24)
        pred, _ = model(bx)
        loss = F.mse_loss(pred, by)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return model, loss.item()


lti_model, lti_mse = train_ssm(selective=False)
sel_model, sel_mse = train_ssm(selective=True)
print(f"[SSM Bất biến thời gian (LTI)] MSE cuối: {lti_mse:.4f}")
print(f"[Selective SSM (Mamba)]        MSE cuối: {sel_mse:.4f}\n")

with torch.no_grad():
    bx, _, _ = sample_selective_copy_batch(128, 24)
    _, delta = sel_model(bx)
    d_mean = delta.mean(dim=-1)
    mask = bx[:, :, 1] == 1.0
    print("Kiểm tra cổng chọn lọc Delta_t của Mamba:")
    print(f"  Trung bình Delta_t tại token TÍN HIỆU (marker=1): {d_mean[mask].mean():.3f}")
    print(f"  Trung bình Delta_t tại token NHIỄU    (marker=0): {d_mean[~mask].mean():.3f}")

[SSM Bất biến thời gian (LTI)] MSE cuối: 0.8421
[Selective SSM (Mamba)]        MSE cuối: 0.0014

Kiểm tra cổng chọn lọc Delta_t của Mamba:
  Trung bình Delta_t tại token TÍN HIỆU (marker=1): 2.841
  Trung bình Delta_t tại token NHIỄU    (marker=0): 0.018

4. Kiến trúc Lai (Hybrid SSM-Transformer) và Mô hình Ngôn ngữ Khuếch tán (Diffusion LMs)

Nếu Mamba vừa có độ phức tạp tuyến tính $\mathcal{O}(L)$ vừa có khả năng nhớ chọn lọc, liệu nó đã có thể thay thế hoàn toàn Softmax Attention trong mọi tác vụ hay chưa? Câu trả lời thực tế là: chưa hoàn toàn, nhưng sự kết hợp giữa cả hai lại tạo ra kiến trúc mạnh nhất hiện nay.

4.1 Điểm mù “Sao chép nguyên văn” của SSM và Kiến trúc Lai Jamba

Hãy hình dung bạn đưa cho mô hình một tài liệu dài 50 trang chứa hàng trăm mã số hợp đồng khác nhau, rồi ở dòng cuối cùng mới hỏi: “Hãy trích xuất nguyên văn số điện thoại nằm ngay sau mã hợp đồng #8492 ở trang 12”.

• Một mô hình SSM / Mamba nén toàn bộ lịch sử đã đọc vào một trạng thái ẩn có kích thước cố định $h_t \in \mathbb{R}^{D \times N}$. Khi chưa biết trước câu hỏi ở cuối văn bản là gì, nó không thể nén nguyên vẹn từng chữ số của cả trăm mã hợp đồng vào một dung lượng bộ nhớ hữu hạn mà không bị mất mát thông tin.
• Ngược lại, Softmax Attention giữ nguyên toàn bộ các vector $(k_j, v_j)$ trong KV Cache — giống như một cuốn sổ mở sẵn cho phép mô hình “ngoái lại nhìn trực tiếp” bất kỳ token nào ở trang 12 với độ phân giải tuyệt đối.

Khám phá này đã dẫn đến xu hướng thiết kế Kiến trúc Lai (Hybrid SSM–Transformer), tiêu biểu là mô hình Jamba của AI21 Labs (Lieber & others, 2024) (và sau đó là Zamba, Samba, Bamba, Nemotron-H). Thay vì dùng 100% lớp Attention hoặc 100% lớp Mamba, Jamba xếp chồng các khối theo tỷ lệ 7 lớp Mamba : 1 lớp Attention, đồng thời tích hợp Sparse MoE vào các lớp FFN!

4.2 Vượt khỏi sinh tuần tự: Mô hình Ngôn ngữ Khuếch tán Mặt nạ (Masked Diffusion LMs)

Cả Transformer, MoE lẫn Mamba đều chia sẻ chung một giả định nền tảng từ Chương 0 đến nay: sinh văn bản tự hồi quy từ trái sang phải (left-to-right autoregressive generation) $P(x_1, \dots, x_L) = \prod_{t=1}^{L} P(x_t \mid x_{<t})$. Cách sinh từng token một này khiến độ trễ suy luận luôn tỉ lệ thuận với số token đầu ra và khó sửa lại các token đã lỡ sinh ở đầu câu.

Liệu ta có thể sinh văn bản theo cách các mô hình Diffusion (như Stable Diffusion, Flux) tạo ra bức ảnh — tức khởi tạo toàn bộ câu bằng các ô trống rồi tinh chỉnh song song toàn bộ câu qua vài chục bước khử nhiễu?

Đó chính là ý tưởng của Mô hình Ngôn ngữ Khuếch tán Mặt nạ (Masked Diffusion Language Models — MDLM) (Sahoo & others, 2024) (cùng các biến thể quy mô lớn như LLaDA hay Mercury):

1. Quá trình thuận (Làm nhiễu — Forward Masking Process): Xuất phát từ câu văn sạch $x_0$ độ dài $L$, ở thời điểm $t \in [0, 1]$, mỗi token độc lập có xác suất $\alpha_t = 1 - t$ được giữ nguyên và xác suất $t$ bị thay thế bằng token mặt nạ [MASK]. Khi $t = 1$, toàn bộ câu chỉ gồm toàn [MASK].
2. Quá trình ngược (Khử nhiễu song song — Reverse Unmasking Process): Một mạng Transformer hai chiều (không dùng mặt nạ nhân quả, giống như BERT ở Chương 0 Mục 10.2!) đọc chuỗi đang bị che một phần $x_t$ và dự đoán phân phối từ vựng gốc cho tất cả các vị trí [MASK] cùng một lúc. Ở mỗi bước khử nhiễu từ $t$ về $s < t$, một tỷ lệ các vị trí [MASK] có độ tự tin cao nhất được “lật mở” đồng thời thành các từ cụ thể, cho đến khi cả câu hoàn chỉnh tại $t = 0$.

Như Sahoo et al. (Sahoo & others, 2024) chứng minh, hàm mục tiêu huấn luyện liên tục theo thời gian của MDLM thực chất là trung bình có trọng số của các hàm mất mát Masked Language Modeling (BERT) ở mọi tỷ lệ che mặt nạ $t \in (0, 1]$:

Bằng cách cho phép lật mở hàng chục token trong một lần chạy mạng duy nhất, các mô hình khuếch tán rời rạc mở ra hướng đi mới cho việc sinh mã nguồn và chỉnh sửa văn bản song song với tốc độ vượt xa giải mã tự hồi quy truyền thống.

4.3 Bảng tổng kết các Kiến trúc Mô hình Ngôn ngữ

Để khép lại Chương 6, bảng dưới đây đặt cạnh nhau 5 họ kiến trúc mà ta đã khảo sát để bạn dễ dàng đối chiếu sự đánh đổi về tính toán, bộ nhớ và năng lực biểu diễn:

Nếu muốn đào sâu hơn nữa, bạn có thể thử sức với các bài tập mở rộng dưới đây:

Tài liệu tham khảo

1. 2017 Shazeer et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR. 1
2. 2020 Katharopoulos et al. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML. 2
3. 2022 Fedus, Zoph & Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR. 1
4. 2022 Gu, Goel & Ré. Efficiently Modeling Long Sequences with Structured State Spaces. ICLR. 3
5. 2023 Gu & Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. COLM 2024 (arXiv 2023). 3
6. 2023 Peng et al. RWKV: Reinventing RNNs for the Transformer Era. EMNLP Findings. 2
7. 2024 Dao & Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML. 3
8. 2024 DeepSeek-AI. DeepSeek-V3 Technical Report. arXiv. 1
9. 2024 Jiang et al. Mixtral of Experts. arXiv. 1
10. 2024 Lieber et al. Jamba: A Hybrid Transformer-Mamba Language Model. arXiv. 4
11. 2024 Sahoo et al. Simple and Effective Masked Diffusion Language Models. NeurIPS. 4