Heaviest-term certified lower bound on exact common information

CHANGELOG.md

@@ -18,6 +18,18 @@ brute-force reference or a closed-form value and has a regression test.
   4-state survey process the 10k and 50k caps took 9 s and 86 s and the 200k cap did not finish in
   15 minutes, so the ladder ran past the survey's hard timeout before any bounds were reported.
 
+### Exact common information: heaviest-term bound
+
+- `heaviest_term_bound` (new module `sofic.generators.exact_ci_bounds`) gives a certified lower bound on
+  `G[X : Y]` by branching on the heaviest term of the decomposition and bounding the remainder by the
+  coverage of its caps, with every remaining atom at most `t/(1 - t)` (`capped_coverage_entropy`, which
+  takes the least concave majorant of the caps). An optimum has at most `2^min - 1` terms (Kumar, Li &
+  El Gamal), so the heaviest weighs at least the reciprocal. `generative_model` runs it after the exact
+  stage when the certified lower is below the estimate and `|S+| + |S-| <= 7`. On the 3 × 3 survey
+  joints `((0,1),(2,·),(0,1))`, `((0,1),(0,2),(0,1))` and `((0,1),(2,1),(0,·))` it certifies `G` to
+  within 1e-5 of the estimate (0.9335, 0.9243, 0.4717; was 0.9118, 0.9056, 0.4436), which closes the
+  first bracket.
+
 ### Generator caps
 
 - `generator_weight_cap` adds the forward analogue of the reverse joint-moment rows: per-state forward

docs/generators/generative_models.rst

@@ -281,6 +281,18 @@ caps (:func:`~sofic.generators.generator_caps.generator_coverage_bound`, the
 ``generator_caps`` stage of the pipeline) exceeds :math:`\op{G}{S^+ : S^-}` on most
 survey processes with positive persistent crypticity.
 
