Age-conditioned semi-Markov router / age_router_experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2from pathlib import Path
  3import numpy as np
  4
  5
  6def survival_from_hazard(h):
  7    h = np.asarray(h, dtype=float)
  8    S = np.ones(len(h) + 1)
  9    for n in range(len(h)):
 10        S[n + 1] = S[n] * (1.0 - h[n])
 11    return S, h * S[:-1]
 12
 13
 14def make_sequence(rng, n, heavy_tail=True):
 15    vals, durations, regime = [], [], 0
 16    while len(vals) < n:
 17        if heavy_tail:
 18            d = int(np.clip(np.floor(2.0 * (1.0 + rng.pareto(1.7))), 2, 80))
 19        else:
 20            d = int(rng.geometric(1 / 12.0))
 21        durations.append(d)
 22        vals.extend([regime] * d)
 23        regime = 1 - regime
 24    return np.asarray(vals[:n], dtype=np.int64), durations
 25
 26
 27def duration_hazard(durations, max_age):
 28    d = np.asarray(durations)
 29    return np.array([
 30        np.mean(d == a + 1) / max(np.mean(d >= a + 1), 1e-12)
 31        for a in range(max_age)
 32    ]).clip(1e-5, 1 - 1e-5)
 33
 34
 35def tokenwise_router(x, noise, rng):
 36    obs = x + rng.normal(0, noise, len(x))
 37    return (obs > 0.5).astype(np.int64)
 38
 39
 40def age_router(x, noise, rng, hazard, max_age, evidence_gain=5.0):
 41    obs = x + rng.normal(0, noise, len(x))
 42    out = np.empty(len(x), dtype=np.int64)
 43    regime = int(obs[0] > 0.5)
 44    age = 0
 45    calls, switches = 1, 0
 46    out[0] = regime
 47    for t in range(1, len(x)):
 48        candidate = int(obs[t] > 0.5)
 49        base = hazard[min(age, max_age - 1)]
 50        disagreement = abs(obs[t] - 0.5) if candidate != regime else 0.0
 51        logit = np.log(base / (1.0 - base)) + evidence_gain * disagreement
 52        p_switch = 1.0 / (1.0 + np.exp(-np.clip(logit, -30, 30)))
 53        if rng.random() < p_switch:
 54            calls += 1
 55            if candidate != regime:
 56                regime = candidate
 57                switches += 1
 58                age = 0
 59            else:
 60                age += 1
 61        else:
 62            age += 1
 63        out[t] = regime
 64    return out, calls, switches
 65
 66
 67def score(pred, x, n):
 68    return {"accuracy": float(np.mean(pred == x)),
 69            "switches": int(np.sum(pred[1:] != pred[:-1])),
 70            "switch_rate": float(np.mean(pred[1:] != pred[:-1])),
 71            "router_calls": int(n)}
 72
 73
 74def evaluate(seed=7, n=30000, noise=0.65):
 75    rng = np.random.default_rng(seed)
 76    x, true_durations = make_sequence(rng, n, heavy_tail=True)
 77    train_rng = np.random.default_rng(seed + 100)
 78    _, train_durations = make_sequence(train_rng, 200000, heavy_tail=True)
 79    max_age = 80
 80    hazards = duration_hazard(train_durations, max_age)
 81    b = tokenwise_router(x, noise, np.random.default_rng(seed + 1))
 82    a, calls, _ = age_router(x, noise, np.random.default_rng(seed + 2), hazards, max_age)
 83    bm, am = score(b, x, n), score(a, x, calls)
 84    h = duration_hazard(true_durations, max_age)
 85    S, p = survival_from_hazard(h)
 86    identity_error = float(abs(np.sum(p) + S[-1] - 1.0))
 87    d = np.asarray(true_durations)
 88    implied_mean = float(np.sum((np.arange(max_age) + 1) * p) + (max_age + 1) * S[-1])
 89    result = {
 90        "settings": {"seed": seed, "n": n, "noise": noise},
 91        "math": {"max_survival_identity_error": identity_error,
 92                 "empirical_mean_duration": float(np.mean(d)),
 93                 "implied_mean_duration_truncated_tail": implied_mean,
 94                 "hazard_first_12": hazards[:12].tolist()},
 95        "baseline": bm, "age_conditioned": am,
 96        "relative_switch_reduction": float(1 - am["switches"] / max(bm["switches"], 1)),
 97        "relative_router_call_reduction": float(1 - calls / n),
 98    }
 99    Path("results.json").write_text(json.dumps(result, indent=2))
100    print(json.dumps(result, indent=2))
101    return result
102
103
104if __name__ == "__main__":
105    evaluate()