Đạo hàm giải tích của Softmax và tích Vector-Jacobian (Softmax VJP)

Mô tả bài toán

Cho vector xác suất sau Softmax p = softmax(z) ∈ R^V với p_i = e^z_i∑_k e^z_k. Ma trận Jacobian của hàm Softmax J ∈ R^{V × V} có các phần tử: J_{i, j} = ∂ p_i∂ z_j = p_i (δ_{i, j} - p_j) = p_i (1 - p_i) & khi i = j
-p_i p_j & khi i ≠ j Trong lan truyền ngược, khi nhận được gradient từ tầng trên g = ∂ L∂ p ∈ R^V, ta cần tính Vector-Jacobian Product (VJP): ∂ L∂ z = J^T g = ∑_i=1^V g_i ∂ p_i∂ z = p ⊙ ( g - ∑_k=1^V g_k p_k ) = p ⊙ (g - g^T p)

Hãy viết hàm softmax_vjp(p: np.ndarray, g: np.ndarray) -> np.ndarray:

• p: vector hoặc ma trận xác suất Softmax (shape (V,) hoặc (B, V)).
• g: vector hoặc ma trận gradient ∂ L∂ p cùng shape với p.
• Tính ∂ L∂ z bằng công thức vector hóa tối ưu O(V) mà không tạo ma trận Jacobian V × V.
• Trả về mảng NumPy cùng shape với kiểu np.float64.

Ràng buộc

• Thời gian chạy tối đa: 6000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.