Kiến trúc Hiệu quả¶

Làm cho mô hình nhanh hơn không chỉ là giảm độ chính xác, mà còn là các kiến trúc thông minh hơn với ít tính toán hơn trên mỗi token. File này bao quát StreamingLLM, attention thưa và tuyến tính, multi-query và grouped-query attention, Mixture of Experts khi suy luận, kiến thức cô đọng (knowledge distillation), tỉa bớt (pruning), và tìm kiếm kiến trúc mạng (NAS)

Lượng tử hóa (file 1) làm cho mỗi phép tính rẻ hơn. File này làm cho có ít phép tính hơn ngay từ đầu. Cả hai bổ sung cho nhau: một mô hình vừa có kiến trúc hiệu quả vừa được lượng tử hóa có thể nhanh gấp 10-100 lần bản gốc.

StreamingLLM: Sinh với Độ dài Vô hạn¶

Transformer tiêu chuẩn lưu tất cả các token trước đó trong cache KV, vốn tăng tuyến tính theo độ dài chuỗi. Đến một điểm nào đó, cache vượt quá bộ nhớ GPU và việc sinh thất bại. StreamingLLM (Xiao et al., 2023) giải quyết điều này bằng một cache KV cuộn (rolling KV-cache) có kích thước cố định.
Quan sát then chốt: một vài token đầu tiên trong chuỗi nhận điểm attention cao không cân xứng bất kể nội dung của chúng là gì. Chúng được gọi là bồn attention (attention sinks). Nếu bạn loại bỏ chúng khỏi cache, phân bố attention sụp đổ và chất lượng sinh suy giảm thảm khốc.
Giải pháp của StreamingLLM: giữ vĩnh viễn một số lượng nhỏ token bồn (sink tokens, 1-4 token đầu tiên) trong cache, cộng thêm một cửa sổ cuộn (rolling window) gồm \(w\) token gần nhất. Tổng kích thước cache là \(\text{sink} + w\), cố định bất kể đã sinh ra bao nhiêu token.

\[\text{Cache} = [\text{token}_0, \text{token}_1, \text{token}_{t-w+1}, \ldots, \text{token}_t]\]

Các bồn attention neo giữ phân bố softmax, và cửa sổ cuộn cung cấp ngữ cảnh gần đây. Điều này cho phép sinh với độ dài vô hạn với bộ nhớ cố định, đánh đổi bằng việc mất quyền truy cập vào ngữ cảnh ở giữa chuỗi.
StreamingLLM hoạt động mà không cần huấn luyện lại với các mô hình đã phát triển bồn attention một cách tự nhiên (hầu hết LLM tiền huấn luyện đều có). Với các mô hình chưa có, thêm một token bồn có thể học duy nhất trong lúc huấn luyện là đủ.

Attention Thưa (Sparse Attention)¶

Self-attention đầy đủ là \(O(n^2)\) với độ dài chuỗi \(n\) vì mỗi token attend đến mọi token khác. Với \(n = 128K\), ma trận attention có \(128K^2 = 16\) tỷ phần tử. Các mẫu attention thưa (sparse attention) giảm điều này bằng cách giới hạn token nào attend đến token nào.

Attention sparsity patterns: full attention is O(n²), sliding window is O(n·w), local+global adds long-range tokens

Sliding window attention (Mistral, Gemma): mỗi token chỉ attend đến \(w\) token trước đó (ví dụ \(w = 4096\)). Attention là \(O(n \cdot w)\) thay vì \(O(n^2)\). Thông tin lan truyền vượt khỏi cửa sổ qua nhiều lớp: sau \(L\) lớp, ngữ cảnh hiệu dụng là \(L \times w\).
Attention cục bộ + toàn cục (Longformer, BigBird): hầu hết token dùng sliding window attention (cục bộ), nhưng một vài token định trước (như [CLS], mỗi token thứ 512) attend đến tất cả token (toàn cục). Điều này nắm bắt cả các mẫu cục bộ lẫn các phụ thuộc tầm xa.
Dilated attention: attend đến mỗi token thứ \(k\) trong một cửa sổ, tạo ra một mẫu thưa phủ một phạm vi lớn hơn với cùng số lượng điểm attention. Tăng hệ số dilation qua các lớp tạo ra một cấu trúc phân cấp tương tự như dilated convolution (chương 8).
Giải pháp thực tế cho các LLM hiện đại là sliding window + full attention xen kẽ: một số lớp dùng sliding window (rẻ, xử lý ngữ cảnh cục bộ), một số lớp dùng full attention (đắt, bắt phụ thuộc tầm xa). Mistral/Mixtral dùng mẫu này.

