The softmax function

Real.softmax s i = exp (s i) / ∑ j, exp (s j), the normalized exponential of a score vector over a finite type: Luce's choice rule with exponential scores ([Luc59]), the multinomial logit of [McF74], and the Boltzmann–Gibbs distribution. An inverse temperature enters by scaling the argument, softmax (α • s). With two alternatives it is Real.sigmoid of the score difference, and Real.logit inverts Real.sigmoid.

softmax s is the density of the exponentially tilted counting measure Measure.count.tilted s; that face — the partition function as mgf, log-sum-exp as cgf — is Core.Probability.SoftmaxTheory.

Main definitions

• Real.softmax — the normalized exponential of a score vector.
• Real.logit — the inverse of Real.sigmoid.

Main results

• Real.sum_softmax, Real.softmax_pos — softmax s is a probability distribution.
• Real.softmax_div_softmax, Real.log_softmax_div_softmax — odds are exponentiated score differences (independence of irrelevant alternatives).
• Real.softmax_le_softmax_iff, Real.softmax_update_strictMono — monotonicity in the scores.
• Real.softmax_add_const — translation invariance.
• Real.softmax_fin_two — two alternatives give Real.sigmoid.
• Real.tendsto_softmax_nhds_zero, Real.tendsto_softmax_atTop, Real.tendsto_softmax_atTop_pi, Real.tendsto_softmax_atBot — softmax in the inverse temperature: uniform at 0, a point mass on a strict maximizer (minimizer) as α → ∞ (α → -∞).
• Real.tendsto_sum_negMulLog_softmax_atTop — its entropy vanishes in the hard limit.
• Real.softmax_sum_apply, Real.sum_softmax_eval_eq — a separable score on a product type gives a product distribution, with coordinate marginals.
• Real.rpow_div_sum_rpow — Luce's power rule is softmax of log-scores.

noncomputable def Real.softmax {ι : Type u_1} [Fintype ι] (s : ι → ℝ) : ι → ℝ

The softmax function softmax s i = exp (s i) / ∑ j, exp (s j). An inverse temperature enters by scaling the argument, softmax (α • s).

Equations

• Real.softmax s i = Real.exp (s i) / ∑ j : ι, Real.exp (s j)

theorem Real.softmax_def {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (i : ι) : softmax s i = exp (s i) / ∑ j : ι, exp (s j)

theorem Real.softmax_add_const {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (c : ℝ) : (softmax fun (i : ι) => s i + c) = softmax s

Softmax is invariant under translating every score by the same constant.

@[simp]

theorem Real.softmax_zero {ι : Type u_1} [Fintype ι] : softmax 0 = fun (x : ι) => (↑(Fintype.card ι))⁻¹

theorem Real.sum_exp_pos {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) : 0 < ∑ j : ι, exp (s j)

theorem Real.softmax_pos {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i : ι) : 0 < softmax s i

theorem Real.softmax_nonneg {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i : ι) : 0 ≤ softmax s i

@[simp]

theorem Real.sum_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) : ∑ i : ι, softmax s i = 1

theorem Real.softmax_le_one {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i : ι) : softmax s i ≤ 1

theorem Real.softmax_div_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i j : ι) : softmax s i / softmax s j = exp (s i - s j)

Odds are exponentiated score differences: independence of irrelevant alternatives.

theorem Real.log_softmax_div_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i j : ι) : log (softmax s i / softmax s j) = s i - s j

theorem Real.softmax_eq_softmax_mul_exp_sub {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i j : ι) : softmax s i = softmax s j * exp (s i - s j)

theorem Real.log_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i : ι) : log (softmax s i) = s i - log (∑ j : ι, exp (s j))

theorem Real.softmax_eq_exp_sub {ι : Type u_1} [Fintype ι] [Nonempty ι] (s : ι → ℝ) (i : ι) : softmax s i = exp (s i - log (∑ j : ι, exp (s j)))

theorem Real.softmax_le_softmax_iff {ι : Type u_1} [Fintype ι] [Nonempty ι] {s : ι → ℝ} {i j : ι} : softmax s i ≤ softmax s j ↔ s i ≤ s j

theorem Real.softmax_lt_softmax_iff {ι : Type u_1} [Fintype ι] [Nonempty ι] {s : ι → ℝ} {i j : ι} : softmax s i < softmax s j ↔ s i < s j

theorem Real.softmax_le_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] {s : ι → ℝ} {i j : ι} : s i ≤ s j → softmax s i ≤ softmax s j

Alias of the reverse direction of Real.softmax_le_softmax_iff.

theorem Real.softmax_lt_softmax {ι : Type u_1} [Fintype ι] [Nonempty ι] {s : ι → ℝ} {i j : ι} : s i < s j → softmax s i < softmax s j

Alias of the reverse direction of Real.softmax_lt_softmax_iff.

theorem Real.softmax_update_strictMono {ι : Type u_1} [Fintype ι] [DecidableEq ι] [Nontrivial ι] (s : ι → ℝ) (i : ι) : StrictMono fun (x : ℝ) => softmax (Function.update s i x) i

softmax (update s i x) i is strictly increasing in the score x.

