Cài đặt Hàm Softmax Ổn định Số học chống Overflow/Underflow

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

• Hiểu hiện tượng tràn số trên (overflow) khi tính e^z với logit lớn (z > 709 trong float64 sinh inf).
• Cài đặt mẹo dịch chuyển cực đại: softmax(z) = softmax(z - max(z)).
• Hỗ trợ tính toán dọc theo trục bất kỳ axis cho tensor đa chiều.

Mô tả bài toán

Trong các mô hình Transformer và Phân loại Đa lớp, vector logits đầu ra của mạng nơ-ron có thể đạt giá trị 1000.0. Một công thức ngây thơ e^z∑ e^z sẽ lập tức tạo ra ∞∞ = NaN, làm sụp đổ toàn bộ quá trình học.

Yêu cầu

Viết hàm stable_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:

• Chuyển đổi logits sang kiểu np.float64.
• Tìm giá trị lớn nhất M dọc theo axis với keepdims=True: M = max(logits, axis).
• Tính số mũ đã dịch chuyển: shift_exp = exp(logits - M).
• Tính tổng các số mũ dọc theo axis với keepdims=True: sum_exp = ∑ shift_exp.
• Chia xác suất: P = shift_expsum_exp.
• Trả về mảng phân phối xác suất cùng shape, trong đó tổng các phần tử trên trục axis luôn bằng 1.0 và không có bất kỳ giá trị NaN hay Inf nào.

import numpy as np

def stable_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    pass

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ả.