Attention Sinks

Attention sink là một token nhận được phần lớn trọng số attention trong nhiều head và layer, mặc dù nội dung của token không liên quan đến truy vấn; thường là token đầu tiên trong chuỗi. Guangxuan Xiao cùng các đồng nghiệp tại MIT, Meta AI và Carnegie Mellon University đã đặt tên cho hiện tượng này trong Các mô hình ngôn ngữ phát trực tuyến hiệu quả với các attention sink, được công bố tại ICLR 2024. Họ cho thấy các mô hình ngôn ngữ lớn (LLMs) như Llama 2, MPT, Falcon và Pythia hoạt động không ổn định khi KV cache dạng trượt loại bỏ các token đầu tiên, và việc giữ lại bốn token đầu tiên giúp khôi phục quá trình tạo sinh ổn định.

Vì sao sink hình thành#

Self-attention chuyển điểm số của mỗi truy vấn thành trọng số bằng softmax, và tổng các trọng số luôn bằng một. Khi không token nào trong ngữ cảnh liên quan đến một head, trọng số vẫn phải được phân bổ đâu đó. Trong mô hình ngôn ngữ nhân quả, các token đầu tiên hiển thị với mọi token xuất hiện sau, vì vậy quá trình huấn luyện biến chúng thành nơi mặc định nhận phần trọng số mà head không cần dùng. Trong phân tích StreamingLLM trên Llama 2, mọi layer trừ hai layer thấp nhất đều tập trung attention vào các token ban đầu trên tất cả head. Logit sink được thêm vào softmax cho thấy hiện tượng này:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Không có sink, head phân bổ toàn bộ trọng số cho sáu giá trị không liên quan và trả về đầu ra có chuẩn 0.83. Với logit sink bằng 3, sink nhận 77% trọng số và đầu ra giảm xuống 0.19, nên head hầu như không làm gì cho truy vấn này. Các model gpt-oss có trọng số mở của OpenAI tích hợp sẵn cơ chế này bằng một logit sink có thể học cho từng head attention, để mỗi head có nơi nhận phần trọng số thừa mà không đóng góp gì vào đầu ra.

Tạo sinh dạng streaming với StreamingLLM#

Model tạo sinh vô hạn sẽ vượt quá bất kỳ cửa sổ ngữ cảnh nào và không thể giữ mọi token trong KV cache; attention cửa sổ thông thường, chỉ giữ các token gần nhất, sẽ thất bại ngay khi các token đầu tiên bị loại bỏ. StreamingLLM giữ cố định key và value của bốn token ban đầu bên cạnh một cửa sổ cuộn chứa các token gần đây, đồng thời gán vị trí theo thứ tự trong cache thay vì văn bản gốc và áp dụng lại embedding vị trí xoay (RoPE) ở mỗi bước:

Cách này duy trì ổn định mô hình hóa ngôn ngữ trên hơn 4 triệu token và chạy nhanh hơn tới 22.2 lần so với cửa sổ trượt tính toán lại cache cho mỗi token mới. Cửa sổ vẫn giới hạn những gì model có thể nhớ: các token bị loại khỏi phần giữa sẽ mất vĩnh viễn, vì vậy chatbot hoặc tác tử AI dạng streaming chỉ nhớ những lượt trao đổi gần nhất. Các tác giả cũng tiền huấn luyện model bằng một token placeholder duy nhất có thể học, đặt ở đầu mỗi mẫu; token này tự đảm nhận vai trò sink.

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

• Attention cửa sổ trượt: Giới hạn mỗi token chỉ nhìn thấy các token lân cận gần đây. Khi áp dụng lúc suy luận cho model được huấn luyện bằng attention đầy đủ, phương pháp này thất bại một khi các sink rời khỏi cửa sổ.
• Suy giảm ngữ cảnh: Độ chính xác giảm khi prompt dài ra nhưng vẫn nằm trong cửa sổ ngữ cảnh.

Sink trong model thị giác#

Transformer thị giác (ViT) thể hiện hành vi tương tự. Trong Vision Transformers Need Registers, các nhà nghiên cứu tại Meta và Inria phát hiện token có chuẩn lớn bất thường trên các mảng nền ít thông tin trong ViT có giám sát và tự giám sát, được tái sử dụng cho tính toán nội bộ. Việc thêm một vài token register vào đầu vào đã đảm nhận vai trò đó, loại bỏ hiện tượng giả khỏi feature map và bản đồ attention, đồng thời cải thiện dự đoán dày đặc.

SAM 2 của Meta, chạy trong gói Ultralytics, phân bổ bộ nhớ video bằng kế hoạch tương tự là giữ anchor cùng một cửa sổ: với mỗi frame mới, memory attention đọc các mục trong ngân hàng bộ nhớ của những frame do người dùng prompt, được giữ trong suốt video, cùng với các mục của 6 frame gần nhất, nhờ đó chi phí attention trên mỗi frame không đổi dù video chạy trong bao lâu.