Skip to content

recall_guard.mia.mcs

recall_guard.mia.mcs

Per-model MCS (Memorization Contamination Score) logistic-regression calibrator.

Implements the mia.mcs component from the honest-model-ranking design. Satisfies Requirements 5.1, 5.2, 5.3, 5.4:

  • MCSCalibrator: frozen dataclass holding the trained LogisticRegression estimator, the canonical feature_order used during training (so predict_proba cannot accidentally feed the classifier a permuted vector), the held-out AUC, and an is_weak flag set when holdout_auc < min_auc.
  • train(model_lm, is_memorized, oos_control, baseline, ref_lm, min_auc=0.6, seed=0): runs the model (and optional reference) on every row, computes MIA features, standardises them against the per-model control baseline, splits a held-out portion via sklearn.model_selection.train_test_split (test_size=0.25, stratify=y, random_state=seed), fits LogisticRegression(class_weight="balanced", solver="liblinear", random_state=seed) on the training half, scores roc_auc_score on the holdout half.
  • MCSCalibrator.predict_proba(features, baseline) -> float, pure: standardises features against baseline, builds the classifier input vector in feature_order, and returns the predict_proba(...)[:, 1] value.

Per-row LM failures (TimeoutError, RuntimeError, ValueError) are skipped with a single WARNING per skip; every other code path is pure. The MemGuard penalty rule consumed downstream is penalized_confidence = raw_confidence * (1 - p_memorized) (Req 5.4: continuous, not threshold-based).

MCSCalibrator dataclass

Per-model logistic-regression calibrator for p(memorized | features).

Attributes:

Name Type Description
model str

The NVIDIA model ID this calibrator was trained for.
classifier LogisticRegression

The fitted sklearn.linear_model.LogisticRegression instance. sklearn estimators are mutable; frozen=True only prevents reassignment of the field reference, which is the design intent.
feature_order list[str]

Canonical order used to flatten the standardised feature dict into the classifier's input vector. Populated at train time and consumed verbatim by :meth:predict_proba so the classifier is never fed a permuted row.
holdout_auc float

ROC-AUC score of the trained classifier on the 25% held-out portion of the labelled IS/OOS corpus. Reported in the manifest and the per-model evaluation result (Req 5.2).
is_weak bool

True iff holdout_auc < min_auc at train time. Surfaced as the weak-calibration warning in top3.md (Req 5.3).
Source code in recall_guard/mia/mcs.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
@dataclass(frozen=True)
class MCSCalibrator:
    """Per-model logistic-regression calibrator for ``p(memorized | features)``.

    Attributes
    ----------
    model:
        The NVIDIA model ID this calibrator was trained for.
    classifier:
        The fitted ``sklearn.linear_model.LogisticRegression`` instance.
        sklearn estimators are mutable; ``frozen=True`` only prevents
        reassignment of the field reference, which is the design intent.
    feature_order:
        Canonical order used to flatten the standardised feature dict
        into the classifier's input vector. Populated at train time and
        consumed verbatim by :meth:`predict_proba` so the classifier is
        never fed a permuted row.
    holdout_auc:
        ROC-AUC score of the trained classifier on the 25% held-out
        portion of the labelled IS/OOS corpus. Reported in the manifest
        and the per-model evaluation result (Req 5.2).
    is_weak:
        ``True`` iff ``holdout_auc < min_auc`` at train time. Surfaced
        as the ``weak-calibration`` warning in ``top3.md`` (Req 5.3).
    """

    model: str
    classifier: LogisticRegression
    feature_order: list[str]
    holdout_auc: float
    is_weak: bool

    def predict_proba(
        self, features: MiaFeatures, baseline: ControlBaseline
    ) -> float:
        """Return the calibrated probability of "memorized" for one record.

        Standardises ``features`` against the model's ``baseline`` and
        feeds the resulting vector to the trained classifier in
        ``self.feature_order``.

        Returns
        -------
        float
            ``p(memorized | features) ∈ [0.0, 1.0]``.

        Raises
        ------
        ValueError
            If any of the four core features standardises to ``None``
            (uncalibrated baseline). A missing ``ref_delta`` does NOT
            raise: the reference feature is optional by contract, so it
            is imputed at the control-baseline mean (standardised 0.0),
            which contributes no memorization evidence either way.
        """
        standardised = standardise(features, baseline)
        if "ref_delta" in self.feature_order and standardised.get("ref_delta") is None:
            standardised = {**standardised, "ref_delta": 0.0}
        row = _row_vector(standardised, self.feature_order)
        # Estimator was trained on a 2-D matrix; predict on a 1-row matrix.
        proba = float(self.classifier.predict_proba(row.reshape(1, -1))[0, 1])
        # sklearn returns values strictly in [0, 1] for LR; clamp defensively
        # against fp64 round-off so the float postcondition holds exactly.
        if proba < 0.0:
            proba = 0.0
        elif proba > 1.0:
            proba = 1.0
        return proba

predict_proba 

predict_proba(features, baseline)

Return the calibrated probability of "memorized" for one record.

