Probes API

The probes module trains a classifier to predict the group from the embeddings. Both probes subclass DetectorBase and do their own train/test split or cross-validation inside fit.

Class reference

SKLearnProbe

SKLearnProbe

SKLearnProbe(
    estimator: BaseEstimator | None = None,
    *,
    metric: MetricName = "f1",
    threshold: float = 0.7,
    average: str = "macro",
    evaluation: EvaluationName = "holdout",
    test_size: float = 0.2,
    cv_folds: int = 5,
    random_state: int = 0
)

Bases: DetectorBase

Shortcut detector based on training a classifier to predict a demographic target.

Idea

• Train a probe classifier to predict a sensitive/demographic attribute y from embeddings X.
• If the probe performs above a user-defined threshold on a metric (e.g., F1), treat this as evidence that embeddings encode the attribute (potential shortcut).

Parameters

estimator: Any scikit-learn estimator supporting fit/predict (optionally predict_proba or decision_function). If None, uses a standardized LogisticRegression. metric: One of: "accuracy", "f1", "precision", "recall", "roc_auc". threshold: Shortcut is detected if metric_value > threshold. average: Averaging strategy for multiclass f1/precision/recall ("macro", "micro", "weighted"). For binary problems, "binary" is used automatically for these metrics. evaluation: "holdout" (train/test split), "cv" (StratifiedKFold cross-validation), or "train" (no splitting). test_size: Used for holdout split. cv_folds: Used for CV. random_state: Reproducibility for splitting.

Fit inputs

embeddings: np.ndarray, shape (n_samples, n_features) target: np.ndarray, shape (n_samples,) Demographic/sensitive attribute labels (e.g., gender).

Source code in shortcut_detect/probes/sklearn_probe.py

def __init__(
    self,
    estimator: BaseEstimator | None = None,
    *,
    metric: MetricName = "f1",
    threshold: float = 0.70,
    average: str = "macro",
    evaluation: EvaluationName = "holdout",
    test_size: float = 0.2,
    cv_folds: int = 5,
    random_state: int = 0,
):
    super().__init__(method="ml_probe")

    if estimator is None:
        # Reasonable default probe: scale then logistic regression
        estimator = Pipeline(
            steps=[
                ("scaler", StandardScaler(with_mean=True, with_std=True)),
                (
                    "clf",
                    LogisticRegression(
                        max_iter=2000,
                        solver="lbfgs",
                        random_state=random_state,
                    ),
                ),
            ]
        )

    self.estimator = estimator
    self.config = MLProbeConfig(
        metric=metric,
        threshold=float(threshold),
        average=average,
        evaluation=evaluation,
        test_size=float(test_size),
        cv_folds=int(cv_folds),
        random_state=int(random_state),
    )

    # Fitted artifacts
    self.estimator_: BaseEstimator | None = None
    self.metric_value_: float | None = None
    self.y_true_eval_: np.ndarray | None = None
    self.y_pred_eval_: np.ndarray | None = None

Methods:

predict

predict(X: ndarray) -> np.ndarray

Predict class labels for embeddings (requires prior fit).

Source code in shortcut_detect/probes/sklearn_probe.py

def predict(self, X: np.ndarray) -> np.ndarray:
    """Predict class labels for embeddings (requires prior fit)."""
    self._ensure_fitted()
    X = np.asarray(X)
    if X.ndim != 2:
        raise ValueError(f"X must be 2D (n_samples, n_features). Got shape={X.shape}.")
    return np.asarray(self.estimator_.predict(X))

TorchProbe

TorchProbe

TorchProbe(
    model: Module,
    loss_fn: Any,
    *,
    optimizer_class: Any = torch.optim.Adam,
    optimizer_kwargs: dict[str, Any] | None = None,
    device: str | None = None,
    metric: MetricName = "accuracy",
    threshold: float = 0.7,
    test_size: float = 0.2,
    random_state: int = 0,
    epochs: int = 10,
    batch_size: int = 128,
    num_workers: int = 0,
    early_stopping: int | None = None,
    use_amp: bool = False,
    verbose: bool = False,
    loader_factory: LoaderFactory | None = None,
    stage_loader_overrides: StageLoaderOverrides = None
)

Bases: DetectorBase

Probe-based shortcut detector using a PyTorch model.

Fits a torch model to predict a demographic target from embeddings and flags a shortcut if the chosen metric exceeds a threshold.

Parameters

