Thủ thuật Gumbel-Max và Gumbel-Softmax vi phân (Gumbel Reparameterization)

Mô tả bài toán

Trong học tăng cường và mô hình sinh rời rạc, thao tác lấy mẫu argmax hoặc sample() từ phân phối Categorical là một hàm bậc thang không thể tính đạo hàm (non-differentiable). Thủ thuật Gumbel-Max và Gumbel-Softmax (Jang et al., 2016) giải quyết triệt để vấn đề này:

• Gumbel Noise: Cho biến ngẫu nhiên đều u ∼ U(0, 1), biến đổi Gumbel chuẩn là:

g = -log(-log(u))

• Gumbel-Max: Lấy argmax_i (z_i + g_i) tương đương chính xác với việc lấy mẫu từ phân phối rời rạc với logits z.
• Gumbel-Softmax (Continuous Relaxation): Xấp xỉ liên tục và khả vi của hàm argmax:

y_i = e^{(z_i + g_i) / τ}∑_k=1^V e^{(z_k + g_k) / τ} trong đó τ > 0 là tham số nhiệt độ. Khi τ → 0, y tiệm cận một vector One-Hot.

Hãy viết hàm gumbel_softmax(logits: np.ndarray, tau: float = 1.0, hard: bool = False, seed: int = 42) -> np.ndarray:

• logits: mảng 1D float64 kích thước (V,).
• tau: nhiệt độ vi phân τ > 0.
• hard: nếu True, vector trả về là One-Hot tại vị trí cực đại (y_hard), nhưng trong đồ thị tính toán được bảo toàn đạo hàm của vector mềm (y_soft) qua mẹo: y = y_hard - y_soft.detach() + y_soft (Straight-Through Estimator).
• seed: seed sinh số ngẫu nhiên Gumbel.
• Trả về vector xác suất y ∈ R^V.

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