scripts/exp_lemma61.py

"""Claim 5 -- independent audit of Lemma 6.1 (dimension-free change of measure).

Lemma 6.1 as printed:  for -beta I <= grad^2 V <= H <= beta I and
pi(x,p) propto exp(-V(x) - ||p||^2/2),

        E_mu[ ||grad V(x)||^2 ] <= tr(H) + beta KL(mu||pi)          (G)
        E_mu[ p^T H p ]         <= tr(H) + beta KL(mu||pi)          (P)

Independent re-derivation.  For any C>0 the Donsker-Varadhan variational
formula gives the *sharp* value

    sup_mu { E_mu[f] - C beta KL(mu||pi) } = C beta log E_pi[ exp(f/(C beta)) ].

Take f = p^T H p with H = beta I_d.  Under pi, p ~ N(0,I_d), so
    C beta log E_pi exp(||p||^2/C) = -(C beta d/2) log(1 - 2/C),   C > 2.
The claimed right-hand side is C tr(H) = C beta d.  Hence the inequality
    E_mu[p^T H p] <= C ( tr(H) + beta KL(mu||pi) )
holds for ALL mu if and only if  -(1/2) log(1-2/C) <= 1, i.e.

    C >= C* := 2 e^2 / (e^2 - 1) = 2.31303528549933...

The identical computation with f = ||grad V||^2 and V = (beta/2)||x||^2 gives the
same C*.  So the printed constant C = 1 is FALSE, the sharp constant is C*, and
the extremal mu is Gaussian with p-covariance e^2 I (resp. x-covariance
e^2/beta I).  Everything below verifies this numerically and probes the
boundary of the assumption grad^2 V <= H.
"""

import numpy as np
from scipy import integrate

import common as C

CSTAR = 2 * np.exp(2) / (np.exp(2) - 1)
res = {
    "C_star_closed_form": float(CSTAR),
    "derivation": "sup_mu {E_mu f - C beta KL} = C beta log E_pi exp(f/(C beta)); "
    "f=p^T H p, H=beta I gives -(C beta d/2) log(1-2/C) <= C beta d "
    "iff C >= 2e^2/(e^2-1)",
}


# ---------------------------------------------------------------------------
# 1. Explicit closed-form counterexample to the printed constant C = 1.
#    mu = N(0, (1/a) I_d) x N(0, s I_d),  V = (a/2)||x||^2, H = a I, beta = a.
# ---------------------------------------------------------------------------
def ratio_p(s):
    """E_mu[p^T H p] / (tr H + beta KL)  for H = beta I, mu_p = N(0, s I)."""
    return 2 * s / (1 + s - np.log(s))


def ratio_g(t):
    """E_mu[||grad V||^2] / (tr H + beta KL) for V=(beta/2)||x||^2,
    mu_x = N(0, t/beta I)."""
    return 2 * t / (1 + t - np.log(t))


ce = []
for d in (1, 10, 1000):
    for beta in (0.5, 1.0, 7.0):
        s = np.exp(2.0)
        lhs = s * beta * d  # E_mu[p^T H p]
        kl = 0.5 * d * (s - 1 - np.log(s))
        rhs = beta * d + beta * kl
        ce.append(
            {
                "d": d,
                "beta": beta,
                "p_cov": float(s),
                "LHS_E_pHp": float(lhs),
                "RHS_trH_plus_beta_KL": float(rhs),
                "ratio": float(lhs / rhs),
                "violates_printed_lemma": bool(lhs > rhs),
            }
        )
res["closed_form_counterexamples_momentum"] = ce
res["sup_ratio_momentum_over_gaussians"] = float(
    max(ratio_p(s) for s in np.geomspace(1.001, 1e4, 200000))
)
res["sup_ratio_gradient_over_gaussians"] = float(
    max(ratio_g(t) for t in np.geomspace(1.001, 1e4, 200000))
)
res["argmax_s_closed_form_e2"] = float(np.exp(2))

