Apply regression target transforms to the unnormalized target [RES-2639]

changelog/1196.changed.md

@@ -0,0 +1 @@
+Regression target transforms (`REGRESSION_Y_PREPROCESS_TRANSFORMS`) now act on the target in its original units and are standardized afterwards, instead of being applied to the already z-normalized target. A transform such as `1_plus_log` therefore does what its name suggests; previously it operated on standardized values, where it is undefined for roughly half of the target. This changes the predictions of the v2, v2.5 and v3 defaults, whose ensembles use `safepower` for half of their estimators; the v2.6 default (`"none"`) and any estimator without a target transform are unaffected.

src/tabpfn/finetuning/data_util.py

@@ -18,6 +18,9 @@
 from tabpfn.architectures.shared.bar_distribution import FullSupportBarDistribution
 from tabpfn.preprocessing.datamodel import FeatureModality, FeatureSchema
 from tabpfn.preprocessing.ensemble import TabPFNEnsemblePreprocessor
+from tabpfn.preprocessing.target_transform import (
+    rebind_target_transform_statistics,
+)
 from tabpfn.utils import infer_random_state, pad_tensors
 
 if TYPE_CHECKING:
@@ -429,6 +432,14 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
                 + train_mean  # Inverse normalization back to raw space
             ).float()
             y_train = y_train_standardized
+            # The configs were built for the statistics of the whole dataset,
+            # but every split standardizes with its own; the target transforms
+            # undo that standardization internally and so have to follow.
+            rebind_target_transform_statistics(
+                [c.target_transform for c in conf],
+                mean=float(train_mean),
+                std=float(train_std),
+            )
         else:
             y_train = y_train_raw
 

src/tabpfn/inference_config.py

@@ -184,6 +184,9 @@ class InferenceConfig:
     it. The preprocessors should be passed as a tuple/list and are then (repeatedly)
     used by the estimators in the ensembles.
 
+    The preprocessor is applied to the target variable in its original units, and
+    its output is standardized afterwards, since that is the scale TabPFN expects.
+
     By default, we use no preprocessing and a power transformation (if we have
     more than one estimator).
 

src/tabpfn/preprocessing/configs.py

@@ -176,7 +176,16 @@ class ClassifierEnsembleConfig(EnsembleConfig):
 
 @dataclass
 class RegressorEnsembleConfig(EnsembleConfig):
-    """Configuration for a regression ensemble member."""
+    """Configuration for a regression ensemble member.
+
+    Attributes:
+        target_transform: Transform mapping the z-normalised target to the
+            target this member is fitted on, and, through
+            ``inverse_transform``, the model's bar-distribution borders back to
+            the z-normalised space. `None` leaves the target z-normalised. See
+            `tabpfn.preprocessing.target_transform` for how the transforms are
+            composed so that they act on the target in its original units.
+    """
 
     target_transform: TransformerMixin | Pipeline | None
 

src/tabpfn/preprocessing/target_transform.py

