routing_signals.py

"""
DOES THE CHEAP MODEL ALREADY KNOW WHEN IT NEEDS HELP?

Run this AFTER cascade.py, on the cascade_probs.npz it saves. No retraining
-- the probability vectors are all that any confidence rule can see, so
every signal below is computable from a file that already exists.

Four signals, three of them from the cheap model's output vector alone:

  MAX PROBABILITY   the highest predicted class probability
  TOP-2 MARGIN      the gap between first and second choice
  NEGATIVE ENTROPY  how peaked the whole distribution is
  RANDOM            a control, because a threshold applied to noise still
                    produces a curve and it is easy to mistake one for a
                    result

The first three are functions of the same ten numbers, so they are likely
to agree closely. That is worth establishing rather than assuming: if they
do agree, no amount of re-deriving statistics from one softmax will help,
and the next thing to try has to see something else.

Each is scored two ways. A RANKING measure -- how well it orders the
examples that genuinely need escalation, those the cheap model gets wrong
and the expensive one gets right. And the ACCURACY CURVE it produces when
the least-confident fraction is escalated, which is what actually matters
because a cascade is judged on accuracy against arithmetic.

The ranking measure is computed without sklearn, as the probability that a
needy example is ranked below a safe one, which is what an ROC AUC is.
"""

import numpy as np
import json


def routing_scores(probs, seed=0):
    """Every signal available from the cheap model's output.

    All in one convention: HIGHER means more likely to be safe without
    escalation, so the least-confident fraction is always the lowest
    scores."""
    p = np.clip(probs, 1e-12, 1.0)
    top2 = np.partition(p, -2, axis=1)[:, -2:]
    rg = np.random.default_rng(seed)
    return {
        "max probability": p.max(1),
        "top-2 margin": top2[:, 1] - top2[:, 0],
        "negative entropy": (p*np.log(p)).sum(1),
        "random (control)": rg.random(len(p)),
    }


def auc(score, need):
    """The probability that a needy example scores below a safe one.

    Computed by ranks rather than by pairs, so it is exact and fast, and
    without sklearn so this runs anywhere."""
    if need.sum() == 0 or need.sum() == len(need):
        return float("nan")
    r = np.argsort(np.argsort(-score))          # 0 = least confident
    n1, n0 = need.sum(), (~need).sum()
    return float((r[need].sum() - n1*(n1-1)/2) / (n1*n0))


