multigeom.geometry.estimator#
Conditional-pair estimator of the multi-scale Fisher information matrix G(x) — Algorithm 1 of the paper.
Given a clean image x ∈ [-1, 1]^{1×64×64}, an encoder f_φ : ℝ^{1×64×64} → ℝ^N,
and a conditional diffusion network ε_θ(x_t, t, r) trained on f_φ-generated
responses with classifier-free guidance, the estimator returns the top-k
eigenpairs of
G(x) = ∫_0^∞ [φ_τ * J_τ](x) dτ (paper Eq. (eq:Gdef))
via Monte-Carlo integration of the line element of Eq. (eq:line) along the
log-SNR axis log τ = log[(1−ᾱ_t)/ᾱ_t]. The integrand reduces to squared
noise-prediction differences after Tweedie's identity and the change of
variable dτ = τ d(log τ) cancel against the 1/τ² factor that Tweedie
contributes (Appendix C.2).
Algorithm 1 (paper) — implemented here verbatim:
Inputs: x ∈ [-1, 1]^{64×64}, ε_θ, f_φ, ᾱ_t schedule, grid {t_j}, σ, N_MC, B
Output: top-k eigenpairs of G(x) ≈ Dᵀ D
1. log τ_j ← log[(1−ᾱ_{t_j})/ᾱ_{t_j}] for j = 1,...,J
2. w_j ← ½ (log τ_{j+1} − log τ_{j−1}) (half-weights at endpoints)
3. for j = 1,...,J:
4. for n = 1,...,N_MC in batches of B:
5. z ~ N(0, I) [B,1,64,64]
6. x_t ← √ᾱ_{t_j} x + √(1−ᾱ_{t_j}) z [B,1,64,64]
7. ε^u ← ε_θ(x_t, t_j, 0) null-token pass
8. x̂ ← (x_t − √(1−ᾱ_{t_j}) ε^u) / √ᾱ_{t_j} Tweedie
9. x̂ ← clip(x̂, −1, 1)
10. μ ← f_φ(x̂) [B,N]
11. ξ_1, ξ_2 ~ N(0, I) [B,N]
12. r_1, r_2 ← μ + σ ξ_1, μ + σ ξ_2 Eq. (eq:obs-model)
13. ε^c_1, ε^c_2 ← ε_θ(x_t, t_j, r_1), ε_θ(x_t, t_j, r_2)
14. Δ ← flatten(ε^c_1 − ε^c_2) [B, 64²]
15. append √(w_j / (2 N_MC)) · Δ to D
16. Stack rows → D ∈ ℝ^{(J·N_MC) × 64²}
17. (U, S, Vᵀ) ← SVD(D)
18. return S²[:k] and rows of Vᵀ[:k] top-k eigenpairs of G(x) ≈ DᵀD
This is the **conditional-pair** estimator (Eq. (eq:cc)). It is *not* the
unconditional-conditional estimator (which has a per-image-dependent bias
documented in Appendix C.2 of the paper).
Classes
|
Top eigenpairs of G(x) returned by compute_g: eigenvalues [k] and eigenvectors [k, P], largest first. |