Pluto-1 private release candidate: dual-encoder ensemble, calibrated Typed Decisions specialist
2778af7 verified | """Calibration: temperature scaling and split-conformal abstention.""" | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| def fit_temperature(logits: np.ndarray, targets: np.ndarray) -> float: | |
| """Fit one temperature by minimizing NLL on held-out logits.""" | |
| x = torch.from_numpy(logits.astype(np.float32)) | |
| y = torch.from_numpy(targets.astype(np.float32)) | |
| log_t = torch.zeros((), requires_grad=True) | |
| opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100) | |
| def closure(): | |
| opt.zero_grad() | |
| t = log_t.exp().clamp(0.02, 50.0) | |
| loss = -(y * F.log_softmax(x / t, dim=-1)).sum(-1).mean() | |
| loss.backward() | |
| return loss | |
| opt.step(closure) | |
| return float(log_t.exp().clamp(0.02, 50.0).item()) | |
| def multiclass_ece(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float: | |
| conf = probs.max(-1) | |
| pred = probs.argmax(-1) | |
| correct = (pred == labels).astype(np.float64) | |
| bins = np.linspace(0.0, 1.0, n_bins + 1) | |
| ece = 0.0 | |
| for lo, hi in zip(bins[:-1], bins[1:]): | |
| m = (conf > lo) & (conf <= hi) | |
| if m.any(): | |
| ece += m.mean() * abs(correct[m].mean() - conf[m].mean()) | |
| return float(ece) | |
| def brier(probs: np.ndarray, labels: np.ndarray) -> float: | |
| onehot = np.zeros_like(probs) | |
| onehot[np.arange(len(labels)), labels] = 1.0 | |
| return float(((probs - onehot) ** 2).sum(-1).mean()) | |
| class ConformalGate: | |
| """Selective predictor with an empirically risk-controlled threshold. | |
| The gate accepts a prediction only when its top probability is at least | |
| ``threshold``. The threshold is the lowest confidence cut whose accepted | |
| set on held-out calibration data has a 95% Wilson upper bound on error at | |
| or below ``alpha``. That gives the widest coverage consistent with the | |
| requested risk on the calibration distribution. | |
| """ | |
| threshold: float | |
| alpha: float | |
| coverage: float | |
| risk: float | |
| n: int | |
| def _wilson_upper(errors: int, total: int, z: float = 1.96) -> float: | |
| if total <= 0: | |
| return 1.0 | |
| p = errors / total | |
| denom = 1.0 + z * z / total | |
| centre = (p + z * z / (2 * total)) / denom | |
| half = z * ((p * (1 - p) / total + z * z / (4 * total * total)) ** 0.5) / denom | |
| return min(1.0, centre + half) | |
| def fit( | |
| cls, | |
| probs: np.ndarray, | |
| labels: np.ndarray, | |
| alpha: float = 0.05, | |
| ) -> "ConformalGate": | |
| pred = probs.argmax(-1) | |
| return cls.fit_scores( | |
| probs.max(-1).astype(np.float64), | |
| (pred == labels), | |
| alpha=alpha, | |
| ) | |
| def fit_scores( | |
| cls, | |
| scores: np.ndarray, | |
| correct: np.ndarray, | |
| alpha: float = 0.05, | |
| ) -> "ConformalGate": | |
| conf = np.asarray(scores, dtype=np.float64) | |
| wrong = (~np.asarray(correct, dtype=bool)).astype(np.int64) | |
| n = len(conf) | |
| order = np.argsort(conf, kind="stable") | |
| conf_s = conf[order] | |
| wrong_s = wrong[order] | |
| suffix_errors = np.concatenate([np.cumsum(wrong_s[::-1])[::-1], [0]]) | |
| best_t = 1.0 | |
| best_cov = 0.0 | |
| best_risk = 0.0 | |
| best_n = 0 | |
| min_accept = max(30, int(0.01 * n)) | |
| for i in range(n): | |
| accepted = n - i | |
| if accepted < min_accept: | |
| break | |
| t = float(conf_s[i]) | |
| if t <= 0.0: | |
| continue | |
| errors = int(suffix_errors[i]) | |
| upper = cls._wilson_upper(errors, accepted) | |
| if upper <= alpha: | |
| best_t = t | |
| best_cov = accepted / n | |
| best_risk = errors / accepted | |
| best_n = accepted | |
| break | |
| return cls( | |
| threshold=best_t, | |
| alpha=alpha, | |
| coverage=float(best_cov), | |
| risk=float(best_risk), | |
| n=int(best_n), | |
| ) | |
| def accepts(self, prob: float) -> bool: | |
| return float(prob) >= self.threshold | |
| def to_dict(self) -> dict: | |
| return { | |
| "threshold": self.threshold, | |
| "alpha": self.alpha, | |
| "coverage": self.coverage, | |
| "risk": self.risk, | |
| "n": self.n, | |
| } | |
| def from_dict(cls, d: dict) -> "ConformalGate": | |
| return cls(**{k: d[k] for k in ("threshold", "alpha", "coverage", "risk", "n")}) | |
| def selective_report( | |
| probs: np.ndarray, labels: np.ndarray, thresholds: list[float] | |
| ) -> list[dict]: | |
| conf = probs.max(-1) | |
| pred = probs.argmax(-1) | |
| wrong = pred != labels | |
| out = [] | |
| for t in thresholds: | |
| accepted = conf >= t | |
| out.append( | |
| { | |
| "threshold": float(t), | |
| "coverage": float(accepted.mean()), | |
| "risk": float(wrong[accepted].mean()) if accepted.any() else None, | |
| "accepted": int(accepted.sum()), | |
| } | |
| ) | |
| return out | |