+Coverage only sees the top weights. :func:`~sofic.generators.exact_ci_bounds.heaviest_term_bound`
+uses the whole decomposition: with heaviest term :math:`T = t\, q r^\top`, any exact decomposition has
+entropy :math:`h(t) + (1 - t) H(\lambda')`, where :math:`\lambda'` decomposes the remainder
+:math:`(P - T)/(1 - t)` with atoms at most :math:`t/(1 - t)`. Branching over boxes of :math:`(q, r, t)`
+and bounding the remainder by certified caps of its entrywise upper envelope gives a lower bound on
+:math:`\op{G}{S^+ : S^-}`; an optimum needs at most :math:`2^{\min(|S^+|, |S^-|)} - 1` terms
+:cite:`Kumar2014`, so :math:`t` is at least the reciprocal. On joints whose optimum leaves two terms
+after the heaviest, the bound is tight. The pipeline runs it (the ``heaviest-term bound`` stage) when
+the certified lower bound is still below the exact-common-information estimate and
+:math:`|S^+| + |S^-| \le 7`. On three survey processes with :math:`3 \times 3` joints it raises the
+lower bound to the estimate, closing ``((0,1),(2,·),(0,1))`` at :math:`0.9335` bits.
+
 :func:`~sofic.generators.generator_caps.certified_generator_weight_cap` removes the LP
 tolerance: every branch-and-bound leaf gets an exact bound from rationalized HiGHS duals or an
 exact Farkas certificate of infeasibility (:mod:`sofic.generators.certificates`, after
@@ -336,6 +348,9 @@ Coupled generators and bounds
 .. autofunction:: sofic.generators.coverage_bound.mu_hyperplane_cap
 .. autofunction:: sofic.generators.generator_caps.generator_weight_cap
 .. autofunction:: sofic.generators.generator_caps.generator_coverage_bound
+.. autofunction:: sofic.generators.exact_ci_bounds.heaviest_term_bound
+.. autofunction:: sofic.generators.exact_ci_bounds.capped_coverage_entropy
+.. autoclass:: sofic.generators.exact_ci_bounds.HeaviestTermBoundResult
 .. autofunction:: sofic.generators.generator_caps.certified_generator_weight_cap
 .. autofunction:: sofic.generators.generator_caps.verify_weight_cap_certificate
 .. autoclass:: sofic.generators.generator_caps.WeightCapCertificate

sofic/generators/exact_ci_bounds.py

@@ -0,0 +1,160 @@
+r"""Certified lower bound on the exact common information by branching on the heaviest term.
+
+Any exact decomposition :math:`P = \sum_w \lambda_w q_w r_w^\top` with heaviest term
+:math:`T = t\, q r^\top` has
+
+.. math:: H(\lambda) = h(t) + (1 - t)\, H(\lambda'),
+
+where :math:`h` is the binary entropy and :math:`\lambda'` an exact decomposition of the remainder
+:math:`(P - T)/(1 - t)` whose atoms are at most :math:`t / (1 - t)`. Over a box of :math:`(q, r, t)` the
+remainder lies entrywise below :math:`\bar R = P - t_{lo}\, q_{lo} r_{lo}^\top`, so (``mu_k`` being monotone)
+its ``k`` heaviest atoms weigh at most :math:`\mu_k(\bar R)/(1 - t_{hi})`; the coverage bound with those caps
+and the atom cap :math:`t_{hi}/(1 - t_{hi})` bounds :math:`H(\lambda')`, and :math:`h(t) + (1 - t) C` is
+concave in :math:`t`, so its minimum over the box is at an endpoint. An optimum needs at most
+:math:`2^{\min(|X|, |Y|)} - 1` terms :cite:`Kumar2014`, so :math:`t \ge 1/(2^{\min} - 1)`. Best-first
+branch and bound over the boxes gives a lower bound valid at any node budget; it is tight wherever the
+remainder of an optimum is covered exactly by its caps (e.g. two remaining terms).
+"""
+
+from __future__ import annotations
+
+import heapq
+import itertools
+import time
+from dataclasses import dataclass
+
+import numpy as np
+
+from sofic.generators.coverage_bound import mu1_vertex, mu_k_spectral_cap
+from sofic.generators.generator_caps import _tighten
+
+Box = tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float, float]
+
+
+@dataclass(frozen=True)
+class HeaviestTermBoundResult:
+    """``value`` bounds :math:`\\op{G}{X : Y}` from below; ``reached`` says whether it met ``target``."""
+
+    value: float
+    nodes: int
+    open_boxes: int
+    reached: bool
+
+
+def _entropy(p: np.ndarray) -> float:
+    p = p[p > 0]
+    return float(-(p * np.log2(p)).sum())
+
+
+def capped_coverage_entropy(caps: list[float], atom: float, n_max: int = 64) -> float:
+    """Lower bound on :math:`H(\\pi)` when the ``k`` heaviest atoms weigh at most ``caps[k-1]`` and every atom
+    at most ``atom``.
+
+    The ``k``-th partial sum of the sorted vector is at most :math:`S_k = \\min(C_k, k\\,\\mathrm{atom}, 1)`, with
+    :math:`C` the least concave majorant of :math:`(0, \\mathrm{caps}, 1, 1, \\ldots)`. :math:`S` is concave, so its
+    increments form a sorted vector that majorizes every feasible one, and entropy is Schur concave.
+    """
+    pts = [0.0] + [min(1.0, float(c)) for c in caps] + [1.0] * max(1, n_max - len(caps))
+    pts = np.maximum.accumulate(np.array(pts))
+    hull = [0]
+    for k in range(1, len(pts)):
+        while len(hull) >= 2:
+            i, j = hull[-2], hull[-1]
+            if (pts[j] - pts[i]) * (k - i) <= (pts[k] - pts[i]) * (j - i):
+                hull.pop()
+            else:
+                break
+        hull.append(k)
+    C = np.interp(np.arange(len(pts)), hull, pts[hull])
+    S = np.minimum(np.minimum(C, atom * np.arange(len(pts))), 1.0)
+    return _entropy(np.diff(S))
+
+
+def _box_bound(P: np.ndarray, box: Box) -> float | None:
+    qlo, _qhi, rlo, _rhi, tlo, thi = box
+    T_lo = tlo * np.outer(qlo, rlo)
+    if np.any(T_lo > P + 1e-15):
+        return None
+    scale = 1.0 - thi
+    if scale <= 0.0:
+        return 0.0
+    Rbar = np.maximum(P - T_lo, 0.0)
+    mu1 = mu1_vertex(Rbar)
+    caps = [min(1.0, (mu1.value if mu1 is not None else float(Rbar.sum())) / scale)]
+    caps += [min(1.0, mu_k_spectral_cap(Rbar, j) / scale) for j in range(2, min(P.shape))]
+    cover = capped_coverage_entropy(caps, thi / scale)
+    h = [_entropy(np.array([t, 1.0 - t])) + (1.0 - t) * cover for t in (tlo, thi)]
+    return min(h)
+
+
+def _split(box: Box, dim: int) -> list[Box]:
+    qlo, qhi, rlo, rhi, tlo, thi = box
+    na = len(qlo)
+    out: list[Box] = []
+    for side in (0, 1):
+        q0, q1, r0, r1, t0, t1 = qlo.copy(), qhi.copy(), rlo.copy(), rhi.copy(), tlo, thi
+        if dim < na:
+            mid = 0.5 * (q0[dim] + q1[dim])
+            if side == 0:
+                q1[dim] = mid
+            else:
+                q0[dim] = mid
+        elif dim < na + len(rlo):
+            j = dim - na
+            mid = 0.5 * (r0[j] + r1[j])
+            if side == 0:
+                r1[j] = mid
+            else:
+                r0[j] = mid
+        else:
+            mid = 0.5 * (t0 + t1)
+            t0, t1 = (t0, mid) if side == 0 else (mid, t1)
+        tq, tr = _tighten(q0, q1), _tighten(r0, r1)
+        if tq is not None and tr is not None:
+            out.append((tq[0], tq[1], tr[0], tr[1], t0, t1))
+    return out
+
+
+def heaviest_term_bound(
+    P: np.ndarray,
+    *,
+    target: float | None = None,
+    max_nodes: int = 200_000,
+    deadline: float | None = None,
+) -> HeaviestTermBoundResult:
+    """Certified lower bound on :math:`\\op{G}{X : Y}` for the joint ``P``; see the module docstring.
+
+    Branches best first on the widest coordinate of the heaviest term's box and stops once every open box
+    is at least ``target``, after ``max_nodes`` boxes, or at ``deadline`` (a :func:`time.monotonic` value).
+    The caps are :func:`~sofic.generators.coverage_bound.mu1_vertex` and the spectral caps, as in the
+    certified coverage bound. The work grows quickly with ``|X| + |Y|``; it is meant for small joints.
+    """
+    P = np.asarray(P, dtype=float)
+    P = P / P.sum()
+    na, nb = P.shape
+    mu1 = mu1_vertex(P)
+    t_hi = min(1.0, mu1.value if mu1 is not None else 1.0)
+    t_lo = 1.0 / (2 ** min(na, nb) - 1)
+    root: Box = (np.zeros(na), np.ones(na), np.zeros(nb), np.ones(nb), t_lo, t_hi)
+    first = _box_bound(P, root)
+    heap: list[tuple[float, int, Box]] = [] if first is None else [(first, 0, root)]
+    order = itertools.count(1)
+    nodes = 0
+    goal = np.inf if target is None else target
+    while heap and nodes < max_nodes and (deadline is None or time.monotonic() < deadline):
+        bound, _, box = heap[0]
+        if bound >= goal:
+            break
+        heapq.heappop(heap)
+        nodes += 1
+        qlo, qhi, rlo, rhi, tlo, thi = box
+        widths = np.concatenate([qhi - qlo, rhi - rlo, [thi - tlo]])
+        for child in _split(box, int(np.argmax(widths))):
+            value = _box_bound(P, child)
+            if value is not None:
+                heapq.heappush(heap, (max(bound, value), next(order), child))
+    value = heap[0][0] if heap else float("inf")
+    return HeaviestTermBoundResult(value, nodes, len(heap), target is not None and value >= target)
+
+
+__all__ = ["HeaviestTermBoundResult", "capped_coverage_entropy", "heaviest_term_bound"]

