Softmax yang stabil

Implementasikan softmax di sepanjang axis terakhir:

\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}

Input bisa berupa vektor (C,) atau batch (N, C). Hasilnya harus benar untuk [1000, 1001, 1002], yang membuat rumus naif menghasilkan nan.

[1, 2, 3]           →  [0.0900, 0.2447, 0.6652]
[1000, 1001, 1002]  →  [0.0900, 0.2447, 0.6652]   // sama!

Dari skor ke probabilitas

Softmax mengubah skor sembarang (logits) menjadi distribusi probabilitas: semuanya positif dan totalnya 1. Skor yang lebih besar mendapat porsi yang lebih besar secara eksponensial.

Invarian terhadap pergeseran

\frac{e^{x_i - c}}{\sum_j e^{x_j - c}} = \frac{e^{-c}\,e^{x_i}}{e^{-c}\sum_j e^{x_j}} = \text{softmax}(x)_i

Dengan $c = \max(x)$ , argumen exp terbesar adalah 0, jadi tidak pernah overflow, dan setidaknya satu suku bernilai 1, jadi penyebutnya tidak pernah 0.

Temperatur

$\text{softmax}(x / T)$ : $T < 1$ membuat distribusi lebih tajam, $T > 1$ lebih rata. Ini adalah tombol “kreativitas” saat sampling dari model bahasa.

Penjelasan

Yang diminta. Softmax di sepanjang axis terakhir untuk input (C,) atau (N, C), aman untuk logits besar.

Ide kunci: geser dulu, hasilnya tidak berubah. Mengurangi konstanta yang sama dari semua logits tidak mengubah softmax, karena faktor $e^{-c}$ muncul di pembilang dan penyebut lalu saling menghapus. Pilih $c = \max(x)$ :

exp terbesar menjadi $e^0 = 1$ , jadi tidak pernah overflow;
penyebutnya minimal 1, jadi tidak pernah 0.

Langkah di kode:

z = x - x.max(axis=-1, keepdims=True). keepdims=True membuat maksimum berbentuk (N, 1) (atau (1,)), sehingga bisa dikurangkan dari setiap baris lewat broadcasting.
e = np.exp(z).
e / e.sum(axis=-1, keepdims=True) membagi setiap baris dengan jumlahnya sendiri.

Cek. [1000, 1001, 1002] digeser menjadi [-2, -1, 0] dan menghasilkan [0.090, 0.245, 0.665], sama persis dengan [1, 2, 3].

Yang dikembalikan: array berbentuk sama dengan input, setiap baris berjumlah 1.

Kesalahan umum: x.max() tanpa axis (memakai maksimum seluruh batch, sehingga baris lain bisa underflow menjadi 0/0); lupa keepdims sehingga bentuknya tidak cocok.