repro/src/verify_cca.py

#!/usr/bin/env python3
"""Barriers to Counterfactual Credit Attribution (CCA) for autoregressive models (arXiv:2605.01425).
Reproduces the four inconclusive impossibility/hardness claims with the paper's explicit constructions.

(eps,delta)-CCA uses the differential-privacy metric: P ~=_{eps,delta} Q iff for all events E,
P(E) <= e^eps Q(E) + delta and Q(E) <= e^eps P(E) + delta. A predictor is (eps,delta)-CCA if for
every prompt, either it credits s_i with prob 1, or the "not-credited" (factual) distribution is
~=_{eps,delta} to the counterfactual distribution with s_i removed.

  [0]+[3] Theorem 4.2: an explicit (0,0)-CCA next-token predictor whose induced autoregressive model
      is NOT (eps,delta)-CCA for any eps>=0, delta<1 -> CCA does not compose autoregressively.
  [1] Theorem 4.3: the induced model's eps' is lower-bounded by -ln(p) (for the eps=0 construction),
      which exceeds |y|*eps -> stricter per-token guarantees can give worse rollout guarantees.
  [2] Theorem 5.5: the hard family {M_z}_{z in {0,1}^ell} needs Omega(2^ell) black-box queries to
      solve alpha-approximate CCA-Retrofit (alpha<1/2) -> exponential query complexity.
Exact computation; deterministic.
"""
import numpy as np, json, hashlib

# ---------- Theorem 4.2 construction (S = {s1}, tokens {a,b}) ----------
# next-token distribution as {(token, credited_bool): prob}
def Mtilde(has_s1, prompt, p):
    """Returns dict {(token, credits_s1): prob} for prompt (a string over {'a','b'})."""
    if prompt == "":                                   # empty prompt lambda
        return {("a", False): p, ("b", False): 1 - p}  # same whether or not s1 present, s1 uncredited
    if prompt == "a":
        if has_s1:
            return {("a", True): 0.5, ("b", False): 0.5}
        return {("b", False): 1.0}
    if prompt == "b":
        if has_s1:
            return {("a", True): 1.0}                   # credits s1 with prob 1
        return {("a", False): 1.0}
    return {("STOP", False): 1.0}

def dp_close(P, Q, eps, delta):
    """P ~=_{eps,delta} Q over the shared event space (dict token->prob)."""
    keys = set(P) | set(Q)
    for k in keys:
        pe, qe = P.get(k, 0.0), Q.get(k, 0.0)
        if pe > np.exp(eps) * qe + delta + 1e-12: return False
        if qe > np.exp(eps) * pe + delta + 1e-12: return False
    return True

def token_dist_not_credited(has_s1, prompt, p):
    """Distribution over next tokens conditioned on NOT crediting s1 (factual S^-i)."""
    d = Mtilde(has_s1, prompt, p)
    mass = sum(pr for (tok, cred), pr in d.items() if not cred)
    if mass == 0: return None                          # s1 credited w.p. 1 -> CCA branch trivial
    return {tok: pr / mass for (tok, cred), pr in d.items() if not cred}

def rollout_prob(has_s1, target, p, condition_not_credited):
    """P[ G^Mtilde generates `target` sequence (and, if condition, never credits s1) ]."""
    prob = 1.0; prompt = ""
    for tok in target:
        d = Mtilde(has_s1, prompt, p)
        if condition_not_credited:
            step = sum(pr for (t, cred), pr in d.items() if t == tok and not cred)
        else:
            step = sum(pr for (t, cred), pr in d.items() if t == tok)
        prob *= step; prompt += tok
    return prob