model: torch.nn.Module that maps embeddings -> logits (classification) or scalar (regression). For classification, return shape (N, C) logits. loss_fn: Loss function (e.g., nn.CrossEntropyLoss()). optimizer_class / optimizer_kwargs: Optimizer configuration. device: "cpu" or "cuda"; defaults to CUDA if available. metric: One of: "accuracy", "f1", "roc_auc", "loss". For multiclass, "f1" uses macro averaging; "roc_auc" only supported for binary. threshold: shortcut_detected is True when metric_value > threshold (except for "loss", where shortcut_detected is True when loss < threshold if you choose to use loss; see notes).

Source code in shortcut_detect/probes/torch_probe.py

def __init__(
    self,
    model: nn.Module,
    loss_fn: Any,
    *,
    optimizer_class: Any = torch.optim.Adam,
    optimizer_kwargs: dict[str, Any] | None = None,
    device: str | None = None,
    metric: MetricName = "accuracy",
    threshold: float = 0.70,
    test_size: float = 0.2,
    random_state: int = 0,
    epochs: int = 10,
    batch_size: int = 128,
    num_workers: int = 0,
    early_stopping: int | None = None,
    use_amp: bool = False,
    verbose: bool = False,
    loader_factory: LoaderFactory | None = None,
    stage_loader_overrides: StageLoaderOverrides = None,
):
    super().__init__(method="torch_probe")

    self.model = model
    self.loss_fn = loss_fn
    self.optimizer_class = optimizer_class
    self.optimizer_kwargs = optimizer_kwargs or {"lr": 1e-3}
    self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

    self.config = TorchProbeConfig(
        metric=metric,
        threshold=float(threshold),
        test_size=float(test_size),
        random_state=int(random_state),
        epochs=int(epochs),
        batch_size=int(batch_size),
        num_workers=int(num_workers),
        early_stopping=early_stopping,
        use_amp=bool(use_amp),
        verbose=bool(verbose),
    )

    self.shortcut_detected_: bool | None = None
    self.metric_value_: float | None = None
    self.n_features_in_: int | None = None

    self._optimizer: torch.optim.Optimizer | None = None
    self._scaler: torch.cuda.amp.GradScaler | None = None
    self.loader_factory = loader_factory
    self.stage_loader_overrides = stage_loader_overrides

Methods:

fit_dataset

fit_dataset(
    dataset: Dataset,
    *,
    val_dataset: Dataset | None = None,
    target_extractor: Callable[[Any], Any] | None = None,
    data_spec: DataSpec | dict[str, Any] | None = None
) -> TorchProbe

Train using map-style or iterable datasets without materializing full arrays.

Source code in shortcut_detect/probes/torch_probe.py

def fit_dataset(
    self,
    dataset: Dataset,
    *,
    val_dataset: Dataset | None = None,
    target_extractor: Callable[[Any], Any] | None = None,
    data_spec: DataSpec | dict[str, Any] | None = None,
) -> TorchProbe:
    """Train using map-style or iterable datasets without materializing full arrays."""
    resolved_spec = resolve_data_spec(data_spec)
    is_iterable = is_iterable_dataset(dataset)

    if is_iterable and val_dataset is None:
        raise ValueError("IterableDataset requires `val_dataset` for TorchProbe.fit_dataset.")

    if val_dataset is None:
        n = safe_len(dataset)
        if n is None or n < 4:
            raise ValueError("Dataset must expose len() and have at least 4 samples.")
        rng = np.random.RandomState(self.config.random_state)
        idx = np.arange(n)
        rng.shuffle(idx)
        test_n = max(1, int(round(self.config.test_size * n)))
        val_idx = idx[:test_n]
        train_idx = idx[test_n:]
        train_dataset = Subset(dataset, train_idx.tolist())
        eval_dataset = Subset(dataset, val_idx.tolist())
    else:
        train_dataset = dataset
        eval_dataset = val_dataset

    train_loader = self._make_loader_from_dataset(
        train_dataset,
        stage="train",
        shuffle=not is_iterable_dataset(train_dataset),
    )
    val_loader = self._make_loader_from_dataset(
        eval_dataset,
        stage="val",
        shuffle=False,
    )
    return self.fit_loaders(
        train_loader,
        val_loader=val_loader,
        target_extractor=target_extractor,
        data_spec=resolved_spec,
    )

fit_loaders

fit_loaders(
    train_loader: DataLoader,
    *,
    val_loader: DataLoader,
    target_extractor: Callable[[Any], Any] | None = None,
    data_spec: DataSpec | dict[str, Any] | None = None
) -> TorchProbe

Train/evaluate from user-provided loaders.

