import Mathlib.Probability.Kernel.Disintegration.StandardBorel import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.FiniteRGOKLError import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ObservationConditionalKernel import AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.StateDependentRGO import AutoSamplingTheory.TechnicalLemmas.Analysis.GibbsGradientMoment import Mathlib.Tactic /-! # One-step KL propagation for the actual retained-state sampler Source: SPHMC arXiv:2609.06906v1 Section 6.3 conditional output argument after (6.5), using the actual Algorithm 3.3 reference update. This is the KL recurrence component, not the numerical observation-error budget in (6.5). The base measure is the real Gibbs law volume.tilted(-V). Its probability property follows from GibbsGradientMoment. StateDependentRGO constructs the state-dependent target T and ideal Gaussian observation H; its joint recovery on a Dirac input yields the actual posterior recovery used internally. Targets are exposed as nested Gibbs tilts. A combined volume density or normalizer formula is not an additional public conclusion. ObservationConditionalKernel supplies the same pre-noise first-hit GD update, full retained reference/count/history, joint observation/state law J and its conditional kernel C. Its actual observation variance tau/(beta+b) is distinct from the ideal observation variance v=(eta+tau)/(beta+b). Conditional precision and center identify the target posterior only for actual-Q-almost every observation and then conditional-C-almost every retained state. No transfer of these conditional versions to ideal H is made; reference randomness remains. All-ENNReal KL mixture and chain inequalities give the remaining-output error, and disintegration folds its conditional average into the actual P-state law. The theorem derives posterior recovery, target alignment, active projection and error measurability. They are not assumptions of the public theorem. Internal J,C,K and their recovery formulas are proof constructions, not extra public conjuncts. The absorbed branch has e=0 and P=Dirac but retains the full remaining-output KL; terminal error is not discarded. Genuine C2 Hessian bounds agree with source (1.1); positive finite dimension and 0 klDiv (L x) (K x)) := by classical have hf : Measurable (fun p : A × B => ENNReal.ofReal (klFun ((Kernel.rnDeriv L K p.1 p.2).toReal))) := by exact ENNReal.measurable_ofReal.comp (continuous_klFun.measurable.comp (Kernel.measurable_rnDeriv L K).ennreal_toReal) have heq (x : A) : klDiv (L x) (K x) = if L x ≪ K x then ∫⁻ y, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K x y).toReal)) ∂K x else ∞ := by rw [klDiv_eq_lintegral_klFun] split_ifs · apply lintegral_congr_ae filter_upwards [Kernel.rnDeriv_eq_rnDeriv_measure (κ := L) (η := K) (a := x)] with y hy rw [hy] · rfl simp_rw [heq] exact Measurable.ite (Kernel.measurableSet_absolutelyContinuous L K) hf.lintegral_kernel_prod_right' measurable_const private theorem conditional_kl_integral {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (μ : Measure A) [IsProbabilityMeasure μ] (L K : Kernel A B) [IsMarkovKernel L] [IsMarkovKernel K] : klDiv (μ ⊗ₘ L) (μ ⊗ₘ K) = ∫⁻ x, klDiv (L x) (K x) ∂μ := by classical by_cases hac : μ ⊗ₘ L ≪ μ ⊗ₘ K · have hfiber := hac.kernel_of_compProd have hwd : (μ ⊗ₘ K).withDensity (fun p => Kernel.rnDeriv L K p.1 p.2) = μ ⊗ₘ L := by rw [← Measure.compProd_withDensity (Kernel.measurable_rnDeriv L K)] apply Measure.compProd_congr filter_upwards [hfiber] with x hx exact Kernel.withDensity_rnDeriv_eq hx have hrn : (μ ⊗ₘ L).rnDeriv (μ ⊗ₘ K) =ᵐ[μ ⊗ₘ K] (fun p => Kernel.rnDeriv L K p.1 p.2) := by rw [← hwd] exact Measure.rnDeriv_withDensity _ (Kernel.measurable_rnDeriv L K) rw [klDiv_eq_lintegral_klFun_of_ac hac] calc _ = ∫⁻ p, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K p.1 p.2).toReal)) ∂(μ ⊗ₘ K) := by apply lintegral_congr_ae filter_upwards [hrn] with p hp rw [hp] _ = ∫⁻ x, ∫⁻ y, ENNReal.ofReal (klFun ((Kernel.rnDeriv L K x y).toReal)) ∂K x ∂μ := by apply Measure.lintegral_compProd exact ENNReal.measurable_ofReal.comp (continuous_klFun.measurable.comp (Kernel.measurable_rnDeriv L K).ennreal_toReal) _ = _ := by apply lintegral_congr_ae filter_upwards [hfiber] with x hx rw [klDiv_eq_lintegral_klFun_of_ac hx] apply lintegral_congr_ae filter_upwards [Kernel.rnDeriv_eq_rnDeriv_measure (κ := L) (η := K) (a := x)] with y hy rw [hy] · rw [klDiv_of_not_ac hac] symm by_contra hfinite have hf := ae_lt_top (measurable_fiber_kl L K) hfinite apply hac apply Measure.AbsolutelyContinuous.compProd_right filter_upwards [hf] with x hx exact (klDiv_ne_top_iff.mp hx.ne).1 private theorem composed_kl_bound {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (P Q : Measure A) [IsProbabilityMeasure P] [IsProbabilityMeasure Q] (L K : Kernel A B) [IsMarkovKernel L] [IsMarkovKernel K] : klDiv (L ∘ₘ P) (K ∘ₘ Q) ≤ klDiv P Q + ∫⁻ x, klDiv (L x) (K x) ∂P := by calc klDiv (L ∘ₘ P) (K ∘ₘ Q) = klDiv (P ⊗ₘ L).snd (Q ⊗ₘ K).snd := by rw [Measure.snd_compProd, Measure.snd_compProd] _ ≤ klDiv (P ⊗ₘ L) (Q ⊗ₘ K) := klDiv_map_le _ _ measurable_snd _ = _ := by rw [klDiv_compProd_eq_add, conditional_kl_integral] private theorem mixture_kl_bound {A B : Type u} [MeasurableSpace A] [MeasurableSpace B] [MeasurableSpace.CountableOrCountablyGenerated A B] (μ : Measure A) [IsProbabilityMeasure μ] (ν : Measure B) [IsProbabilityMeasure ν] (R T : Kernel A B) [IsMarkovKernel R] [IsMarkovKernel T] (hT : ∀ᵐ x ∂μ, T x=ν) : klDiv (R ∘ₘ μ) ν ≤ ∫⁻ x, klDiv (R x) (T x) ∂μ := by let K : Kernel A B := Kernel.const A ν have hc := composed_kl_bound μ μ R K have hconst : K ∘ₘ μ=ν := by change μ.bind (fun _ => ν)=ν rw [Measure.bind_const,measure_univ,one_smul] rw [hconst,klDiv_self,zero_add] at hc refine hc.trans_eq ?_ apply lintegral_congr_ae filter_upwards [hT] with x hx simp only [K,Kernel.const_apply,hx] end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep namespace AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep open MeasureTheory ProbabilityTheory InformationTheory open scoped ENNReal NNReal variable {E S : Type u} [MeasurableSpace E] [MeasurableSpace S] [StandardBorelSpace S] [Nonempty S] [MeasurableSpace.CountableOrCountablyGenerated S E] [MeasurableSpace.CountableOrCountablyGenerated E E] private theorem conditional_step (J : Kernel S (E × S)) [IsMarkovKernel J] (P : Kernel S S) [IsMarkovKernel P] (T H : Kernel S E) [IsMarkovKernel T] [IsMarkovKernel H] (K : Kernel (S × E) E) [IsMarkovKernel K] (R : Kernel S E) [IsMarkovKernel R] (s : S) (hP : P s=J.snd s) (hrecover : (K.comap (Prod.mk s) measurable_prodMk_left) ∘ₘ H s = T s) (halign : ∀ᵐ y ∂J.fst s, ∀ᵐ t ∂J.condKernel (s,y), T t=K (s,y)) : klDiv ((R ∘ₖ P) s) (T s) ≤ klDiv (J.fst s) (H s) + ∫⁻ t, klDiv (R t) (T t) ∂P s := by let C := J.condKernel let LR := (R ∘ₖ C).comap (Prod.mk s) measurable_prodMk_left let KR := K.comap (Prod.mk s) measurable_prodMk_left have hdecomp (f : S → ℝ≥0∞) (hf : Measurable f) : (∫⁻ t, f t ∂P s) = ∫⁻ y, ∫⁻ t, f t ∂C (s,y) ∂J.fst s := by rw [hP,Kernel.snd_apply,lintegral_map hf measurable_snd] conv_lhs => rw [← Kernel.disintegrate J J.condKernel] exact Kernel.lintegral_compProd J.fst J.condKernel s (hf.comp measurable_snd) have hout : LR ∘ₘ J.fst s = (R ∘ₖ P) s := by ext A hA rw [Measure.bind_apply hA LR.aemeasurable,Kernel.comp_apply' R P s hA] rw [hdecomp (fun t => R t A) (R.measurable_coe hA)] apply lintegral_congr intro y change (R ∘ₖ C) (s,y) A = _ exact Kernel.comp_apply' R C (s,y) hA have hc := composed_kl_bound (J.fst s) (H s) LR KR rw [hout,hrecover] at hc have hm : ∀ᵐ y ∂J.fst s, klDiv (LR y) (KR y) ≤ ∫⁻ t, klDiv (R t) (T t) ∂C (s,y) := by filter_upwards [halign] with y hy have hh := mixture_kl_bound (C (s,y)) (K (s,y)) R T hy exact hh calc klDiv ((R ∘ₖ P) s) (T s) ≤ klDiv (J.fst s) (H s) + ∫⁻ y, klDiv (LR y) (KR y) ∂J.fst s := hc _ ≤ klDiv (J.fst s) (H s) + ∫⁻ y, ∫⁻ t, klDiv (R t) (T t) ∂C (s,y) ∂J.fst s := add_le_add le_rfl (lintegral_mono_ae hm) _ = _ := congrArg (fun z => klDiv (J.fst s) (H s)+z) (hdecomp (fun t => klDiv (R t) (T t)) (measurable_fiber_kl R T)).symm end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep open MeasureTheory ProbabilityTheory open scoped NNReal ENNReal noncomputable section namespace AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep variable {E S : Type*} [NormedAddCommGroup E] [InnerProductSpace ℝ E] [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E] [MeasurableSpace S] [MeasurableSingletonClass S] private theorem actual_gibbs_posterior {V : E → ℝ} {α β : ℝ≥0} (hα : 0 < α) (hαβ : α ≤ β) (hV : ContDiff ℝ 2 V) (hHess : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧ fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2) (b a : S → ℝ) (u : S → E) (hb : Measurable b) (ha : Measurable a) (hu : Measurable u) (hb0 : ∀ s, 0 ≤ b s) (ha0 : ∀ s, 0 < a s) : let μ := (volume : Measure E).tilted (fun x => -V x) IsProbabilityMeasure μ ∧ ∃ (T H : Kernel S E) (K : Kernel (S × E) E), IsMarkovKernel T ∧ IsMarkovKernel H ∧ IsMarkovKernel K ∧ (∀ s, T s=μ.tilted (fun x => -(b s/2)*‖x-u s‖^2)) ∧ (∀ s, H s=AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing (T s) (Real.sqrt (a s))) ∧ (∀ s y, K (s,y)=μ.tilted (fun x => -((b s+(a s)⁻¹)/2)* ‖x-(b s+(a s)⁻¹)⁻¹ • (b s • u s+(a s)⁻¹ • y)‖^2)) ∧ ∀ s, (K.comap (Prod.mk s) measurable_prodMk_left) ∘ₘ H s = T s := by intro μ have hμ : IsProbabilityMeasure μ := (AutoSamplingTheory.TechnicalLemmas.Analysis.GibbsGradientMoment.gibbs_gradient_moment hα hαβ hV hHess).1 let := hμ obtain ⟨T,H,B,hT,hH,hB,hTf,hHf,hBpost,hBup,hrec⟩ := AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.StateDependentRGO.state_dependent_recovery μ b a u hb ha hu hb0 ha0 let := hT let := hH let := hB let K := B.snd have hK : IsMarkovKernel K := by dsimp [K]; infer_instance let := hK have hKf (s : S) (y : E) : K (s,y)=μ.tilted (fun x => -((b s+(a s)⁻¹)/2)* ‖x-(b s+(a s)⁻¹)⁻¹ • (b s • u s+(a s)⁻¹ • y)‖^2) := by dsimp only [K] rw [Kernel.snd_apply,hBup,Measure.map_map measurable_snd measurable_prodMk_left] change Measure.map id _ = _ exact Measure.map_id refine ⟨hμ,T,H,K,hT,hH,hK,hTf,hHf,hKf,?_⟩ intro s have he := hrec (Measure.dirac s) inferInstance ext A hA have hset : MeasurableSet (Prod.snd ⁻¹' A : Set (S × E)) := measurable_snd hA have hei := congrArg (fun m : Measure (S × E) => m (Prod.snd ⁻¹' A)) he have hm : Measurable (fun r : S => ∫⁻ y, B (r,y) (Prod.snd ⁻¹' A) ∂H r) := (B.measurable_coe hset).lintegral_kernel_prod_right' rw [Measure.bind_apply hset B.aemeasurable, Measure.lintegral_compProd (B.measurable_coe hset),lintegral_dirac' s hm, Measure.dirac_compProd_apply hset] at hei rw [Measure.bind_apply hA (K.comap (Prod.mk s) measurable_prodMk_left).aemeasurable] change (∫⁻ y, K (s,y) A ∂H s) = T s A simp only [K,Kernel.snd_apply' _ _ hA] exact hei end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep namespace AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep open MeasureTheory ProbabilityTheory InformationTheory open scoped NNReal ENNReal variable {E S : Type u} [NormedAddCommGroup E] [InnerProductSpace ℝ E] [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E] [MeasurableSpace S] local notation "RefState" E:arg => ℝ≥0 × E × E × ℕ × (ℕ → E) private def firstIndex (q : ℕ → S → ℝ) (b : S → ℝ) (s : S) : ℕ := by classical exact if h : ∃ n, q n s ≤ b s then Nat.find h else 0 theorem one_step_kl_error {V : E → ℝ} {α β : ℝ≥0} (hα : 0 < α) (hαβ : α ≤ β) (hV : ContDiff ℝ 2 V) (hHess : ∀ x w : E, (α:ℝ)*‖w‖^2 ≤ fderiv ℝ (fderiv ℝ V) x w w ∧ fderiv ℝ (fderiv ℝ V) x w w ≤ (β:ℝ)*‖w‖^2) (hd : 0 < (Module.finrank ℝ E : ℝ)) (η τ : RefState E → ℝ) (hη : Measurable η) (hτ : Measurable τ) (hη0 : ∀ s, 0 < η s) (hτ0 : ∀ s, 0 < τ s) (M : Kernel (RefState E) E) [IsMarkovKernel M] (threshold : ℝ≥0) : let d : ℝ := Module.finrank ℝ E let v := fun s : RefState E => (η s+τ s)/((β:ℝ)+s.1) let bp := fun s : RefState E => s.1+Real.toNNReal (v s)⁻¹ let obs := fun p : RefState E × (E × E) => p.2.1+Real.sqrt (τ p.1/((β:ℝ)+p.1.1)) • p.2.2 let center := fun p : RefState E × (E × E) => (bp p.1:ℝ)⁻¹ • ((p.1.1:ℝ) • p.1.2.1+(v p.1)⁻¹ • obs p) let F := fun p x => V x+(bp p.1:ℝ)/2*‖x-center p‖^2 let T := fun p x => x-((β:ℝ)+bp p.1)⁻¹ • gradient (F p) x let q := fun n p => ‖gradient (F p) ((T p)^[n] p.2.1)‖^2 let N := firstIndex q (fun p => ((α:ℝ)+bp p.1)*d) let out := fun p => (T p)^[N p] p.2.1 let update := fun p : RefState E × (E × E) => (bp p.1,center p,out p,p.1.2.2.2.1+1,fun n => Nat.casesOn n (obs p) p.1.2.2.2.2) let μ := (volume : Measure E).tilted (fun x => -V x) IsProbabilityMeasure μ ∧ ∃ (T H : Kernel (RefState E) E) (P : Kernel (RefState E) (RefState E)), ∃ hT : IsMarkovKernel T, ∃ hH : IsMarkovKernel H, ∃ hP : IsMarkovKernel P, letI := hT letI := hH letI := hP (∀ s, T s=μ.tilted (fun x => -((s.1:ℝ)/2)*‖x-s.2.1‖^2)) ∧ (∀ s, H s=AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing (T s) (Real.sqrt (v s))) ∧ (∀ s, P s=if threshold ≤ s.1 then Measure.dirac s else ((M s).prod (stdGaussian E)).map (fun z => update (s,z))) ∧ let Q := fun s => AutoSamplingTheory.TechnicalLemmas.Measure.GaussianSmoothing.gaussianSmoothing (M s) (Real.sqrt (τ s/((β:ℝ)+s.1))) let e := fun s => if threshold ≤ s.1 then 0 else klDiv (Q s) (H s) Measurable e ∧ ∀ (R : Kernel (RefState E) E) [IsMarkovKernel R], Measurable (fun s => klDiv (R s) (T s)) ∧ ∀ s, klDiv ((R ∘ₖ P) s) (T s) ≤ e s+∫⁻ t, klDiv (R t) (T t) ∂P s := by classical intro d v bp obs center F Tstep q N out update μ have hβ : 0 < (β:ℝ) := lt_of_lt_of_le (show 0 < (α:ℝ) from hα) (show (α:ℝ) ≤ β from hαβ) have hv (s : RefState E) : 0 < v s := div_pos (add_pos (hη0 s) (hτ0 s)) (add_pos_of_pos_of_nonneg hβ s.1.coe_nonneg) have hvm : Measurable v := by dsimp [v]; fun_prop have hbp (s : RefState E) : (bp s:ℝ)=(s.1:ℝ)+(v s)⁻¹ := by simp only [bp,NNReal.coe_add,Real.coe_toNNReal _ (inv_pos.mpr (hv s)).le] obtain ⟨hμ,T,H,K,hT,hH,hK,hTf,hHf,hKf,hrec⟩ := actual_gibbs_posterior hα hαβ hV hHess (fun s : RefState E => (s.1:ℝ)) v (fun s => s.2.1) (by fun_prop) hvm (by fun_prop) (fun s => s.1.coe_nonneg) hv let := hμ let := hT let := hH let := hK obtain ⟨J,hJ,hJs,hQ,hdis,hsupp,hfact,P,hP,hPs,hactive⟩ := AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ObservationConditionalKernel.observation_conditional_kernel hα hαβ hV hHess hd η τ hη hτ hη0 hτ0 M threshold let := hJ let := hP have href := AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.ReferenceCarryingKernel.reference_carrying_kernel hα hαβ hV hHess hd η τ hη hτ hη0 hτ0 M threshold have hu : Measurable update := href.2.2.2.1 have ho : Measurable obs := by dsimp [obs]; fun_prop have hJs' (s : RefState E) : J s=((M s).prod (stdGaussian E)).map (fun z => (obs (s,z),update (s,z))) := hJs s have hPs' (s : RefState E) : P s=if threshold ≤ s.1 then Measure.dirac s else ((M s).prod (stdGaussian E)).map (fun z => update (s,z)) := hPs s have halign (s : RefState E) : ∀ᵐ y ∂J.fst s, ∀ᵐ t ∂J.condKernel (s,y), T t=K (s,y) := by filter_upwards [hsupp s] with y hy filter_upwards [hy] with t ht rw [hTf,hKf,ht.1,ht.2,hbp] refine ⟨hμ,T,H,P,hT,hH,hP,hTf,hHf,hPs',?_⟩ intro Q e have heq : e=(fun s => if threshold ≤ s.1 then 0 else klDiv (J.fst s) (H s)) := by funext s simp only [e,Q,hQ] have he : Measurable e := by rw [heq] exact Measurable.ite (measurableSet_le measurable_const (by fun_prop)) measurable_const (measurable_fiber_kl J.fst H) refine ⟨he,fun R hR => ?_⟩ let := hR have hmR := measurable_fiber_kl R T refine ⟨hmR,fun s => ?_⟩ by_cases hs : threshold ≤ s.1 · simp only [e,Kernel.comp_apply,hPs',if_pos hs, Measure.dirac_bind R.measurable,lintegral_dirac' s hmR,zero_add,le_refl] · have hPJ : P s=J.snd s := by have hf : Measurable (fun z : E × E => (obs (s,z),update (s,z))) := (ho.prodMk hu).comp (measurable_const.prodMk measurable_id) rw [hPs',if_neg hs,Kernel.snd_apply,hJs',Measure.map_map measurable_snd hf] rfl have hc := conditional_step J P T H K R s hPJ (hrec s) (halign s) simpa only [e,Q,if_neg hs,hQ] using hc end AutoSamplingTheory.ExampleCases.SmoothedPicardHMC.EnhancedKLOneStep