@@ -0,0 +1,147 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+
+"""Target (y) transformation pipelines for regression.
+
+The regressor z-normalises the target before handing it to the ensemble
+preprocessing, because the shared bar-distribution space of the checkpoint is
+defined in that z-normalised space. A target transform such as ``1_plus_log``
+is however only meaningful on the target in its *original* units: applied to
+z-normalised values it operates on a shifted, rescaled target, which is not
+what one would expect (and, for the log-like transforms, mostly produces NaNs
+because roughly half of a z-normalised target is negative).
+
+:func:`wrap_target_transform` therefore composes each transform into a
+three-step pipeline that
+
+1. undoes the regressor's z-normalisation, so the transform sees the target in
+   its original units,
+2. applies the transform, and
+3. standardises the result again, which is the scale the model expects.
+
+Keeping the pipeline's input and output in the z-normalised space means the
+rest of the regressor -- in particular the bar-distribution borders and their
+sanity limits, which are expressed in z-units -- is unaffected:
+``pipeline.inverse_transform`` maps the model's borders straight back into the
+z-normalised space, exactly as an unwrapped transform did.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+
+from tabpfn.preprocessing.steps.utils import make_scaler_safe
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+UNSTANDARDIZE_STEP = "unstandardize_target"
+TARGET_TRANSFORM_STEP = "target_transform"
+STANDARDIZE_STEP = "standardize_target"
+
+
+class UnstandardizeTarget(TransformerMixin, BaseEstimator):
+    """Map a z-normalised target back to its original units.
+
+    ``transform`` undoes a z-normalisation with the given statistics and
+    ``inverse_transform`` re-applies it, so this transformer is the first step
+    of the pipelines built by :func:`wrap_target_transform`.
+
+    Args:
+        mean: Mean that was subtracted by the z-normalisation.
+        std: Standard deviation the z-normalisation divided by.
+    """
+
+    def __init__(self, mean: float = 0.0, std: float = 1.0) -> None:
+        self.mean = mean
+        self.std = std
+
+    def fit(self, X: np.ndarray, y: np.ndarray | None = None) -> UnstandardizeTarget:
+        """Stateless, the statistics are given at construction time."""
+        del X, y
+        return self
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        """Return ``X`` in the original units of the target."""
+        return np.asarray(X) * self.std + self.mean
+
+    def inverse_transform(self, X: np.ndarray) -> np.ndarray:
+        """Return ``X`` z-normalised again."""
+        return (np.asarray(X) - self.mean) / self.std
+
+
+def wrap_target_transform(
+    transform: TransformerMixin | Pipeline,
+    *,
+    mean: float,
+    std: float,
+) -> Pipeline:
+    """Compose a target transform so that it acts on the unnormalised target.
+
+    Args:
+        transform: The target transform to wrap, e.g. one of the presets of
+            :func:`get_all_reshape_feature_distribution_preprocessors`.
+        mean: Mean of the training target, used to undo its z-normalisation.
+        std: Standard deviation of the training target, used to undo its
+            z-normalisation.
+
+    Returns:
+        A pipeline mapping the z-normalised target to the transformed and
+        re-standardised target, whose ``inverse_transform`` maps back into the
+        z-normalised space.
+    """
+    return Pipeline(
+        steps=[
+            (UNSTANDARDIZE_STEP, UnstandardizeTarget(mean=mean, std=std)),
+            (TARGET_TRANSFORM_STEP, transform),
+            # The transform may leave the target on an arbitrary scale (e.g. a
+            # log target), while the model expects a standardised one. The safe
+            # wrapper also keeps non-finite outputs, such as the log of a
+            # non-positive target, from reaching the model.
+            (STANDARDIZE_STEP, make_scaler_safe("standard", StandardScaler())),
+        ],
+    )
+
+
+def rebind_target_transform_statistics(
+    transforms: Iterable[TransformerMixin | Pipeline | None],
+    *,
+    mean: float,
+    std: float,
+) -> None:
+    """Point wrapped target transforms at another z-normalisation, in place.
+
+    Needed when the transforms were built for one z-normalisation but are
+    (re-)fitted on a target normalised with different statistics, as in the
+    fine-tuning data pipeline, which re-splits the dataset and z-normalises
+    with the statistics of every new training split.
+
+    Transforms that are not pipelines from :func:`wrap_target_transform`, such
+    as the ``None`` entries of an unwrapped target transform, are ignored.
+
+    Args:
+        transforms: The target transforms to update.
+        mean: Mean of the z-normalisation the transforms will be fitted on.
+        std: Standard deviation of that z-normalisation.
+    """
+    for transform in transforms:
+        if not isinstance(transform, Pipeline):
+            continue
+        step = transform.named_steps.get(UNSTANDARDIZE_STEP)
+        if isinstance(step, UnstandardizeTarget):
+            step.mean = mean
+            step.std = std
+
+
+__all__ = [
+    "STANDARDIZE_STEP",
+    "TARGET_TRANSFORM_STEP",
+    "UNSTANDARDIZE_STEP",
+    "UnstandardizeTarget",
+    "rebind_target_transform_statistics",
+    "wrap_target_transform",
+]

src/tabpfn/preprocessing/transform.py

@@ -95,7 +95,11 @@ def _transform_labels_one(
 
     Args:
         config: Ensemble config.
-        y_train: The unprocessed labels.
+        y_train: The labels as the estimator handed them over: encoded class
+            indices for classification, and the z-normalised target for
+            regression. Note that the regression target transforms undo that
+            z-normalisation internally, so that they act on the target in its
+            original units (see `tabpfn.preprocessing.target_transform`).
 
     Return: The processed labels.
     """