Source code in shortcut_detect/probes/torch_probe.py

def fit_loaders(
    self,
    train_loader: DataLoader,
    *,
    val_loader: DataLoader,
    target_extractor: Callable[[Any], Any] | None = None,
    data_spec: DataSpec | dict[str, Any] | None = None,
) -> TorchProbe:
    """Train/evaluate from user-provided loaders."""
    resolved_spec = resolve_data_spec(data_spec)
    device = torch.device(self.device)

    train_loss_hist, best_val_loss, best_epoch = self._train_with_loaders(
        train_loader,
        val_loader,
        target_extractor=target_extractor,
    )
    eval_metrics = self._evaluate_loader(
        val_loader,
        device,
        target_extractor=target_extractor,
    )

    metric = self.config.metric
    if metric not in eval_metrics:
        raise ValueError(
            f"Metric '{metric}' not available. Available: {sorted(eval_metrics.keys())}"
        )
    metric_value = float(eval_metrics[metric])
    self.metric_value_ = metric_value

    if metric == "loss":
        shortcut = bool(metric_value < self.config.threshold)
        notes_rule = "shortcut_detected is True when loss < threshold."
        if metric_value <= min(self.config.threshold, 0.35):
            risk_level = "high"
        elif metric_value < self.config.threshold:
            risk_level = "moderate"
        else:
            risk_level = "low"
    else:
        shortcut = bool(metric_value > self.config.threshold)
        notes_rule = f"shortcut_detected is True when {metric} > threshold."
        if metric_value >= max(self.config.threshold, 0.85):
            risk_level = "high"
        elif metric_value >= self.config.threshold:
            risk_level = "moderate"
        else:
            risk_level = "low"

    self.shortcut_detected_ = shortcut
    n_train = (
        resolved_spec.train_size
        if resolved_spec is not None
        else safe_len(getattr(train_loader, "dataset", None))
    )
    n_test = (
        resolved_spec.val_size
        if resolved_spec is not None
        else safe_len(getattr(val_loader, "dataset", None))
    )

    self._set_results(
        shortcut_detected=shortcut,
        risk_level=risk_level,
        metrics={
            "metric": metric,
            "metric_value": metric_value,
            "threshold": self.config.threshold,
            "protocol": "loader",
        },
        notes=("Trained a PyTorch probe model from provided data loaders. " + notes_rule),
        metadata={
            "device": self.device,
            "epochs": self.config.epochs,
            "batch_size": self.config.batch_size,
            "optimizer": getattr(self.optimizer_class, "__name__", str(self.optimizer_class)),
            "optimizer_kwargs": dict(self.optimizer_kwargs),
            "n_train": n_train,
            "n_test": n_test,
        },
        report={
            "protocol": "loader",
            "test_size": self.config.test_size,
            "train_loss_history": train_loss_hist,
            "best_val_loss": best_val_loss,
            "best_epoch": best_epoch,
            "eval_metrics": {k: float(v) for k, v in eval_metrics.items()},
        },
    )
    self._is_fitted = True
    return self

predict

predict(X: ndarray) -> np.ndarray

Predict class labels for embeddings (requires prior fit).

Source code in shortcut_detect/probes/torch_probe.py

def predict(self, X: np.ndarray) -> np.ndarray:
    """Predict class labels for embeddings (requires prior fit)."""
    self._ensure_fitted()
    X_arr = np.asarray(X)
    if X_arr.ndim != 2:
        raise ValueError(f"X must be 2D (n_samples, n_features). Got shape={X_arr.shape}.")
    if self.n_features_in_ is not None and X_arr.shape[1] != self.n_features_in_:
        raise ValueError(
            f"X has {X_arr.shape[1]} features, but the probe was fitted on {self.n_features_in_}."
        )
    device = torch.device(self.device)
    self.model.to(device)
    self.model.eval()
    loader = self._make_loader(
        X_arr,
        np.zeros(X_arr.shape[0], dtype=np.int64),
        stage="predict",
        batch_size=max(64, min(256, int(X_arr.shape[0]))),
        shuffle=False,
        num_workers=0,
    )
    preds_list = []
    with torch.no_grad():
        for xb, _ in loader:
            xb = xb.to(device, non_blocking=True)
            out = self.model(xb)
            if out.dim() == 2 and out.shape[1] > 1:
                preds_list.append(out.argmax(dim=1).detach().cpu().numpy())
            else:
                preds_list.append(out.squeeze(-1).detach().cpu().numpy().astype(np.int64))
    return np.concatenate(preds_list, axis=0)

SKLearnProbe

