"""P061: confidence with sparse known IDs, unknown IDs, and decisive context."""
from collections.abc import Iterator
from functools import partial
import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report
import relflow as rf
PROOF_ID = "P061"
BUDGET = 600
ENTITIES = 4096
def records(*, rows: int, seed: int, novel: bool = False) -> Iterator[dict]:
rng = np.random.default_rng(seed)
identities = rng.integers(0, ENTITIES, rows)
effects = np.random.default_rng(918).choice([-2.5, 2.5], ENTITIES)
values = rng.uniform(-3, 3, rows)
probabilities = 1 / (1 + np.exp(-(2 * values + effects[identities])))
labels = rng.random(rows) < probabilities
for x, identity, label in zip(values, identities, labels, strict=True):
yield {"x": float(x), "entity": f"{'new' if novel else 'id'}-{identity}", "label": "yes" if label else "no"}
def build(unavailable: float) -> rf.Model:
model = rf.Model(
d_model=32,
n_layers=1,
n_heads=4,
dropout=0.0,
batch_size=128,
x=rf.Number,
entity=rf.Category(p_unavailable=unavailable),
label=rf.Category(mask=True),
)
model.optimizer = rf.adamw(learning_rate=0.002)
return model
def probability(model: rf.Model, rows: list[dict]) -> np.ndarray:
predictions = model.predict(pa.Table.from_pylist(rows))["predictions"].to_pylist()
content = [row["label"]["content"] for row in predictions]
return np.asarray([item["probability"] if item["value"] == "yes" else 1 - item["probability"] for item in content])
def score(p: np.ndarray, rows: list[dict], novel: bool) -> dict[str, float]:
x = np.asarray([row["x"] for row in rows])
if novel:
oracle = (1 / (1 + np.exp(-(2 * x - 2.5))) + 1 / (1 + np.exp(-(2 * x + 2.5)))) / 2
else:
effects = np.random.default_rng(918).choice([-2.5, 2.5], ENTITIES)
identity = np.asarray([int(row["entity"].split("-")[1]) for row in rows])
oracle = 1 / (1 + np.exp(-(2 * x + effects[identity])))
y = np.asarray([row["label"] == "yes" for row in rows], dtype=float)
clipped = p.clip(1e-7, 1 - 1e-7)
confidence = np.maximum(p, 1 - p)
return {
"oracle_rmse": float(np.sqrt(np.mean((p - oracle) ** 2))),
"constant_rmse": float(np.sqrt(np.mean((0.5 - oracle) ** 2))),
"nll": float(np.mean(-y * np.log(clipped) - (1 - y) * np.log1p(-clipped))),
"brier": float(np.mean((p - y) ** 2)),
"accuracy": float(np.mean((p > 0.5) == y)),
"ambiguous_confidence": float(confidence[np.abs(x) < 0.5].mean()),
"decisive_confidence": float(confidence[np.abs(x) > 2.5].mean()),
}
def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
metrics, checks = {}, {}
for unavailable in (0.01, 0.2):
lit.seed_everything(seed, workers=True)
model = build(unavailable)
data = rf.SyntheticDataModule(
model=model,
train=partial(records, rows=32768, seed=seed + 1),
validate=partial(records, rows=1024, seed=seed + 2),
seed=seed,
)
trainer = lit.Trainer(
accelerator=accelerator,
devices=1,
max_epochs=-1,
max_steps=BUDGET if steps is None else min(steps, BUDGET),
logger=False,
enable_checkpointing=False,
enable_progress_bar=False,
enable_model_summary=False,
deterministic=True,
num_sanity_val_steps=0,
)
trainer.fit(model, datamodule=data)
model.eval()
vocabulary = rf.Category.vocabulary(model, "/entity")
counts = rf.Category.counts(model, "/entity")
arm = {}
for novel in (False, True):
rows = list(records(rows=8192, seed=seed + 3, novel=novel))
arm["novel" if novel else "known"] = score(probability(model, rows), rows, novel)
rows = list(records(rows=2048, seed=seed + 4, novel=True))
before = probability(model, rows)
renamed = [{**row, "entity": "another-unseen-identity"} for row in rows]
arm["renaming_drift"] = float(np.max(np.abs(before - probability(model, renamed))))
checks[f"{unavailable}: evaluation preserves mapping and exposures"] = vocabulary == rf.Category.vocabulary(
model, "/entity"
) and counts == rf.Category.counts(model, "/entity")
checks[f"{unavailable}: unseen spelling is uninformative"] = arm["renaming_drift"] < 1e-5
model.encode(pa.Table.from_pylist(rows), strata="train")
arm["admission_mean_drift"] = float(np.mean(np.abs(before - probability(model, rows))))
metrics[str(unavailable)] = arm
if unavailable == 0.2:
checks.update(
{
"Known-ID probability RMSE below 0.30": arm["known"]["oracle_rmse"] < 0.30,
"Unknown-ID probability RMSE below 0.10": arm["novel"]["oracle_rmse"] < 0.10,
"Unknown-ID predictions beat the constant baseline": arm["novel"]["oracle_rmse"]
< arm["novel"]["constant_rmse"],
"Ambiguous unknown IDs retain confidence below 0.70": arm["novel"]["ambiguous_confidence"] < 0.70,
"Decisive context retains confidence above 0.90": arm["novel"]["decisive_confidence"] > 0.90,
}
)
return metrics, checks
if __name__ == "__main__":
report(PROOF_ID, run, seed=7611)
