Cài đặt Hàm Softmax và LogSumExp Ổn định Số học Trong AI

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

• Hiểu rõ hiện tượng tràn số mũ (Overflow / Underflow) trong biểu diễn số thực dấu phẩy động 64-bit (float64) khi tính hàm e^x.
• Nắm vững kỹ thuật Trừ Cực Đại (max-subtraction trick) để đưa giá trị số mũ về khoảng ≤ 0.
• Cài đặt hai hàm số toán học cốt lõi trong Transformer và Phân loại: Softmax và LogSumExp.

Mô tả bài toán

Khi tính toán xác suất từ logits đầu ra của mạng nơ-ron, nếu một giá trị logit bằng 1000.0, hàm np.exp(1000.0) sẽ trả về inf, dẫn đến phép chia infinf = NaN.

Hãy cài đặt hai hàm ổn định số học:

• stable_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
• Trừ giá trị lớn nhất dọc theo trục: m = max(logits, axis=axis, keepdims=True).
• Tính e^{logits - m}.
• Chia cho tổng các số mũ: e^{logits - m}Σ e^{logits - m}.
• Tổng xác suất dọc theo trục axis của ma trận kết quả luôn xấp xỉ 1.0.

• stable_log_sum_exp(logits: np.ndarray, axis: int = -1) -> np.ndarray:
• Tính hàm log(Σ_j e^z_j) theo công thức ổn định:

m + log(Σ e^{z - m})

• Trong đó m = max(logits, axis=axis, keepdims=True).
• Loại bỏ chiều rút gọn bằng np.squeeze(..., axis=axis) trước khi trả về.

Ràng buộc

• Hàm phải chạy đúng kể cả khi mảng chứa các giá trị cực lớn (1000.0) hoặc cực nhỏ (-1000.0).
• Không được trả về inf hay nan.
• Thời gian chạy tối đa: 1000ms.
• Giới hạn bộ nhớ: 256MB.