Constructor

SKLearnProbe(
    estimator: BaseEstimator | None = None,
    *,
    metric: str = "f1",
    threshold: float = 0.70,
    average: str = "macro",
    evaluation: str = "holdout",
    test_size: float = 0.2,
    cv_folds: int = 5,
    random_state: int = 0,
)

Parameters

Parameter	Type	Default	Description
estimator	sklearn estimator	standardised LogisticRegression	Any estimator with fit and predict
metric	str	"f1"	"accuracy", "f1", "precision", "recall" or "roc_auc"
threshold	float	0.70	Shortcut when the metric exceeds it
average	str	"macro"	Averaging for multiclass f1, precision, recall
evaluation	str	"holdout"	"holdout", "cv" or "train"
test_size	float	0.2	Held-out share for "holdout"
cv_folds	int	5	Folds for "cv"
random_state	int	0	Seed for the split

Methods

fit()

def fit(embeddings: np.ndarray, target: np.ndarray) -> SKLearnProbe

Splits (or cross-validates), trains the estimator, scores it and fills results_.

predict()

def predict(X: np.ndarray) -> np.ndarray

Labels from the fitted estimator.

get_report(), summary()

The standard detector schema and its one-line summary.

Attributes

Attribute	Type	Description
metric_value_	float	Held-out or cross-validated score
shortcut_detected_	bool	metric_value_ > threshold
estimator_	estimator	Fitted estimator
y_true_eval_, y_pred_eval_	ndarray	Held-out labels and predictions (holdout and train evaluation)
results_	dict	Standard schema; results_["report"] holds fold_scores, mean_score, std_score under evaluation="cv"

Usage

from sklearn.linear_model import LogisticRegression
from shortcut_detect import SKLearnProbe

probe = SKLearnProbe(LogisticRegression(max_iter=1000), metric="accuracy",
                     evaluation="holdout", test_size=0.2, random_state=42)
probe.fit(embeddings, group_labels)

print(f"Accuracy: {probe.metric_value_:.2%}, flagged: {probe.shortcut_detected_}")

---

TorchProbe

Constructor

TorchProbe(
    model: nn.Module,
    loss_fn,
    *,
    optimizer_class=torch.optim.Adam,
    optimizer_kwargs: dict | None = None,   # {"lr": 1e-3}
    device: str | None = None,
    metric: str = "accuracy",
    threshold: float = 0.70,
    test_size: float = 0.2,
    random_state: int = 0,
    epochs: int = 10,
    batch_size: int = 128,
    num_workers: int = 0,
    early_stopping: int | None = None,
    use_amp: bool = False,
    verbose: bool = False,
    loader_factory=None,
    stage_loader_overrides=None,
)

Parameters

Parameter	Type	Default	Description
model	nn.Module	required	The probe network
loss_fn	callable	required	For example nn.CrossEntropyLoss()
optimizer_class	class	torch.optim.Adam	Optimizer
optimizer_kwargs	dict	{"lr": 1e-3}	Optimizer arguments
device	str or None	None	"cuda" when available, else "cpu"
metric	str	"accuracy"	Score reported in metric_value_
threshold	float	0.70	Shortcut when the metric exceeds it
test_size	float	0.2	Held-out share
epochs	int	10	Training epochs
batch_size	int	128	Batch size
early_stopping	int or None	None	Patience in epochs
loader_factory	callable or None	None	Builds a DataLoader per stage
stage_loader_overrides	dict or None	None	Per-stage DataLoader kwargs

Methods

• fit(embeddings, target): in-memory arrays
• fit_dataset(dataset, ...): a map-style or iterable torch.utils.data.Dataset yielding {"embeddings": x, "labels": y}
• fit_loaders(...): pre-built DataLoaders; see the class reference
• predict(X): labels from the trained network

Attributes

Attribute	Type	Description
metric_value_	float	Held-out score
shortcut_detected_	bool	metric_value_ > threshold
n_features_in_	int	Input width seen in fit
results_	dict	Standard schema

Usage

import torch.nn as nn
from shortcut_detect import TorchProbe

class CustomProbe(nn.Module):
    def __init__(self, input_dim, n_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, n_classes),
        )

    def forward(self, x):
        return self.net(x)

probe = TorchProbe(
    model=CustomProbe(512, 3),
    loss_fn=nn.CrossEntropyLoss(),
    device="cuda",
    epochs=50,
)
probe.fit(embeddings, group_labels)
print(probe.metric_value_, probe.shortcut_detected_)

See also

• Probe method guide
• ShortcutDetector API