Sliding Window Attention

Attention cửa sổ trượt (SWA) giới hạn mỗi token chỉ attention trên một số token lân cận cố định—trong model ngôn ngữ nhân quả là w token gần nhất—thay vì toàn bộ chuỗi như self-attention đầy đủ. Khi đó, lượng tính toán attention tăng tuyến tính theo độ dài chuỗi thay vì bậc hai, và model tạo sinh chỉ cần giữ key và value của w token gần nhất trong KV cache. Longformer đưa cửa sổ trượt thành cốt lõi của cơ chế attention vào năm 2020; năm 2023, Mistral 7B của Mistral AI áp dụng kỹ thuật này cho một mô hình ngôn ngữ lớn (LLM) mở được sử dụng rộng rãi, với cửa sổ 4.096 token.

Cửa sổ mở rộng theo độ sâu như thế nào#

Một layer SWA chỉ quan sát w token, nhưng các layer Transformer được xếp chồng lên nhau. Ở layer 2, một token đọc các hidden state vốn đã tóm tắt w token từ layer 1, vì vậy thông tin có thể truyền xa thêm khoảng w − 1 vị trí qua mỗi layer, giống như các phép tích chập xếp chồng mở rộng trường tiếp nhận của mạng thần kinh tích chập (CNN). Trong PyTorch, cửa sổ là một mask attention boolean; nhân mask này với chính nó cho biết thông tin truyền đi xa đến đâu:

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

Với 16 token, cửa sổ tính 58 điểm số mỗi layer thay vì 136; phạm vi tiếp cận của token cuối tăng thêm 3 token mỗi layer, lần lượt từ 4 lên 7, 10 rồi 13. Có thể truyền cùng mask dưới dạng attn_mask cho hàm scaled_dot_product_attention của PyTorch. Với Mistral 7B, 32 layer cùng cửa sổ 4.096 token tạo ra phạm vi attention lý thuyết khoảng 131K token, dù thông tin từ xa phải đi qua nhiều bước mới tới nơi.

Mistral 7B cũng lưu key và value trong một bộ đệm vòng gồm w ô; vị trí i được ghi vào ô i mod w và ghi đè mục cũ nhất. Với chuỗi 32K token, cách này giảm bộ nhớ cache tám lần mà không làm giảm chất lượng; kích thước cache cố định giúp bộ nhớ luôn ổn định, bất kể quá trình tạo sinh kéo dài bao lâu.

Xen kẽ các lớp cục bộ và toàn cục#

Các cửa sổ thuần túy không thể sao chép một dữ kiện từ xa trong một bước, vì vậy các model gần đây kết hợp lớp SWA với lớp toàn cục đầy đủ. Gemma 2 của Google luân phiên chúng theo tỷ lệ một-một với cửa sổ 4.096 token, còn Gemma 3 dùng năm lớp cục bộ cho mỗi lớp toàn cục với cửa sổ chỉ 1.024 token, nhờ đó giữ bộ nhớ đệm KV nhỏ trên cửa sổ ngữ cảnh dài tối đa 128K token. Chỉ các lớp toàn cục lưu key và value cho toàn bộ ngữ cảnh, vì vậy năm trong sáu lớp chỉ lưu 1.024 token mỗi lớp.

Vision Transformer sử dụng cách kết hợp tương tự cho ảnh độ phân giải cao, trong đó mỗi patch là một token. SAM 3 của Meta, chạy trong gói Ultralytics cho phân đoạn ảnh, thực hiện attention trong các cửa sổ patch không chồng lấn kích thước 24 × 24 ở 28 trong số 32 block của bộ mã hóa ảnh, đồng thời thực hiện attention toàn cục trên toàn bộ lưới patch 72 × 72 của ảnh 1.008 pixel tại các block 7, 15, 23 và 31.

SWA và các thuật ngữ liên quan#

• Swin Transformer: Chia ảnh thành các cửa sổ cố định, không chồng lấn, được mọi token trong cửa sổ dùng chung; đồng thời dịch lưới giữa các lớp để các cửa sổ lân cận trao đổi thông tin.
• Attention thưa: Nhóm rộng hơn gồm các mẫu cố định. SWA là thành viên đơn giản nhất, còn các phương pháp khác bổ sung cửa sổ giãn cách, token toàn cục hoặc liên kết ngẫu nhiên.
• Attention sink: Các token đầu tiên của một chuỗi, hấp thụ trọng số attention dư. Khi áp dụng cửa sổ lúc suy luận cho model được huấn luyện bằng attention đầy đủ, việc loại bỏ các token này sẽ làm hỏng quá trình sinh; StreamingLLM giữ chúng trong bộ nhớ đệm bên cạnh cửa sổ.
• Attention biến dạng: Mỗi query thực hiện attention trên một số điểm lấy mẫu mà nó tự học độ lệch, vì vậy vùng lân cận phụ thuộc vào nội dung.
• Flash attention: Kernel GPU dùng cho attention chính xác. Các triển khai của nó chấp nhận kích thước cửa sổ, nên model SWA bỏ qua hoàn toàn các block nằm ngoài cửa sổ.