Standardises features against the model's baseline and feeds the resulting vector to the trained classifier in self.feature_order.

Returns:

Type Description
float

p(memorized | features) ∈ [0.0, 1.0].

Raises:

Type Description
ValueError

If any of the four core features standardises to None (uncalibrated baseline). A missing ref_delta does NOT raise: the reference feature is optional by contract, so it is imputed at the control-baseline mean (standardised 0.0), which contributes no memorization evidence either way.
Source code in recall_guard/mia/mcs.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def predict_proba(
    self, features: MiaFeatures, baseline: ControlBaseline
) -> float:
    """Return the calibrated probability of "memorized" for one record.

    Standardises ``features`` against the model's ``baseline`` and
    feeds the resulting vector to the trained classifier in
    ``self.feature_order``.

    Returns
    -------
    float
        ``p(memorized | features) ∈ [0.0, 1.0]``.

    Raises
    ------
    ValueError
        If any of the four core features standardises to ``None``
        (uncalibrated baseline). A missing ``ref_delta`` does NOT
        raise: the reference feature is optional by contract, so it
        is imputed at the control-baseline mean (standardised 0.0),
        which contributes no memorization evidence either way.
    """
    standardised = standardise(features, baseline)
    if "ref_delta" in self.feature_order and standardised.get("ref_delta") is None:
        standardised = {**standardised, "ref_delta": 0.0}
    row = _row_vector(standardised, self.feature_order)
    # Estimator was trained on a 2-D matrix; predict on a 1-row matrix.
    proba = float(self.classifier.predict_proba(row.reshape(1, -1))[0, 1])
    # sklearn returns values strictly in [0, 1] for LR; clamp defensively
    # against fp64 round-off so the float postcondition holds exactly.
    if proba < 0.0:
        proba = 0.0
    elif proba > 1.0:
        proba = 1.0
    return proba

train 

train(
    model_lm,
    is_memorized,
    oos_control,
    baseline,
    ref_lm,
    min_auc=0.6,
    seed=0,
    max_workers=1,
)

Train the MCS classifier for one model.

Drives the LM over both labelled corpora (in parallel when max_workers > 1), fits a logistic regression on the standardised features, and reports a held-out AUC. Raises ValueError if either class ends up empty after per-row skips.

Source code in recall_guard/mia/mcs.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def train(
    model_lm: NvidiaLM,
    is_memorized: list[EvalRow],
    oos_control: list[EvalRow],
    baseline: ControlBaseline,
    ref_lm: NvidiaLM | None,
    min_auc: float = 0.6,
    seed: int = 0,
    max_workers: int = 1,
) -> MCSCalibrator:
    """Train the MCS classifier for one model.

    Drives the LM over both labelled corpora (in parallel when
    ``max_workers > 1``), fits a logistic regression on the standardised
    features, and reports a held-out AUC. Raises ``ValueError`` if either
    class ends up empty after per-row skips.
    """
    feature_order = _resolve_feature_order(baseline)
    x, y, n_valid_is, n_valid_oos = _gather_train_xy(
        model_lm=model_lm,
        is_memorized=is_memorized,
        oos_control=oos_control,
        baseline=baseline,
        ref_lm=ref_lm,
        feature_order=feature_order,
        max_workers=max_workers,
    )

    # The stratified holdout needs at least one row per class in BOTH the
    # train and holdout halves. Check up front so tiny corpora fail with a
    # clear message instead of an opaque sklearn split error.
    n_total = n_valid_is + n_valid_oos
    n_holdout = math.ceil(_HOLDOUT_FRACTION * n_total)
    if n_holdout < 2 or (n_total - n_holdout) < 2:
        raise ValueError(
            f"mcs.train: {n_total} valid rows "
            f"(n_valid_is={n_valid_is}, n_valid_oos={n_valid_oos}) cannot "
            f"support the stratified {_HOLDOUT_FRACTION:.0%} holdout split "
            f"(holdout would hold {n_holdout} row(s), need >= 2 with both "
            "classes). Provide more calibration rows."
        )

    x_train, x_holdout, y_train, y_holdout = train_test_split(
        x, y,
        test_size=_HOLDOUT_FRACTION,
        random_state=seed,
        stratify=y,
    )

    classifier = LogisticRegression(
        class_weight="balanced",
        solver="liblinear",
        random_state=seed,
    )
    classifier.fit(x_train, y_train)

    holdout_scores = classifier.predict_proba(x_holdout)[:, 1]
    holdout_auc = float(roc_auc_score(y_holdout, holdout_scores))
    is_weak = holdout_auc < float(min_auc)

    logger.info(
        "mcs.train: model=%s n_valid_is=%d n_valid_oos=%d "
        "holdout_auc=%.4f is_weak=%s",
        model_lm.model, n_valid_is, n_valid_oos, holdout_auc, is_weak,
    )

    return MCSCalibrator(
        model=model_lm.model,
        classifier=classifier,
        feature_order=feature_order,
        holdout_auc=holdout_auc,
        is_weak=is_weak,
    )