Deepseek công bố bài báo: Tại sao NSA (Neural Selective Attention) nhanh hơn so với cơ chế Attention truyền thống?
- Vấn đề: Các mô hình ngôn ngữ lớn (LLM) cần xử lý ngữ cảnh dài nhưng cơ chế chú ý truyền thống (Full Attention) đòi hỏi chi phí tính toán rất lớn, đặc biệt với chuỗi dài như 64K tokens.
- Giải pháp: NSA (Native Sparse Attention) được đề xuất với thiết kế phù hợp phần cứng và khả năng huấn luyện toàn diện, giúp duy trì hiệu suất mô hình trong khi giảm chi phí tính toán.
- Hai cải tiến chính:
- Thiết kế thuật toán cân bằng cường độ số học (arithmetic intensity-balanced algorithm) giúp đạt tốc độ cao hơn trên phần cứng hiện đại.
- Hỗ trợ huấn luyện end-to-end giúp giảm chi phí huấn luyện mà không ảnh hưởng đến hiệu suất mô hình.
- Cách tiếp cận: NSA sử dụng chiến lược thưa thớt phân cấp (hierarchical sparse strategy), kết hợp:
- Token compression (nén token): Gom nhóm các token thành các khối nén giúp lưu giữ thông tin quan trọng mà vẫn giảm tải tính toán.
- Token selection (chọn lọc token): Giữ lại các token quan trọng để đảm bảo độ chính xác cục bộ.
- Sliding window attention (cửa sổ trượt): Duy trì thông tin cục bộ bằng cách chỉ tập trung vào các token gần nhất.
- Hiệu suất:
- So với Full Attention, NSA đạt tốc độ nhanh hơn 11.6x khi suy luận trên chuỗi 64K tokens.
- Trong các bài kiểm tra benchmark trên nhiều tác vụ (MMLU, MATH, DROP, HumanEval...), NSA đạt điểm số trung bình cao hơn so với Full Attention trên 7/9 bài kiểm tra.
- Trên bài toán "needle-in-a-haystack" (tìm thông tin trong ngữ cảnh dài), NSA đạt độ chính xác 100% nhờ cơ chế chú ý thưa thớt phân cấp.
- So sánh với các phương pháp Sparse Attention khác:
- Các mô hình Sparse Attention trước đây như H2O, Quest, InfLLM chủ yếu tập trung vào tối ưu hóa suy luận nhưng không hỗ trợ huấn luyện toàn diện.
- NSA vượt trội hơn nhờ tích hợp hoàn chỉnh cả trong quá trình huấn luyện và suy luận, đạt tốc độ cao hơn mà vẫn duy trì hiệu suất mô hình.
- Ứng dụng:
- Hữu ích cho các mô hình ngôn ngữ cần xử lý ngữ cảnh dài như tóm tắt văn bản, phân tích mã nguồn, hội thoại đa vòng, và lập luận chuỗi.
- Phù hợp với phần cứng hiện đại (GPU A100, FlashAttention-2, Triton) để tận dụng tối đa khả năng tính toán.
📌
NSA là một cơ chế chú ý thưa thớt được thiết kế từ đầu để phù hợp với phần cứng hiện đại, mang lại hiệu suất cao hơn Full Attention trong khi giảm đáng kể chi phí tính toán. NSA kết hợp 3 chiến lược nén, chọn lọc, và cửa sổ trượt để đảm bảo tính chính xác trong khi giảm số phép tính cần thiết. Thử nghiệm cho thấy NSA tăng tốc độ suy luận lên đến 11.6 lần trên chuỗi 64K tokens và vượt trội hơn các phương pháp Sparse Attention khác trên các bài toán ngữ cảnh dài. Đây là một bước tiến quan trọng cho các mô hình ngôn ngữ lớn xử lý văn bản dài mà vẫn đảm bảo hiệu suất. 🚀
https://arxiv.org/pdf/2502.11089