1"""MVP verification of robust parameter-update envelopes.
  2
  3The model is f(x) = (a b)x, a two-layer scalar linear neural network.
  4At (a,b)=(1,4), the proposed update to (4,1) has zero endpoint loss for
  5x=1,y=4, but its interpolation has a loss peak at the midpoint.
  6"""
  7import json
  8import math
  9import random
 10from pathlib import Path
 11import numpy as np
 12
 13SEED = 584
 14random.seed(SEED)
 15np.random.seed(SEED)
 16
 17
 18def scalar_loss(a, b, target=4.0):
 19    return (a * b - target) ** 2
 20
 21
 22def path_values(theta0, theta1, grid):
 23    a0, b0 = theta0
 24    a1, b1 = theta1
 25    a = a0 + grid * (a1 - a0)
 26    b = b0 + grid * (b1 - b0)
 27    output = a * b
 28    loss = (output - 4.0) ** 2
 29    return output, loss
 30
 31
 32def envelope_accept(theta0, theta1, loss_budget, K):
 33    grid = np.linspace(0.0, 1.0, K + 1)
 34    _, losses = path_values(theta0, theta1, grid)
 35    violation = float(np.max(losses - loss_budget))
 36    return violation <= 1e-12, violation, float(np.max(losses)), grid[int(np.argmax(losses))]
 37
 38
 39def largest_feasible_fraction(theta0, theta1, loss_budget, K=16):
 40    # Scale the proposed direction, then use bisection on the accepted fraction.
 41    def ok(frac):
 42        end = np.asarray(theta0) + frac * (np.asarray(theta1) - theta0)
 43        return envelope_accept(theta0, end, loss_budget, K)[0]
 44    lo, hi = 0.0, 1.0
 45    if ok(hi):
 46        return 1.0
 47    for _ in range(45):
 48        mid = (lo + hi) / 2.0
 49        if ok(mid):
 50            lo = mid
 51        else:
 52            hi = mid
 53    return lo
 54
 55
 56def run():
 57    theta0 = np.array([1.0, 4.0])
 58    theta1 = np.array([4.0, 1.0])
 59    budget = scalar_loss(*theta0) + 1e-10  # endpoint-only loss-growth budget
 60
 61    # Core numerical claim, with a dense reference grid.
 62    dense = np.linspace(0.0, 1.0, 100001)
 63    outputs, losses = path_values(theta0, theta1, dense)
 64    dense_peak = float(np.max(losses))
 65    peak_s = float(dense[int(np.argmax(losses))])
 66    endpoint_max = max(float(losses[0]), float(losses[-1]))
 67    core = {
 68        "start_loss": float(losses[0]), "end_loss": float(losses[-1]),
 69        "endpoint_max_loss": endpoint_max, "dense_max_loss": dense_peak,
 70        "dense_peak_s": peak_s, "interior_violation": dense_peak - budget,
 71        "analytic_midpoint_output": float(outputs[len(outputs)//2]),
 72    }
 73
 74    # Compare endpoint acceptance against the envelope rule at several K values.
 75    comparisons = {}
 76    for K in (4, 8, 16, 32):
 77        accepted, violation, peak, where = envelope_accept(theta0, theta1, budget, K)
 78        comparisons[str(K)] = {
 79            "endpoint_only_accepts": True,
 80            "envelope_accepts": bool(accepted),
 81            "sampled_max_violation": violation,
 82            "sampled_max_loss": peak,
 83            "sampled_peak_s": where,
 84        }
 85
 86    # An acceptance-transition measurement: robust halving finds the largest
 87    # fraction of this direction whose entire sampled path is feasible.
 88    thresholds = {str(K): largest_feasible_fraction(theta0, theta1, budget, K)
 89                  for K in (8, 16, 32)}
 90
 91    # A small repeated-update proxy: each trial has the same candidate update.
 92    # This isolates the acceptance mechanism rather than confounding it with
 93    # a changing optimizer state.
 94    trials = 50
 95    endpoint_accepted = sum(1 for _ in range(trials)
 96                            if scalar_loss(*theta1) <= budget)
 97    robust_accepted = sum(1 for _ in range(trials)
 98                          if envelope_accept(theta0, theta1, budget, 16)[0])
 99
100    result = {
101        "seed": SEED,
102        "model": "two-layer scalar linear network f(x)=(a*b)x, x=1, target=4",
103        "constraint": "loss(theta(s)) <= loss(theta(0)) + 1e-10",
104        "core_check": core,
105        "K_comparison": comparisons,
106        "largest_feasible_fraction": thresholds,
107        "50_identical_candidate_trials": {
108            "endpoint_only_accepted": endpoint_accepted,
109            "robust_envelope_accepted": robust_accepted,
110        },
111        "interpretation": "Endpoint rule accepts the weight swap, while every envelope grid detects its interior loss spike; K=8,16,32 agree on rejection and similar halving threshold.",
112    }
113    out = Path("robust_envelope_results.json")
114    out.write_text(json.dumps(result, indent=2))
115    print(json.dumps(result, indent=2))
116
117
118if __name__ == "__main__":
119    run()