Attention Tuyến tính và Mô hình Không gian Trạng thái¶

Liệu ta có thể loại bỏ hoàn toàn \(O(n^2)\) attention? Attention tuyến tính và mô hình không gian trạng thái (SSM) xử lý chuỗi trong thời gian \(O(n)\) bằng cách tránh ma trận attention tường minh.
Attention tuyến tính thay thế softmax attention bằng xấp xỉ kernel:

\[\text{Tiêu chuẩn: } O = \text{softmax}(QK^T / \sqrt{d}) V\]

\[\text{Tuyến tính: } O = \phi(Q) (\phi(K)^T V)\]

Bằng cách kết hợp tích \(K^T V\) trước (vốn là \(d \times d\), không phụ thuộc vào độ dài chuỗi), phép tính trở thành \(O(n \cdot d^2)\) thay vì \(O(n^2 \cdot d)\). Với các chuỗi dài mà \(n \gg d\), đây là một tiết kiệm khổng lồ.
RWKV kết hợp ý tưởng từ RNN và transformer. Nó dùng một công thức hồi quy xử lý token tuần tự (như RNN) nhưng có thể song song hóa trong lúc huấn luyện (như transformer). Suy luận có chi phí \(O(1)\) mỗi token (bộ nhớ cố định, không tăng cache KV).
Mamba (Gu & Dao, 2023) là một mô hình không gian trạng thái có chọn lọc (selective state-space model). Nó xử lý chuỗi thông qua các chuyển tiếp trạng thái đã học:

\[h_t = \bar{A} h_{t-1} + \bar{B} x_t, \quad y_t = C h_t\]

trong đó \(\bar{A}\) và \(\bar{B}\) phụ thuộc vào đầu vào (có tính chọn lọc), cho phép Mamba chủ động tập trung vào hoặc bỏ qua các phần của đầu vào. Không giống như các SSM cố định, tính chọn lọc giúp Mamba cạnh tranh được với transformer trên các tác vụ ngôn ngữ trong khi vẫn duy trì độ phức tạp \(O(n)\).
Sự đánh đổi: attention tuyến tính và SSM nhanh hơn cho các chuỗi dài nhưng nhìn chung kém khả năng hơn full attention đối với các tác vụ đòi hỏi truy xuất tầm xa chính xác. Các kiến trúc lai (một số lớp transformer + một số lớp Mamba) thường cho kết quả tốt nhất.

Multi-Query và Grouped-Query Attention¶

Multi-head attention (MHA) tiêu chuẩn (chương 7) sử dụng các phép chiếu \(K\), \(V\) riêng cho mỗi head. Với \(h\) head, điều này đồng nghĩa với \(h\) tensor key và value riêng trong cache KV. Multi-Query Attention (MQA) và Grouped-Query Attention (GQA) giảm điều này.
MQA (Shazeer, 2019): tất cả các head chia sẻ một bộ phép chiếu \(K, V\) duy nhất. Mỗi head vẫn có phép chiếu \(Q\) riêng. Cache KV nhỏ lại với hệ số \(h\) (ví dụ 32 lần với 32 head).
GQA (Ainslie et al., 2023): một điểm trung gian. Các head được nhóm lại, và mỗi nhóm chia sẻ một bộ phép chiếu \(K, V\). Với \(h = 32\) head và \(g = 8\) nhóm, mỗi nhóm 4 head chia sẻ K/V. Cache KV nhỏ lại \(h/g = 4\) lần.

\[\text{MHA: } h \text{ head, } h \text{ bộ K/V} \quad \to \quad \text{GQA: } h \text{ head, } g \text{ bộ K/V} \quad \to \quad \text{MQA: } h \text{ head, } 1 \text{ bộ K/V}\]

MHA vs GQA vs MQA: MHA gives each head its own KV, GQA shares KV across groups, MQA uses a single KV for all heads — dramatically reducing KV-cache size

Hầu hết các LLM hiện đại đều dùng GQA (Llama 2/3, Gemma, Mistral). Nó giảm bộ nhớ cache KV và độ trễ suy luận với tổn thất chất lượng không đáng kể so với MHA.

Multi-head Latent Attention (MLA)¶

MLA (DeepSeek-V2, 2024) đi xa hơn GQA bằng cách nén cache KV vào một không gian ẩn hạng thấp (low-rank latent space). Thay vì lưu trọn vẹn các vector key và value, MLA lưu một vector ẩn nén \(\mathbf{c}_t\) cho mỗi token và tái tạo K/V ngay trong lúc attention:

