diff --git a/scripts/run_failure_analysis.py b/scripts/run_failure_analysis.py index cbfae5c..2fce582 100644 --- a/scripts/run_failure_analysis.py +++ b/scripts/run_failure_analysis.py @@ -129,10 +129,15 @@ def _run_predictions(predictions_path: str, out: str) -> None: def predict(present): if present not in preds: - raise SystemExit(f"no scored condition covers the modality subset {sorted(present)}") + # KeyError, not SystemExit: the harness degrades gracefully when an optional + # subset (the all-masked one Shapley probes) has no scored condition behind it. + raise KeyError(f"no scored condition covers the modality subset {sorted(present)}") return preds[present] - report = analyze_modality_failure(["echo", "ecg"], lvef, ef, predict) + try: + report = analyze_modality_failure(["echo", "ecg"], lvef, ef, predict) + except KeyError as e: + raise SystemExit(e.args[0]) from None _emit( report, out, diff --git a/scripts/validate_harness.py b/scripts/validate_harness.py index 0ac818b..b620900 100644 --- a/scripts/validate_harness.py +++ b/scripts/validate_harness.py @@ -5,6 +5,9 @@ (1) Multi-seed recovery (echo/ECG): does the harness reliably recover planted echo-dominance and the modality-specific silent-failure asymmetry across seeds? (2) N-modality generality (echo/ECG/labs): does it produce a genuine 3x3 complementarity matrix? +(3) Redundancy stress case: two equally strong modalities plus one weak one -- leave-one-out + reads the redundant pair as worthless, Shapley splits the credit between them. +(4) Nonlinear probe: the two-modality recovery repeated with a small MLP instead of Ridge. """ from __future__ import annotations @@ -18,16 +21,17 @@ analyze_modality_failure, make_synthetic_modalities, make_synthetic_multimodal, + masked_mlp_predict_fn, masked_ridge_predict_fn, stratified_train_mask, ) from primed_ai.failure.plots import plot_complementarity_matrix -def _analyze(emb, lvef, ef, seed): +def _analyze(emb, lvef, ef, seed, probe=masked_ridge_predict_fn): train = stratified_train_mask(ef, 0.7, seed=seed + 1000) test = ~train - predict_full = masked_ridge_predict_fn(emb, lvef, train) + predict_full = probe(emb, lvef, train) return analyze_modality_failure( {m: emb[m][test] for m in emb}, lvef[test], @@ -92,6 +96,32 @@ def threemodal(out: str, figdir: str) -> dict: return d +def redundancy(out: str) -> dict: + emb, lvef, ef = make_synthetic_modalities({"m1": 2.0, "m2": 2.0, "weak": 0.5}, n=750, seed=0) + d = _analyze(emb, lvef, ef, 0).to_dict() + c = d["complementarity"] + Path(out).mkdir(parents=True, exist_ok=True) + (Path(out) / "validation_redundancy.json").write_text(json.dumps(d, indent=2)) + print("=== redundancy stress case (m1 == m2 strong, weak) ===") + print("leave-one-out :", c["marginal_value"]) + print("shapley :", c["shapley_value"]) + return d + + +def nonlinear(out: str) -> dict: + emb, lvef, ef, _ = make_synthetic_multimodal(seed=0) + d = _analyze(emb, lvef, ef, 0, probe=masked_mlp_predict_fn).to_dict() + c = d["complementarity"] + Path(out).mkdir(parents=True, exist_ok=True) + (Path(out) / "validation_nonlinear.json").write_text(json.dumps(d, indent=2)) + print("=== nonlinear probe (MLP) on the planted two-modality data ===") + print("leave-one-out :", c["marginal_value"]) + print("shapley :", c["shapley_value"]) + print("winners :", c["per_example_winners"]) + print("silent rates :", {k: v["silent_rate"] for k, v in d["dropout"].items()}) + return d + + def main() -> None: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--seeds", type=int, default=12) @@ -100,6 +130,8 @@ def main() -> None: a = ap.parse_args() multiseed(a.seeds, a.out) threemodal(a.out, a.figdir) + redundancy(a.out) + nonlinear(a.out) if __name__ == "__main__": diff --git a/src/primed_ai/failure/__init__.py b/src/primed_ai/failure/__init__.py index 81a1b3d..0c43652 100644 --- a/src/primed_ai/failure/__init__.py +++ b/src/primed_ai/failure/__init__.py @@ -14,6 +14,7 @@ from .demo import ( make_synthetic_modalities, make_synthetic_multimodal, + masked_mlp_predict_fn, masked_ridge_predict_fn, stratified_train_mask, ) @@ -25,6 +26,7 @@ "CATEGORIES", "make_synthetic_multimodal", "make_synthetic_modalities", + "masked_mlp_predict_fn", "masked_ridge_predict_fn", "stratified_train_mask", ] diff --git a/src/primed_ai/failure/core.py b/src/primed_ai/failure/core.py index b8042a2..eca5d75 100644 --- a/src/primed_ai/failure/core.py +++ b/src/primed_ai/failure/core.py @@ -6,8 +6,9 @@ 1. **Failure taxonomy** -- per-example category (``correct`` / ``imprecise`` / ``critical``) under full-modality inference and under each missing-modality condition. -2. **Complementarity** -- each modality's leave-one-out marginal value, a complementarity - matrix over modality subsets, and per-example attribution ("which modality wins"). +2. **Complementarity** -- each modality's leave-one-out marginal value, its exact Shapley + value over modality coalitions, a complementarity matrix over modality subsets, and + per-example attribution ("which modality wins"). 3. **Dropout profile (loud vs. silent)** -- when a modality is missing at inference, does the model fail *loudly* (output shifts / sits near the decision boundary -> monitorable) or *silently* (confident, stable-looking output that is now clinically wrong)? @@ -21,6 +22,7 @@ from dataclasses import dataclass, field from itertools import combinations +from math import factorial from typing import Callable import numpy as np @@ -124,7 +126,23 @@ def analyze_modality_failure( needed.add(allm - {m}) for a, b in combinations(modalities, 2): needed.add(frozenset({a, b})) - preds = {s: np.asarray(predict_fn(s), dtype=float) for s in needed} + # Exact Shapley attribution needs every coalition, the empty set (all branches + # masked) included. 2^N calls are cheap at the N this harness sees; past the cap + # we keep leave-one-out attribution only rather than blow up predict_fn calls. + shapley_ok = len(modalities) <= 8 + if shapley_ok: + for k in range(len(modalities) + 1): + for combo in combinations(modalities, k): + needed.add(frozenset(combo)) + preds = {} + for s in needed: + try: + preds[s] = np.asarray(predict_fn(s), dtype=float) + except Exception: + if s: + raise + # predict_fn that cannot mask everything: drop Shapley, keep the rest + shapley_ok = False full = preds[allm] err_full = _abs_err(full, y) @@ -157,11 +175,31 @@ def _condition(pred) -> dict: row.append(round(float(_abs_err(preds[subset], y).mean()), 4)) matrix.append(row) + # Exact Shapley value of each modality over the coalition game v(S) = -MAE(f(S)): + # the weighted mean MAE drop from adding m across all coalitions, so positive phi + # is error the modality removes. Unlike leave-one-out, two redundant modalities + # split the credit instead of both reading as worthless. + shapley = None + if shapley_ok: + mae_of = {s: float(_abs_err(p, y).mean()) for s, p in preds.items()} + nm = len(modalities) + shapley = {} + for m in modalities: + rest = [x for x in modalities if x != m] + phi = 0.0 + for k in range(nm): + w = factorial(k) * factorial(nm - k - 1) / factorial(nm) + for combo in combinations(rest, k): + s = frozenset(combo) + phi += w * (mae_of[s] - mae_of[s | {m}]) + shapley[m] = round(phi, 4) + best_solo = min(solo_mae.values()) complementarity = { "full_mae": round(float(err_full.mean()), 4), "solo_mae": solo_mae, "marginal_value": marginal_value, # how much each modality is worth (LOO) + "shapley_value": shapley, # coalition-fair attribution; None past the subset cap "fusion_gain_vs_best_solo": round(best_solo - float(err_full.mean()), 4), "per_example_winners": winners, # which modality "wins" overall "matrix": {"modalities": modalities, "values": matrix}, diff --git a/src/primed_ai/failure/demo.py b/src/primed_ai/failure/demo.py index 7b38f79..07d9868 100644 --- a/src/primed_ai/failure/demo.py +++ b/src/primed_ai/failure/demo.py @@ -50,14 +50,12 @@ def make_synthetic_multimodal( return embeddings, lvef, ef_le_40, planted -def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1.0): - """Train one Ridge model on the full concatenated embeddings; mask absent modalities at +def _masked_predict_fn(model, embeddings: dict, lvef, train_mask): + """Fit ``model`` on the full concatenated embeddings; mask absent modalities at inference by zeroing their columns. This mimics a *single deployed multimodal model* evaluated under missing inputs (no per-condition retraining), which is exactly what the dropout analysis assumes. """ - from sklearn.linear_model import Ridge - mods = list(embeddings) cols, off = {}, 0 for m in mods: @@ -67,7 +65,7 @@ def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1 X = np.concatenate([np.asarray(embeddings[m], dtype=float) for m in mods], axis=1) y = np.asarray(lvef, dtype=float) train_mask = np.asarray(train_mask, dtype=bool) - model = Ridge(alpha=alpha).fit(X[train_mask], y[train_mask]) + model.fit(X[train_mask], y[train_mask]) def predict_fn(present: frozenset) -> np.ndarray: Xm = X.copy() @@ -79,6 +77,26 @@ def predict_fn(present: frozenset) -> np.ndarray: return predict_fn +def masked_ridge_predict_fn(embeddings: dict, lvef, train_mask, alpha: float = 1.0): + """One Ridge model on the full concatenation, absent modalities zeroed at inference.""" + from sklearn.linear_model import Ridge + + return _masked_predict_fn(Ridge(alpha=alpha), embeddings, lvef, train_mask) + + +def masked_mlp_predict_fn(embeddings: dict, lvef, train_mask, seed: int = 0): + """Nonlinear counterpart of :func:`masked_ridge_predict_fn`: one small MLP trained on the + full concatenation. Exists so the harness's recovery checks are not conditional on a + linear probe. + """ + from sklearn.neural_network import MLPRegressor + + model = MLPRegressor( + hidden_layer_sizes=(64,), max_iter=2000, random_state=seed, early_stopping=True + ) + return _masked_predict_fn(model, embeddings, lvef, train_mask) + + def make_synthetic_modalities( strengths: dict, n: int = 600, dim: int = 16, seed: int = 0, exclusive_frac: float = 0.12 ): diff --git a/tests/test_failure.py b/tests/test_failure.py index fe99137..6f4ce8d 100644 --- a/tests/test_failure.py +++ b/tests/test_failure.py @@ -67,6 +67,61 @@ def test_report_is_json_serializable(): assert d["complementarity"]["matrix"]["modalities"] == ["echo", "ecg"] +def test_shapley_efficiency_and_closed_form(): + # exact Shapley over v(S) = -MAE(f(S)): the values must sum to the full-vs-empty MAE + # gap, and at N = 2 each phi has a closed form over the four subset MAEs + emb, lvef, ef, _ = make_synthetic_multimodal(n=600, seed=0) + train = stratified_train_mask(ef, 0.7, seed=12345) + test = ~train + pf = masked_ridge_predict_fn(emb, lvef, train) + report = analyze_modality_failure( + {m: emb[m][test] for m in emb}, lvef[test], ef[test], lambda present: pf(present)[test] + ) + shap = report.complementarity["shapley_value"] + y = report.predictions["__labels__"] + + def mae(key): + return float(np.abs(report.predictions[key] - y).mean()) + + assert sum(shap.values()) == pytest.approx(mae("") - mae("ecg,echo"), abs=1e-3) + closed_echo = 0.5 * ((mae("") - mae("echo")) + (mae("ecg") - mae("ecg,echo"))) + assert shap["echo"] == pytest.approx(closed_echo, abs=1e-3) + assert shap["echo"] > shap["ecg"] # planted dominance survives the attribution change + + +def test_shapley_credits_redundant_modalities_beyond_loo(): + # two equally strong modalities: leave-one-out under-credits each (the twin covers), + # Shapley counts the coalitions where the twin is absent and pays both more + from primed_ai.failure import make_synthetic_modalities + + emb, lvef, ef = make_synthetic_modalities({"m1": 2.0, "m2": 2.0, "weak": 0.5}, n=750, seed=0) + train = stratified_train_mask(ef, 0.7, seed=1000) + test = ~train + pf = masked_ridge_predict_fn(emb, lvef, train) + c = analyze_modality_failure( + {m: emb[m][test] for m in emb}, lvef[test], ef[test], lambda present: pf(present)[test] + ).complementarity + loo, shap = c["marginal_value"], c["shapley_value"] + assert shap["m1"] > loo["m1"] and shap["m2"] > loo["m2"] + assert abs(shap["m1"] - shap["m2"]) < 1.0 # symmetric plants get symmetric credit + assert shap["weak"] < min(shap["m1"], shap["m2"]) + + +def test_shapley_none_when_empty_set_unsupported(): + rng = np.random.default_rng(0) + y = rng.uniform(15, 75, 40) + gate = y <= 40 + + def predict_fn(present: frozenset) -> np.ndarray: + if not present: + raise ValueError("cannot mask everything") + return y + rng.normal(0, len(present), 40) + + report = analyze_modality_failure(["echo", "ecg"], y, gate, predict_fn) + assert report.complementarity["shapley_value"] is None + assert report.complementarity["marginal_value"] # the rest of the report still works + + def test_classify_taxonomy_categories(): y = np.array([35.0, 55.0, 50.0]) pred = np.array(