3. Coding Attention Mechanisms (Part 03)

3.4. Cơ chế self-attention đơn giản có trainable weights

• Điểm khác biệt quan trọng nhất ở phần này so với phần trước là việc các trainable weights matrices sẽ được cập nhật trong quá trình huấn luyện.

• Ta sẽ triển khai qua 2 bước:

  • Bước 1: Viết code chi tiết từng bước (giống phần no weights)

  • Bước 2: Gom thành 1 class Python để import vào LLM architecture (sẽ triển khai ở các phần sau)

Từng bước tính trainable weights

• Theo dõi code minh họa tại 7. Self-attention-trainable.ipynb.

• Trước tiên ở phần này sẽ có thêm 3 weights matrices có thể huấn luyện được: $W_q, W_k, W_v$

• Ba ma trận này dùng để chiếu các token đầu vào đã được embedding $x^{(i)}$ sang 3 loại vector khác nhau: query (q), key (k), value (v).

• 

  • Ở ảnh trên, input token thứ 2 $x^{(2)}$ được chọn làm query input.

  • Vector query $q^{(2)}$ được tính bằng nhân ma trận giữa $x^{(2)}$ và weights matrix $W_q$: $q^{(2)} = x^{(2)} W_q$

  • Tương tự với key $k^{(2)}$ và value $v^{(2)}$: $k^{(2)} = x^{(2)} W_k, v^{(2)} = x^{(2)} W_v$

• Lý do tách ra 3 loại vector:

  • Query (q): vector dùng để đi xét với các vector khác.

  • Key (k): vector dùng để xét khi các vector khác tìm tới mình.

  • Value (v): vector chứa các đặc trưng về ngữ nghĩa.

• 

  • Vì đang tính context vector cho token 2 nên attention scores được tính bằng $q^{(2)}$ với các $k^{(i)}$.

• attention scores $\omega$ được tính bằng dot product tương tự cơ chế self-attention đơn giản ở phần trước. Điểm khác là ở đây được tính giữa vector query và vector key, với query là vector biểu diễn token đang xét và key là các vector biểu diễn các token sẽ được "so sánh" với query.

• 

• Ta sẽ tiến hành tính toán attention weights bằng cách điều chỉnh attention scores với hàm softmax đã thực nghiệm ở phần trước.

• Điểm khác so với phần trước là ta sẽ chuẩn hóa attention scores bằng cách chia cho căn bậc hai của embedding dimension của keys: $\alpha_{2i} = \frac{\omega_{2i}}{\sqrt{d_k}}$

  • với $d_k$ là embedding dimension keys.

• Lý do cho việc scaled dot-product attention này là cải thiện quá trình training bằng cách tránh gradient nhỏ. Ví dụ:

  • Giả sử dot-product là [28, 23, 21, 12, 6, 19] và embedding dim $d_k$ = 1024

  • Không scale: softmax gần như [≈ 1, 0, 0, 0, 0, 0].

  • Scale với $\sqrt{1024}$: softmax thành [0.2211, 0.1891, 0.1776, 0.1341, 0.1112, 0.1669]. $\rightarrow$ Gradient lan truyền tốt hơn.

• 

• Bước cuối cùng là tính context vector như hình ở trên. Tương tự như ở phần trước nhưng thay vì được tạo bằng tổng các input vector nhân với attention weight tương ứng thì ở đây sẽ là tổng các value (v) vector nhân với attention weight tương ứng: $z^{(2)} = \sum_{j=1}^{T} \alpha_{2j} ⋅ v^{(j)}$

Tại sao lại là query, key & value?

• 3 khái niệm này được mượn từ information retrieval & databases, nơi dữ liệu được lưu trữ, tìm kiếm & truy xuất.

• Query (q): dùng để biểu diễn phần tử hiện tại mà mô hình đang muốn hiểu rõ hơn, dùng để so sánh với các phần tử khác xem cần chú ý đến phần tử nào.

• Key (k): dùng để so khớp (match) với query, mỗi phần tử đầu vào đều có 1 key tương ứng.

• Value (v): Tương tự như key-value trong databases. Sau khi mô hình xác định được các key nào liên quan tới query, nó sẽ lấy ra các value tương ứng.

Triển khai Python class cho self-attention

• Theo dõi code minh họa tại 8. Self-attention-class.ipynb.

• 

  • Trong self-attention, ta biến đổi input vector $X$ thông qua 3 ma trận trọng số: $W_q$, $W_k$, $W_v$.

  • Sau đó, tính attention scores dựa trên các vector $Q$ & vector $K$ thu được.

  • attention scores được chuẩn hóa qua hàm softmax để trả về attention weights.

  • Nhân attention weights với các vector $V$ tương ứng để ra được các context vector $Z$.

• Ta có thể cải thiện version SelfAttention_v1 bằng cách sử dụng các tầng nn.Linear của Pytorch. Bản chất của lớp Linear này là thực hiện phép tính $y = xW + b$. Nếu ta tắt bias (tức là $b=0$), thì nó chính xác là $y = xW$.

• nn.Linear khởi tạo weights tiên tiến hơn (ví dụ Xavier initialization hoặc Kaiming initialization), giúp mô hình học ổn định và tránh gradient quá lớn/nhỏ.

• Ngoài ra, nn.Linear có cơ chế tự khởi tạo trọng số $W$ được tối ưu sẵn, thay vì dùng torch.rand().

• Triển khai version SelfAttention_v2 tại 8. Self-attention-class.ipynb.