sofic/generators/generative_search.py

@@ -60,6 +60,7 @@
     targeted_line_split_seeds,
 )
 from sofic.generators.epsilon_machine import EpsilonMachine
+from sofic.generators.exact_ci_bounds import heaviest_term_bound
 from sofic.generators.generator_bounds import (
     flow_bound_result,
     maximal_rectangles,
@@ -75,6 +76,9 @@
 _TARGETED_GAP = 1e-2
 # ``I[S+ : S- | W]`` (bits) above which a channel is not treated as an exact decomposition.
 _SEPARATION_ATOL = 1e-9
+# The heaviest-term bound branches over |S+| + |S-| - 1 coordinates; run it only up to this many causal
+# states (|S+| + |S-|), where it converges within a stage budget.
+_HEAVIEST_TERM_MAX_DIMS = 7
 
 
 def _entropy(p: np.ndarray) -> float:
@@ -583,6 +587,28 @@ def maybe_return_early() -> GenerativeModelResult | None:
                 "cg exceeds g": cg_exceeds_g,
             }
         )
+        if (
+            not wyner
+            and coverage == "certified"
+            and g_estimate is not None
+            and lower < g_estimate - _CERT_TOL
+            and problem.n_plus + problem.n_minus <= _HEAVIEST_TERM_MAX_DIMS
+            and not out_of_time("heaviest-term bound")
+        ):
+            joint_matrix = problem.lifted.plus.T @ (problem.p_pair[:, None] * problem.lifted.minus)
+            # aim inside the certification tolerance of the best upper so far, so a closed bracket certifies
+            aim = min(g_estimate, min((c[0] for c in candidates), default=np.inf)) - 0.5 * _CERT_TOL
+            htb = heaviest_term_bound(joint_matrix, target=aim, deadline=share(0.15))
+            log.append(
+                {
+                    "stage": "heaviest-term bound (certified)",
+                    "value": htb.value,
+                    "nodes": htb.nodes,
+                    "reached": htb.reached,
+                }
+            )
+            if htb.value > lower + 1e-9:
+                lower, lower_source, lower_certified = htb.value, "heaviest-term bound (certified)", True
         if bilevel and n_generating == 0 and optima and not out_of_time("bilevel search"):
             for ent, viol, c in bilevel_search(bidir, optima[:8], deadline=share(0.15))[:4]:
                 log.append({"stage": "bilevel search", "entropy": ent, "kernel violation": viol})