def analyse(path="cascade_probs.npz", cheap_macs=110592, dear_macs=1290240,
            fractions=(0.0, 0.05, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0)):
    d = np.load(path)
    cheap, dear, yte = d["cheap"], d["dear"], d["yte"]
    cc = cheap.argmax(1) == yte
    dc = dear.argmax(1) == yte
    need = ~cc & dc

    print("=" * 78)
    print("DOES THE CHEAP MODEL ALREADY KNOW WHEN IT NEEDS HELP?")
    print("=" * 78)
    print(f"  cheap {cc.mean():.4f} at {cheap_macs:,} multiplies")
    print(f"  dear  {dc.mean():.4f} at {dear_macs:,} multiplies "
          f"({dear_macs/cheap_macs:.1f}x)")
    print(f"\n  an oracle escalating exactly the {need.mean()*100:.1f}% that")
    print(f"  need it reaches {np.where(need, dc, cc).mean():.4f}, at "
          f"{(cheap_macs + need.mean()*dear_macs)/dear_macs:.2f}x the cost")
    print(f"  of always running the expensive model.")
    print(f"\n  and {(cc & ~dc).mean()*100:.1f}% go the other way — the cheap")
    print(f"  model right where the expensive one is wrong, which any")
    print(f"  escalation risks throwing away.")

    scores = routing_scores(cheap)
    print("\n" + "=" * 78)
    print("  HOW WELL DOES EACH SIGNAL RANK THE EXAMPLES THAT NEED HELP?")
    print("=" * 78)
    print(f"  {'signal':>18s} {'AUC':>7s}   (0.5 is chance)")
    aucs = {}
    for nm, s in scores.items():
        aucs[nm] = auc(s, need)
        print(f"  {nm:>18s} {aucs[nm]:7.4f}")

    real = [v for k, v in aucs.items() if "control" not in k]
    print(f"\n  the three output statistics span {min(real):.4f} to "
          f"{max(real):.4f}")
    if max(real) - min(real) < 0.02:
        print("  — which is what should be expected, since all three are")
        print("  functions of the same ten numbers. No further statistic on")
        print("  one softmax will help; a better signal has to see something")
        print("  the output vector does not contain.")

    print("\n" + "=" * 78)
    print("  ACCURACY AGAINST ARITHMETIC")
    print("=" * 78)
    print(f"  {'escalated':>10s} {'cost':>12s} {'vs dear':>8s} " +
          " ".join(f"{k.split()[0][:8]:>9s}" for k in scores))
    best = {}
    for f in fractions:
        cost = cheap_macs + f*dear_macs
        row = []
        for nm, s in scores.items():
            if f == 0:
                a_ = cc.mean()
            elif f == 1:
                a_ = dc.mean()
            else:
                esc = s <= np.quantile(s, f)
                a_ = np.where(esc, dc, cc).mean()
            row.append(a_)
            # 0% and 100% are not cascades: one is the cheap model and the
            # other is the expensive one, and at 100% every signal is
            # identical, so including them makes the comparison vacuous
            if 0 < f < 1 and a_ > best.get(nm, (0, 0))[0]:
                best[nm] = (a_, f)
        print(f"  {f*100:9.0f}% {cost:12,.0f} {dear_macs/cost:7.2f}x " +
              " ".join(f"{v:9.4f}" for v in row))

    print("\n" + "=" * 78)
    print("  READOUT")
    print("=" * 78)
    # the honest comparison is against the CONTROL at the same escalation
    # fraction, not against the best fraction each signal happens to reach
    ctrl_at = {}
    for f in fractions:
        if not (0 < f < 1):
            continue
        s_ = scores["random (control)"]
        ctrl_at[f] = np.where(s_ <= np.quantile(s_, f), dc, cc).mean()
    print(f"  gain over RANDOM escalation at the same fraction:\n")
    print(f"  {'escalated':>10s} " +
          " ".join(f"{k.split()[0][:8]:>9s}" for k in scores
                   if "control" not in k))
    lifts = {}
    for f in fractions:
        if not (0 < f < 1):
            continue
        row = []
        for nm, s_ in scores.items():
            if "control" in nm:
                continue
            a_ = np.where(s_ <= np.quantile(s_, f), dc, cc).mean()
            row.append(a_ - ctrl_at[f])
            lifts.setdefault(nm, []).append(a_ - ctrl_at[f])
        print(f"  {f*100:9.0f}% " + " ".join(f"{v:+9.4f}" for v in row))
    bestlift = max((max(v), k) for k, v in lifts.items())
    print()
    if bestlift[0] < 0.005:
        print("  NO SIGNAL BEATS RANDOM ESCALATION. The cheap model's output")
        print("  does not indicate when it needs help, so the conditional")
        print("  computation opportunity is real and not reachable from the")
        print("  probabilities. What to try next must see the cheap model's")
        print("  REPRESENTATION, or a second cheap model's disagreement —")
        print("  not another statistic on one softmax.")
    else:
        print(f"  THE CHEAP MODEL DOES KNOW, PARTLY. The best signal beats")
        print(f"  random escalation by up to {bestlift[0]:+.4f} at the same")
        print(f"  arithmetic, which is what a usable cascade needs: the rule")
        print(f"  does not have to be right, only better than spending the")
        print(f"  budget at random.")
        print(f"\n  and the three output statistics agree to within "
              f"{max(max(v) for v in lifts.values()) - min(max(v) for v in lifts.values()):.4f},")
        print(f"  so which one is used barely matters.")
    print(f"\n  best single point: ", end="")
    b2 = max(best.items(), key=lambda t: t[1][0])
    cost = cheap_macs + b2[1][1]*dear_macs
    print(f"{b2[0]} at {b2[1][1]*100:.0f}% escalated reaches "
          f"{b2[1][0]:.4f}")
    print(f"  ({dear_macs/cost:.2f}x less arithmetic than always-dear, "
          f"which scores {dc.mean():.4f})")
    json.dump({"auc": aucs, "best": {k: list(v) for k, v in best.items()}},
              open("routing_signals.json", "w"), indent=2)
    print("\n  wrote routing_signals.json")


if __name__ == "__main__":
    analyse()