reproduce.py

#!/usr/bin/env python3
"""Deterministic audit for OpenReview QYA0Q28ssf."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


CLAIMS = [
    "Under finite drift energy assumptions, the intra-generation KL divergence between the generated distribution and the training target is bounded above by (1/2)ε̂ᵢ², the learned-path score error energy (Proposition 3.1).",
    "A matching lower bound on the chi-squared divergence, χ²(p̂^(i+1) ∥ qᵢ) ≥ (1/4)ηᵢε²⋆,ᵢ − C·ε⁴⋆,ᵢ, holds for small score errors (ε²⋆,ᵢ ≤ 1), where ηᵢ ∈ [0,1] is the observability coefficient (Proposition 3.3, Definition 3.2).",
    "Combining the upper and lower bounds yields a two-sided equivalence χ²(p̂^(i+1) ∥ qᵢ) ≍ ε²⋆,ᵢ in the perturbative regime (Theorem 3.4).",
    "When the score-error series ∑ᵢ ε²⋆,ᵢ diverges, the accumulated divergence across generations cannot vanish and has a non-zero floor limsup Dᵢ ≥ αη̄ε̄ / [16(1+(1-α)²)] (Proposition 4.1).",
    "When ∑ᵢ ε²⋆,ᵢ converges, accumulated divergence across generations satisfies D_{N+1} + C_bias ≍ Σᵢ(1-α)^{2(N-i)}ε²⋆,ᵢ + (1-α)^{2(N+1-i₀)}D_{i₀}, decaying geometrically at rate (1-α)² per generation, where α is the fraction of fresh data mixed in each round (Theorem 4.2).",
    "The α-dependent tradeoff between drift and stability is empirically confirmed on a 10D Gaussian mixture and on Fashion-MNIST/CIFAR-10, with low α producing collapse-driven drift and high α maintaining stability (Figure 1).",
]

SOURCE_HASHES = {
    "paper_v1.pdf": "fe979c798cd48a5af02f6c647ecc29b2d6a937841adfa8ceedb492b1c2d81583",
    "source_v1.tar": "729a62da953ae2f9458f5e938ba53a7d6dbd8efbac6f079028fa6af6df0c0a06",
    "Fig1_Accumulation_alpha0.1.pdf": "45eebb04f9604968d8b8287f242306a41ffd0b76a16da9524b935800e1d1c5e7",
    "Fig1_Accumulation_alpha0.5.pdf": "0ae881662d2dd33a959a6c7c0cad78f74d0ea052ba9be66bba44dbc284a53798",
    "Fig1_Accumulation_alpha0.9.pdf": "b4522ecd3f09190dfbbfc3763809df8d650730913ca548e0ca78002e62cd9387",
    "Memory_Heatmap.pdf": "3cf3f1e29336c38af239a264149a755527e1137e8bdb64c57370c721ff8222d9",
    "fashionmnist_memory_alpha_0.1.png": "f790ec1387c2124743147fb70b562525210f4a86d3010d520ee33e8ecefea745",
    "fashionmnist_memory_alpha_0.5.png": "a95679ee2eda087e40e6b31416129bf1ebe20428849cf377431c0c32fa5dd097",
    "fashionmnist_memory_alpha_0.9.png": "1e3924dc05187a1a060ab085d2fce7b42f29efb4b10ec280cb4c75a2c5b423f6",
    "cifar10_observability.pdf": "1d44080768d8ea7a58122ac6938690d2275bc61e9fa0aaa18b5e0e496a094d4d",
}


def sha256(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            value.update(block)
    return value.hexdigest()


def write_csv(path: Path, rows: list[dict]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def gaussian_kl(error_energy: float) -> float:
    """KL(N(mu,I)||N(0,I)) with ||mu||^2=error_energy."""
    return 0.5 * error_energy


def gaussian_chi2(error_energy: float) -> float:
    """chi^2(N(mu,I)||N(0,I)) with ||mu||^2=error_energy."""
    return math.expm1(error_energy)


def discounted_convolution(errors: np.ndarray, alpha: float, initial: float = 0.0) -> np.ndarray:
    beta = (1.0 - alpha) ** 2
    values = np.empty(len(errors), dtype=float)
    current = float(initial)
    for index, error in enumerate(errors):
        current = beta * current + float(error)
        values[index] = current
    return values


def direct_convolution(errors: np.ndarray, alpha: float, initial: float = 0.0) -> np.ndarray:
    beta = (1.0 - alpha) ** 2
    rows = []
    for generation in range(len(errors)):
        powers = beta ** np.arange(generation, -1, -1)
        rows.append(float(np.dot(powers, errors[: generation + 1]) + beta ** (generation + 1) * initial))
    return np.asarray(rows)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", default="outputs")
    parser.add_argument("--source-dir", default="sources")
    args = parser.parse_args()
    out = Path(args.output); src = Path(args.source_dir)
    out.mkdir(parents=True, exist_ok=True)

    source_rows = []
    for filename, expected in SOURCE_HASHES.items():
        actual = sha256(src / filename)
        source_rows.append({"file": filename, "expected_sha256": expected, "actual_sha256": actual, "match": expected == actual})

    source_claim_rows = [
        {"claim": 1, "source": "arXiv v1 Proposition 3.1, Eq. 16", "status": "supported", "scope": "A1 finite learned-path energy and A2 true Girsanov martingale; KL direction is learned endpoint to training target"},
        {"claim": 2, "source": "arXiv v1 Definition 3.2 and Proposition 3.3, Eqs. 17/20", "status": "supported", "scope": "A1-A4, i>=i0, epsilon_star^2<=1; C depends on regularity/observability constants"},
        {"claim": 3, "source": "arXiv v1 Theorem 3.4", "status": "supported", "scope": "A1-A4, positive observability and tail effects; equivalence constants are not one"},
        {"claim": 4, "source": "arXiv v1 Proposition 4.1", "status": "literal_conflation_unsupported", "scope": "non-summability yields non-summable drift/cannot vanish; the displayed positive floor additionally requires a uniform per-generation error floor"},
        {"claim": 5, "source": "arXiv v1 Theorem 4.2, Eq. 23", "status": "supported", "scope": "A1-A5, positive observability, small summable errors, and C_bias; asymp hides constants"},
        {"claim": 6, "source": "arXiv v1 Figures 1,5-10 and Appendix G", "status": "supported", "scope": "10D GMM is quantitative; Fashion/CIFAR divergence uses learned-feature Gaussian proxies and is source evidence, not rerun here"},
    ]

    # Claim 1: a Gaussian shift realizes equality in the path-KL identity.
    kl_rows = []
    for dimension in [1, 2, 5, 10]:
        for amplitude in [0.02, 0.05, 0.1, 0.2, 0.4]:
            energy = dimension * amplitude * amplitude
            kl = gaussian_kl(energy)
            kl_rows.append({
                "dimension": dimension, "coordinate_shift": amplitude,
                "learned_path_error_energy": f"{energy:.12e}", "endpoint_kl": f"{kl:.12e}",
                "half_energy_upper": f"{0.5*energy:.12e}", "upper_bound_holds": kl <= 0.5*energy + 1e-15,
                "ratio_to_half_energy": f"{kl/(0.5*energy):.12f}",
            })

    # Claims 2-3: exact observability for deterministic Brownian integrands.
    # eta=(integral e)^2 / integral e^2 on [0,1]. Each base profile has unit energy.
    profiles = [
        ("constant", 1.0),
        ("linear_increasing", 0.75),
        ("mean_zero_time_control", 0.0),
    ]
    observability_rows = []
    sandwich_rows = []
    for profile, eta in profiles:
        for scale in [0.02, 0.05, 0.1, 0.2, 0.35, 0.5]:
            path_energy = scale * scale
            endpoint_energy = eta * path_energy
            chi2 = gaussian_chi2(endpoint_energy)
            observability_rows.append({
                "profile": profile, "scale": scale, "path_energy": f"{path_energy:.12e}",
                "observability_eta": f"{eta:.12f}", "endpoint_shift_energy": f"{endpoint_energy:.12e}",
                "endpoint_chi2": f"{chi2:.12e}", "eta_in_unit_interval": 0.0 <= eta <= 1.0,
            })
            if eta > 0:
                leading = eta * path_energy
                sandwich_rows.append({
                    "profile": profile, "scale": scale, "path_energy": f"{path_energy:.12e}",
                    "observability_eta": f"{eta:.12f}", "chi2": f"{chi2:.12e}",
                    "chi2_over_eta_energy": f"{chi2/leading:.12f}",
                    "one_quarter_leading_bound": f"{0.25*leading:.12e}",
                    "leading_lower_holds_without_remainder": chi2 >= 0.25*leading,
                    "chi2_over_path_energy": f"{chi2/path_energy:.12f}",
                })

    # Claim 4: keep the two distinct premises of Proposition 4.1 distinct.
    generation_count = 6000
    constant_error = np.full(generation_count, 0.02)
    harmonic_error = 0.02 / (np.arange(generation_count, dtype=float) + 1.0)
    persistent_rows = []
    for alpha in [0.1, 0.5, 0.9]:
        beta = (1-alpha)**2
        constant_d = discounted_convolution(constant_error, alpha)
        harmonic_d = discounted_convolution(harmonic_error, alpha)
        eta_floor = 0.4
        source_floor = alpha * eta_floor * 0.02 / (16.0*(1.0+beta))
        persistent_rows.append({
            "alpha": alpha, "beta": f"{beta:.12f}",
            "constant_error_series_diverges": True, "constant_error_has_positive_floor": True,
            "constant_final_divergence": f"{constant_d[-1]:.12e}", "constant_exact_limit": f"{0.02/(1-beta):.12e}",
            "source_floor_with_eta_0_4": f"{source_floor:.12e}", "constant_exceeds_source_floor": constant_d[-1] >= source_floor,
            "harmonic_error_partial_sum": f"{harmonic_error.sum():.12e}", "harmonic_series_diverges_as_horizon_grows": True,
            "harmonic_error_has_positive_floor": False, "harmonic_final_divergence": f"{harmonic_d[-1]:.12e}",
            "harmonic_tail_tends_to_zero": harmonic_d[-1] < harmonic_d[generation_count//2],
            "literal_claim_floor_follows_from_divergence_alone": False,
        })

    # Claim 5: summable errors and the exact geometric convolution identity.
    summable_errors = 0.08 / np.square(np.arange(1, 241, dtype=float))
    accumulation_rows = []
    memory_rows = []
    for alpha in [0.0, 0.1, 0.5, 0.9]:
        beta = (1-alpha)**2
        recurrence = discounted_convolution(summable_errors, alpha, initial=0.07)
        direct = direct_convolution(summable_errors, alpha, initial=0.07)
        for generation in [0, 1, 2, 5, 10, 20, 50, 100, 239]:
            accumulation_rows.append({
                "alpha": alpha, "beta": f"{beta:.12f}", "generation": generation,
                "error_energy": f"{summable_errors[generation]:.12e}",
                "recurrence_value": f"{recurrence[generation]:.12e}", "direct_convolution": f"{direct[generation]:.12e}",
                "identity_residual": f"{abs(recurrence[generation]-direct[generation]):.3e}",
            })
        if beta == 0:
            half_life = 0.0
        elif beta == 1:
            half_life = math.inf
        else:
            half_life = math.log(0.5)/math.log(beta)
        memory_rows.append({
            "alpha": alpha, "geometric_rate_beta": f"{beta:.12f}",
            "one_generation_retention": f"{beta:.12f}",
            "five_generation_retention": f"{beta**5:.12e}",
            "twenty_generation_retention": f"{beta**20:.12e}",
            "memory_half_life_generations": "inf" if math.isinf(half_life) else f"{half_life:.12f}",
            "final_discounted_divergence": f"{recurrence[-1]:.12e}",
        })

    # Claim 6 source assets: record the exact qualitative statements they depict.
    experiment_rows = [
        {"dataset": "10D five-component Gaussian mixture", "source_assets": "Fig1_Accumulation_alpha0.1/0.5/0.9.pdf; Memory_Heatmap.pdf", "fresh_data_levels": "0.1,0.5,0.9", "reported_outcome": "low alpha accumulates drift; high alpha has short geometric memory", "evidence_type": "pinned source experiment"},
        {"dataset": "Fashion-MNIST", "source_assets": "fashionmnist_memory_alpha_0.1/0.5/0.9.png", "fresh_data_levels": "0.1,0.5,0.9", "reported_outcome": "memory heatmaps show longer persistence at low alpha", "evidence_type": "pinned source experiment"},
        {"dataset": "CIFAR-10", "source_assets": "cifar10_observability.pdf", "fresh_data_levels": "0.1,0.5,0.9", "reported_outcome": "observability stabilizes around 0.25-0.35; feature-space proxy is used", "evidence_type": "pinned source experiment"},
    ]

    gates = {
        "all_source_hashes_match": all(row["match"] for row in source_rows),
        "all_six_claims_mapped": len(source_claim_rows) == 6,
        "claim1_kl_upper_all_20_panels": all(row["upper_bound_holds"] for row in kl_rows),
        "claim1_gaussian_path_identity_exact": all(abs(float(row["ratio_to_half_energy"])-1.0) < 1e-12 for row in kl_rows),
        "claim2_eta_unit_interval": all(row["eta_in_unit_interval"] for row in observability_rows),
        "claim2_lower_leading_term_all_observable_panels": all(row["leading_lower_holds_without_remainder"] for row in sandwich_rows),
        "claim2_unobservable_control_has_energy": any(row["profile"] == "mean_zero_time_control" and float(row["path_energy"]) > 0 and float(row["endpoint_chi2"]) == 0 for row in observability_rows),
        "claim3_two_sided_ratio_finite_positive": all(0.70 <= float(row["chi2_over_path_energy"]) <= 1.30 for row in sandwich_rows),
        "claim3_small_error_ratio_near_observability": all(1.0 <= float(row["chi2_over_eta_energy"]) <= 1.04 for row in sandwich_rows if float(row["path_energy"]) <= 0.04),
        "claim4_constant_floor_panels_pass": all(row["constant_exceeds_source_floor"] for row in persistent_rows),
        "claim4_literal_conflation_detected": all(not row["literal_claim_floor_follows_from_divergence_alone"] for row in persistent_rows),
        "claim4_harmonic_nonfloor_control": all(row["harmonic_tail_tends_to_zero"] and not row["harmonic_error_has_positive_floor"] for row in persistent_rows),
        "claim5_summable_error_series": float(summable_errors.sum()) < 0.14,
        "claim5_convolution_identity": max(float(row["identity_residual"]) for row in accumulation_rows) < 1e-14,
        "claim5_geometric_rates_exact": all(abs(float(row["geometric_rate_beta"]) - (1-float(row["alpha"]))**2) < 1e-12 for row in memory_rows),
        "claim5_fresh_data_shortens_memory": float(memory_rows[1]["twenty_generation_retention"]) > float(memory_rows[2]["twenty_generation_retention"]) > float(memory_rows[3]["twenty_generation_retention"]),
        "claim5_no_fresh_data_destructive_control": memory_rows[0]["memory_half_life_generations"] == "inf" and float(memory_rows[0]["twenty_generation_retention"]) == 1.0,
        "claim6_all_three_source_datasets_pinned": {row["dataset"] for row in experiment_rows} == {"10D five-component Gaussian mixture", "Fashion-MNIST", "CIFAR-10"},
        "claim6_three_alpha_levels_pinned": all(row["fresh_data_levels"] == "0.1,0.5,0.9" for row in experiment_rows),
    }

    write_csv(out / "source_pins.csv", source_rows)
    write_csv(out / "source_claim_audit.csv", source_claim_rows)
    write_csv(out / "gaussian_kl_upper.csv", kl_rows)
    write_csv(out / "observability_controls.csv", observability_rows)
    write_csv(out / "two_sided_sandwich.csv", sandwich_rows)
    write_csv(out / "persistent_error_audit.csv", persistent_rows)
    write_csv(out / "summable_accumulation.csv", accumulation_rows)
    write_csv(out / "fresh_data_memory.csv", memory_rows)
    write_csv(out / "source_experiments.csv", experiment_rows)

    figure, axes = plt.subplots(2, 2, figsize=(12, 8))
    for profile in ["constant", "linear_increasing"]:
        rows = [row for row in sandwich_rows if row["profile"] == profile]
        axes[0, 0].loglog([float(row["path_energy"]) for row in rows], [float(row["chi2"]) for row in rows], "o-", label=profile)
    axes[0, 0].set(title="Observable Gaussian endpoint χ²", xlabel="path energy", ylabel="χ²"); axes[0, 0].legend()

    alphas = [0.1, 0.5, 0.9]
    for alpha in alphas:
        axes[0, 1].plot(np.arange(80), discounted_convolution(np.full(80, 0.02), alpha), label=f"alpha={alpha}")
    axes[0, 1].set(title="Persistent error floor", xlabel="generation", ylabel="discounted divergence"); axes[0, 1].legend()

    for alpha in [0.1, 0.5, 0.9]:
        axes[1, 0].semilogy(np.arange(80), discounted_convolution(summable_errors[:80], alpha, 0.07), label=f"alpha={alpha}")
    axes[1, 0].set(title="Summable errors: fresh data shortens memory", xlabel="generation", ylabel="discounted divergence"); axes[1, 0].legend()

    beta_values = [(1-alpha)**2 for alpha in alphas]
    ages = np.arange(21)
    for alpha, beta in zip(alphas, beta_values):
        axes[1, 1].semilogy(ages, np.maximum(beta**ages, 1e-20), "o-", label=f"alpha={alpha}")
    axes[1, 1].set(title="Exact memory kernel (1-alpha)^(2m)", xlabel="error age m", ylabel="retained weight"); axes[1, 1].legend()
    figure.tight_layout(); figure.savefig(out / "diffusion_collapse_audit.png", dpi=170); plt.close(figure)

    result = {
        "paper": {"openreview_id": "QYA0Q28ssf", "arxiv_id": "2602.16601v1", "title": "Quantifying Error Propagation and Model Collapse in Diffusion Models"},
        "claims": CLAIMS,
        "scope": {
            "claim4": "Literal live claim conflates Proposition 4.1(i)'s non-summability premise with part (ii)'s additional uniform error-floor premise; the displayed positive floor is not supported by divergence of the series alone.",
            "gaussian_audit": "Exact Gaussian endpoint/path identities and deterministic memory recurrences test the formulas but do not replace the paper's stochastic-process proofs.",
            "experiments": "10D GMM, Fashion-MNIST, and CIFAR-10 conclusions are pinned source evidence; image divergence uses the source's learned-feature Gaussian proxy and is not an independent rerun.",
        },
        "diagnostics": {"summable_error_total_240": float(summable_errors.sum()), "harmonic_partial_sum_6000": float(harmonic_error.sum()), "constant_error": 0.02},
        "gates": gates,
        "all_gates_pass": all(gates.values()),
    }
    (out / "results.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    if not result["all_gates_pass"]:
        raise SystemExit(f"failed gates: {[name for name, passed in gates.items() if not passed]}")
    print(f"PASS {sum(gates.values())}/{len(gates)} gates")


if __name__ == "__main__":
    main()