\[\mathbf{c}_t = W_{\text{compress}} \cdot [\mathbf{k}_t; \mathbf{v}_t], \quad \mathbf{k}_t = W_K^{\text{up}} \cdot \mathbf{c}_t, \quad \mathbf{v}_t = W_V^{\text{up}} \cdot \mathbf{c}_t\]

Vector nén \(\mathbf{c}_t\) nhỏ hơn nhiều so với K và V gốc cộng lại. DeepSeek-V2 đạt giảm 93.3% kích thước cache KV so với MHA, vượt qua cả MQA, trong khi vẫn duy trì chất lượng ngang MHA.
Sự đánh đổi: tái tạo K/V từ ẩn thêm một chi phí tính toán nhỏ cho mỗi thao tác attention. Nhưng vì decode LLM bị giới hạn bởi băng thông bộ nhớ (không phải tính toán), đây là một lợi thế ròng: ít bộ nhớ cần nạp hơn > tốn thêm chút tính toán mỗi token.

Flash Attention¶

Flash Attention (Dao et al., 2022, được trình bày chi tiết trong chương 16 file 05) không phải là một thay đổi kiến trúc mà là một tối ưu hóa hiện thực thuộc về mọi cuộc thảo luận về attention hiệu quả. Nó tính toán attention đầy đủ chính xác với:
- Bộ nhớ O(n) thay vì O(n²) (ma trận attention không bao giờ được tạo ra trong HBM).
- Nhanh gấp 2-4 lần so với attention chuẩn (bằng cách giữ dữ liệu trong SRAM qua tiling và online softmax).
- Không mất mát chất lượng — output giống hệt về mặt toán học với attention chuẩn.
Flash Attention hiện là hiện thực attention mặc định trong PyTorch (torch.nn.functional.scaled_dot_product_attention), JAX, và tất cả các framework suy luận chính. Nếu bạn đang chạy attention trong năm 2024+, gần như chắc chắn bạn đang dùng Flash Attention.

Ring Attention¶

Ring Attention (Liu et al., 2023) phân bố tính toán attention qua nhiều thiết bị cho các chuỗi quá dài không thể chứa trên bộ nhớ của một GPU, thậm chí với Flash Attention.
Ý tưởng: phân hoạch chuỗi trên \(N\) thiết bị. Mỗi thiết bị giữ Q, K, V của \(n/N\) token. Các thiết bị được sắp xếp trong một vòng tròn. Trong mỗi bước:
1. Mỗi thiết bị tính attention cục bộ (Q của nó với K/V cục bộ).
2. Mỗi thiết bị gửi khối K/V của nó đến thiết bị kế tiếp trong vòng.
3. Mỗi thiết bị nhận K/V từ thiết bị trước và tính attention với chúng.
4. Sau \(N\) bước, mọi thiết bị đã attend đến mọi khối K/V.
Việc truyền thông được chồng lấn với tính toán (overlapped): trong khi tính attention trên khối K/V hiện tại, khối tiếp theo đang được truyền. Điều này ẩn độ trễ truyền thông gần như hoàn toàn.
Ring Attention cho phép cửa sổ ngữ cảnh triệu token bằng cách phân bố cache KV trên một vòng GPU. Bộ nhớ mỗi thiết bị là O(n/N), làm cho các chuỗi dài tùy ý trở nên khả thi (chỉ bị giới hạn bởi số lượng thiết bị).

Mixture of Experts khi Suy luận¶

Các mô hình MoE (chương 7) chỉ kích hoạt một phần tham số của chúng trên mỗi token (thường 2 trên 8 expert). Khi suy luận, thách thức đặc biệt là caching chuyên gia (expert caching): tất cả các expert phải ở trong bộ nhớ (vì bất kỳ token nào cũng có thể được định tuyến đến bất kỳ expert nào), nhưng chỉ có 2 expert hoạt động mỗi token.
Với mô hình Mixtral 8x7B: tổng tham số = 47B (8 × 7B expert, nhưng có các thành phần chia sẻ). Tham số hoạt động mỗi token ≈ 13B (2 expert + các lớp chia sẻ). Mô hình có chất lượng ngang LLM-70B với chi phí suy luận ngang LLM-13B, nhưng đòi hỏi 47B tham số trong bộ nhớ.
Expert offloading: đối với các triển khai bị giới hạn bởi bộ nhớ GPU, giữ các expert không hoạt động trên CPU hoặc SSD và nạp chúng khi cần. Cách này hiệu quả vì định tuyến token đủ dự đoán được để nạp trước các expert có khả năng được dùng.
Expert caching: duy trì một cache LRU của các expert được dùng gần đây trong bộ nhớ GPU. Nếu cùng một expert được kích hoạt liên tục (phổ biến với dữ liệu cùng miền), tỷ lệ cache hit sẽ cao.