tests/test_exact_ci_bounds.py

@@ -0,0 +1,63 @@
+"""Certified heaviest-term lower bound on the exact common information."""
+
+from __future__ import annotations
+
+import numpy as np
+from hypothesis import given, settings
+from hypothesis import strategies as st
+
+from sofic.generators.coupled_generators import _decompose
+from sofic.generators.exact_ci_bounds import capped_coverage_entropy, heaviest_term_bound
+
+
+def _h(p: float) -> float:
+    return float(-p * np.log2(p) - (1 - p) * np.log2(1 - p))
+
+
+def test_capped_coverage_is_below_a_feasible_vector_the_greedy_misses() -> None:
+    # caps (0.5, 0.6, 1.0): the sorted vector (0.34, 0.26, 0.26, 0.14) is feasible with S_3 = 0.86 > 0.7
+    y = np.array([0.34, 0.26, 0.26, 0.14])
+    assert capped_coverage_entropy([0.5, 0.6, 1.0], 1.0) <= float(-(y * np.log2(y)).sum())
+
+
+def test_dsbs_reaches_its_exact_common_information() -> None:
+    eps = 0.45
+    D = np.array([[1 - eps, eps], [eps, 1 - eps]]) / 2
+    G = _h(1 / 11)  # Kumar-Li-El Gamal Prop. 8 at eps = 9/20: weights (10/11, 1/11)
+    result = heaviest_term_bound(D, target=G - 1e-4, max_nodes=50_000)
+    assert result.reached
+    assert result.value <= G + 1e-9
+
+
+@settings(max_examples=8, deadline=None)
+@given(st.integers(0, 2**32 - 1))
+def test_bound_is_below_every_exact_decomposition(seed: int) -> None:
+    rng = np.random.default_rng(seed)
+    P = rng.dirichlet(np.ones(9)).reshape(3, 3)
+    result = heaviest_term_bound(P, max_nodes=2_000)
+    for k in (3, 4, 5):
+        lam, A, B, residual = _decompose(
+            P, rng.dirichlet(np.ones(k)), rng.dirichlet(np.ones(3), size=k), rng.dirichlet(np.ones(3), size=k)
+        )
+        if residual <= 1e-10:
+            p = lam[lam > 1e-12]
+            assert result.value <= float(-(p * np.log2(p)).sum()) + 1e-9
+
+
+def test_survey_case_closes_against_its_exact_common_information() -> None:
+    """((0,1),(2,·),(0,1)): the bound rises to within 0.01 of the best exact decomposition (0.93349 bits)."""
+    P = np.array([[0.107, 0.3019, 0.1449], [0.0781, 0.1449, 0.0], [0.0379, 0.107, 0.0781]])
+    P = P / P.sum()
+    best = np.inf
+    rng = np.random.default_rng(0)
+    for k in (2, 3, 4):
+        for _ in range(20):
+            lam, A, B, residual = _decompose(
+                P, rng.dirichlet(np.ones(k)), rng.dirichlet(np.ones(3), size=k), rng.dirichlet(np.ones(3), size=k)
+            )
+            if residual <= 1e-10:
+                p = lam[lam > 1e-12]
+                best = min(best, float(-(p * np.log2(p)).sum()))
+    result = heaviest_term_bound(P, target=best - 1e-2, max_nodes=100_000)
+    assert result.reached
+    assert best - 1e-2 <= result.value <= best + 1e-9