Độ Lệch Tuyến Tính Trong Chú Ý (Attention with Linear Biases - ALiBi)

Mục tiêu kiến thức

• ALiBi (Press et al., 2022) loại bỏ hoàn toàn positional embeddings ở input. Thay vào đó, nó cộng trực tiếp một ma trận phạt tuyến tính tĩnh vào ma trận attention scores:

Softmax(Q K^T + m · B)

• Ma trận khoảng cách B ∈ R^{S × S}:

B_{i, j} = -|i - j| (hoặc B_{i, j} = -(i - j) với i ≥ j)

• Hệ số góc m của từng head:

Với H heads (thường H = 2^k), tập các hệ số góc là dãy cấp số nhân: m ∈ { 2^{-8H · 1}, 2^{-8H · 2}, …, 2^{-8H · H} } Nếu H không phải lũy thừa của 2, ta tính theo lũy thừa gần nhất.

• ALiBi cho phép mô hình huấn luyện ở độ dài 1024 token có thể ngoại suy mượt mà lên 2048, 4096 token mà không bị sụp đổ hàm mất mát.

Yêu cầu

Viết hàm:

def compute_alibi_bias(num_heads: int, seq_len: int) -> torch.Tensor:
    pass

• Trả về tensor shape (num_heads, seq_len, seq_len) kiểu float32.
• Với mỗi head h, giá trị tại (i, j) bằng -m_h · (i - j) nếu i ≥ j và -∞ nếu i < j (đảm bảo tính nhân quả Causal Masking).

Ràng buộc

• Thời gian chạy tối đa: 2000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.