Giải mã DeltaNet: Từ cơ chế KDA đến ứng dụng thực tế trong Qwen và Kimi

Một lưu ý về ký hiệu: bài viết này mặc định sử dụng ký hiệu bra-ket vì (theo quan điểm lấy cảm hứng từ cơ học lượng tử của tôi) nó làm cho các hình dạng trong phần dẫn xuất này trở nên rất rõ ràng. Tùy chọn chuyển đổi ký hiệu toán học ở trên sẽ viết lại mọi phương trình bằng cách sử dụng các vectơ in đậm thông thường và các phép chuyển vị tường minh. Trong chế độ bra-ket, ∣q⟩ là một vectơ cột, ⟨k∣ là một vectơ hàng, ⟨k∣q⟩ là một số, và ∣v⟩⟨k∣ là một ma trận. Các vectơ mặc định hướng sang phải, trong khi các khóa (keys) hướng sang trái khi được viết vào trạng thái linear-attention. Chúng ta làm việc với một đầu chú ý (attention head) nhân quả và các vectơ giá trị thực, giả định rằng các khóa của DeltaNet đã được chuẩn hóa, và để trạng thái ánh xạ từ không gian khóa sang không gian giá trị.

Các biến thể linear attention hiện đại rất phức tạp, và thoạt nhìn không dễ để thấy chúng được thiết kế để đạt được mục đích gì. Để tham khảo, đây là phương trình cập nhật trạng thái cho Kimi Delta Attention (KDA):

S~t=St−1Diag⁡(αt), ∣v^t⟩=S~t∣kt⟩, ∣et⟩=βt(∣vt⟩−∣v^t⟩), St=S~t+∣et⟩⟨kt∣, ∣ot⟩=St(dk−1/2∣qt⟩)

Lý do khiến chúng khó hiểu là vì đây là biến thể mới nhất trong một họ các biến thể linear attention đã được phát triển trong vài năm qua, và sự phức tạp của chúng đã tăng lên không thể tránh khỏi, khiến cho từ bên ngoài, các biến thể mới nhất có vẻ khó tiếp cận.

Trong bài viết này, chúng ta sẽ đi qua họ các biến thể linear attention DeltaNet, trong đó có hai biến thể được sử dụng bởi các dòng mô hình Qwen và Kimi mới nhất, đồng thời chỉ ra cách bạn có thể đi đến các phương trình tương tự bằng cách khẳng định những điều đơn giản về trạng thái ẩn của mình.

Đó là lộ trình mà chúng ta sẽ thực hiện:

softmax attention → linear attention → DeltaNet → Gated DeltaNet → KDA

Chỉ sau khi dẫn xuất xong KDA, chúng ta mới chuyển sang các chương trình Triton theo kiểu đệ quy và theo khối (chunkwise) để thực thi nó.

1. Bắt đầu với quadratic attention

Đối với một truy vấn (query) tại token t, ordinary causal softmax attention là:

ati=exp⁡ ⁣(s⟨ki∣qt⟩)∑j≤texp⁡ ⁣(s⟨kj∣qt⟩), s=dk−1/2, ∣ot⟩=∑i≤tati∣vi⟩.

Mỗi trọng số attention là một đại lượng vô hướng. Nó đo lường sự tương đồng giữa một khóa và một truy vấn, sau đó softmax biến đổi tất cả các điểm số cho truy vấn đó thành một phân phối. Đầu ra là tổng có trọng số của các vectơ giá trị.

Trên một chuỗi có độ dài T, có T^2 cặp khóa-truy vấn. Trong quá trình suy luận tự hồi quy (autoregressive inference), chúng ta có thể lưu vào bộ nhớ đệm (cache) các khóa và giá trị thay vì tính toán lại chúng, nhưng bộ nhớ đệm vẫn tăng theo chuỗi và mỗi truy vấn mới vẫn phải kiểm tra toàn bộ lịch sử.

Trở ngại đối với việc sắp xếp lại phép tính này chính là softmax. Mẫu số của nó phụ thuộc đồng thời vào truy vấn hiện tại và mọi khóa trước đó. Vì vậy, tạm thời hãy loại bỏ nó.

