RegressionGaussian.lean

1 import KomorebiLSM.RegressionGeometry
2 import Mathlib.Analysis.SpecialFunctions.Gaussian.GaussianIntegral
3 import Mathlib.MeasureTheory.Measure.Haar.NormedSpace
4 import Mathlib.MeasureTheory.Integral.Pi
5 import Mathlib.Analysis.Matrix.PosDef
6 import Mathlib.MeasureTheory.Measure.Lebesgue.EqHaar
7 import Mathlib.MeasureTheory.Integral.Bochner.ContinuousLinearMap
8 import Mathlib.Tactic
9 
10 /-! Actual Gaussian quadratic integrals for the ridge-mixture argument. -/
11 set_option autoImplicit false
12 namespace Komorebi.LSM.RegressionConfidence
13 open MeasureTheory
14 open Matrix
15 open scoped BigOperators ENNReal ComplexOrder
16 
17 /-- Completing the square gives the actual one-dimensional Gaussian tilt. -/
18 theorem integral_exp_quadratic (b c : ℝ) (hb : 0 < b) :
19     (∫ x : ℝ,Real.exp (-b*x^2+c*x)) =
20       Real.sqrt (Real.pi/b)*Real.exp (c^2/(4*b)) := by
21   have hid (x : ℝ) : -b*x^2+c*x = -b*(x+(-c/(2*b)))^2+c^2/(4*b) := by
22     field_simp
23     ring
24   simp_rw [hid,Real.exp_add]
25   rw [integral_mul_const,integral_add_right_eq_self (fun x : ℝ => Real.exp (-b*x^2))
26     (-c/(2*b)),integral_gaussian]
27 
28 /-- The completed-square density is genuinely Bochner integrable. -/
29 theorem integrable_exp_quadratic (b c : ℝ) (hb : 0 < b) :
30     Integrable (fun x : ℝ => Real.exp (-b*x^2+c*x)) := by
31   have hid (x : ℝ) : -b*x^2+c*x = -b*(x+(-c/(2*b)))^2+c^2/(4*b) := by
32     field_simp
33     ring
34   simp_rw [hid,Real.exp_add]
35   exact ((integrable_exp_neg_mul_sq hb).comp_add_right (-c/(2*b))).mul_const _
36 
37 /-- The actual finite-product quadratic Gaussian integral. -/
38 theorem integral_exp_diagonal {I : Type*} [Fintype I] (b c : I → ℝ)
39     (hb : ∀ i,0 < b i) :
40     (∫ x : I → ℝ,Real.exp (-(∑ i,b i*(x i)^2)+∑ i,c i*x i)) =
41       ∏ i,Real.sqrt (Real.pi/b i)*Real.exp ((c i)^2/(4*b i)) := by
42   simp_rw [← Finset.sum_neg_distrib,← Finset.sum_add_distrib,Real.exp_sum]
43   rw [integral_fintype_prod_volume_eq_prod (f := fun i (x : ℝ) =>
44     Real.exp (-(b i*x^2)+c i*x))]
45   congr with i
46   simpa only [neg_mul] using integral_exp_quadratic (b i) (c i) (hb i)
47 
48 /-- Strict positivity of the computed integral supplies its Bochner integrability. -/
49 theorem integrable_exp_diagonal {I : Type*} [Fintype I] (b c : I → ℝ)
50     (hb : ∀ i,0 < b i) :
51     Integrable (fun x : I → ℝ => Real.exp (-(∑ i,b i*(x i)^2)+∑ i,c i*x i)) := by
52   have hp : 0 < ∫ x : I → ℝ,Real.exp (-(∑ i,b i*(x i)^2)+∑ i,c i*x i) := by
53     rw [integral_exp_diagonal b c hb]
54     apply Finset.prod_pos
55     intro i _
56     exact mul_pos (Real.sqrt_pos.2 (div_pos Real.pi_pos (hb i))) (Real.exp_pos _)
57   by_contra hi
58   rw [integral_undef hi] at hp
59   exact lt_irrefl 0 hp
60 
61 /-- Real unitary coordinate changes preserve the actual product Lebesgue integral. -/
62 theorem integral_unitary_mulVec {I : Type*} [Fintype I] [DecidableEq I]
63     (U : unitary (Matrix I I ℝ)) (f : (I → ℝ) → ℝ) :
64     (∫ x : I → ℝ,f ((U : Matrix I I ℝ) *ᵥ x)) = ∫ x : I → ℝ,f x := by
65   let A : Matrix I I ℝ := U
66   have hleft : A.transpose*A=1 := by
67     simpa only [A,Matrix.star_eq_conjTranspose,Matrix.conjTranspose_eq_transpose_of_trivial]
68       using Unitary.coe_star_mul_self U
69   have hright : A*A.transpose=1 := by
70     have h := Unitary.coe_mul_star_self U
71     change (U : Matrix I I ℝ)*star (U : Matrix I I ℝ)=1 at h
72     simpa only [A,Matrix.star_eq_conjTranspose,Matrix.conjTranspose_eq_transpose_of_trivial]
73       using h
74   have hdet : A.det^2=1 := by
75     have h := congrArg Matrix.det hleft
76     simpa only [Matrix.det_mul,Matrix.det_transpose,Matrix.det_one,pow_two] using h
77   have habs : |A.det|=1 := by
78     nlinarith [sq_abs A.det,abs_nonneg A.det]
79   have hd : LinearMap.det A.toLin' ≠ 0 := by
80     rw [LinearMap.det_toLin']
81     intro hz
82     simp [hz] at hdet
83   have hmap := Measure.map_linearMap_addHaar_pi_eq_smul_addHaar hd
84     (volume : Measure (I → ℝ))
85   simp only [LinearMap.det_toLin',abs_inv,habs,inv_one,ENNReal.ofReal_one,one_smul] at hmap
86   have hpres : MeasurePreserving (fun x => A *ᵥ x) (volume : Measure (I → ℝ)) volume :=
87     ⟨A.toLin'.continuous_of_finiteDimensional.measurable,hmap⟩
88   let e : (I → ℝ) ≃ₗ[ℝ] (I → ℝ) :=
89     { A.toLin' with
90       invFun := fun x => A.transpose *ᵥ x
91       left_inv := fun x => by
92         change A.transpose *ᵥ (A *ᵥ x)=x
93         rw [mulVec_mulVec,hleft,Matrix.one_mulVec]
94       right_inv := fun x => by
95         change A *ᵥ (A.transpose *ᵥ x)=x
96         rw [mulVec_mulVec,hright,Matrix.one_mulVec] }
97   exact hpres.integral_comp e.toContinuousLinearEquiv.toHomeomorph.measurableEmbedding f
98 
99 /-- The centered positive-definite Gaussian integral is computed from the
100 actual spectral decomposition, including its determinant factor. -/
101 theorem integral_exp_posDef {I : Type*} [Fintype I] [DecidableEq I]
102     (V : Matrix I I ℝ) (hV : V.PosDef) :
103     (∫ x : I → ℝ,Real.exp (-(x ⬝ᵥ (V *ᵥ x))/2)) =
104       Real.sqrt ((2*Real.pi)^Fintype.card I / V.det) := by
105   let U := hV.isHermitian.eigenvectorUnitary
106   let eigen := hV.isHermitian.eigenvalues
107   have hd : (U : Matrix I I ℝ).transpose*V*(U : Matrix I I ℝ)=diagonal eigen := by
108     simpa only [U,eigen,Unitary.conjStarAlgAut_star_apply,Matrix.star_eq_conjTranspose,
109       Matrix.conjTranspose_eq_transpose_of_trivial,RCLike.ofReal_real_eq_id,Function.id_comp]
110       using hV.isHermitian.conjStarAlgAut_star_eigenvectorUnitary
111   have hq (x : I → ℝ) :
112       ((U : Matrix I I ℝ) *ᵥ x) ⬝ᵥ (V *ᵥ ((U : Matrix I I ℝ) *ᵥ x)) =
113         ∑ i,eigen i*(x i)^2 := by
114     calc
115       _ = x ⬝ᵥ ((U : Matrix I I ℝ).transpose *ᵥ
116           (V *ᵥ ((U : Matrix I I ℝ) *ᵥ x))) := by
117         rw [dotProduct_transpose_mulVec,dotProduct_comm]
118       _ = x ⬝ᵥ (((U : Matrix I I ℝ).transpose*V*(U : Matrix I I ℝ)) *ᵥ x) := by
119         rw [mulVec_mulVec,mulVec_mulVec]
120       _ = _ := by
121         rw [hd]
122         simp [dotProduct,Matrix.mulVec_diagonal,pow_two,mul_comm,mul_assoc]
123   rw [← integral_unitary_mulVec U (fun x => Real.exp (-(x ⬝ᵥ (V *ᵥ x))/2))]
124   simp_rw [hq]
125   have heigen (i : I) : 0 < eigen i := hV.eigenvalues_pos i
126   have hi := integral_exp_diagonal (fun i => eigen i/2) (fun _ => 0)
127     (fun i => div_pos (heigen i) (by norm_num))
128   have heq (x : I → ℝ) :
129       -(∑ i,(eigen i/2)*(x i)^2)+(∑ i,(0:ℝ)*x i) = -(∑ i,eigen i*(x i)^2)/2 := by
130     simp only [zero_mul,Finset.sum_const_zero,add_zero,neg_div]
131     congr 1
132     rw [Finset.sum_div]
133     apply Finset.sum_congr rfl
134     intro i _
135     ring
136   simp_rw [heq,zero_pow (by decide : (2:ℕ)≠0),zero_div,Real.exp_zero,mul_one] at hi
137   rw [hi]
138   calc
139     (∏ i,Real.sqrt (Real.pi/(eigen i/2))) = ∏ i,Real.sqrt ((2*Real.pi)/eigen i) := by
140       apply Finset.prod_congr rfl
141       intro i _
142       congr 1
143       field_simp
144     _ = Real.sqrt (∏ i,(2*Real.pi)/eigen i) :=
145       (Real.sqrt_prod Finset.univ (fun i _ => (div_pos (by positivity) (heigen i)).le)).symm
146     _ = _ := by
147       congr 1
148       rw [Finset.prod_div_distrib,Finset.prod_const]
149       exact congrArg (fun t : ℝ => (2*Real.pi)^Fintype.card I/t)
150         hV.isHermitian.det_eq_prod_eigenvalues.symm
151 
152 /-- Completing the square uses the actual nonsingular inverse of the
153 positive-definite matrix. -/
154 theorem quadratic_complete_square {I : Type*} [Fintype I] [DecidableEq I]
155     (V : Matrix I I ℝ) (hV : V.PosDef) (s q : I → ℝ) :
156     -(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q =
157       -((q-V⁻¹ *ᵥ s) ⬝ᵥ (V *ᵥ (q-V⁻¹ *ᵥ s)))/2+(s ⬝ᵥ (V⁻¹ *ᵥ s))/2 := by
158   let a := V⁻¹ *ᵥ s
159   have hinv : V *ᵥ a=s := by
160     dsimp [a]
161     rw [mulVec_mulVec,Matrix.mul_nonsing_inv V (isUnit_iff_ne_zero.mpr (ne_of_gt hV.det_pos)),
162       Matrix.one_mulVec]
163   have hsymm : V.transpose=V := by
164     simpa only [Matrix.IsHermitian,Matrix.conjTranspose_eq_transpose_of_trivial] using hV.isHermitian
165   have hcross : a ⬝ᵥ (V *ᵥ q)=s ⬝ᵥ q := by
166     have h := dotProduct_transpose_mulVec V a q
167     rw [hsymm,hinv] at h
168     exact h.trans (dotProduct_comm _ _)
169   change -(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q =
170     -((q-a) ⬝ᵥ (V *ᵥ (q-a)))/2+(s ⬝ᵥ a)/2
171   simp only [Matrix.mulVec_sub,sub_dotProduct,dotProduct_sub]
172   rw [hinv,hcross,dotProduct_comm q s,dotProduct_comm a s]
173   ring
174 
175 /-- The actual multivariate Gaussian exponential tilt, with its exact
176 determinant and inverse-quadratic factors. -/
177 theorem integral_exp_posDef_tilt {I : Type*} [Fintype I] [DecidableEq I]
178     (V : Matrix I I ℝ) (hV : V.PosDef) (s : I → ℝ) :
179     (∫ q : I → ℝ,Real.exp (-(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q)) =
180       Real.sqrt ((2*Real.pi)^Fintype.card I / V.det)*
181         Real.exp ((s ⬝ᵥ (V⁻¹ *ᵥ s))/2) := by
182   simp_rw [quadratic_complete_square V hV s,Real.exp_add]
183   rw [integral_mul_const]
184   have ht := integral_add_right_eq_self (μ := (volume : Measure (I → ℝ)))
185     (fun q : I → ℝ => Real.exp (-(q ⬝ᵥ (V *ᵥ q))/2)) (-(V⁻¹ *ᵥ s))
186   simp only [← sub_eq_add_neg] at ht
187   rw [ht,integral_exp_posDef V hV]
188 
189 theorem integrable_exp_posDef_tilt {I : Type*} [Fintype I] [DecidableEq I]
190     (V : Matrix I I ℝ) (hV : V.PosDef) (s : I → ℝ) :
191     Integrable (fun q : I → ℝ => Real.exp (-(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q)) := by
192   have hp : 0 < ∫ q : I → ℝ,Real.exp (-(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q) := by
193     rw [integral_exp_posDef_tilt V hV s]
194     exact mul_pos (Real.sqrt_pos.2 (div_pos (by positivity) hV.det_pos)) (Real.exp_pos _)
195   by_contra hi
196   rw [integral_undef hi] at hp
197   exact lt_irrefl 0 hp
198 
199 section GaussianPrior
200 variable {I : Type*} [Fintype I] [DecidableEq I]
201 
202 /-- The partition function computed by the actual Gaussian integral. -/
203 noncomputable def gaussianPartition (V : Matrix I I ℝ) : ℝ :=
204   Real.sqrt ((2*Real.pi)^Fintype.card I/V.det)
205 
206 theorem gaussianPartition_pos {V : Matrix I I ℝ} (hV : V.PosDef) :
207     0 < gaussianPartition V :=
208   Real.sqrt_pos.2 (div_pos (by positivity) hV.det_pos)
209 
210 /-- The genuine centered Gaussian density with precision gamma times the identity. -/
211 noncomputable def gaussianDensity (γ : ℝ) (q : I → ℝ) : ℝ :=
212   (gaussianPartition (γ • (1 : Matrix I I ℝ)))⁻¹ *
213     Real.exp (-(q ⬝ᵥ ((γ • (1 : Matrix I I ℝ)) *ᵥ q))/2)
214 
215 theorem gaussianDensity_nonneg (γ : ℝ) (q : I → ℝ) : 0 ≤ gaussianDensity γ q :=
216   mul_nonneg (inv_nonneg.2 (Real.sqrt_nonneg _)) (Real.exp_nonneg _)
217 
218 theorem gaussianDensity_measurable (γ : ℝ) : Measurable (gaussianDensity (I := I) γ) := by
219   unfold gaussianDensity
220   fun_prop
221 
222 theorem gaussianDensity_integrable (γ : ℝ) (hγ : 0 < γ) :
223     Integrable (gaussianDensity (I := I) γ) := by
224   unfold gaussianDensity
225   have hV : (γ • (1 : Matrix I I ℝ)).PosDef := Matrix.PosDef.one.smul hγ
226   have hi := integrable_exp_posDef_tilt (γ • (1 : Matrix I I ℝ)) hV (0 : I → ℝ)
227   simpa only [zero_dotProduct,add_zero] using hi.const_mul
228     (gaussianPartition (γ • (1 : Matrix I I ℝ)))⁻¹
229 
230 theorem integral_gaussianDensity (γ : ℝ) (hγ : 0 < γ) :
231     (∫ q : I → ℝ,gaussianDensity γ q) = 1 := by
232   have hV : (γ • (1 : Matrix I I ℝ)).PosDef := Matrix.PosDef.one.smul hγ
233   unfold gaussianDensity
234   rw [integral_const_mul,integral_exp_posDef _ hV]
235   exact inv_mul_cancel₀ (ne_of_gt (gaussianPartition_pos hV))
236 
237 /-- The actual normalized Gaussian mixing measure, constructed by density. -/
238 noncomputable def gaussianPrior (γ : ℝ) : Measure (I → ℝ) :=
239   volume.withDensity (fun q => ENNReal.ofReal (gaussianDensity γ q))
240 
241 theorem gaussianPrior_probability (γ : ℝ) (hγ : 0 < γ) :
242     IsProbabilityMeasure (gaussianPrior (I := I) γ) := by
243   constructor
244   rw [gaussianPrior,withDensity_apply _ MeasurableSet.univ,Measure.restrict_univ,
245     ← ofReal_integral_eq_lintegral_ofReal (gaussianDensity_integrable γ hγ)
246       (ae_of_all _ (gaussianDensity_nonneg γ)),integral_gaussianDensity γ hγ]
247   exact ENNReal.ofReal_one
248 
249 theorem gaussianPrior_integral (γ : ℝ) (g : (I → ℝ) → ℝ) :
250     (∫ q,g q∂gaussianPrior γ)=∫ q,gaussianDensity γ q*g q := by
251   rw [gaussianPrior,integral_withDensity_eq_integral_toReal_smul
252     (gaussianDensity_measurable γ).ennreal_ofReal (ae_of_all _ fun _ => ENNReal.ofReal_lt_top)]
253   apply integral_congr_ae
254   exact ae_of_all _ (fun q => by
255     change (ENNReal.ofReal (gaussianDensity γ q)).toReal • g q = gaussianDensity γ q*g q
256     rw [ENNReal.toReal_ofReal (gaussianDensity_nonneg γ q),smul_eq_mul])
257 
258 theorem gaussianPartition_ratio (V W : Matrix I I ℝ) (hV : V.PosDef) (hW : W.PosDef) :
259     gaussianPartition V/gaussianPartition W=Real.sqrt (W.det/V.det) := by
260   unfold gaussianPartition
261   rw [← Real.sqrt_div (div_nonneg (by positivity) hV.det_pos.le)]
262   congr 1
263   have hc : (2*Real.pi)^Fintype.card I ≠ 0 := by positivity
264   field_simp [hc,ne_of_gt hV.det_pos,ne_of_gt hW.det_pos]
265 
266 /-- Mixing a quadratic exponential against the constructed Gaussian prior
267 gives the exact determinant ratio; this is an identity of actual integrals. -/
268 theorem gaussianPrior_tilt (γ : ℝ) (hγ : 0 < γ) (V : Matrix I I ℝ)
269     (hV : V.PosDef) (s : I → ℝ) :
270     (∫ q : I → ℝ,Real.exp (s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2)∂gaussianPrior γ) =
271       Real.sqrt ((γ • (1 : Matrix I I ℝ)).det/V.det)*
272         Real.exp ((s ⬝ᵥ (V⁻¹ *ᵥ s))/2) := by
273   rw [gaussianPrior_integral]
274   have he (q : I → ℝ) : gaussianDensity γ q*
275       Real.exp (s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2) =
276       (gaussianPartition (γ • (1 : Matrix I I ℝ)))⁻¹ *
277         Real.exp (-(q ⬝ᵥ (V *ᵥ q))/2+s ⬝ᵥ q) := by
278     unfold gaussianDensity
279     rw [mul_assoc,← Real.exp_add]
280     congr 2
281     simp only [Matrix.sub_mulVec,dotProduct_sub]
282     ring
283   simp_rw [he]
284   rw [integral_const_mul,integral_exp_posDef_tilt V hV s]
285   change (gaussianPartition (γ • (1 : Matrix I I ℝ)))⁻¹*
286     (gaussianPartition V*Real.exp ((s ⬝ᵥ (V⁻¹ *ᵥ s))/2)) = _
287   calc
288     _ = (gaussianPartition V/gaussianPartition (γ • (1 : Matrix I I ℝ)))*
289         Real.exp ((s ⬝ᵥ (V⁻¹ *ᵥ s))/2) := by ring
290     _ = _ := by rw [gaussianPartition_ratio V _ hV (Matrix.PosDef.one.smul hγ)]
291 
292 theorem gaussianPrior_tilt_integrable (γ : ℝ) (hγ : 0 < γ) (V : Matrix I I ℝ)
293     (hV : V.PosDef) (s : I → ℝ) :
294     Integrable (fun q : I → ℝ => Real.exp (s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2))
295       (gaussianPrior γ) := by
296   have h0 : (γ • (1 : Matrix I I ℝ)).PosDef := Matrix.PosDef.one.smul hγ
297   have hp : 0 < ∫ q : I → ℝ,Real.exp (s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2)
298       ∂gaussianPrior γ := by
299     rw [gaussianPrior_tilt γ hγ V hV s]
300     exact mul_pos (Real.sqrt_pos.2 (div_pos h0.det_pos hV.det_pos)) (Real.exp_pos _)
301   by_contra hi
302   rw [integral_undef hi] at hp
303   exact lt_irrefl 0 hp
304 
305 /-- The exact self-normalized mixture formula for the actual ridge Gram matrix. -/
306 theorem gaussian_ridge_mixture (γ σ : ℝ) (hγ : 0 < γ) (hσ : 0 < σ)
307     (φ : ℕ → I → ℝ) (n : ℕ) (score : I → ℝ) :
308     (∫ q : I → ℝ,Real.exp ((q ⬝ᵥ score)/σ-(q ⬝ᵥ (designGram φ n *ᵥ q))/2)
309       ∂gaussianPrior γ) =
310       Real.sqrt ((γ • (1 : Matrix I I ℝ)).det/(ridgeMatrix γ φ n).det)*
311         Real.exp ((weightedNorm (ridgeMatrix γ φ n)⁻¹ score)^2/(2*σ^2)) := by
312   let V := ridgeMatrix γ φ n
313   let s : I → ℝ := σ⁻¹ • score
314   have hV := ridgeMatrix_posDef γ hγ φ n
315   have hp (q : I → ℝ) : s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2 =
316       (q ⬝ᵥ score)/σ-(q ⬝ᵥ (designGram φ n *ᵥ q))/2 := by
317     simp only [V,ridgeMatrix,add_sub_cancel_left,s,smul_dotProduct,smul_eq_mul]
318     rw [dotProduct_comm score q]
319     ring
320   have he : (s ⬝ᵥ (V⁻¹ *ᵥ s))/2 = (weightedNorm V⁻¹ score)^2/(2*σ^2) := by
321     rw [weightedNorm_sq hV.inv.posSemidef]
322     dsimp [s]
323     simp only [Matrix.mulVec_smul,smul_dotProduct,dotProduct_smul,smul_eq_mul]
324     field_simp [ne_of_gt hσ]
325     ring
326   have hi := gaussianPrior_tilt γ hγ V hV s
327   simp_rw [hp] at hi
328   rw [he] at hi
329   exact hi
330 
331 theorem gaussian_ridge_mixture_integrable (γ σ : ℝ) (hγ : 0 < γ)
332     (φ : ℕ → I → ℝ) (n : ℕ) (score : I → ℝ) :
333     Integrable (fun q : I → ℝ =>
334       Real.exp ((q ⬝ᵥ score)/σ-(q ⬝ᵥ (designGram φ n *ᵥ q))/2)) (gaussianPrior γ) := by
335   let V := ridgeMatrix γ φ n
336   let s : I → ℝ := σ⁻¹ • score
337   have hp (q : I → ℝ) : s ⬝ᵥ q-(q ⬝ᵥ ((V-γ • 1) *ᵥ q))/2 =
338       (q ⬝ᵥ score)/σ-(q ⬝ᵥ (designGram φ n *ᵥ q))/2 := by
339     simp only [V,ridgeMatrix,add_sub_cancel_left,s,smul_dotProduct,smul_eq_mul]
340     rw [dotProduct_comm score q]
341     ring
342   have hi := gaussianPrior_tilt_integrable γ hγ V (ridgeMatrix_posDef γ hγ φ n) s
343   simpa only [hp] using hi
344 
345 end GaussianPrior
346 
347 end Komorebi.LSM.RegressionConfidence

Back to source controls