Give each regression member its own target pipeline [RES-2639]

changelog/1198.changed.md

@@ -0,0 +1 @@
+Internal refactor of the regression target handling, in preparation for changing when the target transform is applied. Each ensemble member now owns one invertible pipeline from the target, in its original units, to the target the model is fitted on, instead of the estimator z-normalizing the target up front. Model inputs are unchanged; predictions can differ in the last few float32 digits, since a transformed member's bar-distribution borders are now mapped back through its own pipeline. Fixes a latent inconsistency for row-subsampled estimators, whose target was normalized with the whole training set's statistics while their borders were mapped back as if it had been their own.

src/tabpfn/finetuning/data_util.py

@@ -66,6 +66,10 @@ class RegressorBatch:
         znorm_space_bardist: Bar distribution in z-normalized target space.
         X_query_raw: Original unprocessed test features.
         y_query_raw: Original unprocessed test targets.
+        y_train_mean: Mean of this split's training target, defining the frame
+            the two bar distributions relate by. The regressor needs it to map
+            an estimator's borders back out of its target pipeline.
+        y_train_std: Standard deviation of this split's training target.
     """
 
     X_context: list[torch.Tensor]
@@ -79,6 +83,8 @@ class RegressorBatch:
     znorm_space_bardist: FullSupportBarDistribution
     X_query_raw: torch.Tensor
     y_query_raw: torch.Tensor
+    y_train_mean: float = 0.0
+    y_train_std: float = 1.0
 
 
 @dataclass
@@ -423,12 +429,13 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
                 train_std = eps
 
             y_test_standardized = (y_test_raw - train_mean) / train_std
-            y_train_standardized = (y_train_raw - train_mean) / train_std
             raw_space_bardist_ = FullSupportBarDistribution(
                 znorm_space_bardist_.borders * train_std
                 + train_mean  # Inverse normalization back to raw space
             ).float()
-            y_train = y_train_standardized
+            # The members standardize the target themselves, each with the
+            # statistics of the split it is fitted on.
+            y_train = y_train_raw
         else:
             y_train = y_train_raw
 
@@ -498,6 +505,8 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
                 znorm_space_bardist=znorm_space_bardist_,
                 X_query_raw=x_test_raw,
                 y_query_raw=y_test_raw,
+                y_train_mean=float(train_mean),
+                y_train_std=float(train_std),
             )
 
         return ClassifierBatch(
@@ -664,6 +673,10 @@ def meta_dataset_collator(
         znorm_space_bardist=first_item.znorm_space_bardist,
         X_query_raw=_collate_tensor_field(batch, "X_query_raw", padding_val),
         y_query_raw=_collate_tensor_field(batch, "y_query_raw", padding_val),
+        # Taken from the first item, like the bar distributions above and for
+        # the same reason: they belong together as one frame.
+        y_train_mean=first_item.y_train_mean,
+        y_train_std=first_item.y_train_std,
     )
 
 

src/tabpfn/finetuning/finetuned_regressor.py

@@ -357,10 +357,15 @@ def _should_skip_batch(self, batch: RegressorBatch) -> bool:  # type: ignore[ove
 
     @override
     def _setup_batch(self, batch: RegressorBatch) -> None:  # type: ignore[override]
-        """Set up bar distribution for this batch."""
+        """Set up bar distribution and target frame for this batch."""
         self.finetuned_estimator_.raw_space_bardist_ = batch.raw_space_bardist
         self.finetuned_estimator_.bardist_ = batch.znorm_space_bardist
         self._bardist_loss = batch.znorm_space_bardist
+        # `fit_from_preprocessed` never sees the target, so the frame its
+        # members' pipelines were fitted in has to come from the batch. The
+        # forward pass needs it to map an estimator's borders back.
+        self.finetuned_estimator_.y_train_mean_ = batch.y_train_mean
+        self.finetuned_estimator_.y_train_std_ = batch.y_train_std
 
     @override
     def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor:  # type: ignore[override]

src/tabpfn/preprocessing/target_transform.py

@@ -0,0 +1,106 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+
+"""Target (y) transformation pipelines for regression.
+
+Every regression ensemble member owns one invertible map from the target, in
+its original units, to the target the model is fitted on:
+
+* ``transform`` standardizes the target -- optionally after reshaping it with
+  one of the `REGRESSION_Y_PREPROCESS_TRANSFORMS` presets -- because the
+  checkpoint's bar distribution is defined for a standardized target.
+* ``inverse_transform`` maps the model's bar-distribution borders back into
+  the original units of the target.
+
+The estimator owns the affine frame the ensemble is aggregated in --
+``y_train_mean_`` and ``y_train_std_``, the same frame ``raw_space_bardist_``
+decodes from -- and applies it to the borders after the inverse. That split
+keeps these pipelines self-contained: the statistics of a member's own target
+are learned in ``fit``, so nothing has to be baked into a pipeline when it is
+built, and nothing has to be rebound when it is refitted on a different split.
+
+Aggregating in the frame of the checkpoint rather than in the original units
+of the target is deliberate: the bar distribution has thousands of borders,
+and `translate_probs_across_borders` resolves positions within a bucket in
+float32. For a target with a large offset (say 1e8 with a standard deviation
+of 1e2) the borders in original units are spaced far below float32 resolution
+at that magnitude, and would collapse into duplicates.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.pipeline import Pipeline
+
+if TYPE_CHECKING:
+    from sklearn.base import TransformerMixin as Transformer
+
+STANDARDIZE_STEP = "standardize_target"
+TARGET_TRANSFORM_STEP = "target_transform"
+
+
+class StandardizeTarget(TransformerMixin, BaseEstimator):
+    """Z-normalise the target, learning the statistics in ``fit``.
+
+    `sklearn.preprocessing.StandardScaler` would do, but this keeps the
+    arithmetic of the estimator's own z-normalisation -- `np.mean` and `np.std`
+    plus an epsilon -- so that a member without a target transform is fitted on
+    exactly the target the regressor used to compute before preprocessing.
+
+    Attributes:
+        mean_: Mean of the target seen in `fit`.
+        std_: Standard deviation of that target, plus `EPSILON`.
+    """
+
+    EPSILON = 1e-20
+    """Guards against a division by zero for a (near-)constant target, which
+    `TabPFNRegressor.fit` rejects before it reaches the model anyway."""
+
+    def fit(self, X: np.ndarray, y: np.ndarray | None = None) -> StandardizeTarget:
+        """Learn the mean and standard deviation of the target ``X``."""
+        del y
+        self.mean_ = float(np.mean(X))
+        self.std_ = float(np.std(X)) + self.EPSILON
+        return self
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        """Return the standardized target."""
+        return (np.asarray(X) - self.mean_) / self.std_
+
+    def inverse_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 make_target_transform(transform: Transformer | Pipeline | None) -> Pipeline:
+    """Build the target pipeline of one ensemble member.
+
+    Args:
+        transform: The preset to reshape the target with, e.g. one of
+            :func:`get_all_reshape_feature_distribution_preprocessors`, or
+            None to only standardize the target.
+
+    Returns:
+        A pipeline mapping the target in its original units to the target the
+        model is fitted on, whose ``inverse_transform`` maps back.
+    """
+    if transform is None:
+        return Pipeline(steps=[(STANDARDIZE_STEP, StandardizeTarget())])
+    return Pipeline(
+        steps=[
+            # The preset reshapes the standardized target, as it always has;
+            # the ordering of these two steps is what RES-2639 changes.
+            (STANDARDIZE_STEP, StandardizeTarget()),
+            (TARGET_TRANSFORM_STEP, transform),
+        ],
+    )
+
+
+__all__ = [
+    "STANDARDIZE_STEP",
+    "TARGET_TRANSFORM_STEP",
+    "StandardizeTarget",
+    "make_target_transform",
+]

