mixle.task.calibrate module

CalibratedTaskModel wraps a task model in conformal answer sets.

A distilled student classifies by argmax over a softmax, but the softmax value alone is not a coverage guarantee. Conformal prediction adds the serving contract: on a held-out calibration set it learns a score threshold (mixle.inference.conformal.conformal_label_threshold()) such that the prediction set covers the true label with probability >= 1 - alpha under the usual exchangeability assumption AND the condition that the underlying model’s behavior is unchanged since calibration – the threshold ranks the calibration-time model’s scores, so an in-place retrain/update voids the statement exactly like distribution shift does, and just as silently (STAT-RR22-08/RR23-09). The density gate mitigates shift; drifted traffic or a changed model calls for recalibration.

The decision rule the cascade and the cost model consume:

  • singleton set -> answer locally. NOTE: the 1 - alpha guarantee is MARGINAL – over the whole population, the set contains the truth with probability >= 1 - alpha. It is NOT a conditional guarantee on the answered slice: error among singleton answers can be far higher (an adversarial review measured 9% marginal miscoverage alongside 47% answered-slice error at alpha=0.10). Measure answered-slice agreement per deployment (report() does) rather than reading it off alpha;

  • empty or multi-label set -> escalate to the expensive teacher/frontier (genuinely ambiguous).

When the answered slice itself needs a risk guarantee, calibrate it directly: CalibratedTaskModel.calibrate_selective() certifies thresholds from a PRE-SPECIFIED grid (1001 evenly spaced values, fixed in code, never derived from the sample) with an exact Clopper-Pearson binomial test per grid point at Bonferroni level delta / 1001, and serves at the loosest certified threshold. That is Learn-Then-Test risk control with explicit multiple-testing correction: with probability at least 1 - delta over an i.i.d. calibration draw, the error rate among the answers the model actually gives is at most alpha. It exists because the marginal guarantee above is NOT that statement – the adversarial statistical review measured 9% marginal miscoverage alongside 47.4% answered-slice error on a population built to split them.

escalation_rate is the empirical p_escalate used by the cost model. Conformal coverage is marginal, and a softmax still cannot see true OOD; a generative-density gate (mixle.task.density) covers that residual. Calibration persists in the artifact, so a loaded model decides identically in a fresh process.

selective_risk_threshold(top_probabilities, correct, *, alpha, delta=0.05, min_answered=10)[source]

Choose a confidence threshold whose ANSWERED-SLICE error is controlled, not just marginal.

Learn-Then-Test risk control with explicit multiple-testing correction (Angelopoulos et al., arXiv:2110.01052): every candidate threshold in the pre-specified grid _SELECTIVE_GRID is tested against the calibration sample with an exact Clopper-Pearson binomial upper bound on its answered-slice error rate, each test at Bonferroni level delta / len(grid). By the union bound, with probability at least 1 - delta over the calibration draw, EVERY threshold that passes its test simultaneously has true answered-slice error at most alpha – which is what licenses returning the loosest certified threshold after seeing all the results. Returns None when no grid threshold with at least min_answered answered rows certifies: “this model cannot guarantee that risk; escalate everything”.

Guarantee, stated exactly: if calibration rows are drawn i.i.d. from the deployment distribution, then with probability at least 1 - delta over the calibration draw, the deployment error rate among rows answered at the returned threshold (answer iff top-class probability >= tau) is at most alpha. The conditional-Binomial step needs i.i.d. sampling, not just exchangeability. The correction’s price is conservatism: certifying a zero-error slice needs about log(delta/1001)/log(1-alpha) answered rows (~94 at alpha=0.10, delta=0.05, versus ~29 uncorrected), and tau is quantized to 0.001.

Parameters:
Return type:

float | None

class CalibratedTaskModel(task, *, alpha=0.1, qhat=None, tau=None, density_gate=None)[source]

Bases: object

A TaskModel plus a conformal threshold: predicts label sets and decides answer-vs-escalate.

THE THRESHOLD IS BOUND TO THE PREDICTOR THAT PRODUCED THE CALIBRATION SCORES (STAT-RR22-08). qhat ranks nonconformity scores computed by the contained task model at calibration time; every coverage/selective-risk statement is conditional on that model’s behavior remaining unchanged at serving. Mutating the contained model in place (retraining its raw weights, swapping adapter state) silently voids the guarantee – an exact-wheel probe that mutated the contained model flipped prediction sets from correct to wrong with coverage 1.0 -> 0.0 while qhat sat unchanged – and no runtime check here can detect an in-place mutation. Every library flow that changes the model (improve() and friends) recalibrates in the same operation; if you touch the model yourself, call calibrate(...) again before serving.

Parameters:

  • task (TaskModel)

  • alpha (float)

  • qhat (float | None)

  • tau (float | None)

  • density_gate (Any)

property labels: list[str]

Return labels in the probability-vector order used by the adapter.

calibrate(texts, teacher_labels)[source]

Set the conformal threshold from held-out (texts, teacher_labels) for 1 - alpha set coverage.

Parameters:
Return type:

CalibratedTaskModel

predict_sets(texts)[source]

Conformal label set per input (the classes whose score clears the calibrated threshold).

Impossible structured evidence stays an empty prediction set so the calibrated serving contract escalates it; malformed adapter probabilities raise (see _proba()).

Parameters:

texts (Sequence[Any])

Return type:

list[list[str]]

predict_set(text)[source]

Return the conformal label set for one input.

Parameters:

text (Any)

Return type:

list[str]

calibrate_selective(texts, teacher_labels, *, delta=0.05, min_answered=10)[source]

Calibrate the ANSWERED-SLICE risk gate (see selective_risk_threshold()).

Runs after calibrate() (both gates then apply) or instead of it (serving answers the argmax label whenever the top-class probability clears tau), and applies the same fail-closed input validation as calibrate(). When no threshold can control the risk, tau becomes inf and everything escalates – an honest refusal, never a silent downgrade to the marginal-only contract.

Parameters:
Return type:

CalibratedTaskModel

decide(text)[source]

Return the label if the input is a confident, in-distribution singleton, else ESCALATE (None).

Parameters:

text (Any)

Return type:

Any

batch_decide(texts)[source]

Return local labels or ESCALATE for a batch of inputs.

Parameters:

texts (Sequence[Any])

Return type:

list[Any]

escalation_rate(texts)[source]

Empirical p_escalate – the fraction of inputs escalated (ambiguous set or, if gated, OOD).

Parameters:

texts (Sequence[Any])

Return type:

float

save(path)[source]

Persist the model, the full calibration (alpha, qhat, tau), and any density gate in the artifact.

qhat and tau can each legitimately be +inf (qhat: too little calibration data to admit any confident singleton; tau: calibrate_selective could not control the risk). Those are real, callable thresholds, so both persist as the JSON-safe sentinel "inf" and reload back to float('inf'). The persistence contract is decision identity INCLUDING abstentions: the re-review’s STAT-R1 caught tau being dropped here, which silently downgraded a selectively calibrated artifact to the marginal-only rule and answered a row it had abstained on.

Parameters:

path (str)

Return type:

str

classmethod load(path, *, device='cpu')[source]

Rebuild a calibrated model (with its density gate, if any) from an artifact; decisions match exactly.

Parameters:
Return type:

CalibratedTaskModel