src/tabpfn/regressor.py

@@ -89,6 +89,7 @@
 from tabpfn.preprocessing.steps import (
     get_all_reshape_feature_distribution_preprocessors,
 )
+from tabpfn.preprocessing.target_transform import wrap_target_transform
 from tabpfn.utils import (
     DevicesSpecification,
     convert_batch_of_cat_ix_to_schema,
@@ -854,6 +855,12 @@ def _initialize_dataset_preprocessing(
         Handle the preprocessing of the input (X and y). We also return the
         BarDistribution here, since it is vital for computing the standardized
         target variable in the DatasetCollectionWithPreprocessing class.
+
+        Also sets ``y_train_mean_``/``y_train_std_``, the statistics of the
+        z-normalisation every caller applies to the target before the ensemble
+        preprocessing. They are needed here already because the target
+        transforms are composed with that z-normalisation, see
+        :func:`wrap_target_transform`.
         """
         X, y, feature_names, n_features, _ = ensure_compatible_fit_inputs(
             X,
@@ -887,6 +894,8 @@ def _initialize_dataset_preprocessing(
         self.inferred_feature_schema_ = feature_schema
         self.ordinal_encoder_ = ordinal_encoder
 
+        self.y_train_mean_, self.y_train_std_ = _target_znorm_statistics(y)
+
         # TODO: Introduce regressor target transformer that also keeps track of
         # target name
         possible_target_transforms = get_all_reshape_feature_distribution_preprocessors(
@@ -897,11 +906,24 @@ def _initialize_dataset_preprocessing(
         for (
             y_target_preprocessor
         ) in self.inference_config_.REGRESSION_Y_PREPROCESS_TRANSFORMS:
-            if y_target_preprocessor is not None:
-                preprocessor = possible_target_transforms[y_target_preprocessor]
-            else:
-                preprocessor = None
-            target_preprocessors.append(preprocessor)
+            # Nothing to transform, so the target stays z-normalised. `"none"`
+            # is the identity transform and therefore equivalent to no
+            # transform at all; composing it would only add floating-point
+            # noise to the target.
+            if y_target_preprocessor in (None, "none"):
+                target_preprocessors.append(None)
+                continue
+            # The target reaching the transform is z-normalised, so the
+            # transform is composed with the inverse of that normalisation and
+            # a re-standardisation, to make it act on the target in its
+            # original units.
+            target_preprocessors.append(
+                wrap_target_transform(
+                    possible_target_transforms[y_target_preprocessor],
+                    mean=self.y_train_mean_,
+                    std=self.y_train_std_,
+                )
+            )
 
         preprocessor_configs = self.inference_config_.PREPROCESS_TRANSFORMS
         self.n_estimators_ = scale_n_estimators_for_feature_coverage(
@@ -1185,10 +1207,9 @@ 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)
 
-        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
-        self.y_train_mean_ = mean.item()
+        # `y_train_mean_`/`y_train_std_` were set by
+        # `_initialize_dataset_preprocessing`, which needs them to compose the
+        # target transforms with this z-normalisation.
         y = (y - self.y_train_mean_) / self.y_train_std_
         self._rebuild_raw_space_bardist()
 
@@ -2141,6 +2162,16 @@ def to(self, device: DevicesSpecification) -> None:
             self.raw_space_bardist_.to(self.devices_[0])
 
 
+def _target_znorm_statistics(y: np.ndarray) -> tuple[float, float]:
+    """Return the mean and standard deviation used to z-normalise the target.
+
+    The standard deviation carries a small epsilon so that a (near-)constant
+    target cannot cause a division by zero. A truly constant target is caught
+    by `fit` and never reaches the model.
+    """
+    return float(np.mean(y)), float(np.std(y)) + 1e-20
+
+
 def _logits_to_output(
     *,
     output_type: str,