# ---------------------------------------------------------------------------
# 2. Randomised search over Gaussian mu (arbitrary mean AND covariance) and
#    diagonal quadratic V, for both inequalities.  Exact formulas.
#      V = 1/2 sum a_i x_i^2, H = diag(hh) with a_i <= hh_i <= beta.
#      mu = N(mx, Sx) x N(mp, Sp) diagonal.
# ---------------------------------------------------------------------------
rng = np.random.default_rng(C.SEED)
worstG, worstP = 0.0, 0.0
argG, argP = None, None
nviolG = nviolP = 0
NTRIAL = 200000
for _ in range(NTRIAL):
    d = int(rng.integers(1, 6))
    beta = float(np.exp(rng.uniform(-1.5, 1.5)))
    a = beta * np.exp(rng.uniform(np.log(1e-3), 0.0, d))
    hh = np.minimum(beta, a * np.exp(rng.uniform(0, 2.0, d)))  # a <= H <= beta
    trH = hh.sum()
    sx = np.exp(rng.uniform(-3, 4, d))
    sp = np.exp(rng.uniform(-3, 4, d))
    mx = rng.normal(0, np.exp(rng.uniform(-2, 2)), d)
    mp = rng.normal(0, np.exp(rng.uniform(-2, 2)), d)
    kl = 0.5 * np.sum(a * sx + a * mx**2 - 1 - np.log(a * sx)) + 0.5 * np.sum(
        sp + mp**2 - 1 - np.log(sp)
    )
    rhs = trH + beta * kl
    lhsG = float(np.sum(a * a * sx + a * a * mx**2))
    lhsP = float(np.sum(hh * sp + hh * mp**2))
    if lhsG / rhs > worstG:
        worstG, argG = lhsG / rhs, dict(
            d=d,
            beta=beta,
            a=a.tolist(),
            sx=sx.tolist(),
            mx=mx.tolist(),
            lhs=lhsG,
            rhs=float(rhs),
        )
    if lhsP / rhs > worstP:
        worstP, argP = lhsP / rhs, dict(
            d=d,
            beta=beta,
            hh=hh.tolist(),
            sp=sp.tolist(),
            mp=mp.tolist(),
            lhs=lhsP,
            rhs=float(rhs),
        )
    nviolG += lhsG > rhs
    nviolP += lhsP > rhs
res["random_gaussian_search"] = {
    "n_trials": NTRIAL,
    "seed": C.SEED,
    "worst_ratio_gradient": float(worstG),
    "worst_ratio_momentum": float(worstP),
    "n_violations_of_printed_lemma_gradient": int(nviolG),
    "n_violations_of_printed_lemma_momentum": int(nviolP),
    "violation_rate_gradient": float(nviolG / NTRIAL),
    "violation_rate_momentum": float(nviolP / NTRIAL),
    "worst_case_gradient": argG,
    "worst_case_momentum": argP,
    "exceeds_C_star": bool(max(worstG, worstP) > CSTAR + 1e-9),
}


# ---------------------------------------------------------------------------
# 3. Non-Gaussian mu and NON-QUADRATIC, non-convex V (the lemma only asks
#    -beta I <= grad^2 V <= H).   1-d quadrature, V(x) = a x^2/2 + c cos(w x).
#    grad^2 V = a + c w^2 cos(wx) in [a - c w^2, a + c w^2]; require
#    a - c w^2 >= -beta and H := a + c w^2 <= beta.
#    mu_x is an arbitrary density on a grid (tilted / bimodal / heavy mixture).
# ---------------------------------------------------------------------------
def audit_nonquadratic(a, c, w, beta, mu_logdens, L=40.0, n=200001):
    x = np.linspace(-L, L, n)
    V = a * x**2 / 2 + c * np.cos(w * x)
    dV = a * x - c * w * np.sin(w * x)
    logpi = -V
    logpi -= np.log(np.trapezoid(np.exp(logpi), x))
    lm = mu_logdens(x)
    lm -= np.log(np.trapezoid(np.exp(lm), x))
    mu = np.exp(lm)
    kl = np.trapezoid(mu * (lm - logpi), x)
    lhs = np.trapezoid(mu * dV**2, x)
    H = a + c * w * w
    return float(lhs), float(H), float(kl), float(lhs / (H + beta * kl))


