Thiên Lệch Tích Tụ, Phương Sai Triệt Tiêu: Tối Ưu Huấn Luyện AI Với Stochastic Rounding

Trong tính toán số học, đặc biệt là trong huấn luyện trí tuệ nhân tạo (AI), việc quản lý độ chính xác của số thực là một bài toán cân bằng khó khăn giữa hiệu năng và tài nguyên bộ nhớ. Một bài viết mới đây đã chỉ ra rằng cách chúng ta xử lý các lỗi làm tròn nhỏ nhoi — dù để chúng tích tụ hay để chúng triệt tiêu lẫn nhau — có thể tạo ra sự khác biệt to lớn giữa việc huấn luyện thành công hay thất bại khi sử dụng các định dạng độ chính xác thấp như BF16.

Làm tròn đến giá trị gần nhất (Round-to-nearest) tạo ra cùng một lỗi làm tròn mỗi lần. Ngược lại, làm tròn ngẫu nhiên (Stochastic rounding) tạo ra các lỗi khác nhau mỗi lần, nhưng trung bình lại bằng không. Khi cùng một lỗi lặp đi lặp lại, nó sẽ tích tụ. Khi các lỗi có giá trị trung bình bằng không, chúng sẽ triệt tiêu lẫn nhau một phần.

Ví dụ về tích tụ sai số

Hãy tưởng tượng bạn cộng 0.001 vào 1.0 một nghìn lần trong định dạng BF16. Nếu sử dụng phương pháp làm tròn đến giá trị gần nhất, kết quả sẽ không bao giờ thay đổi khỏi 1.0. Lý do là mọi lần cập nhật đều rơi vào vị trí gần với 1.0 hơn so với giá trị có thể biểu diễn tiếp theo, nên nó luôn bị làm tròn ngược lại 1.0. Tuy nhiên, nếu sử dụng làm tròn ngẫu nhiên, kết quả cuối cùng sẽ đạt được 2.0. Mỗi lần cập nhật sẽ có xác suất làm tròn lên tỷ lệ thuận với vị trí của nó trong khoảng làm tròn. Về mặt kỳ vọng, tổng phép cộng là chính xác.

So sánh tham số đồ chơi

Về mặt toán học, qua $n$ bước, các lỗi có thiên lệch (biased errors) sẽ tăng trưởng theo cấp độ $O(n)$, nhưng các lỗi không thiên lệch (unbiased errors) chỉ tăng trưởng theo $O(\sqrt{n})$. Phương sai (variance) khuếch tán như một bước đi ngẫu nhiên, tăng lên theo từng bước, nhưng nó tăng chậm hơn nhiều so với thiên lệch. Trong các quá trình chạy dài với các cập nhật nhỏ, sự khác biệt này là tất cả.

Thử nghiệm thực tế

Một thí nghiệm đã được thực hiện với một mạng MLP (Multi-Layer Perceptron) nhỏ, được huấn luyện trên nhiệm vụ hồi quy teacher-student sử dụng trình tối ưu hóa HeavyBall’s AdamW. Tất cả các cấu hình đều lưu trữ tham số ở định dạng BF16. Thí nghiệm thay đổi chế độ làm tròn cho trạng thái của trình tối ưu hóa.

Kết quả cho thấy BF16 kết hợp với Stochastic Rounding (SR) — sử dụng 6 byte cho tham số, mômen bậc nhất và bậc hai — đạt được hiệu năng khớp với trạng thái FP32 (sử dụng 10 byte). Các cập nhật trên từng bước có thể ồn hơn, nhưng độ ồn này không có thiên lệch và sẽ bị triệt tiêu trong quá trình huấn luyện.

Ngược lại, BF16 kết hợp với Round-to-Nearest (RNE) — cũng dùng 6 byte — lại bị dậm chân tại chỗ (plateau) với mức mất mát (loss) cao hơn nhiều bậc magnitude. Cùng một lỗi lặp lại liên tục khiến quá trình huấn luyện bị đình trệ.

Kết luận

Làm tròn ngẫu nhiên thay thế phương pháp làm tròn gần nhất ngay trong nhân của trình tối ưu hóa mà không cần thêm bộ nhớ hay băng thông nào cả. Nếu loại bỏ thiên lệch, 6 byte có thể hoạt động ngang ngửa 10 byte. Nhưng nếu giữ lại thiên lệch, 6 byte sẽ ngay lập tức đụng trần hiệu năng.

Đây là một phát hiện quan trọng đối với việc tối ưu hóa mô hình AI trên phần cứng hiện đại, cho phép giảm thiểu đáng kể lượng bộ nhớ VRAM cần thiết mà không đánh đổi quá nhiều về độ chính xác của mô hình.