1.1 Loại bỏ softmax

Để rõ ràng, hãy hấp thụ hằng số tỷ lệ s vào truy vấn. Phiên bản attention tối giản có chủ đích khi đó là:

∣ot⟩=∑i≤t⟨ki∣qt⟩∣vi⟩.

Tích vô hướng vô hướng có thể di chuyển sang bên phải:

∣ot⟩=∑i≤t∣vi⟩⟨ki∣qt⟩=(∑i≤t∣vi⟩⟨ki∣)∣qt⟩.

Mọi thứ phụ thuộc vào quá khứ giờ đây có thể được thu thập vào một ma trận có kích thước cố định V×K:

St=∑i≤t∣vi⟩⟨ki∣

và attention trở thành một phép ghi đệ quy theo sau bởi một phép đọc:

St=St−1+∣vt⟩⟨kt∣, ∣ot⟩=St∣qt⟩.

Đồng nhất thức

(∣v⟩⟨k∣)∣q⟩=⟨k∣q⟩∣v⟩

chính là toàn bộ thủ thuật. Tích ngoài là một ma trận; tích trong là một số. Chúng ta không còn lưu trữ mọi khóa và giá trị trong quá khứ nữa. Chúng ta lưu trữ tổng các tích ngoài của chúng trong trạng thái St có kích thước cố định.

Điều này là tuyến tính theo độ dài chuỗi thay vì bậc hai: quét các token một lần, cập nhật cùng một trạng thái dv×dk tại mỗi bước. Chúng ta đã trả giá cho hiệu quả đó bằng cách loại bỏ sự chuẩn hóa và tính chọn lọc của softmax. Các phương pháp linear-attention tinh vi hơn sử dụng các bản đồ đặc trưng (feature maps) và bộ chuẩn hóa, nhưng dạng không trang trí này phơi bày vấn đề bộ nhớ vốn là động lực thúc đẩy DeltaNet.

1.2 Phép cộng không phải là phép gán

Giả sử chúng ta viết một cặp ∣vt⟩⟨kt∣ và ngay lập tức truy vấn trạng thái mới với cùng khóa đó:

St∣kt⟩=(St−1+∣vt⟩⟨kt∣)∣kt⟩=St−1∣kt⟩+∣vt⟩⟨kt∣kt⟩⏟1=St−1∣kt⟩+∣vt⟩.

Phép ghi không làm cho bộ nhớ trả về ∣vt⟩. Nó cộng ∣vt⟩ vào bất cứ thứ gì mà bộ nhớ đã trả về trước đó.

Nếu trạng thái cũ đã tạo ra giá trị chính xác, phép ghi cộng dồn làm cho trạng thái mới tạo ra gấp đôi giá trị đó. Tổng quát hơn, các khóa không trực giao với nhau, vì vậy mỗi lần ghi có thể gây nhiễu cho các lần ghi trước đó. Linear attention đã cung cấp cho chúng ta một bộ nhớ liên kết nhỏ gọn, nhưng cập nhật của nó hoạt động giống như += trong khi thứ chúng ta muốn lại gần với = hơn.

2. DeltaNet: ghi sai số, không phải giá trị

DeltaNet thay thế phép ghi linear-attention vô điều kiện bằng một hiệu chỉnh theo quy tắc delta. Có hai cách hữu ích để dẫn xuất nó.

2.1 Dẫn xuất một: yêu cầu rằng phép ghi có thể được đọc lại

Trước khi ghi token t, hãy hỏi bộ nhớ xem nó hiện đang liên kết gì với khóa mới:

∣v^t⟩=St−1∣kt⟩.

Nếu chúng ta muốn bộ nhớ trả về ∣vt⟩, chúng ta không nên cộng toàn bộ giá trị. Chúng ta chỉ nên cộng phần chênh lệch:

∣vt⟩−∣v^t⟩.

Giới thiệu một cường độ ghi đã học βt∈[0,1] và định nghĩa