theorem Real.softmax_lt_softmax_of_single_lt {ι : Type u_1} [Fintype ι] [DecidableEq ι] [Nontrivial ι] {s s' : ι → ℝ} {i : ι} (hlt : s i < s' i) (heq : ∀ (j : ι), j ≠ i → s' j = s j) : softmax s i < softmax s' i

Raising one score while holding the others fixed raises its probability.

Limits in the inverse temperature

softmax (α • s) is continuous in the inverse temperature α: uniform at α = 0, concentrating on a strict maximizer as α → ∞ and on a strict minimizer as α → -∞, with its entropy vanishing in the limit.

theorem Real.softmax_eq_inv_sum_exp_sub {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (i : ι) : softmax s i = (∑ j : ι, exp (s j - s i))⁻¹

theorem Real.softmax_le_exp_sub {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (i : ι) [Nonempty ι] (j : ι) : softmax s i ≤ exp (s i - s j)

theorem Real.continuous_softmax_smul {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (i : ι) [Nonempty ι] : Continuous fun (α : ℝ) => softmax (α • s) i

theorem Real.tendsto_softmax_nhds_zero {ι : Type u_1} [Fintype ι] (s : ι → ℝ) (i : ι) [Nonempty ι] : Filter.Tendsto (fun (α : ℝ) => softmax (α • s) i) (nhds 0) (nhds (↑(Fintype.card ι))⁻¹)

At inverse temperature 0, softmax is uniform.

theorem Real.tendsto_softmax_atTop {ι : Type u_1} [Fintype ι] {s : ι → ℝ} {i : ι} (h : ∀ (j : ι), j ≠ i → s j < s i) : Filter.Tendsto (fun (α : ℝ) => softmax (α • s) i) Filter.atTop (nhds 1)

As the inverse temperature grows, softmax concentrates on a strict maximizer.

theorem Real.tendsto_softmax_atTop_of_lt {ι : Type u_1} [Fintype ι] {s : ι → ℝ} {i j : ι} (h : s i < s j) : Filter.Tendsto (fun (α : ℝ) => softmax (α • s) i) Filter.atTop (nhds 0)

As the inverse temperature grows, a strictly dominated alternative vanishes.

theorem Real.tendsto_softmax_atTop_pi {ι : Type u_1} [Fintype ι] {s : ι → ℝ} {i : ι} [DecidableEq ι] (h : ∀ (j : ι), j ≠ i → s j < s i) : Filter.Tendsto (fun (α : ℝ) => softmax (α • s)) Filter.atTop (nhds (Pi.single i 1))

Softmax converges to the point mass on a strict maximizer.

theorem Real.tendsto_softmax_atBot {ι : Type u_1} [Fintype ι] {s : ι → ℝ} {i : ι} (h : ∀ (j : ι), j ≠ i → s i < s j) : Filter.Tendsto (fun (α : ℝ) => softmax (α • s) i) Filter.atBot (nhds 1)

As the inverse temperature tends to -∞, softmax concentrates on a strict minimizer.

theorem Real.tendsto_sum_negMulLog_softmax_atTop {ι : Type u_1} [Fintype ι] {s : ι → ℝ} {i : ι} (h : ∀ (j : ι), j ≠ i → s j < s i) : Filter.Tendsto (fun (α : ℝ) => ∑ j : ι, (softmax (α • s) j).negMulLog) Filter.atTop (nhds 0)

The entropy of softmax vanishes as it concentrates on a strict maximizer.

Two alternatives and the logit

theorem Real.softmax_fin_two (s : Fin 2 → ℝ) : softmax s 0 = (s 0 - s 1).sigmoid

With two alternatives, softmax is the logistic function of the score difference.

noncomputable def Real.logit (p : ℝ) : ℝ

The logit function log (p / (1 - p)), the inverse of Real.sigmoid.

Equations

• p.logit = Real.log (p / (1 - p))

@[simp]

theorem Real.logit_sigmoid (x : ℝ) : x.sigmoid.logit = x

theorem Real.sigmoid_logit {p : ℝ} (hp : p ∈ Set.Ioo 0 1) : p.logit.sigmoid = p

theorem Real.logit_softmax_fin_two (s : Fin 2 → ℝ) : (softmax s 0).logit = s 0 - s 1

With two alternatives, the log-odds are the score difference.

theorem Real.rpow_div_sum_rpow {ι : Type u_1} [Fintype ι] [Nonempty ι] {f : ι → ℝ} (hf : ∀ (i : ι), 0 < f i) (α : ℝ) (i : ι) : f i ^ α / ∑ j : ι, f j ^ α = softmax (α • fun (j : ι) => log (f j)) i

Luce's power rule f i ^ α / ∑ j, f j ^ α is the softmax of the scaled log-scores.

Product types

A separable score s f = ∑ i, c i (f i) on assignments f : ι → V gives a softmax that factorizes into the coordinate softmaxes, so marginalizing the joint distribution at coordinate i recovers softmax (c i).

theorem Real.softmax_sum_apply {ι : Type u_1} [Fintype ι] {V : Type u_2} [DecidableEq ι] [Fintype V] (c : ι → V → ℝ) (f : ι → V) : softmax (fun (g : ι → V) => ∑ i : ι, c i (g i)) f = ∏ i : ι, softmax (c i) (f i)

The softmax of a separable score is the product of the coordinate softmaxes.

theorem Real.sum_softmax_eval_eq {ι : Type u_1} [Fintype ι] {V : Type u_2} [DecidableEq ι] [Fintype V] [DecidableEq V] [Nonempty V] (c : ι → V → ℝ) (i : ι) (v : V) : ∑ f : ι → V with f i = v, softmax (fun (g : ι → V) => ∑ j : ι, c j (g j)) f = softmax (c i) v

Marginalizing the softmax of a separable score at coordinate i recovers the coordinate softmax softmax (c i).