Causal mask · mặt nạ nhân quả tam giác dưới và tính bất biến quá khứ

1Định nghĩa

Causal mask (mặt nạ nhân quả) là quy tắc chặn thông tin theo chiều thời gian bên trong khối Attention của các mô hình Decoder-only: tại bất kỳ vị trí thứ i nào trong chuỗi, token hiện tại chỉ được phép đọc thông tin từ chính nó và các vị trí đứng trước (j <= i), tuyệt đối không được đọc các vị trí đứng sau (j > i). Về mặt toán học, trước khi đưa bảng điểm tương quan Query × Keyᵀ vào hàm softmax, hệ thống cộng một ma trận tam giác dưới mà phần tam giác trên (vị trí tương lai) mang giá trị âm vô cùng (-∞) để xác suất chú ý sau softmax tại các vị trí tương lai bằng đúng 0.

Trong hệ thống huấn luyện và suy luận, causal mask đóng hai vai trò sống còn: lúc huấn luyện, nó cho phép đưa cả câu mẫu dài T token vào tính forward song song chỉ trong một lượt (kết hợp teacher forcing) mà không sợ token đứng trước nhìn trộm đáp án ở vị trí đứng sau; lúc suy luận, chính tính chất "token mới thêm vào cuối chuỗi không bao giờ làm thay đổi biểu diễn của các token cũ" giúp hệ thống lưu cố định các vector Key/Value quá khứ vào KV cache ở pha prefill và tái sử dụng nguyên vẹn qua từng bước decode.

Ví dụ đời thường: Causal mask giống như tờ bìa che đáp án có khoét khung trượt từng dòng khi bạn tự làm bài tập đọc hiểu: khi đang làm câu hỏi ở dòng thứ 5, tờ bìa chỉ để lộ từ dòng 1 đến dòng 5 và che kín hoàn toàn các dòng từ 6 trở đi ở phía dưới — nhờ vậy dù cả trang giấy đã in sẵn trước mặt, bạn vẫn phải tự suy luận từng bước đúng như khi câu chuyện đang diễn ra theo thời gian thực.

2Luồng

flowchart TD
  step1["Tính ma trận điểm số Query × Keyᵀ"]
  step2["Gán -∞ cho vùng tương lai j #gt; i trước softmax"]
  step3["Bỏ qua tính toán vùng bị che trong kernel"]
  step1 --> step2
  step2 --> step3

Luồng mô tả cách causal mask triệt tiêu trọng số chú ý về tương lai và giúp kernel bỏ qua nửa trên của ma trận attention.

3Giải pháp

Trong các kernels hiện đại như FlashAttention (với cờ causal=True), hệ thống không tạo một ma trận mặt nạ T × T đầy đủ trên bộ nhớ HBM rồi mới cộng vào điểm số: kernel tự kiểm tra chỉ số hàng/cột của từng khối (tile) và bỏ qua hoàn toàn việc nhân ma trận ở các khối nằm hẳn phía trên đường chéo chính, giúp giảm gần một nửa số phép tính attention ở pha prefill.

Ở pha decode chuẩn (khi mỗi bước chỉ đưa vào 1 token mới ở vị trí cuối cùng và đọc lại KV cache của các vị trí trước đó), mọi token nằm trong KV cache của request đều là quá khứ hợp lệ (j <= i), nên bước decode 1 token không cần áp lại mặt nạ tam giác trừ khi dùng speculative decoding xác minh nhiều token nháp cùng lúc (tree/block causal mask).

Cần phân biệt rõ `causal mask` (mặt nạ trong phép tính Attention ngăn nhìn token tương lai) với `loss mask` trong bước tạo nhãn hoặc SFT (mặt nạ ở hàm mất mát dùng để bỏ qua không chấm điểm phần câu hỏi hoặc padding).

4Profiling

• Kiểm tra rò rỉ tương lai (causality leak test)

  Thay đổi một token ở vị trí k > i trong chuỗi đầu vào rồi chạy forward: logits tại mọi vị trí 1..i phải giữ nguyên tuyệt đối ở mức bit/tolerance số học.

• Tốc độ prefill khi bật cờ is_causal trong kernel

  So sánh thời gian chạy và lượng HBM tiêu thụ giữa việc truyền ma trận mask T × T tường minh với việc bật cờ is_causal=True trong FlashAttention.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Lower-triangular mask

Ma trận mặt nạ tam giác dưới giữ nguyên điểm số ở vị trí j <= i và gán -∞ ở vị trí j > i trước khi tính softmax.

Past invariance (tính bất biến quá khứ)

Nhờ causal mask, vector ẩn và cặp Key/Value tại vị trí i chỉ phụ thuộc vào các token 1..i; nối thêm token i+1 vào cuối chuỗi không làm đổi KV đã tính của các vị trí 1..i.

Block-diagonal causal mask (sequence packing)

Khi ghép nhiều văn bản ngắn vào chung một chuỗi huấn luyện (packing), mặt nạ vừa chặn tương lai vừa chặn các văn bản khác nhau nhìn chéo sang nhau.

7Bài liên quan

8Nguồn