Tính self-attention bằng một ví dụ nhỏ

Giả sử một vị trí cần trộn thông tin từ hai vị trí trước đó. Điểm tương đồng lần lượt là 2 và 0. Nếu biến chúng thành trọng số bằng softmax, vị trí thứ nhất nhận khoảng 88% đóng góp. Con số ấy đến từ phép tính, không phải mô hình đã đánh dấu một từ là “quan trọng” theo nghĩa của con người.

Trước khi học, bạn nên hiểu token, vector, phép nhân ma trận và xác suất tổng bằng một. Bài này xét một đầu attention, bỏ qua dropout để có thể tính tay. Các số đều là ví dụ giả định.

Query, key và value có công việc khác nhau

Với ma trận biểu diễn đầu vào X, ba phép chiếu tạo Q = XWq, K = XWk và V = XWv. Query dùng để hỏi cần lấy thông tin nào; key tham gia tính mức phù hợp; value chứa thông tin được trộn. Đây là cách diễn giải vai trò tính toán, không phải ba vùng lưu nghĩa cố định.

Giả sử query q = [2, 0], hai key là k1 = [1, 0], k2 = [0, 1]. Tích vô hướng cho điểm thô [2, 0]. Nếu chiều key là 2, scaled dot-product attention chia các điểm cho căn bậc hai của 2, được xấp xỉ [1,414; 0].

Phép chia giúp kiểm soát độ lớn điểm khi số chiều tăng dưới các giả định thống kê thích hợp. Nó không bảo đảm mọi mô hình luôn có cùng phân bố điểm. Trọng số chiếu vẫn được học qua gradient từ mục tiêu huấn luyện.

Softmax tạo một phép trộn có trọng số

Softmax biến điểm s_i thành exp(s_i) chia tổng exp(s_j). Với điểm đã chia ở trên, hai trọng số xấp xỉ 0,804 và 0,196. Chú ý chúng khác 0,881 và 0,119 của điểm chưa chia: thứ tự phép tính ảnh hưởng kết quả.

Chọn v1 = [10, 0] và v2 = [0, 10]. Đầu ra attention là 0,804 nhân v1 cộng 0,196 nhân v2, tức [8,04; 1,96]. Đó là một vector mới; attention không đơn giản trả lại vị trí có điểm lớn nhất.

Quy tắc đầy đủ của một đầu là softmax(QK^T / sqrt(d_k) + M)V. Ma trận M chứa mặt nạ. Công trình Attention Is All You Need trình bày scaled dot-product attention và multi-head attention trong Transformer cho dịch máy. Ví dụ số ở đây chỉ giải thích phép toán, không tái lập thí nghiệm của bài báo.

Mặt nạ quyết định những vị trí được nhìn

Khi dự đoán tự hồi quy, vị trí hiện tại không được nhìn token tương lai trong dữ liệu huấn luyện. Causal mask đặt điểm của các vị trí ấy thành âm vô cực trước softmax; trọng số tương ứng trở thành không. Padding mask xử lý phần đệm của batch, có mục đích khác.

Giả sử thêm key thứ ba có điểm rất cao nhưng thuộc tương lai. Nếu quên mặt nạ, loss huấn luyện có thể đẹp vì mô hình nhìn được đáp án. Khi sinh thật, đáp án chưa tồn tại và chất lượng sụp xuống. Đây là lỗi rò rỉ dữ liệu qua kiến trúc, không thể sửa chỉ bằng thêm dữ liệu.

Một query được so với các key để tạo trọng số. Attention dùng những trọng số đó để tổng hợp các value thành biểu diễn mới của vị trí đang xét.

↔ Vuốt ngang để xem toàn bộ sơ đồ.

1. 01 / Query

   Vị trí đang xét có một vector truy vấn, gọi là query.

2. 02 / Key

   Query được so với các vector key ở những vị trí mà mô hình được phép tham chiếu.

3. 03 / Trọng số

   Điểm tương hợp được chuẩn hóa thành trọng số attention cho từng vị trí.

4. 04 / Tổng hợp

   Các vector value được cộng theo trọng số đó, tạo biểu diễn có thông tin từ ngữ cảnh.

Sơ đồ giản lược một đầu attention, bỏ qua kích thước vector và phép chiếu. Các đường nối không phải lời giải thích đầy đủ cho quyết định của mô hình.

Sơ đồ trên giúp theo dõi việc kết hợp ngữ cảnh. Nó không biểu diễn đầy đủ các ma trận, các đầu song song hay causal mask của ví dụ này. Multi-head attention lặp phép tính với những phép chiếu khác nhau, ghép đầu ra rồi chiếu tiếp; không nên gán sẵn mỗi đầu cho một loại ngữ pháp.

Chi phí và giới hạn diễn giải

Với chuỗi dài n, ma trận điểm có n nhân n phần tử trên mỗi đầu. Tăng chiều dài gấp đôi làm số cặp tăng bốn lần. Một kernel tối ưu có thể tránh lưu toàn bộ ma trận trung gian, nhưng phép attention dày vẫn phải xử lý các cặp theo cấu hình của nó.

Trọng số lớn cũng không đủ chứng minh nguyên nhân của một quyết định cuối cùng. Value, residual connection, các lớp sau và nhiều đầu khác đều tham gia. Muốn kiểm tra vai trò một phần đầu vào, cần thiết kế phép can thiệp và quan sát đầu ra thay đổi, thay vì chỉ nhìn màu trên bản đồ attention.

Bài tập tự kiểm bằng ba vector

Giữ q và các key, đổi v2 thành [10, 10]. Tính lại đầu ra mà không tính lại softmax. Bạn đạt yêu cầu khi ra xấp xỉ [10; 1,96] và giải thích được vì sao thay value không đổi trọng số trong phép tính này.

Tiếp theo, che vị trí thứ nhất bằng mask. Đầu ra phải bằng v2, vì vị trí còn lại nhận trọng số một. Nếu kết quả vẫn pha v1, hãy kiểm tra mask được cộng trước hay sau softmax. Đọc tiếp vị trí và RoPE để hiểu thứ tự đi vào attention như thế nào.