The softmax function

This file defines the softmax function Real.softmax s i = exp (s i) / ∑ j, exp (s j), which normalizes the exponentials of a score vector over a finite type. It is Luce's choice rule with exponential scores, McFadden's multinomial logit, and the Boltzmann–Gibbs distribution. An inverse temperature enters by scaling the argument, as in softmax (α • s). With two alternatives softmax is Real.sigmoid of the score difference.

Log-sum-exp x ↦ log (∑ i, exp (x i)), the logarithm of the normalizer, is convex, and its derivative is the softmax expectation of the derivative of the scores.

Main definitions

• Real.softmax: the normalized exponential of a score vector.

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, Real.softmax_const: translation invariance, and the uniform distribution for constant scores.
• Real.convexOn_log_sum_exp, HasFDerivAt.log_sum_exp: log-sum-exp is convex, and its derivative is the softmax-weighted derivative of the scores.
• Real.tendsto_softmax_atTop: softmax concentrates on a strict maximizer as the inverse temperature grows.
• Real.softmax_eq_sigmoid_of_univ_eq_pair, Real.softmax_fin_two: two alternatives give Real.sigmoid.
• Real.softmax_sum_apply: a separable score on a product type gives a product distribution.

References

• R. D. Luce, Individual Choice Behavior: A Theoretical Analysis (1959)
• D. McFadden, Conditional logit analysis of qualitative choice behavior (1974)

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

The softmax function sends a score vector s to i ↦ exp (s i) / ∑ j, exp (s j). An inverse temperature enters by scaling the argument, as in softmax (α • s).

Equations

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

theorem Real.softmax_apply {ι : 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 ι))⁻¹

@[simp]

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

Equal scores give the uniform distribution.

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)

The odds of two alternatives are their exponentiated score difference, whatever the other alternatives are.

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

Convexity

theorem Real.convexOn_log_sum_exp {ι : Type u_1} [Fintype ι] [Nonempty ι] : ConvexOn ℝ Set.univ fun (s : ι → ℝ) => log (∑ i : ι, exp (s i))

Log-sum-exp, the logarithm of softmax's normalizer, is convex. Hölder's inequality with exponents a⁻¹ and b⁻¹ bounds ∑ i, exp (a * s i + b * t i) by (∑ i, exp (s i)) ^ a * (∑ i, exp (t i)) ^ b.

The derivative of log-sum-exp

theorem HasFDerivAt.log_sum_exp {ι : Type u_1} [Fintype ι] {E : Type u_2} [NormedAddCommGroup E] [NormedSpace ℝ E] [Nonempty ι] {f : ι → E → ℝ} {f' : ι → E →L[ℝ] ℝ} {x : E} (hf : ∀ (i : ι), HasFDerivAt (f i) (f' i) x) : HasFDerivAt (fun (y : E) => Real.log (∑ i : ι, Real.exp (f i y))) (∑ i : ι, Real.softmax (fun (x_1 : ι) => f x_1 x) i • f' i) x

The derivative of log-sum-exp of a family of scores is the softmax-weighted sum of their derivatives.

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

The limit of low temperature

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

Two alternatives

theorem Real.softmax_eq_sigmoid_of_univ_eq_pair {ι : Type u_1} [Fintype ι] [DecidableEq ι] {i j : ι} (hij : i ≠ j) (h : Finset.univ = {i, j}) (s : ι → ℝ) : softmax s i = (s i - s j).sigmoid

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

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

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

Product types

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 ∑ i, c i (g i) on assignments g : ι → V is the product of the coordinate softmaxes.