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

changelog/1198.breaking.md

@@ -0,0 +1 @@
+Previously saved fitted regressors (`.tabpfn_fit` archives and pickled estimators) are incompatible with the updated target preprocessing and must be refitted. Custom fine-tuning loops must pass `y_train_mean` and `y_train_std` to `fit_from_preprocessed`.

changelog/1198.changed.md

@@ -0,0 +1 @@
+Regression ensembles now normalize targets using each member’s training rows; without subsampling, only floating-point differences are expected.

src/tabpfn/finetuning/data_util.py

@@ -18,6 +18,7 @@
 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 StandardizeTarget
 from tabpfn.utils import infer_random_state, pad_tensors
 
 if TYPE_CHECKING:
@@ -66,6 +67,8 @@ 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: Training-target mean defining the shared aggregation frame.
+        y_train_std: Standard deviation of this split's training target.
     """
 
     X_context: list[torch.Tensor]
@@ -79,6 +82,8 @@ class RegressorBatch:
     znorm_space_bardist: FullSupportBarDistribution
     X_query_raw: torch.Tensor
     y_query_raw: torch.Tensor
+    y_train_mean: float
+    y_train_std: float
 
 
 Batch = TypeVar("Batch", ClassifierBatch, RegressorBatch)
@@ -366,7 +371,7 @@ def __len__(self) -> int:
         return len(self.configs)
 
     @override
-    def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa: C901, PLR0912
+    def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa: C901
         """Retrieves, splits, and preprocesses the dataset config at the index.
 
         Performs train/test splitting and applies potentially multiple
@@ -438,22 +443,10 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
         # This however makes also sense when considering that
         # this attribute changes on every dataset
         if is_regression_task:
-            train_mean = np.mean(y_train_raw)
-            train_std = np.std(y_train_raw)
-
-            eps = 1e-8
-            if train_std < eps:
-                warnings.warn(
-                    f"Target variable has constant or near-constant values "
-                    f"(std={train_std:.2e}). Adding epsilon={eps} to prevent "
-                    f"division by zero in standardization.",
-                    UserWarning,
-                    stacklevel=2,
-                )
-                train_std = eps
+            frame = StandardizeTarget().fit(y_train_raw)
+            train_mean, train_std = frame.mean_, frame.std_
 
             y_test_standardized = (y_test_raw - train_mean) / train_std
-            y_train_standardized = (y_train_raw - train_mean) / train_std
             znorm_borders = znorm_space_bardist_.borders
             raw_space_bardist_ = (
                 FullSupportBarDistribution(
@@ -463,9 +456,6 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
                 .float()
                 .to(znorm_borders.device)
             )
-            y_train = y_train_standardized
-        else:
-            y_train = y_train_raw
 
         num_columns = x_train_raw.shape[1]
         feature_schema = FeatureSchema.from_only_categorical_indices(
@@ -480,7 +470,7 @@ def __getitem__(self, index: int) -> ClassifierBatch | RegressorBatch:  # noqa:
         )
         ensemble_members = ensemble_preprocessor.fit_transform_ensemble_members(
             X_train=x_train_raw,
-            y_train=y_train,
+            y_train=y_train_raw,
         )
         X_trains_preprocessed = [m.X_train for m in ensemble_members]
         y_trains_preprocessed = [m.y_train for m in ensemble_members]
@@ -535,6 +525,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(
@@ -700,6 +692,9 @@ 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),
+        # The target frame belongs to the bar distributions above.
+        y_train_mean=first_item.y_train_mean,
+        y_train_std=first_item.y_train_std,
     )
 
 

src/tabpfn/finetuning/finetuned_base.py

@@ -1175,6 +1175,14 @@ def _ddp_broadcast_primary_metric(metric: float) -> float:
                     rank_batch.cat_indices,
                     rank_batch.configs,
                     performance_options=finetuning_performance_options,
+                    **(
+                        {
+                            "y_train_mean": rank_batch.y_train_mean,
+                            "y_train_std": rank_batch.y_train_std,
+                        }
+                        if isinstance(rank_batch, RegressorBatch)
+                        else {}
+                    ),
                 )
 
                 if using_ddp:

src/tabpfn/finetuning/finetuned_regressor.py

@@ -40,15 +40,15 @@
 def _targets_in_estimator_space(
     y_Q: torch.Tensor, config: RegressorEnsembleConfig
 ) -> torch.Tensor:
-    """Map z-scored targets into the space the estimator predicts in.
+    """Map raw targets into the space the estimator predicts in.
 
     Members with a target transform (the default ensemble's "safepower") saw
     transformed context targets, so their loss must use transformed targets too.
     """
     if config.target_transform is None:
-        return y_Q
+        raise ValueError("A fitted target pipeline is required for regression loss.")
     y_np = config.target_transform.transform(y_Q.cpu().numpy().reshape(-1, 1))
-    return torch.as_tensor(y_np.ravel(), dtype=y_Q.dtype, device=y_Q.device)
+    return torch.as_tensor(y_np.ravel(), device=y_Q.device)
 
 
 def _compute_regression_loss(  # noqa: C901
@@ -368,7 +368,7 @@ 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 the bar distributions 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
@@ -406,8 +406,11 @@ def _forward_with_loss(self, batch: RegressorBatch) -> torch.Tensor:  # type: ig
         logits_BQL = logits_QBEL.permute(1, 2, 0, 3).reshape(B * E, Q, L)
 
         targets_BQ = torch.stack(
-            [_targets_in_estimator_space(y_query_batch[0], c[0]) for c in batch.configs]
-        ).to(self.device)
+            [
+                _targets_in_estimator_space(batch.y_query_raw[0], c[0])
+                for c in batch.configs
+            ]
+        ).to(device=self.device, dtype=y_query_batch.dtype)
 
         return _compute_regression_loss(
             logits_BQL=logits_BQL,

src/tabpfn/preprocessing/target_transform.py

@@ -0,0 +1,72 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+
+"""Invertible regression target pipelines.
+
+Each member maps raw targets to model inputs and inversely maps borders to raw
+units. The estimator maps those borders into a shared standardized space for
+aggregation, avoiding float32 precision loss for targets with large offsets.
+"""
+
+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):
+    """Standardize using fitted ``mean_`` and ``std_`` (population std + epsilon).
+
+    Keep NumPy arithmetic rather than StandardScaler to preserve exact model
+    inputs when the member sees the full training target.
+    """
+
+    EPSILON = 1e-20
+    """Prevent division by zero for constant targets."""
+
+    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:
+    """Standardize raw targets, then apply the optional preset.
+
+    The inverse returns values in the target's original units.
+    """
+    if transform is None:
+        return Pipeline(steps=[(STANDARDIZE_STEP, StandardizeTarget())])
+    return Pipeline(
+        steps=[
+            # The preset reshapes the standardized target.
+            (STANDARDIZE_STEP, StandardizeTarget()),
+            (TARGET_TRANSFORM_STEP, transform),
+        ],
+    )
+
+
+__all__ = [
+    "STANDARDIZE_STEP",
+    "TARGET_TRANSFORM_STEP",
+    "StandardizeTarget",
+    "make_target_transform",
+]