src/tabpfn/regressor.py

@@ -89,6 +89,10 @@
 from tabpfn.preprocessing.steps import (
     get_all_reshape_feature_distribution_preprocessors,
 )
+from tabpfn.preprocessing.target_transform import (
+    StandardizeTarget,
+    make_target_transform,
+)
 from tabpfn.utils import (
     DevicesSpecification,
     convert_batch_of_cat_ix_to_schema,
@@ -893,15 +897,18 @@ def _initialize_dataset_preprocessing(
             num_examples=y.shape[0],  # Use length of validated y
             random_state=random_state,  # Use the provided rng
         )
-        target_preprocessors: list[TransformerMixin | Pipeline | None] = []
-        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)
+        # Every member standardizes the target itself, so the members are handed
+        # the target in its original units and no statistic of it is needed here.
+        target_preprocessors: list[TransformerMixin | Pipeline | None] = [
+            make_target_transform(
+                None
+                if y_target_preprocessor is None
+                else possible_target_transforms[y_target_preprocessor]
+            )
+            for y_target_preprocessor in (
+                self.inference_config_.REGRESSION_Y_PREPROCESS_TRANSFORMS
+            )
+        ]
 
         preprocessor_configs = self.inference_config_.PREPROCESS_TRANSFORMS
         self.n_estimators_ = scale_n_estimators_for_feature_coverage(
@@ -1185,11 +1192,11 @@ 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 = (y - self.y_train_mean_) / self.y_train_std_
+        # The frame the ensemble is aggregated in and `raw_space_bardist_`
+        # decodes from. The target itself is handed to the preprocessing in its
+        # original units; every member standardizes it with its own statistics.
+        self.y_train_mean_ = float(np.mean(y))
+        self.y_train_std_ = float(np.std(y)) + StandardizeTarget.EPSILON
         self._rebuild_raw_space_bardist()
 
         self._build_ensemble_preprocessor_and_executor(
@@ -1689,6 +1696,9 @@ def predict_batched(  # noqa: C901, PLR0912
         items: list[RegressorBatch] = []
         fused_index: list[int] = []
         raw_space_bardists: list[FullSupportBarDistribution] = []
+        # Per dataset, like `raw_space_bardists`: the frame its members' target
+        # pipelines were fitted in, needed to map their borders back.
+        znorm_frames: list[tuple[float, float]] = []
         znorm_borders: torch.Tensor | None = None
 
         for idx, (X, y, X_test) in enumerate(
@@ -1706,6 +1716,7 @@ def predict_batched(  # noqa: C901, PLR0912
 
             # Rebuilt fresh on every fit, so each dataset keeps its own object.
             raw_space_bardists.append(worker.raw_space_bardist_)
+            znorm_frames.append((worker.y_train_mean_, worker.y_train_std_))
             # Fixed by the checkpoint, so identical across datasets.
             if znorm_borders is None:
                 znorm_borders = worker.znorm_space_bardist_.borders.clone()
@@ -1795,11 +1806,14 @@ def predict_batched(  # noqa: C901, PLR0912
                 batch.X_query, autocast=worker.use_autocast_, task_type="regression"
             ):
                 for fused_pos in range(len(fused_index)):
+                    znorm_mean, znorm_std = znorm_frames[fused_pos]
                     contribution = worker._translate_batched_logits(
                         output=output[:, fused_pos, :],
                         config=configs_for_est[fused_pos],
                         znorm_borders=znorm_borders,
                         std_borders=std_borders,
+                        znorm_mean=znorm_mean,
+                        znorm_std=znorm_std,
                     )
                     previous = accumulated[fused_pos]
                     accumulated[fused_pos] = (
@@ -1830,11 +1844,15 @@ def _translate_batched_logits(
         config: RegressorEnsembleConfig,
         znorm_borders: torch.Tensor,
         std_borders: np.ndarray,
+        znorm_mean: float,
+        znorm_std: float,
     ) -> torch.Tensor:
         """Map one estimator's output for one dataset onto the shared borders.
 
         Same border translation as :meth:`predict`, for a single
-        (estimator, dataset) pair of the fused forward.
+        (estimator, dataset) pair of the fused forward. `znorm_mean` and
+        `znorm_std` belong to this dataset, not to `self`: the worker is refitted
+        per dataset, so its own attributes only ever hold the last one's.
         """
         out_d = output.float()
         if self.softmax_temperature != 1:
@@ -1846,6 +1864,8 @@ def _translate_batched_logits(
             logit_cancel_mask, descending_borders, borders_t = transform_borders_one(
                 std_borders,
                 target_transform=config.target_transform,
+                znorm_mean=znorm_mean,
+                znorm_std=znorm_std,
                 repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM,
             )
             if descending_borders:
@@ -1978,6 +1998,9 @@ def _iter_forward_executor(
                 # the transformation done to the borders for a given output is dependant
                 # upon the target_transform of the config.
                 if config_for_ensemble.target_transform is None:
+                    # No target pipeline: the target was handed to the model
+                    # already in the frame of the checkpoint's borders, as the
+                    # differentiable-input path does.
                     borders_t = std_borders.copy()
                     logit_cancel_mask = None
                     descending_borders = False
@@ -1986,6 +2009,8 @@ def _iter_forward_executor(
                         transform_borders_one(
                             std_borders,
                             target_transform=config_for_ensemble.target_transform,
+                            znorm_mean=self.y_train_mean_,
+                            znorm_std=self.y_train_std_,
                             repair_nan_borders_after_transform=self.inference_config_.FIX_NAN_BORDERS_AFTER_TARGET_TRANSFORM,
                         )
                     )

src/tabpfn/utils.py

@@ -465,13 +465,25 @@ def transform_borders_one(
     borders: np.ndarray,
     target_transform: TransformerMixin | Pipeline,
     *,
+    znorm_mean: float,
+    znorm_std: float,
     repair_nan_borders_after_transform: bool,
 ) -> tuple[npt.NDArray[np.bool_] | None, bool, np.ndarray]:
     """Transforms the borders used for the bar distribution for regression.
 
+    `target_transform.inverse_transform` returns the borders in the original
+    units of the target, which are then mapped into the z-normalised frame the
+    ensemble is aggregated in. The sanity limits below are applied in that
+    frame, where they mean "this many standard deviations of the target": in
+    the original units they would reject every border of any target whose
+    magnitude happens to exceed them.
+
     Args:
         borders: The borders to transform.
         target_transform: The target transformer to use.
+        znorm_mean: Mean of the training target, defining the frame the borders
+            are returned in.
+        znorm_std: Standard deviation of the training target.
         repair_nan_borders_after_transform:
             Whether to repair any borders that are NaN after the transformation.
 
@@ -482,7 +494,15 @@ def transform_borders_one(
         descending_borders: Whether the borders are descending after transformation
         borders_t: The transformed borders themselves.
     """
-    borders_t = target_transform.inverse_transform(borders.reshape(-1, 1)).squeeze()  # type: ignore
+    # In float64: mapping out of the member's standardization and back into the
+    # shared frame is a round trip whose two halves nearly cancel, and in
+    # float32 that cancellation costs about six digits of the borders. The
+    # result is narrowed back to the dtype of the borders that came in, so the
+    # aggregation and the decode keep the precision they are used to.
+    borders_t = target_transform.inverse_transform(  # type: ignore
+        borders.reshape(-1, 1).astype(np.float64)
+    ).squeeze()
+    borders_t = ((borders_t - znorm_mean) / znorm_std).astype(borders.dtype)
 
     logit_cancel_mask: npt.NDArray[np.bool_] | None = None
     if repair_nan_borders_after_transform:

tests/test_finetuning_regressor.py

@@ -518,4 +518,15 @@ def test_regressor_dataset_and_collator_batches_type(
         assert batch.X_query_raw.shape[0] == 1
         assert batch.y_query_raw.shape[0] == 1
         assert batch.y_query.shape[0] == 1
+
+        # The frame relating the two bar distributions. `fit_from_preprocessed`
+        # never sees the target, so the regressor can only get it from here, and
+        # the forward pass needs it to map an estimator's borders back.
+        assert batch.y_train_std > 0.0
+        np.testing.assert_allclose(
+            batch.raw_space_bardist.borders.cpu().numpy(),
+            batch.znorm_space_bardist.borders.cpu().numpy() * batch.y_train_std
+            + batch.y_train_mean,
+            rtol=1e-5,
+        )
         break