From e8f10a89de52b1790cff3a195498b96a9672daa4 Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 12 Nov 2026 06:10:28 +0100 Subject: [PATCH 1/3] feat: Thompson sampling for exponential families, the Jeffreys prior, and Gaussian Thompson sampling as an instance Co-Authored-By: Claude Code --- LMLExplore/Bandits/AsymptoticOptimality.lean | 2 +- .../Bandits/ThompsonSampling/Basic.lean | 233 ++++++++++++++++++ .../Gaussian.lean} | 37 ++- .../ExponentialFamily/Gaussian.lean | 99 +++++++- .../ExponentialFamily/Posterior.lean | 180 ++++++++++++++ LMLExplore/Probability/Moments/Basic.lean | 21 ++ 6 files changed, 568 insertions(+), 4 deletions(-) create mode 100644 LMLExplore/Bandits/ThompsonSampling/Basic.lean rename LMLExplore/Bandits/{ThompsonSampling.lean => ThompsonSampling/Gaussian.lean} (94%) create mode 100644 LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean diff --git a/LMLExplore/Bandits/AsymptoticOptimality.lean b/LMLExplore/Bandits/AsymptoticOptimality.lean index 9d5331fc..433a39e6 100644 --- a/LMLExplore/Bandits/AsymptoticOptimality.lean +++ b/LMLExplore/Bandits/AsymptoticOptimality.lean @@ -9,7 +9,7 @@ public import LMLExplore.Analysis.Asymptotics.Lemmas public import LMLExplore.Bandits.AsymptoticallyOptimalUCB public import LMLExplore.Bandits.KLUCB public import LMLExplore.Bandits.LaiRobbins -public import LMLExplore.Bandits.ThompsonSampling +public import LMLExplore.Bandits.ThompsonSampling.Gaussian /-! # Consistency and asymptotic optimality of UCB-type algorithms and Thompson sampling diff --git a/LMLExplore/Bandits/ThompsonSampling/Basic.lean b/LMLExplore/Bandits/ThompsonSampling/Basic.lean new file mode 100644 index 00000000..33a0dc43 --- /dev/null +++ b/LMLExplore/Bandits/ThompsonSampling/Basic.lean @@ -0,0 +1,233 @@ +/- +Copyright (c) 2026 Rémy Degenne. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Rémy Degenne +-/ +module + +public import LeanMachineLearning.SequentialLearning.SumRewards +public import LMLExplore.Bandits.PerturbedLeader +public import LMLExplore.Probability.Distributions.ExponentialFamily.Posterior +public import LMLExplore.SequentialLearning.FiniteActions + +/-! +# Thompson sampling for exponential families + +Let every arm `a` have a natural exponential family `ρ a` (a measure on `ℝ` with the statistic +`id`, log-partition function `ψ_a = cgf id (ρ a)` and natural parameters `Θ_a`, the interior of +`integrableExpSet id (ρ a)`) and a prior `π a` on its natural parameter, a measure on `ℝ` that +may be infinite. Thompson sampling [KKM13, Section 2] samples at every round, independently for +every arm `a` pulled `s ≥ 1` times with sum of rewards `S`, a natural parameter `θ_a` from the +posterior `(ρ a).expFamilyPosterior id (π a) s S`, the probability measure proportional to +`exp (θ S - s ψ_a(θ)) π_a(dθ)` on `Θ_a`, and takes its mean `ψ_a'(θ_a)` as the perturbation of +`a`. The perturbation of an arm that has not been pulled, or whose posterior normalizing constant +is `0` or infinite (an improper posterior), is `⊤`. It then plays the arm of largest perturbation. + +This is a follow-the-perturbed-leader policy (`perturbedLeaderAlgorithm`) whose perturbation +laws, in `EReal`, depend on the history only through the number of pulls and the sum of rewards +of each arm, so [LS, Theorem 36.2] applies to it in the reward-stack form of +`PerturbedLeader.lintegral_pullCount_le`. With the Jeffreys prior `(ρ a).jeffreysPrior id` it is +the algorithm of [KKM13]; for the Gaussian family of variance `σ²` the Jeffreys prior is a +constant multiple of Lebesgue measure, and Thompson sampling is Gaussian Thompson sampling +`gaussianTS σ2` (`gaussianTS_eq_expFamilyTS`). + +## Main definitions + +* `LMLExplore.Bandits.expFamilyTSStackLaw ρ π s S`: the law of the perturbation of an arm of + family `ρ` and prior `π` that has been pulled `s` times with sum of rewards `S`: `δ_⊤` if + `s = 0` or the posterior is `0`, and otherwise the image of the posterior by the mean map + `ψ'`. +* `LMLExplore.Bandits.expFamilyTSLaw ρ π n h a`: the law of the perturbation of arm `a` after the + history `h` of length `n`. +* `LMLExplore.Bandits.expFamilyTS ρ π`: Thompson sampling for the families `ρ` with priors `π`, the + algorithm `perturbedLeaderAlgorithm` with perturbation laws `expFamilyTSLaw ρ π`. + +## Main statements + +* `expFamilyTSStackLaw_of_ne_zero`: the law of the perturbation of a pulled arm with a proper + posterior is the image of the posterior by the mean map. +* `measurable_expFamilyTSStackLaw`, `measurable_expFamilyTSLaw`: measurability of the + perturbation laws in the sum of rewards and in the history. +* `expFamilyTSLaw_history`: along a sequence of rounds, the perturbation law in terms of LML's + `pullCount` and `sumRewards`. +* `ExpFamilyTS.lintegral_pullCount_le`: [LS, Theorem 36.2] for Thompson sampling against a + stochastic bandit, in the reward-stack model. + +## References + +* [KKM13] N. Korda, E. Kaufmann and R. Munos, *Thompson sampling for 1-dimensional exponential + family bandits*, NeurIPS 2013. +* [LS] T. Lattimore and Cs. Szepesvári, *Bandit Algorithms*, Cambridge University Press, 2020, + Chapter 36. +-/ + +public section + +open MeasureTheory ProbabilityTheory Learning Bandits + +open scoped ENNReal + +namespace LMLExplore.Bandits + +section Law + +variable {α : Type*} [DecidableEq α] {ρ π : Measure ℝ} {s : ℕ} {S : ℝ} + +open scoped Classical in +/-- The law of the perturbation of an arm in Thompson sampling for the natural exponential family +of `ρ` with prior `π` on the natural parameter, as a function of the number of pulls `s` of the +arm and of its sum of rewards `S`: the Dirac mass at `⊤` if `s = 0` or if the posterior +`ρ.expFamilyPosterior id π s S` is `0`, and otherwise the image of the posterior by the mean map +`θ ↦ ψ'(θ)`, with `ψ = cgf id ρ`. This is the perturbation law of the reward-stack model, in which +the sum of rewards after `s` pulls is the sum of the first `s` rewards of the stack. -/ +noncomputable def expFamilyTSStackLaw (ρ π : Measure ℝ) (s : ℕ) (S : ℝ) : Measure EReal := + if s = 0 ∨ ρ.expFamilyPosterior id π s S = 0 then Measure.dirac ⊤ + else (ρ.expFamilyPosterior id π s S).map fun θ ↦ ((deriv (cgf id ρ) θ : ℝ) : EReal) + +open scoped Classical in +/-- The defining equation of `expFamilyTSStackLaw`. -/ +lemma expFamilyTSStackLaw_def (ρ π : Measure ℝ) (s : ℕ) (S : ℝ) : + expFamilyTSStackLaw ρ π s S = + if s = 0 ∨ ρ.expFamilyPosterior id π s S = 0 then Measure.dirac ⊤ + else (ρ.expFamilyPosterior id π s S).map fun θ ↦ ((deriv (cgf id ρ) θ : ℝ) : EReal) := + (rfl) + +/-- The perturbation of an arm that has not been pulled is `⊤`. -/ +@[simp] +lemma expFamilyTSStackLaw_zero (ρ π : Measure ℝ) (S : ℝ) : + expFamilyTSStackLaw ρ π 0 S = Measure.dirac ⊤ := by + simp [expFamilyTSStackLaw_def] + +/-- The perturbation of an arm whose posterior is `0` (its normalizing constant is `0` or +infinite) is `⊤`. -/ +lemma expFamilyTSStackLaw_of_expFamilyPosterior_eq_zero (h : ρ.expFamilyPosterior id π s S = 0) : + expFamilyTSStackLaw ρ π s S = Measure.dirac ⊤ := by + simp [expFamilyTSStackLaw_def, h] + +/-- The perturbation of an arm that has been pulled `s ≠ 0` times with sum of rewards `S`, and +whose posterior is not `0`, is the mean `ψ'(θ)` of a natural parameter `θ` drawn from the +posterior. -/ +lemma expFamilyTSStackLaw_of_ne_zero (hs : s ≠ 0) (h : ρ.expFamilyPosterior id π s S ≠ 0) : + expFamilyTSStackLaw ρ π s S = + (ρ.expFamilyPosterior id π s S).map fun θ ↦ ((deriv (cgf id ρ) θ : ℝ) : EReal) := by + simp [expFamilyTSStackLaw_def, hs, h] + +instance : IsProbabilityMeasure (expFamilyTSStackLaw ρ π s S) := by + by_cases h : s = 0 ∨ ρ.expFamilyPosterior id π s S = 0 + · rcases h with h | h + · rw [h, expFamilyTSStackLaw_zero] + infer_instance + · rw [expFamilyTSStackLaw_of_expFamilyPosterior_eq_zero h] + infer_instance + · push Not at h + have := Measure.isProbabilityMeasure_expFamilyPosterior_of_ne_zero h.2 + rw [expFamilyTSStackLaw_of_ne_zero h.1 h.2] + exact (Measure.isProbabilityMeasure_map_iff (by fun_prop)).2 inferInstance + +/-- For an s-finite family `ρ` and an s-finite prior `π`, the perturbation law of an arm pulled `s` +times is a measurable function of its sum of rewards. -/ +@[fun_prop] +lemma measurable_expFamilyTSStackLaw (ρ π : Measure ℝ) [SFinite ρ] [SFinite π] (s : ℕ) : + Measurable fun S ↦ expFamilyTSStackLaw ρ π s S := by + classical + have hpost := Measure.measurable_expFamilyPosterior (ρ := ρ) aemeasurable_id π s + have hset : MeasurableSet {S | s = 0 ∨ ρ.expFamilyPosterior id π s S = 0} := by + simp_rw [← Measure.measure_univ_eq_zero] + exact (MeasurableSet.const _).union + ((Measure.measurable_coe MeasurableSet.univ).comp hpost (measurableSet_singleton 0)) + simp_rw [expFamilyTSStackLaw_def] + exact Measurable.ite hset measurable_const + ((Measure.measurable_map _ (by fun_prop)).comp hpost) + +/-- The law of the perturbation of arm `a` in Thompson sampling for the families `ρ` with priors +`π`, at Lean round `n` for the history `h` of the first `n` rounds: the perturbation law +`expFamilyTSStackLaw (ρ a) (π a)` at the number of pulls and the sum of rewards of `a`. -/ +noncomputable def expFamilyTSLaw (ρ π : α → Measure ℝ) (n : ℕ) (h : Hist Unit α ℝ n) (a : α) : + Measure EReal := + expFamilyTSStackLaw (ρ a) (π a) (pullCount' n h a) (sumRewards' n h a) + +/-- The defining equation of `expFamilyTSLaw`. -/ +lemma expFamilyTSLaw_def (ρ π : α → Measure ℝ) (n : ℕ) (h : Hist Unit α ℝ n) (a : α) : + expFamilyTSLaw ρ π n h a = + expFamilyTSStackLaw (ρ a) (π a) (pullCount' n h a) (sumRewards' n h a) := + (rfl) + +instance {ρ π : α → Measure ℝ} {n : ℕ} {h : Hist Unit α ℝ n} {a : α} : + IsProbabilityMeasure (expFamilyTSLaw ρ π n h a) := by + rw [expFamilyTSLaw_def] + infer_instance + +/-- For s-finite families and priors, the perturbation law of an arm at round `n` is a measurable +function of the history: it depends on it through the number of pulls and the sum of rewards of +the arm. -/ +@[fun_prop] +lemma measurable_expFamilyTSLaw [MeasurableSpace α] [MeasurableSingletonClass α] + (ρ π : α → Measure ℝ) [∀ a, SFinite (ρ a)] [∀ a, SFinite (π a)] (n : ℕ) (a : α) : + Measurable fun h : Hist Unit α ℝ n ↦ expFamilyTSLaw ρ π n h a := + (measurable_from_prod_countable_right + (f := fun q : ℕ × ℝ ↦ expFamilyTSStackLaw (ρ a) (π a) q.1 q.2) + (measurable_expFamilyTSStackLaw (ρ a) (π a))).comp + ((measurable_pullCount' n a).prodMk (measurable_sumRewards' n a)) + +/-- Along a sequence of rounds `(O, A, R)`, the perturbation law of an arm at round `n` is the +perturbation law `expFamilyTSStackLaw` at its number of pulls `pullCount A a n ω` and its sum of +rewards `sumRewards A R a n ω` before round `n`. -/ +lemma expFamilyTSLaw_history {Ω : Type*} (O : ℕ → Ω → Unit) (A : ℕ → Ω → α) (R : ℕ → Ω → ℝ) + (ρ π : α → Measure ℝ) (n : ℕ) (ω : Ω) (a : α) : + expFamilyTSLaw ρ π n (history O A R n ω) a = + expFamilyTSStackLaw (ρ a) (π a) (pullCount A a n ω) (sumRewards A R a n ω) := by + rw [expFamilyTSLaw_def, ← pullCount_eq_pullCount', ← sumRewards_eq_sumRewards'] + +end Law + +variable {α : Type*} [Fintype α] [Nonempty α] [DecidableEq α] [MeasurableSpace α] + [MeasurableSingletonClass α] + +/-- **Thompson sampling for exponential families** [KKM13, Section 2]: for natural exponential +families `ρ a` and priors `π a` on their natural parameters, the follow-the-perturbed-leader +algorithm whose perturbation of an arm `a` pulled `s ≥ 1` times with sum of rewards `S` is the mean +`ψ_a'(θ_a)` of a natural parameter `θ_a` drawn from the posterior +`(ρ a).expFamilyPosterior id (π a) s S`, and `⊤` for unpulled arms and improper posteriors. -/ +noncomputable def expFamilyTS (ρ π : α → Measure ℝ) [∀ a, SFinite (ρ a)] [∀ a, SFinite (π a)] : + Algorithm Unit α ℝ := + perturbedLeaderAlgorithm (fun n x a ↦ expFamilyTSLaw ρ π n x.1 a) + (fun n a ↦ (measurable_expFamilyTSLaw ρ π n a).comp measurable_fst) fun _ _ _ ↦ inferInstance + +/-- Thompson sampling for exponential families is the follow-the-perturbed-leader algorithm with +perturbation laws `expFamilyTSLaw ρ π`. -/ +lemma expFamilyTS_def (ρ π : α → Measure ℝ) [∀ a, SFinite (ρ a)] [∀ a, SFinite (π a)] : + expFamilyTS ρ π = + perturbedLeaderAlgorithm (fun n x a ↦ expFamilyTSLaw ρ π n x.1 a) + (fun n a ↦ (measurable_expFamilyTSLaw ρ π n a).comp measurable_fst) + fun _ _ _ ↦ inferInstance := + (rfl) + +namespace ExpFamilyTS + +variable {Ω : Type*} {mΩ : MeasurableSpace Ω} {P : Measure Ω} {O : ℕ → Ω → Unit} + {A : ℕ → Ω → α} {R : ℕ → Ω → ℝ} {ρ π : α → Measure ℝ} [∀ a, SFinite (ρ a)] + [∀ a, SFinite (π a)] + +/-- **[LS, Theorem 36.2] for Thompson sampling for exponential families.** In a run of Thompson +sampling against the stochastic bandit `Environment.bandit ν`, for arms `i ≠ j`, a level `c` and +every horizon `n`, +`E[T_i(n)] ≤ 1 + E[∑_{s < n} (G_{j s}⁻¹ - 1)] + ∑_{s < n} P(G_{i s} > 1 / n)`, where `G_{a s}` is +the probability that the perturbation of arm `a` after its first `s` rewards in the reward-stack +model `streamMeasure ν` exceeds `c`. -/ +theorem lintegral_pullCount_le [IsProbabilityMeasure P] {ν : Kernel α ℝ} [IsMarkovKernel ν] + (h : IsAlgEnvSeq O A R (expFamilyTS ρ π) (Environment.bandit ν) P) {i j : α} (hij : i ≠ j) + (c : EReal) (n : ℕ) : + ∫⁻ ω, (pullCount A i n ω : ℝ≥0∞) ∂P ≤ + 1 + ∑ s ∈ Finset.range n, ∫⁻ x, + (expFamilyTSStackLaw (ρ j) (π j) s (∑ k ∈ Finset.range s, x k j) (Set.Ioi c))⁻¹ - 1 + ∂(streamMeasure ν) + + ∑ s ∈ Finset.range n, streamMeasure ν {x | (n : ℝ≥0∞)⁻¹ < + expFamilyTSStackLaw (ρ i) (π i) s (∑ k ∈ Finset.range s, x k i) (Set.Ioi c)} := + -- the perturbation law of an arm is a function of its number of pulls and its sum of rewards + PerturbedLeader.lintegral_pullCount_le h (fun a ↦ expFamilyTSStackLaw (ρ a) (π a)) + (fun a ↦ measurable_expFamilyTSStackLaw (ρ a) (π a)) + (fun t x a ↦ expFamilyTSLaw_def ρ π t x.1 a) hij c n + +end ExpFamilyTS + +end LMLExplore.Bandits diff --git a/LMLExplore/Bandits/ThompsonSampling.lean b/LMLExplore/Bandits/ThompsonSampling/Gaussian.lean similarity index 94% rename from LMLExplore/Bandits/ThompsonSampling.lean rename to LMLExplore/Bandits/ThompsonSampling/Gaussian.lean index 5272910e..f1eb5fc3 100644 --- a/LMLExplore/Bandits/ThompsonSampling.lean +++ b/LMLExplore/Bandits/ThompsonSampling/Gaussian.lean @@ -11,6 +11,7 @@ public import LMLExplore.Bandits.Concentration public import LMLExplore.Bandits.Instances public import LMLExplore.Bandits.PerturbedLeader public import LMLExplore.Bandits.Regret +public import LMLExplore.Bandits.ThompsonSampling.Basic public import LMLExplore.Probability.Distributions.Gaussian.Tail public import LMLExplore.SequentialLearning.FiniteActions @@ -26,8 +27,12 @@ Gaussian samples, and "pull each arm once" is the behaviour of the policy itself separate initialization phase. It is the limit of Thompson sampling with independent Gaussian priors on the means when the prior -variances tend to infinity. It differs from LML's `Bandits.tsAlgorithm`, which samples from the -posterior of a prior over environments. +variances tend to infinity. For `σ² ≠ 0` it is the instance of Thompson sampling for exponential +families `expFamilyTS` with the Gaussian family of variance `σ²` and its Jeffreys prior, a constant +multiple of Lebesgue measure on the natural parameter: the posterior of the natural parameter +`θ` of an arm pulled `s` times with sum of rewards `S` is `𝒩(S / (s σ²), 1 / (s σ²))` (for the +family generated by `𝒩(0, σ²)`), and its mean `σ² θ` has law `𝒩(S / s, σ² / s)`. It differs +from LML's `Bandits.tsAlgorithm`, which samples from the posterior of a prior over environments. [LS] writes the variance `1 / t` of the update; the intended variance after `s` samples of the arm is `1 / s`, here `σ² / s`. @@ -43,6 +48,9 @@ is `1 / s`, here `σ² / s`. ## Main statements +* `gaussianTSStackLaw_eq_expFamilyTSStackLaw`, `gaussianTS_eq_expFamilyTS`: for `σ² ≠ 0`, the + perturbation laws of Gaussian Thompson sampling are those of Thompson sampling for the Gaussian + family of variance `σ²` with its Jeffreys prior, and the two algorithms are equal. * `gaussianTSLaw_history`: along a sequence of rounds, the perturbation law in terms of LML's `pullCount` and `empMean`. * `GaussianTS.pullCount_action_eq_zero`: along a run, while some arm has not been pulled, the arm @@ -465,6 +473,31 @@ lemma gaussianTS_def (σ2 : ℝ≥0) : fun _ _ _ ↦ inferInstance := (rfl) +/-- For `σ² ≠ 0`, the perturbation law of Gaussian Thompson sampling of an arm pulled `s` times +with sum of rewards `y` is that of Thompson sampling for the Gaussian family of variance `σ²`, +generated by any `𝒩(m, σ²)`, with its Jeffreys prior. -/ +theorem gaussianTSStackLaw_eq_expFamilyTSStackLaw (m : ℝ) {σ2 : ℝ≥0} (hσ2 : σ2 ≠ 0) (s : ℕ) + (y : ℝ) : + gaussianTSStackLaw σ2 s y = + expFamilyTSStackLaw (gaussianReal m σ2) ((gaussianReal m σ2).jeffreysPrior id) s y := by + rcases eq_or_ne s 0 with rfl | hs + · simp + have hpost := expFamilyPosterior_gaussianReal_jeffreysPrior m hσ2 hs y + rw [gaussianTSStackLaw_of_ne_zero hs, + expFamilyTSStackLaw_of_ne_zero hs (by rw [hpost]; exact IsProbabilityMeasure.ne_zero _), + ← map_deriv_cgf_expFamilyPosterior_gaussianReal_jeffreysPrior m hσ2 hs, + Measure.map_map (by fun_prop) (measurable_deriv _), Function.comp_def] + +/-- For `σ² ≠ 0`, Gaussian Thompson sampling is Thompson sampling for exponential families with the +Gaussian family of variance `σ²` on every arm, generated by any `𝒩(m a, σ²)`, and its Jeffreys +prior. -/ +theorem gaussianTS_eq_expFamilyTS (m : α → ℝ) {σ2 : ℝ≥0} (hσ2 : σ2 ≠ 0) : + gaussianTS σ2 = expFamilyTS (fun a ↦ gaussianReal (m a) σ2) + (fun a ↦ (gaussianReal (m a) σ2).jeffreysPrior id) := by + rw [gaussianTS_def, expFamilyTS_def] + congr with n x a + rw [gaussianTSLaw_def, expFamilyTSLaw_def, gaussianTSStackLaw_eq_expFamilyTSStackLaw (m a) hσ2] + namespace GaussianTS variable {σ2 : ℝ≥0} {Ω : Type*} {mΩ : MeasurableSpace Ω} {env : Environment Unit α ℝ} diff --git a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean index 060a7335..c701631f 100644 --- a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean +++ b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean @@ -7,6 +7,7 @@ module public import LMLExplore.Probability.Distributions.ExponentialFamily.Gamma public import LMLExplore.Probability.Distributions.ExponentialFamily.Invariance +public import LMLExplore.Probability.Distributions.ExponentialFamily.Posterior public import LMLExplore.Probability.Distributions.ExponentialFamily.UpperConfidence public import LMLExplore.Probability.Distributions.Gaussian.Real public import TauCeti.Probability.Distributions.Tilted @@ -28,6 +29,10 @@ parametrization is `d(μ, μ') = (μ - μ')² / (2 v)`. Its extension `dbar` to expression, and the upper confidence level of KL-UCB is the Hoeffding-type level `U(x, c) = x + √(2 v c)` of asymptotically optimal UCB (T. Lattimore and Cs. Szepesvári, *Bandit Algorithms*, Cambridge University Press, 2020, Algorithm 6). +The variance function `ψ'' = v` is constant, so the Jeffreys prior is `√v` times Lebesgue +measure, and the posterior of the natural parameter after `s ≠ 0` observations of sum `S` is the +Gaussian law `𝒩((S - s m) / (s v), 1 / (s v))`, whose image by the mean map is `𝒩(S / s, v / s)`: +Thompson sampling with the Jeffreys prior samples the mean of an arm from `𝒩(μ̂, v / s)`. Gaussian arms with a known variance are the reference model of bandit identification: with `d(μ, μ') = (μ - μ')² / (2 v)`, the transportation costs, the characteristic time and the @@ -75,6 +80,15 @@ Fixed variance, statistic `id`: `v ≠ 0`, the upper confidence level is `U(x, c) = x + √(2 v c)` for finite `c`, and `U(x, ∞) = ∞`. `LMLExplore.expFamilyUpper_ofReal_gaussianReal` is the finite case for a real budget `c ≥ 0`. +* `LMLExplore.iteratedDeriv_two_cgf_gaussianReal`, `LMLExplore.jeffreysPrior_gaussianReal`: the + variance function of the family is the constant `v`, and its Jeffreys prior is `√v` times + Lebesgue measure. +* `LMLExplore.expFamilyPosterior_gaussianReal_volume`, + `LMLExplore.expFamilyPosterior_gaussianReal_jeffreysPrior`: for `v ≠ 0` and `s ≠ 0` + observations of sum `S`, the posterior of the natural parameter for the flat prior and for the + Jeffreys prior is `𝒩((S - s m) / (s v), 1 / (s v))`. +* `LMLExplore.map_deriv_cgf_expFamilyPosterior_gaussianReal_jeffreysPrior`: the image of that + posterior by the mean map `ψ'` is `𝒩(S / s, v / s)`. Fixed mean, statistic `x ↦ (x - m) ^ 2`: @@ -98,7 +112,7 @@ public section open MeasureTheory ProbabilityTheory Real Set -open scoped NNReal +open scoped ENNReal NNReal namespace LMLExplore @@ -235,6 +249,89 @@ theorem expFamilyUpper_top_gaussianReal (m : ℝ) (hv : v ≠ 0) (x : ℝ) : exact Measure.le_expFamilyUpper_of_expFamilyRate_le (mem_expFamilyMeans_gaussianReal m hv _) ((le_max_left x y).trans (le_add_of_nonneg_right zero_le_one)) le_top +/-- The variance function of the exponential family of `𝒩(m, v)` is constant: the second +derivative of its log-partition function is `v` at every natural parameter. -/ +@[simp] +theorem iteratedDeriv_two_cgf_gaussianReal (m : ℝ) (v : ℝ≥0) (η : ℝ) : + iteratedDeriv 2 (cgf id (gaussianReal m v)) η = v := by + rw [iteratedDeriv_succ, iteratedDeriv_one, + show deriv (cgf id (gaussianReal m v)) = fun η ↦ m + v * η from + funext (deriv_cgf_gaussianReal m v)] + simp + +/-- The Jeffreys prior of the exponential family of `𝒩(m, v)` is `√v` times Lebesgue measure. -/ +theorem jeffreysPrior_gaussianReal (m : ℝ) (v : ℝ≥0) : + (gaussianReal m v).jeffreysPrior id = ENNReal.ofReal √v • volume := by + simp [Measure.jeffreysPrior_def, integrableExpSet_id_gaussianReal, withDensity_const] + +/-- For `v ≠ 0` and `s ≠ 0` observations whose sum is `S`, the posterior of the natural parameter +of the exponential family of `𝒩(m, v)` for the flat prior (Lebesgue measure) is the Gaussian law +`𝒩((S - s m) / (s v), 1 / (s v))`. -/ +theorem expFamilyPosterior_gaussianReal_volume (m : ℝ) (hv : v ≠ 0) {s : ℕ} (hs : s ≠ 0) + (S : ℝ) : + (gaussianReal m v).expFamilyPosterior id volume s S = + gaussianReal ((S - s * m) / (s * v)) (s * v)⁻¹ := by + have hsv : (s : ℝ≥0) * v ≠ 0 := mul_ne_zero (Nat.cast_ne_zero.2 hs) hv + have ha : (0 : ℝ) < s * v := by positivity + set μ₀ : ℝ := (S - s * m) / (s * v) + -- the unnormalized posterior is a constant multiple `C` of `𝒩(μ₀, 1 / (s v))`, by completing the + -- square `θ S - s (m θ + v θ² / 2) = s v μ₀² / 2 - s v (θ - μ₀)² / 2` + set C : ℝ≥0∞ := ENNReal.ofReal (√(2 * π * (s * v : ℝ)⁻¹) * exp (s * v * μ₀ ^ 2 / 2)) + have hC0 : C ≠ 0 := by + simp only [C, ne_eq, ENNReal.ofReal_eq_zero, not_le] + positivity + have hdens : (fun θ ↦ ENNReal.ofReal (exp (θ * S - s * cgf id (gaussianReal m v) θ))) = + C • gaussianPDF μ₀ ((s : ℝ≥0) * v)⁻¹ := by + ext θ + have h2 : 0 < 2 * π * (s * v : ℝ)⁻¹ := by positivity + rw [cgf_gaussianReal HasLaw.id, Pi.smul_apply, smul_eq_mul, gaussianPDF, gaussianPDFReal, + ← ENNReal.ofReal_mul (by positivity)] + congr 1 + push_cast + rw [mul_mul_mul_comm, mul_inv_cancel₀ (Real.sqrt_pos.2 h2).ne', one_mul, ← exp_add] + congr 1 + simp only [μ₀] + field_simp + ring + rw [Measure.expFamilyPosterior_def, integrableExpSet_id_gaussianReal, interior_univ, hdens, + withDensity_smul' C _ ENNReal.ofReal_ne_top, ← gaussianReal_of_var_ne_zero _ (inv_ne_zero hsv), + ProbabilityTheory.cond, Measure.restrict_univ, smul_apply, measure_univ, smul_eq_mul, + mul_one, smul_smul, ENNReal.inv_mul_cancel hC0 ENNReal.ofReal_ne_top, one_smul] + +/-- For `v ≠ 0` and `s ≠ 0` observations whose sum is `S`, the posterior of the natural parameter +of the exponential family of `𝒩(m, v)` for its Jeffreys prior is the Gaussian law +`𝒩((S - s m) / (s v), 1 / (s v))`: the Jeffreys prior is a constant multiple of the flat prior. -/ +theorem expFamilyPosterior_gaussianReal_jeffreysPrior (m : ℝ) (hv : v ≠ 0) {s : ℕ} (hs : s ≠ 0) + (S : ℝ) : + (gaussianReal m v).expFamilyPosterior id ((gaussianReal m v).jeffreysPrior id) s S = + gaussianReal ((S - s * m) / (s * v)) (s * v)⁻¹ := by + have hv' : (0 : ℝ) < v := by positivity + rw [jeffreysPrior_gaussianReal, Measure.expFamilyPosterior_smul + (by simpa using Real.sqrt_pos.2 hv') ENNReal.ofReal_ne_top, + expFamilyPosterior_gaussianReal_volume m hv hs] + +/-- For `v ≠ 0` and `s ≠ 0` observations whose sum is `S`, the image of the posterior of the +natural parameter of the exponential family of `𝒩(m, v)` for its Jeffreys prior by the mean map +`ψ'` is the Gaussian law `𝒩(S / s, v / s)`: the mean of a natural parameter drawn from the +posterior is distributed as `𝒩(μ̂, v / s)`, with `μ̂ = S / s` the empirical mean. -/ +theorem map_deriv_cgf_expFamilyPosterior_gaussianReal_jeffreysPrior (m : ℝ) (hv : v ≠ 0) {s : ℕ} + (hs : s ≠ 0) (S : ℝ) : + ((gaussianReal m v).expFamilyPosterior id ((gaussianReal m v).jeffreysPrior id) s S).map + (deriv (cgf id (gaussianReal m v))) = gaussianReal (S / s) (v / s) := by + have hv' : (v : ℝ) ≠ 0 := NNReal.coe_ne_zero.2 hv + have hs' : (s : ℝ) ≠ 0 := Nat.cast_ne_zero.2 hs + rw [expFamilyPosterior_gaussianReal_jeffreysPrior m hv hs, + show deriv (cgf id (gaussianReal m v)) = (fun x ↦ x + m) ∘ (fun x ↦ (v : ℝ) * x) from + funext fun η ↦ by rw [deriv_cgf_gaussianReal, Function.comp_apply, add_comm], + ← Measure.map_map (by fun_prop) (by fun_prop), gaussianReal_map_const_mul, + gaussianReal_map_add_const] + congr 1 + · field_simp + ring + · ext + push_cast + field_simp + /-! ### Fixed mean: the statistic `x ↦ (x - m) ^ 2` -/ section SubSq diff --git a/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean b/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean new file mode 100644 index 00000000..50d4e9cc --- /dev/null +++ b/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean @@ -0,0 +1,180 @@ +/- +Copyright (c) 2026 Rémy Degenne. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Rémy Degenne +-/ +module + +public import Mathlib.Probability.ConditionalProbability +public import LMLExplore.Probability.Distributions.ExponentialFamily.Basic +public import LMLExplore.Probability.Moments.Basic + +/-! +# Posteriors of the natural parameter of an exponential family, and the Jeffreys prior + +Let `ρ` be a measure on `𝓧` and `T : 𝓧 → ℝ` a statistic, with log-partition function +`ψ = cgf T ρ` and natural parameters `Θ`, the interior of `integrableExpSet T ρ`. After `s` +independent observations `x₁, …, x_s` from the member `ρ_θ = ρ.expFamily T θ` of the family, the +density of their joint law with respect to `ρ^{⊗ s}` is +`∏ᵢ exp (θ T(xᵢ) - ψ(θ)) = exp (θ S - s ψ(θ))`, with `S = ∑ᵢ T(xᵢ)` the sum of the statistics. It +depends on the observations only through `s` and `S`. For a prior `ν` on the natural parameter, +the posterior is therefore the probability measure proportional to `exp (θ S - s ψ(θ)) ν(dθ)` on +`Θ`, the conjugate update of [KKM13, Section 2]. It is defined for every measure `ν` on `ℝ`, which +may be infinite (an improper prior, such as Lebesgue measure): the posterior is a probability +measure when the normalizing constant `∫_Θ exp (θ S - s ψ(θ)) dν(θ)` is positive and finite, and +`0` otherwise. Mathlib's `ProbabilityTheory.posterior` is the posterior kernel of a finite prior, +determined only up to null sets of observations; the conjugate formula here is defined at every +value of the sufficient statistic `(s, S)` and for improper priors, as a sampling rule requires. + +The Jeffreys prior of the family is `√ψ''(θ) dθ` on `Θ`, the square root of the Fisher +information of the natural parameter times Lebesgue measure [KKM13, Section 2]. It is +σ-finite but usually infinite: for the Gaussian family of variance `v`, `ψ'' = v` and the Jeffreys +prior is a constant multiple of Lebesgue measure. + +These posteriors are the sampling laws of Thompson sampling for exponential families. + +## Main definitions + +* `MeasureTheory.Measure.expFamilyPosterior ρ T ν s S`: the posterior of the natural parameter for + the prior `ν`, after `s` observations whose statistics sum to `S`. +* `MeasureTheory.Measure.jeffreysPrior ρ T`: the Jeffreys prior `√ψ''(θ) dθ` on `Θ`. + +## Main statements + +* `MeasureTheory.Measure.expFamilyPosterior_apply`: the posterior of a measurable set `B` is + `∫_{Θ ∩ B} exp (θ S - s ψ(θ)) dν(θ) / ∫_Θ exp (θ S - s ψ(θ)) dν(θ)`. +* `MeasureTheory.Measure.isProbabilityMeasure_expFamilyPosterior`, + `MeasureTheory.Measure.expFamilyPosterior_eq_zero_iff`: the posterior is a probability measure + if the normalizing constant is positive and finite, and `0` otherwise; + `MeasureTheory.Measure.isProbabilityMeasure_expFamilyPosterior_of_ne_zero`: a posterior that is + not `0` is a probability measure. +* `MeasureTheory.Measure.expFamilyPosterior_smul`: the posterior does not change when the prior is + multiplied by a positive finite constant. +* `MeasureTheory.Measure.measurable_expFamilyPosterior`: for an s-finite measure `ρ`, a measurable + statistic and an s-finite prior, the posterior is a measurable function of `S`. + +## References + +* [KKM13] N. Korda, E. Kaufmann and R. Munos, *Thompson sampling for 1-dimensional exponential + family bandits*, NeurIPS 2013. +-/ + +public section + +open ProbabilityTheory Real Set + +open scoped ENNReal + +namespace MeasureTheory.Measure + +variable {𝓧 : Type*} {m𝓧 : MeasurableSpace 𝓧} {ρ : Measure 𝓧} {T : 𝓧 → ℝ} {ν : Measure ℝ} + {s : ℕ} {S : ℝ} + +/-- The posterior of the natural parameter of the exponential family of `ρ` and `T`, for the prior +`ν`, after `s` observations whose statistics sum to `S`: the measure +`exp (θ S - s ψ(θ)) ν(dθ)` with `ψ = cgf T ρ`, conditioned on the interior of the natural +parameters `integrableExpSet T ρ` (that is, restricted to it and normalized). It is a probability +measure if the normalizing constant is positive and finite +(`isProbabilityMeasure_expFamilyPosterior`) and `0` otherwise (`expFamilyPosterior_eq_zero_iff`). +-/ +noncomputable def expFamilyPosterior (ρ : Measure 𝓧) (T : 𝓧 → ℝ) (ν : Measure ℝ) (s : ℕ) + (S : ℝ) : Measure ℝ := + (ν.withDensity fun θ ↦ ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)))[| + interior (integrableExpSet T ρ)] + +/-- The defining equation of `expFamilyPosterior`. -/ +lemma expFamilyPosterior_def (ρ : Measure 𝓧) (T : 𝓧 → ℝ) (ν : Measure ℝ) (s : ℕ) (S : ℝ) : + ρ.expFamilyPosterior T ν s S = + (ν.withDensity fun θ ↦ ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)))[| + interior (integrableExpSet T ρ)] := + (rfl) + +instance : IsZeroOrProbabilityMeasure (ρ.expFamilyPosterior T ν s S) := by + rw [expFamilyPosterior_def] + infer_instance + +/-- The posterior of a measurable set `B` is the ratio +`∫_{Θ ∩ B} exp (θ S - s ψ(θ)) dν(θ) / ∫_Θ exp (θ S - s ψ(θ)) dν(θ)`, with `Θ` the interior of +`integrableExpSet T ρ`. -/ +lemma expFamilyPosterior_apply {B : Set ℝ} (hB : MeasurableSet B) : + ρ.expFamilyPosterior T ν s S B = + (∫⁻ θ in interior (integrableExpSet T ρ), ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν)⁻¹ + * ∫⁻ θ in interior (integrableExpSet T ρ) ∩ B, + ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν := by + have hΘ : MeasurableSet (interior (integrableExpSet T ρ)) := isOpen_interior.measurableSet + rw [expFamilyPosterior_def, cond_apply hΘ, withDensity_apply _ hΘ, + withDensity_apply _ (hΘ.inter hB)] + +/-- The posterior is a probability measure if its normalizing constant +`∫_Θ exp (θ S - s ψ(θ)) dν(θ)` is positive and finite. -/ +lemma isProbabilityMeasure_expFamilyPosterior + (h0 : ∫⁻ θ in interior (integrableExpSet T ρ), + ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν ≠ 0) + (htop : ∫⁻ θ in interior (integrableExpSet T ρ), + ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν ≠ ∞) : + IsProbabilityMeasure (ρ.expFamilyPosterior T ν s S) := by + have hΘ : MeasurableSet (interior (integrableExpSet T ρ)) := isOpen_interior.measurableSet + rw [expFamilyPosterior_def] + exact cond_isProbabilityMeasure_of_finite (by rwa [withDensity_apply _ hΘ]) + (by rwa [withDensity_apply _ hΘ]) + +/-- A posterior that is not `0` is a probability measure. -/ +lemma isProbabilityMeasure_expFamilyPosterior_of_ne_zero (h : ρ.expFamilyPosterior T ν s S ≠ 0) : + IsProbabilityMeasure (ρ.expFamilyPosterior T ν s S) := + ⟨IsZeroOrProbabilityMeasure.measure_univ.resolve_left (by rwa [Measure.measure_univ_eq_zero])⟩ + +/-- The posterior is `0` exactly when its normalizing constant `∫_Θ exp (θ S - s ψ(θ)) dν(θ)` is +`0` or infinite, for example for an improper prior and too few observations. -/ +lemma expFamilyPosterior_eq_zero_iff : + ρ.expFamilyPosterior T ν s S = 0 ↔ + ∫⁻ θ in interior (integrableExpSet T ρ), + ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν = ∞ ∨ + ∫⁻ θ in interior (integrableExpSet T ρ), + ENNReal.ofReal (exp (θ * S - s * cgf T ρ θ)) ∂ν = 0 := by + rw [expFamilyPosterior_def, cond_eq_zero, + withDensity_apply _ isOpen_interior.measurableSet] + +/-- The posterior does not depend on the normalization of the prior: multiplying the prior by a +constant `c ∉ {0, ∞}` does not change the posterior. -/ +lemma expFamilyPosterior_smul {c : ℝ≥0∞} (hc0 : c ≠ 0) (hc : c ≠ ∞) : + ρ.expFamilyPosterior T (c • ν) s S = ρ.expFamilyPosterior T ν s S := by + ext B hB + rw [expFamilyPosterior_apply hB, expFamilyPosterior_apply hB, Measure.restrict_smul, + Measure.restrict_smul, lintegral_smul_measure, lintegral_smul_measure, smul_eq_mul, + smul_eq_mul, ENNReal.mul_inv (Or.inl hc0) (Or.inl hc), mul_mul_mul_comm, + ENNReal.inv_mul_cancel hc0 hc, one_mul] + +/-- For an s-finite measure `ρ`, an almost-everywhere measurable statistic `T` and an s-finite +prior `ν`, the posterior after `s` observations is a measurable function of the sum `S` of their +statistics. -/ +@[fun_prop] +lemma measurable_expFamilyPosterior [SFinite ρ] (hT : AEMeasurable T ρ) (ν : Measure ℝ) + [SFinite ν] (s : ℕ) : + Measurable fun S ↦ ρ.expFamilyPosterior T ν s S := by + refine Measure.measurable_of_measurable_coe _ fun B hB ↦ ?_ + simp_rw [expFamilyPosterior_apply hB] + have hf : Measurable fun p : ℝ × ℝ ↦ ENNReal.ofReal (exp (p.2 * p.1 - s * cgf T ρ p.2)) := by + have := measurable_cgf hT + fun_prop + exact (hf.lintegral_prod_right' (ν := ν.restrict _)).inv.mul + (hf.lintegral_prod_right' (ν := ν.restrict _)) + +/-- The Jeffreys prior of the exponential family of `ρ` and `T`: the measure `√ψ''(θ) dθ` on the +interior of the natural parameters `integrableExpSet T ρ`, with `ψ = cgf T ρ` and `ψ''(θ)` the +variance of `T` under the member of natural parameter `θ`, the Fisher information of the natural +parameter [KKM13, Section 2]. It is σ-finite and usually infinite. -/ +noncomputable def jeffreysPrior (ρ : Measure 𝓧) (T : 𝓧 → ℝ) : Measure ℝ := + (volume.restrict (interior (integrableExpSet T ρ))).withDensity + fun θ ↦ ENNReal.ofReal √(iteratedDeriv 2 (cgf T ρ) θ) + +/-- The defining equation of `jeffreysPrior`. -/ +lemma jeffreysPrior_def (ρ : Measure 𝓧) (T : 𝓧 → ℝ) : + ρ.jeffreysPrior T = (volume.restrict (interior (integrableExpSet T ρ))).withDensity + fun θ ↦ ENNReal.ofReal √(iteratedDeriv 2 (cgf T ρ) θ) := + (rfl) + +instance : SigmaFinite (ρ.jeffreysPrior T) := by + rw [jeffreysPrior_def] + infer_instance + +end MeasureTheory.Measure diff --git a/LMLExplore/Probability/Moments/Basic.lean b/LMLExplore/Probability/Moments/Basic.lean index 7c497403..8ae85da4 100644 --- a/LMLExplore/Probability/Moments/Basic.lean +++ b/LMLExplore/Probability/Moments/Basic.lean @@ -5,6 +5,7 @@ Authors: Rémy Degenne -/ module +public import Mathlib.MeasureTheory.Integral.Prod public import Mathlib.Probability.Moments.Basic /-! @@ -18,6 +19,9 @@ variables, shared by the concentration inequalities of the library. * `ProbabilityTheory.iIndepFun.mgf_sum_eq_prod_of_mem`: the moment-generating function of a sum of independent random variables over a `Finset` is the product of their moment-generating functions, assuming measurability only on the `Finset`. +* `ProbabilityTheory.measurable_mgf`, `ProbabilityTheory.measurable_cgf`: for an s-finite measure, + the moment-generating and cumulant generating functions are measurable functions of their + argument. * `ProbabilityTheory.measureReal_ge_le_exp_of_mgf_le`: Chernoff's bound `P(ε ≤ Y) ≤ exp (-t ε + K)` when `exp (t Y)` is integrable, `mgf Y μ t ≤ exp K` and `t ≥ 0`. * `MeasureTheory.measureReal_abs_ge_le_add`: the union bound over both tails, @@ -58,6 +62,23 @@ theorem iIndepFun.mgf_sum_eq_prod_of_mem {X : ι → Ω → ℝ} {s : Finset ι} congr with ω rw [Finset.sum_apply, Finset.sum_coe_sort s fun i ↦ X i ω] +/-- The moment-generating function of an almost-everywhere measurable random variable under an +s-finite measure is a measurable function of its argument. -/ +@[fun_prop] +theorem measurable_mgf [SFinite μ] {X : Ω → ℝ} (hX : AEMeasurable X μ) : Measurable (mgf X μ) := by + have h_eq : mgf X μ = mgf (hX.mk X) μ := funext fun _ ↦ mgf_congr hX.ae_eq_mk + rw [h_eq] + have h : Measurable fun p : ℝ × Ω ↦ exp (p.1 * hX.mk X p.2) := + measurable_exp.comp (measurable_fst.mul (hX.measurable_mk.comp measurable_snd)) + exact (h.stronglyMeasurable.integral_prod_right' (ν := μ)).measurable + +/-- The cumulant generating function of an almost-everywhere measurable random variable under an +s-finite measure is a measurable function of its argument. -/ +@[fun_prop] +theorem measurable_cgf [SFinite μ] {X : Ω → ℝ} (hX : AEMeasurable X μ) : Measurable (cgf X μ) := by + unfold cgf + exact (measurable_mgf hX).log + /-- **Chernoff's bound**: if `exp (t Y)` is integrable and the moment-generating function of `Y` at `t ≥ 0` is at most `exp K`, then `P(ε ≤ Y) ≤ exp (-t ε + K)`. -/ theorem measureReal_ge_le_exp_of_mgf_le [IsFiniteMeasure μ] {Y : Ω → ℝ} {t ε K : ℝ} From 09364c9fc8c5e0ba8df2ee1e6a8c92a3041e0c37 Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 12 Nov 2026 06:30:33 +0100 Subject: [PATCH 2/3] refactor: factor conditioning under scaling and the Gaussian mean map out of the posterior proofs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ProbabilityTheory.cond_smul ((c • μ)[|s] = μ[|s] for c ∉ {0, ∞}) and LMLExplore.deriv_cgf_gaussianReal_eq (the function-level Gaussian mean map), and use them in expFamilyPosterior_smul, expFamilyPosterior_gaussianReal_volume, iteratedDeriv_two_cgf_gaussianReal and map_deriv_cgf_expFamilyPosterior_gaussianReal_jeffreysPrior. Co-Authored-By: Claude Code --- .../Probability/ConditionalProbability.lean | 39 +++++++++++++++++++ .../ExponentialFamily/Gaussian.lean | 20 ++++++---- .../ExponentialFamily/Posterior.lean | 9 ++--- 3 files changed, 54 insertions(+), 14 deletions(-) create mode 100644 LMLExplore/Probability/ConditionalProbability.lean diff --git a/LMLExplore/Probability/ConditionalProbability.lean b/LMLExplore/Probability/ConditionalProbability.lean new file mode 100644 index 00000000..cdab34da --- /dev/null +++ b/LMLExplore/Probability/ConditionalProbability.lean @@ -0,0 +1,39 @@ +/- +Copyright (c) 2026 Rémy Degenne. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Rémy Degenne +-/ +module + +public import Mathlib.Probability.ConditionalProbability + +/-! +# Conditional probability measures + +Complements to Mathlib's `ProbabilityTheory.cond`. + +## Main statements + +* `ProbabilityTheory.cond_smul`: conditioning a measure on a set does not change when the measure + is multiplied by a constant `c ∉ {0, ∞}`. +-/ + +public section + +open MeasureTheory + +open scoped ENNReal + +namespace ProbabilityTheory + +variable {Ω : Type*} {mΩ : MeasurableSpace Ω} + +/-- Conditioning a measure on a set does not change when the measure is multiplied by a constant +`c ∉ {0, ∞}`. -/ +lemma cond_smul (μ : Measure Ω) (s : Set Ω) {c : ℝ≥0∞} (hc0 : c ≠ 0) (hc : c ≠ ∞) : + (c • μ)[|s] = μ[|s] := by + rw [cond, cond, Measure.restrict_smul, smul_apply, smul_eq_mul, smul_smul, + ENNReal.mul_inv (Or.inl hc0) (Or.inl hc), mul_right_comm, ENNReal.inv_mul_cancel hc0 hc, + one_mul] + +end ProbabilityTheory diff --git a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean index c701631f..7514d883 100644 --- a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean +++ b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean @@ -63,7 +63,8 @@ Fixed variance, statistic `id`: * `LMLExplore.expFamily_gaussianReal`: the member of natural parameter `η` of the exponential family of `𝒩(m, v)` is `𝒩(m + v η, v)`. -* `LMLExplore.deriv_cgf_gaussianReal`: the mean map of the family is `η ↦ m + v η`. +* `LMLExplore.deriv_cgf_gaussianReal`, `LMLExplore.deriv_cgf_gaussianReal_eq`: the mean map of the + family is `η ↦ m + v η`. * `LMLExplore.not_ae_eq_const_gaussianReal`: for `v ≠ 0`, the identity is not almost everywhere constant under `𝒩(m, v)`, the nondegeneracy assumption of the mean parametrization. * `LMLExplore.expFamilyMeans_gaussianReal`, `LMLExplore.mem_expFamilyMeans_gaussianReal`, @@ -141,6 +142,11 @@ theorem deriv_cgf_gaussianReal (m : ℝ) (v : ℝ≥0) (η : ℝ) : rw [h.deriv] ring +/-- The mean map of the exponential family of `𝒩(m, v)` is the function `η ↦ m + v η`. -/ +theorem deriv_cgf_gaussianReal_eq (m : ℝ) (v : ℝ≥0) : + deriv (cgf id (gaussianReal m v)) = fun η ↦ m + v * η := + funext (deriv_cgf_gaussianReal m v) + /-- For `v ≠ 0`, the identity is not almost everywhere constant under `𝒩(m, v)`: the Gaussian law has no atoms. This is the nondegeneracy assumption of the mean parametrization of the exponential family of `𝒩(m, v)`. -/ @@ -254,9 +260,7 @@ derivative of its log-partition function is `v` at every natural parameter. -/ @[simp] theorem iteratedDeriv_two_cgf_gaussianReal (m : ℝ) (v : ℝ≥0) (η : ℝ) : iteratedDeriv 2 (cgf id (gaussianReal m v)) η = v := by - rw [iteratedDeriv_succ, iteratedDeriv_one, - show deriv (cgf id (gaussianReal m v)) = fun η ↦ m + v * η from - funext (deriv_cgf_gaussianReal m v)] + rw [iteratedDeriv_succ, iteratedDeriv_one, deriv_cgf_gaussianReal_eq] simp /-- The Jeffreys prior of the exponential family of `𝒩(m, v)` is `√v` times Lebesgue measure. -/ @@ -295,8 +299,7 @@ theorem expFamilyPosterior_gaussianReal_volume (m : ℝ) (hv : v ≠ 0) {s : ℕ ring rw [Measure.expFamilyPosterior_def, integrableExpSet_id_gaussianReal, interior_univ, hdens, withDensity_smul' C _ ENNReal.ofReal_ne_top, ← gaussianReal_of_var_ne_zero _ (inv_ne_zero hsv), - ProbabilityTheory.cond, Measure.restrict_univ, smul_apply, measure_univ, smul_eq_mul, - mul_one, smul_smul, ENNReal.inv_mul_cancel hC0 ENNReal.ofReal_ne_top, one_smul] + cond_smul _ _ hC0 ENNReal.ofReal_ne_top, cond_univ] /-- For `v ≠ 0` and `s ≠ 0` observations whose sum is `S`, the posterior of the natural parameter of the exponential family of `𝒩(m, v)` for its Jeffreys prior is the Gaussian law @@ -321,8 +324,9 @@ theorem map_deriv_cgf_expFamilyPosterior_gaussianReal_jeffreysPrior (m : ℝ) (h have hv' : (v : ℝ) ≠ 0 := NNReal.coe_ne_zero.2 hv have hs' : (s : ℝ) ≠ 0 := Nat.cast_ne_zero.2 hs rw [expFamilyPosterior_gaussianReal_jeffreysPrior m hv hs, - show deriv (cgf id (gaussianReal m v)) = (fun x ↦ x + m) ∘ (fun x ↦ (v : ℝ) * x) from - funext fun η ↦ by rw [deriv_cgf_gaussianReal, Function.comp_apply, add_comm], + deriv_cgf_gaussianReal_eq, + show (fun η ↦ m + v * η) = (fun x ↦ x + m) ∘ (fun x ↦ (v : ℝ) * x) from + funext fun η ↦ add_comm _ _, ← Measure.map_map (by fun_prop) (by fun_prop), gaussianReal_map_const_mul, gaussianReal_map_add_const] congr 1 diff --git a/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean b/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean index 50d4e9cc..c19b7036 100644 --- a/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean +++ b/LMLExplore/Probability/Distributions/ExponentialFamily/Posterior.lean @@ -5,7 +5,7 @@ Authors: Rémy Degenne -/ module -public import Mathlib.Probability.ConditionalProbability +public import LMLExplore.Probability.ConditionalProbability public import LMLExplore.Probability.Distributions.ExponentialFamily.Basic public import LMLExplore.Probability.Moments.Basic @@ -138,11 +138,8 @@ lemma expFamilyPosterior_eq_zero_iff : constant `c ∉ {0, ∞}` does not change the posterior. -/ lemma expFamilyPosterior_smul {c : ℝ≥0∞} (hc0 : c ≠ 0) (hc : c ≠ ∞) : ρ.expFamilyPosterior T (c • ν) s S = ρ.expFamilyPosterior T ν s S := by - ext B hB - rw [expFamilyPosterior_apply hB, expFamilyPosterior_apply hB, Measure.restrict_smul, - Measure.restrict_smul, lintegral_smul_measure, lintegral_smul_measure, smul_eq_mul, - smul_eq_mul, ENNReal.mul_inv (Or.inl hc0) (Or.inl hc), mul_mul_mul_comm, - ENNReal.inv_mul_cancel hc0 hc, one_mul] + rw [expFamilyPosterior_def, expFamilyPosterior_def, withDensity_smul_measure, + cond_smul _ _ hc0 hc] /-- For an s-finite measure `ρ`, an almost-everywhere measurable statistic `T` and an s-finite prior `ν`, the posterior after `s` observations is a measurable function of the sum `S` of their From 8d5088b5a73c947b6c00312c37318d0e23031168 Mon Sep 17 00:00:00 2001 From: Remy Degenne Date: Thu, 12 Nov 2026 06:55:11 +0100 Subject: [PATCH 3/3] docs: state the nonzero-variance hypothesis in the Gaussian Thompson sampling module prose Co-Authored-By: Claude Code --- LMLExplore/Bandits/ThompsonSampling/Basic.lean | 2 +- .../Distributions/ExponentialFamily/Gaussian.lean | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/LMLExplore/Bandits/ThompsonSampling/Basic.lean b/LMLExplore/Bandits/ThompsonSampling/Basic.lean index 33a0dc43..c6008fee 100644 --- a/LMLExplore/Bandits/ThompsonSampling/Basic.lean +++ b/LMLExplore/Bandits/ThompsonSampling/Basic.lean @@ -27,7 +27,7 @@ This is a follow-the-perturbed-leader policy (`perturbedLeaderAlgorithm`) whose laws, in `EReal`, depend on the history only through the number of pulls and the sum of rewards of each arm, so [LS, Theorem 36.2] applies to it in the reward-stack form of `PerturbedLeader.lintegral_pullCount_le`. With the Jeffreys prior `(ρ a).jeffreysPrior id` it is -the algorithm of [KKM13]; for the Gaussian family of variance `σ²` the Jeffreys prior is a +the algorithm of [KKM13]; for the Gaussian family of variance `σ² ≠ 0` the Jeffreys prior is a constant multiple of Lebesgue measure, and Thompson sampling is Gaussian Thompson sampling `gaussianTS σ2` (`gaussianTS_eq_expFamilyTS`). diff --git a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean index 7514d883..5d7ed455 100644 --- a/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean +++ b/LMLExplore/Probability/Distributions/ExponentialFamily/Gaussian.lean @@ -30,9 +30,10 @@ expression, and the upper confidence level of KL-UCB is the Hoeffding-type level `U(x, c) = x + √(2 v c)` of asymptotically optimal UCB (T. Lattimore and Cs. Szepesvári, *Bandit Algorithms*, Cambridge University Press, 2020, Algorithm 6). The variance function `ψ'' = v` is constant, so the Jeffreys prior is `√v` times Lebesgue -measure, and the posterior of the natural parameter after `s ≠ 0` observations of sum `S` is the -Gaussian law `𝒩((S - s m) / (s v), 1 / (s v))`, whose image by the mean map is `𝒩(S / s, v / s)`: -Thompson sampling with the Jeffreys prior samples the mean of an arm from `𝒩(μ̂, v / s)`. +measure. For `v ≠ 0`, the posterior of the natural parameter after `s ≠ 0` observations of sum +`S` is the Gaussian law `𝒩((S - s m) / (s v), 1 / (s v))`, whose image by the mean map is +`𝒩(S / s, v / s)`: Thompson sampling with the Jeffreys prior samples the mean of an arm from +`𝒩(μ̂, v / s)`. Gaussian arms with a known variance are the reference model of bandit identification: with `d(μ, μ') = (μ - μ')² / (2 v)`, the transportation costs, the characteristic time and the