diff --git a/changelog/1261.added.md b/changelog/1261.added.md new file mode 100644 index 000000000..9789cfb9a --- /dev/null +++ b/changelog/1261.added.md @@ -0,0 +1 @@ +Add `prediction_scaling` to `TabPFNClassifier` and `TabPFNRegressor` to rescale predictions toward the training target distribution. Modes: `"none"`, `"balanced"` (classification only, the former `balance_probabilities`), `"sampler"` (undo the label shift introduced by `SAMPLE_SUBSAMPLING_METHOD="majority_downsample"`, computed from the context and training class shares at no extra cost), and `"holdout"` (fit the correction on held-out rows, sharing the holdout with `tuning_config` when set). The default `"auto"` applies `"sampler"` under `majority_downsample` and nothing otherwise, so predictions are unchanged for every other configuration. diff --git a/changelog/1261.deprecated.md b/changelog/1261.deprecated.md new file mode 100644 index 000000000..28a4dc7a4 --- /dev/null +++ b/changelog/1261.deprecated.md @@ -0,0 +1 @@ +`TabPFNClassifier(balance_probabilities=True)` is deprecated in favor of `prediction_scaling="balanced"`. It still works and produces identical probabilities, but now emits a `DeprecationWarning`. diff --git a/src/tabpfn/classifier.py b/src/tabpfn/classifier.py index aae4404fc..85143486c 100644 --- a/src/tabpfn/classifier.py +++ b/src/tabpfn/classifier.py @@ -62,6 +62,7 @@ from tabpfn.inference_tuning import ( ClassifierEvalMetrics, ClassifierTuningConfig, + TuningConfig, find_optimal_classification_thresholds, find_optimal_temperature, get_tuning_splits, @@ -73,6 +74,16 @@ prepend_cache_path, save_fitted_tabpfn_model, ) +from tabpfn.prediction_scaling import ( + PredictionScaling, + PredictionScalingMode, + apply_class_weights, + balanced_class_weights, + context_class_prior, + fit_holdout_class_weights, + resolve_prediction_scaling, + sampler_class_weights, +) from tabpfn.preprocessing import ( ClassifierEnsembleConfig, EnsembleConfig, @@ -241,6 +252,13 @@ class TabPFNClassifier(ClassifierMixin, BaseEstimator): """The ensemble configurations used during fit. Stored for reuse in prompt tuning.""" + prediction_scaling_: PredictionScalingMode + """The prediction scaling mode in effect after resolving `"auto"`.""" + + prediction_scaling_weights_: np.ndarray | None + """Per-class weights applied to the averaged probabilities, or `None` when no + rescaling is applied.""" + def __init__( # noqa: PLR0913 self, *, @@ -248,6 +266,7 @@ def __init__( # noqa: PLR0913 auto_scale_n_estimators: bool = True, categorical_features_indices: Sequence[int] | None = None, softmax_temperature: float | Literal["auto"] = "auto", + prediction_scaling: PredictionScaling = "auto", balance_probabilities: bool = False, average_before_softmax: bool = False, model_path: str @@ -346,13 +365,39 @@ def __init__( # noqa: PLR0913 combined with a `SOFTMAX_TEMPERATURE` in `inference_config`, which is the other way of naming one. + prediction_scaling: + Post-processing that rescales the predicted class probabilities + toward a reference class prior, as opposed to `softmax_temperature`, + which changes their confidence. Each mode multiplies the averaged + probabilities by a per-class weight vector and renormalizes; the + modes differ in where the weights come from. + + - `"none"`: no rescaling. + - `"balanced"`: divide by the class prior so the predicted prior + moves toward uniform. Helps metrics that are insensitive to class + imbalance (balanced accuracy, macro ROC AUC). Without row + subsampling this is the training prior; under + `SAMPLE_SUBSAMPLING_METHOD="majority_downsample"` it is the + context prior the model actually saw. + - `"sampler"`: undo the label shift introduced by row subsampling by + multiplying with training prior over context prior. Exact under the + label-shift assumption and free. A no-op when the sampler did not + change the prior. + - `"holdout"`: fit the weights on held-out rows so the mean predicted + probability per class matches the observed class frequency. Costs + one extra fit per tuning fold; shares the holdout with + `tuning_config` when that is set. Not available with + `differentiable_input=True`. + - `"auto"` (default): `"sampler"` when row subsampling shifted the + prior, otherwise `"none"`. Existing behavior is unchanged for every + configuration that does not use `"majority_downsample"`. + + Rescaling by a per-class factor is monotone within each class, so for + binary tasks it never changes the ranking of the positive class. + balance_probabilities: - Whether to balance the probabilities based on the class distribution - in the training data. This can help to improve predictive performance - when the classes are highly imbalanced and the metric of interest is - insensitive to class imbalance (e.g., balanced accuracy, balanced log - loss, roc-auc macro ovo, etc.). This is only applied when predicting - during a post-processing step. + Deprecated alias for `prediction_scaling="balanced"`. Will be removed + in a future release. average_before_softmax: Only used if `n_estimators > 1`. Whether to average the predictions of @@ -562,6 +607,7 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this self.auto_scale_n_estimators = auto_scale_n_estimators self.categorical_features_indices = categorical_features_indices self.softmax_temperature = softmax_temperature + self.prediction_scaling = prediction_scaling self.balance_probabilities = balance_probabilities self.average_before_softmax = average_before_softmax self.model_path = model_path @@ -890,6 +936,11 @@ def _get_tuning_classifier(self, **overwrite_kwargs: Any) -> TabPFNClassifier: # Fit on the already-expanded array, where a declared column may # have moved down past an expanded date or text column. "categorical_features_indices": self.categorical_features_indices_, + # Holdout rows must be scored without any rescaling: the holdout mode + # learns the whole correction from them, and the free modes are + # applied by this estimator, not the clone. + "prediction_scaling": "none", + "balance_probabilities": False, } params.update(forced) @@ -897,6 +948,69 @@ def _get_tuning_classifier(self, **overwrite_kwargs: Any) -> TabPFNClassifier: return TabPFNClassifier(**params) + def _resolve_prediction_scaling(self, *, y_encoded: np.ndarray) -> None: + """Resolve `prediction_scaling` and compute the free weight vectors. + + Requires `class_counts_` and `ensemble_preprocessor_`. The holdout mode + leaves the weights unset here; they are fitted from held-out rows in + `_maybe_calibrate_temperature_and_tune_decision_thresholds`. + """ + requested = self.prediction_scaling + if self.balance_probabilities: + if PredictionScalingMode(requested) not in ( + PredictionScalingMode.AUTO, + PredictionScalingMode.BALANCED, + ): + raise ValueError( + "balance_probabilities=True conflicts with " + f"prediction_scaling={requested!r}. Drop balance_probabilities; " + "it is a deprecated alias for prediction_scaling='balanced'." + ) + warnings.warn( + "balance_probabilities is deprecated and will be removed in a future " + "release; pass prediction_scaling='balanced' instead.", + DeprecationWarning, + stacklevel=3, + ) + requested = "balanced" + + preprocessor = self.ensemble_preprocessor_ + shifted = preprocessor.sampler_shifted_prior + self.prediction_scaling_ = resolve_prediction_scaling( + requested, + task_type="classifier", + sampler_shifted_prior=shifted, + ) + + context_prior = ( + context_class_prior( + y_encoded, preprocessor.subsample_row_indices, self.n_classes_ + ) + if shifted + else None + ) + # The differentiable-input path does not go through the label encoder + # and therefore has no `class_counts_`; the encoded labels carry them. + train_counts = getattr(self, "class_counts_", None) + if train_counts is None: + train_counts = np.bincount( + np.asarray(y_encoded).astype(np.int64), minlength=self.n_classes_ + ) + mode = self.prediction_scaling_ + if mode == PredictionScalingMode.BALANCED: + self.prediction_scaling_weights_ = balanced_class_weights( + train_counts, context_prior + ) + elif mode == PredictionScalingMode.SAMPLER: + self.prediction_scaling_weights_ = ( + sampler_class_weights(train_counts, context_prior) + if context_prior is not None + else None + ) + else: + # NONE, or HOLDOUT until the holdout fit fills the weights in. + self.prediction_scaling_weights_ = None + @config_context(transform_output="default") # type: ignore def fit(self, X: XType, y: YType) -> Self: """Fit the model. @@ -931,8 +1045,6 @@ def fit(self, X: XType, y: YType) -> Self: ) self.ensemble_configs_ = ensemble_configs - self._maybe_calibrate_temperature_and_tune_decision_thresholds(X=X, y=y) - self.ensemble_preprocessor_ = TabPFNEnsemblePreprocessor( configs=ensemble_configs, n_samples=X.shape[0], @@ -957,6 +1069,12 @@ def fit(self, X: XType, y: YType) -> Self: task_type=self.estimator_type, ) + # The free scaling modes need the context prior, which the row sampler + # fixed above; the holdout mode is fitted inside the tuning step below, + # which also applies whatever weights are already known. + self._resolve_prediction_scaling(y_encoded=y) + self._maybe_calibrate_temperature_and_tune_decision_thresholds(X=X, y=y) + self.executor_ = create_inference_engine( fit_mode=self.fit_mode, X_train=X, @@ -1053,6 +1171,25 @@ def fit_from_preprocessed( return self + def _prediction_scaling_may_apply(self) -> bool: + """Whether `prediction_scaling` can produce non-trivial weights. + + `"auto"` and `"sampler"` only act when the row sampler shifts the prior, + so they stay usable in batched prediction unless that sampler is + configured. + """ + mode = PredictionScalingMode(self.prediction_scaling) + if mode == PredictionScalingMode.NONE: + return False + if mode in (PredictionScalingMode.BALANCED, PredictionScalingMode.HOLDOUT): + return True + config = self.get_inference_config() + return ( + config.SUBSAMPLE_SAMPLES is not None + and SampleSubsamplingMethod(config.SAMPLE_SUBSAMPLING_METHOD) + == SampleSubsamplingMethod.MAJORITY_DOWNSAMPLE + ) + def predict_proba_batched( # noqa: C901, PLR0912 self, X_train_list: list[XType], @@ -1087,8 +1224,9 @@ def predict_proba_batched( # noqa: C901, PLR0912 ValueError: If the input lists have unequal or zero length, the datasets do not all share the same set of classes, or the training (or test) arrays do not all share one shape. - NotImplementedError: If ``balance_probabilities`` or ``tuning_config`` - is configured on the estimator — their state is per-dataset and + NotImplementedError: If an active ``prediction_scaling`` (or the + deprecated ``balance_probabilities``) or ``tuning_config`` is + configured on the estimator — their state is per-dataset and cannot be applied correctly across a shared batch. Score those datasets individually with ``predict_proba``. Also raised for ``inference_precision=torch.float64``, which the fused forward @@ -1118,10 +1256,12 @@ def predict_proba_batched( # noqa: C901, PLR0912 # estimator but their fitted state (thresholds, class counts, calibrated # temperature) is per-dataset; applying the last dataset's state across the # whole batch would be wrong, so they are not supported here. - if self.balance_probabilities: + if self.balance_probabilities or self._prediction_scaling_may_apply(): raise NotImplementedError( - "predict_proba_batched does not support balance_probabilities=True; " - "score datasets individually with predict_proba." + "predict_proba_batched does not support prediction_scaling " + f"({self.prediction_scaling!r}) or balance_probabilities=True; the " + "class weights are fitted per dataset. Score datasets individually " + "with predict_proba, or pass prediction_scaling='none'." ) if self.tuning_config is not None: raise NotImplementedError( @@ -1322,6 +1462,14 @@ def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Sel y_train=y, task_type=self.estimator_type, ) + requested_scaling = PredictionScalingMode(self.prediction_scaling) + if requested_scaling == PredictionScalingMode.HOLDOUT: + raise ValueError( + "prediction_scaling='holdout' is not supported with " + "differentiable_input=True; use 'auto', 'none', 'balanced', or " + "'sampler'." + ) + self._resolve_prediction_scaling(y_encoded=y.detach().cpu().float().numpy()) self.executor_ = InferenceEngineCachePreprocessing( X_train=X, @@ -1355,10 +1503,9 @@ def _maybe_calibrate_temperature_and_tune_decision_thresholds( # `softmax_temperature_` is already resolved by `_initialize_model_variables`. self.tuned_classification_thresholds_ = None - tuning_config_resolved = resolve_tuning_config( - tuning_config=self.tuning_config, - num_samples=X.shape[0], - config_cls=ClassifierTuningConfig, + needs_holdout_scaling = self._needs_holdout_prediction_scaling() + tuning_config_resolved = self._resolve_tuning_config_for_fit( + num_samples=X.shape[0], force_holdout=needs_holdout_scaling ) if tuning_config_resolved is None: if self.eval_metric_ is ClassifierEvalMetrics.F1: @@ -1429,6 +1576,9 @@ def _maybe_calibrate_temperature_and_tune_decision_thresholds( # whatever the checkpoint declared. self.softmax_temperature_ = calibrated_softmax_temperature + if needs_holdout_scaling: + self._fit_holdout_prediction_scaling(holdout_raw_logits, holdout_y_true) + if tuning_config_resolved.tune_decision_thresholds: holdout_probas = ( self.logits_to_probabilities(holdout_raw_logits) @@ -1445,6 +1595,50 @@ def _maybe_calibrate_temperature_and_tune_decision_thresholds( ) self.tuned_classification_thresholds_ = tuned_classification_thresholds + def _resolve_tuning_config_for_fit( + self, *, num_samples: int, force_holdout: bool + ) -> TuningConfig | None: + """Resolve `tuning_config`, creating a default one when a holdout is + needed for `prediction_scaling="holdout"` but no tuning was requested. + """ + resolved = resolve_tuning_config( + tuning_config=self.tuning_config, + num_samples=num_samples, + config_cls=ClassifierTuningConfig, + ) + if resolved is None and force_holdout: + resolved = ClassifierTuningConfig().resolve(num_samples=num_samples) + return resolved + + def _needs_holdout_prediction_scaling(self) -> bool: + return ( + getattr(self, "prediction_scaling_", None) == PredictionScalingMode.HOLDOUT + ) + + def _fit_holdout_prediction_scaling( + self, + holdout_raw_logits: np.ndarray, + holdout_y_true: np.ndarray, + ) -> None: + """Fit `prediction_scaling_weights_` so the mean predicted probability per + class on the holdout rows matches the holdout class frequency. + + Runs after temperature calibration, so the weights are fitted on the + probabilities that will actually be produced at predict time. + """ + holdout_probas = ( + self.logits_to_probabilities(holdout_raw_logits) + .float() + .detach() + .cpu() + .numpy() + ) + self.prediction_scaling_weights_ = fit_holdout_class_weights( + holdout_probas=holdout_probas, + holdout_y_true=holdout_y_true, + n_classes=self.n_classes_, + ) + def _compute_holdout_validation_data( self, X: XType, @@ -1654,7 +1848,6 @@ def logits_to_probabilities_fn( raw_logits=raw_logits, softmax_temperature=softmax_temperature, average_before_softmax=self.average_before_softmax, - balance_probabilities=self.balance_probabilities, ) .float() .detach() @@ -1706,12 +1899,19 @@ def _apply_softmax(self, logits: torch.Tensor) -> torch.Tensor: return torch.nn.functional.softmax(logits, dim=-1) def _apply_balancing(self, probas: torch.Tensor) -> torch.Tensor: - """Applies class balancing to a probability tensor.""" + """Applies class balancing by training class counts to a probability tensor.""" counts = getattr(self, "class_counts_", None) if counts is None: return probas return balance_probas_by_class_counts(probas, counts) + def _apply_prediction_scaling(self, probas: torch.Tensor) -> torch.Tensor: + """Applies the fitted `prediction_scaling_weights_` to a probability tensor.""" + weights = getattr(self, "prediction_scaling_weights_", None) + if weights is None: + return probas + return apply_class_weights(probas, weights) + def logits_to_probabilities( self, raw_logits: np.ndarray | torch.Tensor, @@ -1729,6 +1929,10 @@ def logits_to_probabilities( softmax_temperature: Optional override for temperature scaling. average_before_softmax: Optional override for averaging order. balance_probabilities: Optional override for probability balancing. + When `None`, the fitted `prediction_scaling_weights_` are applied, + which already cover a `"balanced"` prediction scaling. Passing + `True` forces plain balancing by training class counts and + skips the fitted weights; `False` skips both. Returns: Probabilities with shape (n_samples, n_classes). @@ -1748,11 +1952,8 @@ def logits_to_probabilities( if average_before_softmax is None else average_before_softmax ) - use_balance = ( - self.balance_probabilities - if balance_probabilities is None - else balance_probabilities - ) + use_fitted_weights = balance_probabilities is None + use_balance = bool(balance_probabilities) steps: list[Callable[[torch.Tensor], torch.Tensor]] = [] @@ -1777,7 +1978,9 @@ def apply_temp(t: torch.Tensor) -> torch.Tensor: f"Expected logits with 2 or more dims, got {raw_logits.ndim}" ) - if use_balance: + if use_fitted_weights: + steps.append(self._apply_prediction_scaling) + elif use_balance: steps.append(self._apply_balancing) output = raw_logits diff --git a/src/tabpfn/prediction_scaling.py b/src/tabpfn/prediction_scaling.py new file mode 100644 index 000000000..a21e03524 --- /dev/null +++ b/src/tabpfn/prediction_scaling.py @@ -0,0 +1,261 @@ +# Copyright (c) Prior Labs GmbH 2026. +"""Post-hoc rescaling of predictions toward a reference target distribution. + +TabPFN predicts in-context, so the prior it expresses is the one it sees in its +context rows. Two things pull that away from the training data: class balancing +requested by the user, and row subsampling that changes the class mix per +estimator (``SAMPLE_SUBSAMPLING_METHOD="majority_downsample"``). Every mode in +this module is one multiplicative correction of the predicted distribution; they +differ only in where the factors come from. + +Classification: a per-class weight vector applied to the averaged probabilities +and renormalized. Regression: a per-bucket weight on the bar distribution for the +sampler correction, and an affine map of the raw-space borders for the holdout +level correction. +""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING, Literal + +import numpy as np +import torch + +if TYPE_CHECKING: + from tabpfn.architectures.shared.bar_distribution import ( + FullSupportBarDistribution, + ) + +PredictionScaling = Literal["auto", "none", "balanced", "sampler", "holdout"] + + +class PredictionScalingMode(str, Enum): + """Source of the prediction scaling factors.""" + + AUTO = "auto" + NONE = "none" + BALANCED = "balanced" + SAMPLER = "sampler" + HOLDOUT = "holdout" + + +HOLDOUT_WEIGHT_FIT_ITERATIONS = 200 +"""Fixed-point iterations for matching mean predicted class probabilities to the +holdout class frequencies. The map contracts quickly; this is a generous cap.""" + +HOLDOUT_WEIGHT_FIT_TOLERANCE = 1e-8 + + +def resolve_prediction_scaling( + prediction_scaling: PredictionScaling | PredictionScalingMode, + *, + task_type: Literal["classifier", "regressor"], + sampler_shifted_prior: bool, +) -> PredictionScalingMode: + """Resolve ``"auto"`` and validate the mode against the task. + + ``"auto"`` becomes ``"sampler"`` when the row sampler changed the target prior + of the context, and ``"none"`` otherwise. ``"balanced"`` has no meaning for a + continuous target and is rejected for regressors. + """ + mode = PredictionScalingMode(prediction_scaling) + if mode == PredictionScalingMode.AUTO: + return ( + PredictionScalingMode.SAMPLER + if sampler_shifted_prior + else PredictionScalingMode.NONE + ) + if task_type == "regressor" and mode == PredictionScalingMode.BALANCED: + raise ValueError( + "prediction_scaling='balanced' is only defined for classification. " + "Use 'none', 'sampler', 'holdout', or 'auto' for regression." + ) + return mode + + +# --------------------------------------------------------------------------- # +# Classification +# --------------------------------------------------------------------------- # + + +def context_class_prior( + y_encoded: np.ndarray, + row_indices: list[np.ndarray] | None, + n_classes: int, +) -> np.ndarray: + """Class prior the estimators see in their context, averaged over estimators. + + With no row subsampling every estimator sees the full training set, so the + context prior is the training prior. + """ + y_encoded = np.asarray(y_encoded).astype(np.int64, copy=False) + if row_indices is None: + counts = np.bincount(y_encoded, minlength=n_classes).astype(np.float64) + return counts / counts.sum() + priors = [] + for idx in row_indices: + counts = np.bincount(y_encoded[idx], minlength=n_classes).astype(np.float64) + priors.append(counts / counts.sum()) + return np.mean(priors, axis=0) + + +def sampler_class_weights( + train_class_counts: np.ndarray, + context_prior: np.ndarray, +) -> np.ndarray: + """Weights that undo a label shift between context and training data. + + Under label shift ``p_train(c | x) ∝ p_context(c | x) * π_train(c) / π_context(c)``. + Classes absent from the context cannot be corrected and keep weight one. + """ + train_counts = np.asarray(train_class_counts, dtype=np.float64) + train_prior = train_counts / train_counts.sum() + context_prior = np.asarray(context_prior, dtype=np.float64) + weights = np.ones_like(train_prior) + present = context_prior > 0 + weights[present] = train_prior[present] / context_prior[present] + return weights + + +def balanced_class_weights( + train_class_counts: np.ndarray, + context_prior: np.ndarray | None, +) -> np.ndarray: + """Weights that move the predicted prior toward uniform. + + Balancing divides by the prior the model actually expresses. Without row + subsampling that is the training prior, which reproduces the historical + ``balance_probabilities`` behavior exactly. When the sampler shifted the + context prior, dividing by the context prior balances relative to what the + model saw, which is the sampler correction composed with plain balancing. + """ + train_counts = np.asarray(train_class_counts, dtype=np.float64) + prior = ( + train_counts / train_counts.sum() + if context_prior is None + else np.asarray(context_prior, dtype=np.float64) + ) + weights = np.ones_like(prior) + present = prior > 0 + weights[present] = 1.0 / prior[present] + return weights + + +def apply_class_weights( + probas: torch.Tensor, + weights: np.ndarray | torch.Tensor, +) -> torch.Tensor: + """Multiply class probabilities by per-class weights and renormalize.""" + w = torch.as_tensor(np.asarray(weights), dtype=probas.dtype, device=probas.device) + scaled = probas * w + return scaled / scaled.sum(dim=-1, keepdim=True) + + +def fit_holdout_class_weights( + holdout_probas: np.ndarray, + holdout_y_true: np.ndarray, + n_classes: int, +) -> np.ndarray: + """Weights that make the mean predicted probability per class match the + holdout class frequency. + + Solves ``mean_i normalize(p_i * w)_c = freq_c`` for ``w`` by fixed-point + iteration ``w_c <- w_c * freq_c / mean_i normalize(p_i * w)_c``. The solution + is unique up to a common scale; the result is normalized so its + prediction-weighted mean is one. Classes without holdout rows keep weight one. + """ + probas = np.asarray(holdout_probas, dtype=np.float64) + y_true = np.asarray(holdout_y_true).astype(np.int64, copy=False) + freq = np.bincount(y_true, minlength=n_classes).astype(np.float64) + freq /= freq.sum() + present = freq > 0 + + weights = np.ones(n_classes, dtype=np.float64) + for _ in range(HOLDOUT_WEIGHT_FIT_ITERATIONS): + scaled = probas * weights + scaled /= scaled.sum(axis=1, keepdims=True) + mean_pred = scaled.mean(axis=0) + update = np.ones_like(weights) + update[present] = freq[present] / np.maximum(mean_pred[present], 1e-12) + weights *= update + if np.max(np.abs(update - 1.0)) < HOLDOUT_WEIGHT_FIT_TOLERANCE: + break + return weights / weights[present].mean() + + +# --------------------------------------------------------------------------- # +# Regression +# --------------------------------------------------------------------------- # + + +def majority_value_shares( + y: np.ndarray, + row_indices: list[np.ndarray] | None, +) -> tuple[float, float, float]: + """The most frequent target value and its share in training and in context. + + Returns ``(value, train_share, context_share)``; the context share is averaged + over estimators and equals the training share without row subsampling. + """ + y = np.asarray(y, dtype=np.float64) + values, counts = np.unique(y, return_counts=True) + value = float(values[np.argmax(counts)]) + train_share = float(counts.max() / len(y)) + if row_indices is None: + return value, train_share, train_share + context_share = float(np.mean([(y[idx] == value).mean() for idx in row_indices])) + return value, train_share, context_share + + +def sampler_bucket_log_weights( + bardist: FullSupportBarDistribution, + *, + majority_value: float, + train_share: float, + context_share: float, +) -> torch.Tensor: + """Per-bucket log weights that undo a label shift on one target value. + + The bucket holding the majority value is reweighted by + ``train_share / context_share`` and every other bucket by the ratio of the + complements, so a predicted distribution that was calibrated to the context + is moved back to the training prior. Adding the result to the aggregated + log-probabilities before the bar distribution's ``log_softmax`` applies it. + """ + if not (0 < context_share < 1) or not (0 < train_share < 1): + raise ValueError( + "Sampler prediction scaling needs the majority target value to cover " + f"part of both the training data and the context, got " + f"train_share={train_share}, context_share={context_share}." + ) + borders = bardist.borders.detach() + idx = int( + bardist.map_to_bucket_idx( + torch.tensor([majority_value], dtype=borders.dtype, device=borders.device) + ).item() + ) + log_weights = torch.full( + (bardist.num_bars,), + float(np.log((1.0 - train_share) / (1.0 - context_share))), + dtype=torch.float32, + ) + log_weights[idx] = float(np.log(train_share / context_share)) + return log_weights + + +def fit_holdout_level_scaling( + holdout_pred_mean: np.ndarray, + holdout_y_true: np.ndarray, +) -> tuple[float, float]: + """Affine map ``y -> scale * y + shift`` aligning the predicted level to the + holdout target mean. + + Multiplicative when both means are positive, which is the natural correction + for nonnegative targets such as claim amounts; an additive shift otherwise. + """ + pred_mean = float(np.mean(holdout_pred_mean)) + true_mean = float(np.mean(holdout_y_true)) + if pred_mean > 0 and true_mean > 0: + return true_mean / pred_mean, 0.0 + return 1.0, true_mean - pred_mean diff --git a/src/tabpfn/preprocessing/ensemble.py b/src/tabpfn/preprocessing/ensemble.py index dd973ce37..b6f497a1b 100644 --- a/src/tabpfn/preprocessing/ensemble.py +++ b/src/tabpfn/preprocessing/ensemble.py @@ -245,6 +245,7 @@ def __init__( # noqa: PLR0913 task_type=task_type, ) + self.sample_subsampling_method_ = resolved_sample_subsampling_method self.subsample_row_indices = _get_subsample_indices_for_estimators( subsample_samples=subsample_samples, num_estimators=len(self.configs), @@ -255,6 +256,19 @@ def __init__( # noqa: PLR0913 task_type=task_type, ) + @property + def sampler_shifted_prior(self) -> bool: + """True when row subsampling changed the target prior of every context. + + Only ``majority_downsample`` does this by design; ``balanced`` and + ``stratified`` preserve the training proportions in expectation. + """ + return ( + self.subsample_row_indices is not None + and self.sample_subsampling_method_ + == SampleSubsamplingMethod.MAJORITY_DOWNSAMPLE + ) + def any_estimator_uses_gpu_svd(self) -> bool: """True if any ensemble estimator will run SVD on the GPU. diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index 061b47fa3..7cc50e6a0 100644 --- a/src/tabpfn/regressor.py +++ b/src/tabpfn/regressor.py @@ -75,6 +75,14 @@ prepend_cache_path, save_fitted_tabpfn_model, ) +from tabpfn.prediction_scaling import ( + PredictionScaling, + PredictionScalingMode, + fit_holdout_level_scaling, + majority_value_shares, + resolve_prediction_scaling, + sampler_bucket_log_weights, +) from tabpfn.preprocessing import ( EnsembleConfig, FeatureSubsamplingMethod, @@ -284,6 +292,7 @@ def __init__( # noqa: PLR0913 categorical_features_indices: Sequence[int] | None = None, softmax_temperature: float | Literal["auto"] = "auto", average_before_softmax: bool = False, + prediction_scaling: PredictionScaling = "auto", model_path: str | Path | list[str] @@ -390,6 +399,32 @@ def __init__( # noqa: PLR0913 - If `False`, the softmax function is applied to each set of logits. Then, we average the resulting probabilities of each forward pass. + prediction_scaling: + Post-processing that rescales the predicted target distribution + toward the training data, as opposed to `softmax_temperature`, which + changes its spread. + + - `"none"`: no rescaling. + - `"sampler"`: undo the label shift introduced by + `SAMPLE_SUBSAMPLING_METHOD="majority_downsample"`, which + over-represents the rare target values in every context. The + probability mass of the bucket holding the most frequent target + value is reweighted by its training share over its context share, + every other bucket by the ratio of the complements. Exact under the + label-shift assumption and free; a no-op when the sampler did not + change the prior. + - `"holdout"`: fit an affine map of the predicted distribution on + held-out rows so the mean prediction matches the held-out target + mean. Multiplicative when both means are positive, additive + otherwise. Costs one extra fit per tuning fold; shares the holdout + with `tuning_config` when that is set. Not available with + `differentiable_input=True`. + - `"auto"` (default): `"sampler"` when row subsampling shifted the + prior, otherwise `"none"`. Existing behavior is unchanged for every + configuration that does not use `"majority_downsample"`. + + `"balanced"` is a classification-only mode and is rejected here. + model_path: The path to the TabPFN model file, i.e., the pre-trained weights. @@ -585,6 +620,7 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this self.categorical_features_indices = categorical_features_indices self.softmax_temperature = softmax_temperature self.average_before_softmax = average_before_softmax + self.prediction_scaling = prediction_scaling self.model_path = model_path self.device = device self.ignore_pretraining_limits = ignore_pretraining_limits @@ -698,6 +734,17 @@ def model_(self) -> Architecture: ) return self.models_[0] + prediction_scaling_: PredictionScalingMode + """The prediction scaling mode in effect after resolving `"auto"`.""" + + prediction_scaling_log_weights_: torch.Tensor | None + """Per-bucket log weights added to the aggregated log-probabilities for the + `"sampler"` mode, or `None`.""" + + prediction_scaling_affine_: tuple[float, float] + """`(scale, shift)` applied to the raw-space borders for the `"holdout"` mode; + `(1.0, 0.0)` otherwise.""" + @property def norm_bardist_(self) -> FullSupportBarDistribution: """WARNING: DEPRECATED. Please use `raw_space_bardist_` instead. @@ -787,9 +834,12 @@ def _rebuild_raw_space_bardist(self) -> None: ``y_train_std_`` must already be set as Python floats. """ borders = self.znorm_space_bardist_.borders.detach() - self.raw_space_bardist_ = FullSupportBarDistribution( - borders * self.y_train_std_ + self.y_train_mean_, - ).float() + raw_borders = borders * self.y_train_std_ + self.y_train_mean_ + # The holdout prediction scaling is an affine map of the target axis. + scale, shift = getattr(self, "prediction_scaling_affine_", (1.0, 0.0)) + if (scale, shift) != (1.0, 0.0): + raw_borders = raw_borders * scale + shift + self.raw_space_bardist_ = FullSupportBarDistribution(raw_borders).float() def _build_ensemble_preprocessor_and_executor( self, @@ -1044,6 +1094,9 @@ def _get_tuning_regressor(self, **overwrite_kwargs: Any) -> TabPFNRegressor: # Fit on the already-expanded array, where a declared column may # have moved down past an expanded date or text column. "categorical_features_indices": self.categorical_features_indices_, + # Holdout rows are scored without rescaling: the holdout mode learns + # the whole correction from them. + "prediction_scaling": "none", } params.update(forced) @@ -1051,6 +1104,58 @@ def _get_tuning_regressor(self, **overwrite_kwargs: Any) -> TabPFNRegressor: return TabPFNRegressor(**params) + def _reset_prediction_scaling(self) -> None: + """Set the prediction scaling attributes to their no-op values.""" + self.prediction_scaling_ = PredictionScalingMode.NONE + self.prediction_scaling_log_weights_ = None + self.prediction_scaling_affine_ = (1.0, 0.0) + + def _resolve_prediction_scaling(self, *, y_raw: np.ndarray) -> None: + """Resolve `prediction_scaling` and compute the sampler bucket weights. + + Requires `ensemble_preprocessor_` and `raw_space_bardist_`. The holdout + mode's affine map is fitted earlier, in + `_maybe_calibrate_ensemble_temperature`, because it has to be in place + before the raw-space borders are built. + """ + shifted = self.ensemble_preprocessor_.sampler_shifted_prior + self.prediction_scaling_ = resolve_prediction_scaling( + self.prediction_scaling, + task_type="regressor", + sampler_shifted_prior=shifted, + ) + if self.prediction_scaling_ == PredictionScalingMode.SAMPLER and shifted: + value, train_share, context_share = majority_value_shares( + y_raw, self.ensemble_preprocessor_.subsample_row_indices + ) + self.prediction_scaling_log_weights_ = sampler_bucket_log_weights( + self.raw_space_bardist_, + majority_value=value, + train_share=train_share, + context_share=context_share, + ) + else: + self.prediction_scaling_log_weights_ = None + + def _prediction_scaling_may_apply(self) -> bool: + """Whether `prediction_scaling` can change predictions for this config. + + `"auto"` and `"sampler"` only act when the row sampler shifts the prior, + so they stay usable in batched prediction unless that sampler is + configured. + """ + mode = PredictionScalingMode(self.prediction_scaling) + if mode == PredictionScalingMode.NONE: + return False + if mode in (PredictionScalingMode.BALANCED, PredictionScalingMode.HOLDOUT): + return True + config = self.get_inference_config() + return ( + config.SUBSAMPLE_SAMPLES is not None + and SampleSubsamplingMethod(config.SAMPLE_SUBSAMPLING_METHOD) + == SampleSubsamplingMethod.MAJORITY_DOWNSAMPLE + ) + def fit_from_preprocessed( self, X_preprocessed: list[torch.Tensor], @@ -1200,6 +1305,15 @@ def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Sel self.y_train_mean_ = y_mean.detach().item() self.y_train_std_ = y_std.detach().item() y = (y_float - y_mean) / y_std + if ( + PredictionScalingMode(self.prediction_scaling) + == PredictionScalingMode.HOLDOUT + ): + raise ValueError( + "prediction_scaling='holdout' is not supported with " + "differentiable_input=True; use 'auto', 'none', or 'sampler'." + ) + self._reset_prediction_scaling() self._rebuild_raw_space_bardist() # Force sequential preprocessing: with differentiable input X carries @@ -1214,6 +1328,9 @@ def fit_with_differentiable_input(self, X: torch.Tensor, y: torch.Tensor) -> Sel n_preprocessing_jobs=1, inference_mode=False, ) + # `y` is z-normalized by now; the sampler weights locate the majority + # value on the raw-space borders, so they need the original target. + self._resolve_prediction_scaling(y_raw=y_float.detach().cpu().float().numpy()) return self @@ -1234,6 +1351,7 @@ def fit(self, X: XType, y: YType) -> Self: # that the constant-target fit, which returns before calibration, still # exposes the attribute. self.ensemble_softmax_temperature_ = 1.0 + self._reset_prediction_scaling() if self.differentiable_input: raise ValueError( @@ -1284,6 +1402,7 @@ def fit(self, X: XType, y: YType) -> Self: # tuning regressor derives its own mean/std from its own training split. self._maybe_calibrate_ensemble_temperature(X=X, y=y) + y_raw = np.asarray(y, dtype=np.float64) mean, std = np.mean(y), np.std(y) # TODO: y_train_std_ and y_train_mean_ don't seem to be used anywhere else. self.y_train_std_ = std.item() + 1e-20 @@ -1301,6 +1420,7 @@ def fit(self, X: XType, y: YType) -> Self: # TODO: Standard fit usually uses inference_mode=True, before it was enabled inference_mode=True, ) + self._resolve_prediction_scaling(y_raw=y_raw) return self @@ -1471,10 +1591,21 @@ def _maybe_calibrate_ensemble_temperature(self, X: XType, y: YType) -> None: num_samples=X.shape[0], config_cls=RegressorTuningConfig, ) + needs_holdout_scaling = ( + PredictionScalingMode(self.prediction_scaling) + == PredictionScalingMode.HOLDOUT + ) + if tuning_config_resolved is None and needs_holdout_scaling: + # The holdout scaling mode needs held-out predictions even when no + # tuning was requested; use the tuning defaults for the split. + tuning_config_resolved = RegressorTuningConfig().resolve( + num_samples=X.shape[0] + ) if tuning_config_resolved is None: return - if not tuning_config_resolved.calibrate_temperature: + calibrate = tuning_config_resolved.calibrate_temperature + if not calibrate and not needs_holdout_scaling: return holdout_folds = self._compute_holdout_validation_data( @@ -1484,12 +1615,47 @@ def _maybe_calibrate_ensemble_temperature(self, X: XType, y: YType) -> None: n_folds=int(tuning_config_resolved.tuning_n_folds), ) - # Falls back to the current no-op temperature if every fold was - # dropped, e.g. because each training split had a constant target. - self.ensemble_softmax_temperature_ = find_regression_optimal_temperature( - holdout_folds=holdout_folds, - metric_name=self.eval_metric_, - current_default_temperature=self.ensemble_softmax_temperature_, + if calibrate: + # Falls back to the current no-op temperature if every fold was + # dropped, e.g. because each training split had a constant target. + self.ensemble_softmax_temperature_ = find_regression_optimal_temperature( + holdout_folds=holdout_folds, + metric_name=self.eval_metric_, + current_default_temperature=self.ensemble_softmax_temperature_, + ) + + if needs_holdout_scaling: + self._fit_holdout_prediction_scaling(holdout_folds) + + def _fit_holdout_prediction_scaling( + self, + holdout_folds: list[ + tuple[torch.Tensor, FullSupportBarDistribution, torch.Tensor] + ], + ) -> None: + """Fit `prediction_scaling_affine_` so the mean prediction on the holdout + rows matches the holdout target mean. + + Uses the calibrated ensemble temperature, so the level is fitted on the + distribution that will actually be produced at predict time. Leaves the + no-op map in place when no fold was usable. + """ + temperature = self.ensemble_softmax_temperature_ + pred_means: list[np.ndarray] = [] + y_true: list[np.ndarray] = [] + with torch.no_grad(): + for logits, raw_space_bardist, y_holdout in holdout_folds: + if logits.shape[0] == 0: + continue + pred_means.append( + raw_space_bardist.mean(logits / temperature).float().cpu().numpy() + ) + y_true.append(y_holdout.float().cpu().numpy()) + if not pred_means: + return + self.prediction_scaling_affine_ = fit_holdout_level_scaling( + holdout_pred_mean=np.concatenate(pred_means), + holdout_y_true=np.concatenate(y_true), ) def _compute_holdout_validation_data( @@ -1681,7 +1847,11 @@ def _reduce_accumulated_logits( temperature = getattr(self, "ensemble_softmax_temperature_", 1.0) if temperature != 1.0: logits = logits / temperature - + # The sampler prediction scaling reweights buckets; the bar distribution's + # log_softmax renormalizes afterwards. + log_weights = getattr(self, "prediction_scaling_log_weights_", None) + if log_weights is not None: + logits = logits + log_weights.to(device=logits.device, dtype=logits.dtype) return logits def predict_batched( # noqa: C901, PLR0912 @@ -1753,6 +1923,13 @@ def predict_batched( # noqa: C901, PLR0912 # dataset's own holdout, so there is no single temperature to apply to a # shared batch. Mirrors the same guard in # `TabPFNClassifier.predict_proba_batched`. + if self._prediction_scaling_may_apply(): + raise NotImplementedError( + "predict_batched does not support prediction_scaling " + f"({self.prediction_scaling!r}); the correction is fitted per " + "dataset. Score datasets individually with predict, or pass " + "prediction_scaling='none'." + ) if self.tuning_config is not None: raise NotImplementedError( "predict_batched does not support tuning_config (ensemble " diff --git a/tests/test_prediction_scaling.py b/tests/test_prediction_scaling.py new file mode 100644 index 000000000..0d7231a6f --- /dev/null +++ b/tests/test_prediction_scaling.py @@ -0,0 +1,488 @@ +# Copyright (c) Prior Labs GmbH 2026. +"""Tests for `prediction_scaling`: the factor sources and the estimator wiring.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from tabpfn import TabPFNClassifier, TabPFNRegressor +from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution +from tabpfn.prediction_scaling import ( + PredictionScalingMode, + apply_class_weights, + balanced_class_weights, + context_class_prior, + fit_holdout_class_weights, + fit_holdout_level_scaling, + majority_value_shares, + resolve_prediction_scaling, + sampler_bucket_log_weights, + sampler_class_weights, +) +from tabpfn.utils import balance_probas_by_class_counts + +# --------------------------------------------------------------------------- # +# Factor sources +# --------------------------------------------------------------------------- # + + +def test__resolve_prediction_scaling__auto_follows_sampler_shift(): + assert ( + resolve_prediction_scaling( + "auto", task_type="classifier", sampler_shifted_prior=True + ) + == PredictionScalingMode.SAMPLER + ) + assert ( + resolve_prediction_scaling( + "auto", task_type="regressor", sampler_shifted_prior=False + ) + == PredictionScalingMode.NONE + ) + + +def test__resolve_prediction_scaling__balanced_rejected_for_regressor(): + with pytest.raises(ValueError, match="only defined for classification"): + resolve_prediction_scaling( + "balanced", task_type="regressor", sampler_shifted_prior=False + ) + + +def test__context_class_prior__averages_over_estimators(): + y = np.array([0] * 8 + [1] * 2) + indices = [np.array([0, 1, 8]), np.array([2, 3, 4, 9])] + prior = context_class_prior(y, indices, n_classes=2) + np.testing.assert_allclose(prior, [(2 / 3 + 3 / 4) / 2, (1 / 3 + 1 / 4) / 2]) + np.testing.assert_allclose(context_class_prior(y, None, 2), [0.8, 0.2]) + + +def test__sampler_class_weights__restore_training_prior_under_label_shift(): + """Probabilities calibrated to a shifted context are moved back to the + training prior, and the correction is exact for a Bayes-optimal model. + """ + rng = np.random.default_rng(0) + train_counts = np.array([950, 50]) + context_prior = np.array([0.5, 0.5]) + # Class-conditional likelihoods for a batch of rows. + likelihood = rng.random((1000, 2)) + p_context = likelihood * context_prior + p_context /= p_context.sum(axis=1, keepdims=True) + p_train = likelihood * (train_counts / train_counts.sum()) + p_train /= p_train.sum(axis=1, keepdims=True) + + weights = sampler_class_weights(train_counts, context_prior) + corrected = apply_class_weights(torch.tensor(p_context), weights).numpy() + np.testing.assert_allclose(corrected, p_train, atol=1e-12) + + +def test__sampler_class_weights__class_missing_from_context_keeps_weight_one(): + weights = sampler_class_weights(np.array([90, 10]), np.array([1.0, 0.0])) + np.testing.assert_allclose(weights, [0.9, 1.0]) + + +def test__balanced_class_weights__match_legacy_balancing_without_subsampling(): + counts = np.array([700, 200, 100]) + probas = torch.tensor([[0.6, 0.3, 0.1], [0.2, 0.5, 0.3]]) + legacy = balance_probas_by_class_counts(probas, counts) + weighted = apply_class_weights(probas, balanced_class_weights(counts, None)) + torch.testing.assert_close(weighted, legacy) + + +def test__balanced_class_weights__use_context_prior_when_shifted(): + weights = balanced_class_weights(np.array([900, 100]), np.array([0.5, 0.5])) + np.testing.assert_allclose(weights, [2.0, 2.0]) + + +def test__apply_class_weights__renormalizes_and_keeps_binary_ranking(): + rng = np.random.default_rng(1) + p1 = rng.random(50) + probas = torch.tensor(np.stack([1 - p1, p1], axis=1)) + out = apply_class_weights(probas, np.array([0.2, 3.0])) + torch.testing.assert_close(out.sum(dim=1), torch.ones(50, dtype=out.dtype)) + assert np.array_equal(np.argsort(out[:, 1].numpy()), np.argsort(p1)) + + +def test__fit_holdout_class_weights__matches_mean_prediction_to_frequency(): + rng = np.random.default_rng(2) + y = np.array([0] * 900 + [1] * 100) + # Over-confident in the minority: mean predicted positive rate ~0.4. + logits = rng.normal(size=(1000, 2)) + logits[:, 1] += 0.5 + probas = np.exp(logits) / np.exp(logits).sum(axis=1, keepdims=True) + weights = fit_holdout_class_weights(probas, y, n_classes=2) + scaled = probas * weights + scaled /= scaled.sum(axis=1, keepdims=True) + np.testing.assert_allclose(scaled.mean(axis=0), [0.9, 0.1], atol=1e-6) + + +def test__majority_value_shares__zero_inflated_target(): + y = np.array([0.0] * 90 + list(range(1, 11))) + indices = [np.array(list(range(10)) + list(range(90, 100)))] + value, train_share, context_share = majority_value_shares(y, indices) + assert value == 0.0 + assert train_share == 0.9 + assert context_share == 0.5 + assert majority_value_shares(y, None)[2] == 0.9 + + +def test__sampler_bucket_log_weights__reweights_spike_bucket_only(): + bardist = FullSupportBarDistribution(torch.linspace(-1.0, 9.0, 11)) + log_w = sampler_bucket_log_weights( + bardist, majority_value=0.0, train_share=0.9, context_share=0.5 + ) + spike = int(bardist.map_to_bucket_idx(torch.tensor([0.0])).item()) + assert log_w.shape == (10,) + np.testing.assert_allclose(log_w[spike].item(), np.log(0.9 / 0.5)) + others = torch.cat([log_w[:spike], log_w[spike + 1 :]]) + np.testing.assert_allclose(others.numpy(), np.log(0.1 / 0.5)) + + # A uniform prediction moves most of its mass onto the spike bucket, so + # the predicted mean drops toward the majority value. + uniform = torch.zeros(1, 10) + assert bardist.mean(uniform + log_w).item() < bardist.mean(uniform).item() + + +def test__sampler_bucket_log_weights__rejects_degenerate_shares(): + bardist = FullSupportBarDistribution(torch.linspace(-1.0, 9.0, 11)) + with pytest.raises(ValueError, match="cover part of both"): + sampler_bucket_log_weights( + bardist, majority_value=0.0, train_share=0.9, context_share=1.0 + ) + + +def test__fit_holdout_level_scaling__multiplicative_then_additive(): + assert fit_holdout_level_scaling(np.array([2.0, 4.0]), np.array([1.0, 2.0])) == ( + 0.5, + 0.0, + ) + scale, shift = fit_holdout_level_scaling( + np.array([-1.0, 1.0]), np.array([1.0, 3.0]) + ) + assert (scale, shift) == (1.0, 2.0) + + +# --------------------------------------------------------------------------- # +# Classifier wiring +# --------------------------------------------------------------------------- # + + +def _imbalanced_classification( + seed: int = 0, n_majority: int = 270, n_minority: int = 30 +) -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + X = rng.normal(size=(n_majority + n_minority, 4)) + y = np.array([0] * n_majority + [1] * n_minority) + X[y == 1] += 1.0 + return X, y + + +def _downsampling_classifier(**kwargs) -> TabPFNClassifier: + return TabPFNClassifier( + n_estimators=2, + device="cpu", + random_state=0, + inference_config={ + "SUBSAMPLE_SAMPLES": 90, + "SAMPLE_SUBSAMPLING_METHOD": "majority_downsample", + }, + **kwargs, + ) + + +def test__classifier__auto_resolves_to_sampler_under_majority_downsample(): + X, y = _imbalanced_classification() + clf = _downsampling_classifier().fit(X, y) + assert clf.prediction_scaling_ == PredictionScalingMode.SAMPLER + assert clf.prediction_scaling_weights_ is not None + # Minority rows are all in context (30 of 90), so its weight is its + # training share over its context share: 0.1 / (1/3). + np.testing.assert_allclose(clf.prediction_scaling_weights_[1], 0.1 / (30 / 90)) + np.testing.assert_allclose(clf.prediction_scaling_weights_[0], 0.9 / (60 / 90)) + + +def test__classifier__sampler_scaling_restores_prior_and_keeps_ranking(): + X, y = _imbalanced_classification() + scaled = _downsampling_classifier().fit(X, y).predict_proba(X) + raw = _downsampling_classifier(prediction_scaling="none").fit(X, y).predict_proba(X) + # Same ranking of the positive class, lower mean positive probability. + assert np.array_equal(np.argsort(scaled[:, 1]), np.argsort(raw[:, 1])) + assert scaled[:, 1].mean() < raw[:, 1].mean() + assert abs(scaled[:, 1].mean() - 0.1) < abs(raw[:, 1].mean() - 0.1) + + +def test__classifier__auto_is_none_without_prior_shift(): + X, y = _imbalanced_classification() + clf = TabPFNClassifier(n_estimators=2, device="cpu", random_state=0).fit(X, y) + assert clf.prediction_scaling_ == PredictionScalingMode.NONE + assert clf.prediction_scaling_weights_ is None + + +def test__classifier__scaling_is_a_fixed_per_row_function(): + """The weights are fitted once and applied row by row, so the scaled output + equals the unscaled output with the fitted weights applied. This is what makes + the correction independent of the batch being predicted; the forward pass + itself is not bit-identical across batch sizes on every platform, so the two + are compared within one batch. + """ + X, y = _imbalanced_classification() + scaled = _downsampling_classifier().fit(X, y) + raw = _downsampling_classifier(prediction_scaling="none").fit(X, y) + expected = apply_class_weights( + torch.tensor(raw.predict_proba(X[:20])), scaled.prediction_scaling_weights_ + ).numpy() + np.testing.assert_allclose(scaled.predict_proba(X[:20]), expected, atol=1e-6) + + +def test__classifier__balance_probabilities_is_deprecated_alias(): + X, y = _imbalanced_classification() + with pytest.warns(DeprecationWarning, match="prediction_scaling='balanced'"): + legacy = TabPFNClassifier( + n_estimators=2, device="cpu", random_state=0, balance_probabilities=True + ).fit(X, y) + new = TabPFNClassifier( + n_estimators=2, device="cpu", random_state=0, prediction_scaling="balanced" + ).fit(X, y) + assert legacy.prediction_scaling_ == PredictionScalingMode.BALANCED + np.testing.assert_allclose(legacy.predict_proba(X), new.predict_proba(X), atol=1e-6) + + +def test__classifier__balance_probabilities_conflicts_with_other_mode(): + X, y = _imbalanced_classification() + clf = TabPFNClassifier( + n_estimators=2, + device="cpu", + random_state=0, + balance_probabilities=True, + prediction_scaling="sampler", + ) + with pytest.raises(ValueError, match="conflicts with"): + clf.fit(X, y) + + +def test__classifier__balanced_matches_legacy_output_without_subsampling(): + X, y = _imbalanced_classification() + clf = TabPFNClassifier( + n_estimators=2, device="cpu", random_state=0, prediction_scaling="balanced" + ).fit(X, y) + raw = TabPFNClassifier(n_estimators=2, device="cpu", random_state=0).fit(X, y) + expected = balance_probas_by_class_counts( + torch.tensor(raw.predict_proba(X)), clf.class_counts_ + ).numpy() + np.testing.assert_allclose(clf.predict_proba(X), expected, atol=1e-6) + + +def test__classifier__holdout_mode_fits_weights_without_tuning_config(): + X, y = _imbalanced_classification(n_majority=450, n_minority=50) + clf = _downsampling_classifier(prediction_scaling="holdout").fit(X, y) + assert clf.prediction_scaling_ == PredictionScalingMode.HOLDOUT + assert clf.prediction_scaling_weights_ is not None + assert clf.prediction_scaling_weights_.shape == (2,) + # The holdout weights push the mean positive probability toward the base + # rate; without them the downsampled context predicts far too many positives. + raw = _downsampling_classifier(prediction_scaling="none").fit(X, y) + assert abs(clf.predict_proba(X)[:, 1].mean() - 0.1) < abs( + raw.predict_proba(X)[:, 1].mean() - 0.1 + ) + + +def test__classifier__holdout_rejected_with_differentiable_input(): + X, y = _imbalanced_classification() + clf = TabPFNClassifier( + n_estimators=2, + device="cpu", + random_state=0, + differentiable_input=True, + prediction_scaling="holdout", + ) + with pytest.raises(ValueError, match="not supported with"): + clf.fit_with_differentiable_input( + torch.tensor(X, dtype=torch.float32), torch.tensor(y) + ) + + +def test__classifier__differentiable_input_applies_sampler_scaling(): + X, y = _imbalanced_classification() + clf = _downsampling_classifier(differentiable_input=True) + clf.fit_with_differentiable_input( + torch.tensor(X, dtype=torch.float32), torch.tensor(y) + ) + assert clf.prediction_scaling_ == PredictionScalingMode.SAMPLER + assert clf.prediction_scaling_weights_ is not None + + +@pytest.mark.parametrize("mode", ["balanced", "holdout"]) +def test__classifier__predict_proba_batched_rejects_active_scaling(mode: str): + X, y = _imbalanced_classification() + clf = TabPFNClassifier( + n_estimators=2, device="cpu", random_state=0, prediction_scaling=mode + ) + with pytest.raises(NotImplementedError, match="prediction_scaling"): + clf.predict_proba_batched([X], [y], [X[:5]]) + + +def test__classifier__predict_proba_batched_allows_inactive_auto(): + X, y = _imbalanced_classification() + clf = TabPFNClassifier(n_estimators=2, device="cpu", random_state=0) + out = clf.predict_proba_batched([X], [y], [X[:5]]) + assert len(out) == 1 + assert out[0].shape[0] == 5 + + +# --------------------------------------------------------------------------- # +# Regressor wiring +# --------------------------------------------------------------------------- # + + +def _zero_inflated_regression( + seed: int = 0, n_zeros: int = 270, n_nonzero: int = 30 +) -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + X = rng.normal(size=(n_zeros + n_nonzero, 4)) + y = np.concatenate([np.zeros(n_zeros), rng.exponential(size=n_nonzero) + 0.5]) + X[y > 0] += 1.0 + return X, y + + +def _downsampling_regressor(**kwargs) -> TabPFNRegressor: + return TabPFNRegressor( + n_estimators=2, + device="cpu", + random_state=0, + inference_config={ + "SUBSAMPLE_SAMPLES": 90, + "SAMPLE_SUBSAMPLING_METHOD": "majority_downsample", + }, + **kwargs, + ) + + +def test__regressor__auto_resolves_to_sampler_under_majority_downsample(): + X, y = _zero_inflated_regression() + reg = _downsampling_regressor().fit(X, y) + assert reg.prediction_scaling_ == PredictionScalingMode.SAMPLER + assert reg.prediction_scaling_log_weights_ is not None + assert reg.prediction_scaling_log_weights_.shape == ( + reg.raw_space_bardist_.num_bars, + ) + assert reg.prediction_scaling_affine_ == (1.0, 0.0) + + +def test__regressor__sampler_scaling_lowers_predicted_level(): + X, y = _zero_inflated_regression() + scaled = _downsampling_regressor().fit(X, y).predict(X) + raw = _downsampling_regressor(prediction_scaling="none").fit(X, y).predict(X) + assert scaled.mean() < raw.mean() + assert abs(scaled.mean() - y.mean()) < abs(raw.mean() - y.mean()) + + +def test__regressor__sampler_scaling_applies_to_every_output_type(): + X, y = _zero_inflated_regression() + scaled = _downsampling_regressor().fit(X, y) + raw = _downsampling_regressor(prediction_scaling="none").fit(X, y) + s_full = scaled.predict(X[:10], output_type="full") + r_full = raw.predict(X[:10], output_type="full") + assert not np.allclose(s_full["median"], r_full["median"]) + assert not np.allclose(s_full["quantiles"][0], r_full["quantiles"][0]) + + +def test__regressor__auto_is_none_without_prior_shift(): + X, y = _zero_inflated_regression() + reg = TabPFNRegressor(n_estimators=2, device="cpu", random_state=0).fit(X, y) + assert reg.prediction_scaling_ == PredictionScalingMode.NONE + assert reg.prediction_scaling_log_weights_ is None + + +def test__regressor__scaling_is_a_fixed_per_row_function(): + """The bucket log weights are fitted once and added row by row: within one + batch the scaled aggregated logits are the unscaled ones plus the weights. + See the classifier counterpart for why this is compared within a batch. + """ + X, y = _zero_inflated_regression() + scaled = _downsampling_regressor().fit(X, y) + raw = _downsampling_regressor(prediction_scaling="none").fit(X, y) + scaled_logits = scaled.predict(X[:20], output_type="full")["logits"] + raw_logits = raw.predict(X[:20], output_type="full")["logits"] + # Buckets with no mass carry -inf in both tensors; compare the finite ones. + finite = torch.isfinite(raw_logits) & torch.isfinite(scaled_logits) + assert finite.float().mean() > 0.9 + delta = (scaled_logits - raw_logits).cpu() + expected = scaled.prediction_scaling_log_weights_.to(delta.dtype).expand_as(delta) + torch.testing.assert_close( + delta[finite.cpu()], expected[finite.cpu()], atol=1e-5, rtol=1e-5 + ) + + +def test__regressor__balanced_rejected(): + X, y = _zero_inflated_regression() + with pytest.raises(ValueError, match="only defined for classification"): + TabPFNRegressor( + n_estimators=2, device="cpu", random_state=0, prediction_scaling="balanced" + ).fit(X, y) + + +def test__regressor__holdout_mode_fits_affine_map_and_rescales_borders(): + X, y = _zero_inflated_regression(n_zeros=450, n_nonzero=50) + reg = _downsampling_regressor(prediction_scaling="holdout").fit(X, y) + assert reg.prediction_scaling_ == PredictionScalingMode.HOLDOUT + scale, shift = reg.prediction_scaling_affine_ + assert (scale, shift) != (1.0, 0.0) + assert reg.prediction_scaling_log_weights_ is None + # The raw-space borders carry the affine map, so every output type follows. + expected = ( + reg.znorm_space_bardist_.borders * reg.y_train_std_ + reg.y_train_mean_ + ) * scale + shift + torch.testing.assert_close( + reg.raw_space_bardist_.borders, expected.float(), atol=1e-4, rtol=1e-5 + ) + raw = _downsampling_regressor(prediction_scaling="none").fit(X, y) + assert abs(reg.predict(X).mean() - y.mean()) < abs(raw.predict(X).mean() - y.mean()) + + +def test__regressor__differentiable_input_applies_sampler_scaling_in_raw_units(): + """The differentiable path z-normalizes `y` before fitting; the sampler weights + must still locate the spike on the raw-space borders and match `fit()`. + """ + X, y = _zero_inflated_regression() + reference = _downsampling_regressor().fit(X, y) + reg = _downsampling_regressor(differentiable_input=True) + reg.fit_with_differentiable_input( + torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32) + ) + assert reg.prediction_scaling_ == PredictionScalingMode.SAMPLER + assert reg.prediction_scaling_log_weights_ is not None + spike = int(reg.raw_space_bardist_.map_to_bucket_idx(torch.tensor([0.0])).item()) + log_w = reg.prediction_scaling_log_weights_ + assert log_w.argmax().item() == spike + torch.testing.assert_close(log_w, reference.prediction_scaling_log_weights_) + + +def test__regressor__holdout_rejected_with_differentiable_input(): + X, y = _zero_inflated_regression() + reg = TabPFNRegressor( + n_estimators=2, + device="cpu", + random_state=0, + differentiable_input=True, + prediction_scaling="holdout", + ) + with pytest.raises(ValueError, match="not supported with"): + reg.fit_with_differentiable_input( + torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32) + ) + + +def test__regressor__predict_batched_rejects_active_scaling(): + X, y = _zero_inflated_regression() + reg = _downsampling_regressor() + with pytest.raises(NotImplementedError, match="prediction_scaling"): + reg.predict_batched([X], [y], [X[:5]]) + + +def test__regressor__predict_batched_allows_inactive_auto(): + X, y = _zero_inflated_regression() + reg = TabPFNRegressor(n_estimators=2, device="cpu", random_state=0) + out = reg.predict_batched([X], [y], [X[:5]]) + assert len(out) == 1 + assert len(out[0]) == 5