Block Sparse Attention with Log-Linear Complexity

Nhóm nghiên cứu đề xuất PISA, một cơ chế chú ý thưa theo khối sử dụng chiến lược lựa chọn pyramid Top-K để giảm chi phí tính toán từ bậc hai xuống bậc log-linear. PISA xây dựng một hệ thống phân cấp khóa từ thô đến mịn thông qua việc gom nhóm và áp dụng điểm số LogSumExp trên các tập ứng viên bị chặn. Họ phát triển các nhân Triton nhận thức phần cứng cho cả huấn luyện và suy luận để hợp nhất định tuyến phân cấp và chấm điểm mà không cần hiện thực hóa ma trận điểm số query-key. Phương pháp này đạt kết quả tương đương trên các bài kiểm tra suy luận thông thường và tốt hơn trên các tác vụ truy xuất so với baseline.

Through pooling, we construct O(log N) levels of keys, yielding an overall complexity of O(Nlog N), where N denotes the sequence length.