Kiến thức cô đọng (Knowledge Distillation)¶

Cô đọng (distillation, chương 6) huấn luyện một mô hình "học sinh" nhỏ để bắt chước một mô hình "giáo viên" lớn. Học sinh học từ các dự đoán mềm (phân bố xác suất trên các lớp) của giáo viên, vốn chứa nhiều thông tin hơn là chỉ nhãn cứng (hard labels) đơn thuần.

\[\mathcal{L} = \alpha \cdot \text{KL}(p_{\text{teacher}}^{T} \| p_{\text{student}}^{T}) + (1 - \alpha) \cdot \mathcal{L}_{\text{CE}}(y, p_{\text{student}})\]

trong đó \(T\) là nhiệt độ (temperature, \(T\) cao hơn làm mềm phân bố, bộc lộ sự bất định của giáo viên) và \(\alpha\) cân bằng giữa mất mát cô đọng (distillation loss) và mất mát entropy chéo chuẩn.
Đối với LLM: cô đọng được dùng để tạo ra các mô hình nhỏ, nhanh từ các mô hình lớn, có năng lực. GPT-4 → một học sinh 7B nắm bắt phần lớn hành vi của GPT-4 cho một tác vụ cụ thể. Học sinh có thể rẻ hơn 10-100 lần để phục vụ.
Cô đọng theo tác vụ cụ thể: chỉ cô đọng trên dữ liệu liên quan đến tác vụ triển khai của bạn. Một mô hình 7B được cô đọng trên Q&A y khoa từ một giáo viên 70B có thể vượt trội mô hình 70B trên chính tác vụ đó (vì dung lượng giới hạn của học sinh hoàn toàn tập trung vào miền mục tiêu).

Tỉa bớt (Pruning)¶

Tỉa bớt loại bỏ các trọng số không cần thiết (đặt chúng về 0), giảm kích thước và tính toán của mô hình.
Tỉa bớt không cấu trúc (unstructured pruning, dựa trên độ lớn): loại bỏ các trọng số riêng lẻ có giá trị tuyệt đối nhỏ nhất. Điều này tạo ra một ma trận trọng số thưa (sparse). Đơn giản và hiệu quả cho nén, nhưng phần cứng hiện tại (GPU) không thể tăng tốc các phép toán thưa một cách hiệu quả trừ khi độ thưa tuân theo các mẫu cụ thể.
Tỉa bớt có cấu trúc (structured pruning): loại bỏ toàn bộ đơn vị — các head attention, nơ-ron MLP, hoặc lớp. Điều này tạo ra một mô hình dày đặc (dense) nhỏ hơn, có thể tăng tốc trực tiếp trên phần cứng tiêu chuẩn. Sự đánh đổi là độ hạt thô hơn (loại bỏ cả một head có thể loại bỏ cả trọng số có ích và vô ích cùng nhau).
Độ thưa 2:4 (NVIDIA Ampere+): một mẫu độ thưa được hỗ trợ bởi phần cứng, trong đó 2 trên mỗi 4 trọng số bằng 0. Các sparse Tensor Cores của GPU bỏ qua các phép nhân với số không, đạt tốc độ tăng ~2x. Đây là mẫu thưa duy nhất có tăng tốc phần cứng thực tế hiện nay.
Giả thuyết Ticket Trúng Số (Lottery Ticket Hypothesis, Frankle & Carlin, 2019): trong một mạng khởi tạo ngẫu nhiên, tồn tại một mạng con ("vé trúng số") có thể được huấn luyện riêng lẻ để đạt hiệu suất ngang mạng đầy đủ. Việc tìm ra các mạng con này (bằng cách huấn luyện, tỉa bớt và tua lại) rất tốn kém, nhưng insight này thúc đẩy nghiên cứu về tỉa bớt.

Tìm kiếm kiến trúc mạng (NAS)¶