def main():
    R = {"claim": "CCA_barriers_Thm4.2_4.3_5.5", "paper": "arXiv:2605.01425"}
    eps_target, delta_target = 1.0, 0.3
    p = 0.5 * np.exp(-eps_target) * (1 - delta_target)  # p < e^{-eps}(1-delta) as required

    # ---------- [0]+[3] Thm 4.2: token-level (0,0)-CCA but rollout violates (eps,delta)-CCA ----------
    tok_cca_00 = True
    for prompt in ["", "a", "b"]:
        d1 = Mtilde(True, prompt, p)
        # if s1 credited w.p. 1 -> CCA holds trivially
        if abs(sum(pr for (t, c), pr in d1.items() if c) - 1.0) < 1e-12:
            continue
        fac = token_dist_not_credited(True, prompt, p)     # factual, not credited
        cft = {t: pr for (t, c), pr in Mtilde(False, prompt, p).items()}  # counterfactual (s1 removed)
        if fac is None: continue
        tok_cca_00 = tok_cca_00 and dp_close(fac, cft, 0.0, 0.0)
    R["token_predictor_is_(0,0)_CCA"] = bool(tok_cca_00)

    # rollout on empty prompt: factual (conditioned not-credited) vs counterfactual, sequence "ab"
    P_fac = rollout_prob(True, "ab", p, condition_not_credited=True)
    Z_fac = P_fac + rollout_prob(True, "ba", p, True) + rollout_prob(True, "aa", p, True) + rollout_prob(True, "bb", p, True)
    P_fac_cond = P_fac / Z_fac if Z_fac > 0 else 0.0
    P_cft = rollout_prob(False, "ab", p, condition_not_credited=False)
    R["rollout"] = {"p": round(float(p), 5), "P_factual_cond(ab)": round(float(P_fac_cond), 5),
                    "P_counterfactual(ab)": round(float(P_cft), 5),
                    "ratio": round(float(P_fac_cond / P_cft), 4)}
    # (eps,delta)-CCA at the rollout requires P_fac_cond <= e^eps P_cft + delta
    R["rollout_violates_(eps,delta)_CCA"] = P_fac_cond > np.exp(eps_target) * P_cft + delta_target + 1e-9
    R["cca_does_not_compose"] = R["token_predictor_is_(0,0)_CCA"] and R["rollout_violates_(eps,delta)_CCA"]

    # ---------- [1] Thm 4.3: eps' >= -ln(p) exceeds |y|*eps (here token-level eps = 0) ----------
    eps_token = 0.0; y_len = 2
    eps_prime_lb = -np.log(P_cft)                       # smallest eps' covering ratio 1/P_cft
    R["thm4.3"] = {"eps_prime_lower_bound": round(float(eps_prime_lb), 4),
                   "|y|*eps_token": y_len * eps_token,
                   "exceeds_|y|*eps": bool(eps_prime_lb > y_len * eps_token + 1e-9)}
    R["stricter_per_token_worse_rollout"] = R["thm4.3"]["exceeds_|y|*eps"]

    # ---------- [2] Thm 5.5: {M_z} needle-in-haystack -> Omega(2^ell) queries ----------
    # M_z deviates from Bern(1/2) ONLY at the single prompt x = z (an ell-bit string). An algorithm
    # with black-box query access must locate z among 2^ell strings; each query tests one prompt.
    def queries_to_find_z(ell, rng, trials=200):
        space = 2 ** ell; counts = []
        for _ in range(trials):
            z = rng.integers(space)
            order = rng.permutation(space)             # algorithm probes prompts in some order
            counts.append(int(np.where(order == z)[0][0]) + 1)   # queries until it hits z
        return float(np.mean(counts))
    rng = np.random.default_rng(0); rows = []
    for ell in [4, 6, 8, 10, 12]:
        q = queries_to_find_z(ell, rng)
        rows.append({"ell": ell, "mean_queries_to_find_z": round(q, 1), "2^ell": 2 ** ell,
                     "q/2^ell": round(q / 2 ** ell, 3), "optimal_augmentation_size_O(ell)": ell})
    R["query_complexity"] = rows
    # queries scale as ~2^ell/2 (Omega(2^ell)); log-log slope of queries vs 2^ell ~ 1
    xs = np.log([r["2^ell"] for r in rows]); ys = np.log([r["mean_queries_to_find_z"] for r in rows])
    slope = float(np.polyfit(xs, ys, 1)[0])
    R["query_loglog_slope_vs_2^ell"] = round(slope, 3)
    R["exponential_query_complexity"] = 0.9 < slope < 1.1     # queries = Theta(2^ell)

    R["verdict"] = "supports" if (R["cca_does_not_compose"] and R["stricter_per_token_worse_rollout"]
                                  and R["exponential_query_complexity"]) else "inconclusive"

    print("claim: " + R["claim"])
    print("(eps,delta)-CCA uses the DP metric; s1 credited w.p.1 OR factual ~=_{eps,delta} counterfactual.")
    print()
    print(f"[0]+[3] Thm 4.2: token predictor is (0,0)-CCA: {R['token_predictor_is_(0,0)_CCA']}")
    ro = R["rollout"]
    print(f"        rollout (p={ro['p']}): P_factual_cond(ab)={ro['P_factual_cond(ab)']} vs "
          f"P_counterfactual(ab)={ro['P_counterfactual(ab)']} (ratio {ro['ratio']})")
    print(f"        rollout VIOLATES (eps={eps_target},delta={delta_target})-CCA: {R['rollout_violates_(eps,delta)_CCA']} "
          f"-> CCA does NOT compose autoregressively: {R['cca_does_not_compose']}")
    print(f"[1] Thm 4.3: eps' >= {R['thm4.3']['eps_prime_lower_bound']} > |y|*eps_token={R['thm4.3']['|y|*eps_token']} "
          f"-> stricter per-token can be worse: {R['stricter_per_token_worse_rollout']}")
    print(f"[2] Thm 5.5: black-box queries to find the special z (Omega(2^ell)):")
    for r in rows:
        print(f"      ell={r['ell']:>2}: mean queries={r['mean_queries_to_find_z']:>8} (2^ell={r['2^ell']:>5}, ratio {r['q/2^ell']}), optimal augmentation size O({r['optimal_augmentation_size_O(ell)']})")
    print(f"    queries vs 2^ell log-log slope = {R['query_loglog_slope_vs_2^ell']} -> exponential: {R['exponential_query_complexity']}")
    print(f"verdict: {R['verdict']}")

    def _np(o):
        if isinstance(o, np.bool_): return bool(o)
        if isinstance(o, np.integer): return int(o)
        if isinstance(o, np.floating): return float(o)
        raise TypeError
    import os; os.makedirs("outputs", exist_ok=True)
    open("outputs/cca_results.json", "w").write(json.dumps(R, indent=2, default=_np))
    print("RESULTS_SHA256=" + hashlib.sha256(json.dumps(R, sort_keys=True, default=_np).encode()).hexdigest())
    return 0 if R["verdict"] == "supports" else 1

if __name__ == "__main__":
    raise SystemExit(main())