nq = []
rng2 = np.random.default_rng(C.SEED + 1)
worst_nq = 0.0
for _ in range(4000):
    a = float(np.exp(rng2.uniform(-2, 0.5)))
    w = float(np.exp(rng2.uniform(-1, 1.2)))
    c = float(rng2.uniform(0, 1)) * a / (w * w) * 0.9  # keeps H <= 2a
    beta = a + c * w * w  # tightest beta = H
    if a - c * w * w < -beta:
        continue
    kind = int(rng2.integers(0, 4))
    s = float(np.exp(rng2.uniform(-1.5, 2.5)))
    mshift = float(rng2.normal(0, 3))
    if kind == 0:
        f = lambda x, s=s, m=mshift: -((x - m) ** 2) / (2 * s)
    elif kind == 1:  # bimodal
        f = lambda x, s=s, m=mshift: np.logaddexp(
            -((x - m) ** 2) / (2 * s), -((x + m) ** 2) / (2 * s)
        )
    elif kind == 2:  # exponential tilt of pi
        th = float(rng2.normal(0, 1.5))
        f = lambda x, th=th, a=a, c=c, w=w: -(a * x**2 / 2 + c * np.cos(w * x)) + th * x
    else:  # heavy-tailed (Student-t like)
        nu = float(rng2.uniform(2.5, 8))
        f = lambda x, s=s, nu=nu: -0.5 * (nu + 1) * np.log1p(x * x / (nu * s))
    lhs, H, kl, r = audit_nonquadratic(a, c, w, beta, f)
    if not np.isfinite(r):
        continue
    nq.append(r)
    worst_nq = max(worst_nq, r)
nq = np.array(nq)
res["nonquadratic_nonconvex_V_quadrature"] = {
    "n_cases": int(len(nq)),
    "worst_ratio": float(worst_nq),
    "n_violating_printed_lemma": int((nq > 1.0).sum()),
    "n_violating_C_star": int((nq > CSTAR).sum()),
    "seed": C.SEED + 1,
    "V": "a x^2/2 + c cos(w x)  (non-convex where a - c w^2 < 0)",
}

# ---------------------------------------------------------------------------
# 4. Boundary audit: what happens when the assumption grad^2 V <= H is violated
#    (H too small)?  The corrected inequality must fail.
# ---------------------------------------------------------------------------
bd = []
for shrink in (1.0, 0.7, 0.4, 0.2, 0.1):
    a, beta, d = 1.0, 1.0, 5
    H = shrink * a  # shrink<1 violates grad^2 V <= H
    s = np.exp(2.0)
    lhs = d * a * a * s / a  # E||grad V||^2 with mu_x = N(0, s/a I)
    kl = 0.5 * d * (s - 1 - np.log(s))
    rhs = CSTAR * (d * H + beta * kl)
    bd.append(
        {
            "H_over_hessian": shrink,
            "LHS": float(lhs),
            "C_star_RHS": float(rhs),
            "holds": bool(lhs <= rhs),
        }
    )
res["boundary_audit_H_below_hessian"] = bd

# ---------------------------------------------------------------------------
# 5. Does the corrected constant C* hold universally?  Check the DV certificate
#    numerically:  C* beta log E_pi exp(f/(C* beta)) <= C* tr(H)  for f = p^T H p.
# ---------------------------------------------------------------------------
cert = []
for d in (1, 3, 10):
    for frac in (1.0, 0.9, 0.5, 0.25, 0.1):
        beta = 1.0
        hh = np.full(d, frac * beta)
        trH = hh.sum()
        lam = hh / (CSTAR * beta)
        assert np.all(2 * lam < 1)
        sup = CSTAR * beta * np.sum(-0.5 * np.log(1 - 2 * lam))
        cert.append(
            {
                "d": d,
                "H_over_beta": frac,
                "DV_sup": float(sup),
                "C_star_trH": float(CSTAR * trH),
                "holds": bool(sup <= CSTAR * trH + 1e-12),
            }
        )
res["DV_certificate_for_C_star"] = cert
res["C_star_is_sharp"] = bool(abs(cert[0]["DV_sup"] - cert[0]["C_star_trH"]) < 1e-9)

res["verdict"] = (
    "Lemma 6.1 is FALSE as printed (constant 1); the sharp constant is "
    "C* = 2e^2/(e^2-1) = 2.3130. The paper's downstream use is via '<~' so the "
    "proof strategy is unaffected."
)
C.dump("lemma61", res)
for k in (
    "C_star_closed_form",
    "sup_ratio_momentum_over_gaussians",
    "sup_ratio_gradient_over_gaussians",
    "C_star_is_sharp",
    "verdict",
):
    print(k, "=", res[k])
print(
    "random search:",
    res["random_gaussian_search"]["worst_ratio_gradient"],
    res["random_gaussian_search"]["worst_ratio_momentum"],
    "violations",
    res["random_gaussian_search"]["n_violations_of_printed_lemma_gradient"],
    res["random_gaussian_search"]["n_violations_of_printed_lemma_momentum"],
    "exceeds C*",
    res["random_gaussian_search"]["exceeds_C_star"],
)
print("non-quadratic:", res["nonquadratic_nonconvex_V_quadrature"])