NAS tự động hóa thiết kế kiến trúc bằng cách tìm kiếm trên một không gian các kiến trúc khả thi để tìm ra kiến trúc tối đa hóa độ chính xác dưới các ràng buộc phần cứng (độ trễ, bộ nhớ, năng lượng).
EfficientNet (chương 8) được tìm ra bởi NAS: quy tắc tỷ lệ hợp chất (cân bằng độ sâu, độ rộng, độ phân giải) nổi lên từ tìm kiếm, chứ không phải từ trực giác con người.
Cho hiệu quả suy luận, NAS có thể tìm các kiến trúc tối ưu cho các mục tiêu phần cứng cụ thể: "tìm một mô hình có độ trễ <5ms trên iPhone Neural Engine và độ chính xác >80% trên ImageNet." Không gian tìm kiếm bao gồm các loại lớp, độ rộng, hàm kích hoạt và mẫu attention.
Mạng Once-for-all (một lần cho mọi nhu cầu) huấn luyện một mạng quá tham số duy nhất và trích xuất các mạng con cho các mục tiêu triển khai khác nhau. Một lần huấn luyện tạo ra các mô hình cho cloud GPU, mobile GPU và CPU, mỗi loại tối ưu cho mục tiêu của nó.

Bài tập lập trình (dùng CoLab hoặc notebook)¶

Hiện thực sliding window attention và so sánh mức sử dụng bộ nhớ với full attention.

import jax
import jax.numpy as jnp

def full_attention(Q, K, V):
    """Standard O(n^2) attention."""
    scores = Q @ K.T / jnp.sqrt(Q.shape[-1])
    weights = jax.nn.softmax(scores, axis=-1)
    return weights @ V

def sliding_window_attention(Q, K, V, window_size=128):
    """Sliding window attention: each token attends to window_size previous tokens."""
    n = Q.shape[0]
    d = Q.shape[-1]
    output = jnp.zeros_like(Q)

    for i in range(n):
        start = max(0, i - window_size + 1)
        k_window = K[start:i+1]
        v_window = V[start:i+1]
        scores = Q[i] @ k_window.T / jnp.sqrt(d)
        weights = jax.nn.softmax(scores)
        output = output.at[i].set(weights @ v_window)

    return output

n, d = 512, 64
key = jax.random.PRNGKey(0)
Q = jax.random.normal(key, (n, d))
K = jax.random.normal(jax.random.PRNGKey(1), (n, d))
V = jax.random.normal(jax.random.PRNGKey(2), (n, d))

print(f"Full attention memory:    O(n^2) = {n*n} entries")
print(f"Window (w=128) memory:   O(n*w) = {n*128} entries")
print(f"Reduction: {n*n / (n*128):.1f}x")

So sánh kích thước cache KV cho MHA, GQA và MQA. Cho thấy tại sao GQA là điểm cân bằng thực tế.

def kv_cache_size(n_heads, n_kv_heads, d_head, seq_len, bytes=2):
    """KV-cache size in MB."""
    return 2 * n_kv_heads * d_head * seq_len * bytes / 1e6

n_heads = 32
d_head = 128
seq_len = 32768

mha = kv_cache_size(n_heads, n_heads, d_head, seq_len)       # 32 KV heads
gqa = kv_cache_size(n_heads, 8, d_head, seq_len)              # 8 KV heads
mqa = kv_cache_size(n_heads, 1, d_head, seq_len)              # 1 KV head

print(f"MHA (32 KV heads): {mha:.0f} MB per layer")
print(f"GQA (8 KV heads):  {gqa:.0f} MB per layer ({mha/gqa:.0f}x smaller)")
print(f"MQA (1 KV head):   {mqa:.0f} MB per layer ({mha/mqa:.0f}x smaller)")

Mô phỏng tỉa bớt có cấu trúc bằng cách loại bỏ các head attention ít quan trọng nhất từ một lớp attention ngẫu nhiên và đo sự thay đổi output.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
n_heads, seq_len, d_head = 8, 64, 32

# Random multi-head attention output (one per head)
head_outputs = jax.random.normal(key, (n_heads, seq_len, d_head))

# Full output: concatenate all heads
full_output = head_outputs.reshape(seq_len, n_heads * d_head)

# Importance: measure each head's contribution by its norm
head_norms = jnp.linalg.norm(head_outputs, axis=(1, 2))
print("Head importance (by norm):", jnp.round(head_norms, 2))

# Prune least important heads
for n_keep in [8, 6, 4, 2]:
    top_heads = jnp.argsort(head_norms)[-n_keep:]
    pruned = head_outputs[top_heads].reshape(seq_len, n_keep * d_head)

    # Pad to original size for comparison (zero out pruned heads)
    full_pruned = jnp.zeros_like(head_outputs)
    full_pruned = full_pruned.at[top_heads].set(head_outputs[top_heads])
    full_pruned = full_pruned.reshape(seq_len, n_heads * d_head)

    error = jnp.linalg.norm(full_output - full_pruned) / jnp.linalg.norm(full_output)
    print(f"Keep {n_keep}/{n_heads} heads: relative error = {error:.4f}, "
          f"memory = {n_keep/n_heads:.0%}")