diff --git a/changelog/1196.changed.md b/changelog/1196.changed.md new file mode 100644 index 000000000..e2ce8110f --- /dev/null +++ b/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. diff --git a/src/tabpfn/finetuning/data_util.py b/src/tabpfn/finetuning/data_util.py index 34734f3a7..9b0350b33 100644 --- a/src/tabpfn/finetuning/data_util.py +++ b/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 diff --git a/src/tabpfn/inference_config.py b/src/tabpfn/inference_config.py index b56befa83..9c23aef07 100644 --- a/src/tabpfn/inference_config.py +++ b/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). diff --git a/src/tabpfn/preprocessing/configs.py b/src/tabpfn/preprocessing/configs.py index cdb7af874..10db02c1b 100644 --- a/src/tabpfn/preprocessing/configs.py +++ b/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 diff --git a/src/tabpfn/preprocessing/target_transform.py b/src/tabpfn/preprocessing/target_transform.py new file mode 100644 index 000000000..6b90922d2 --- /dev/null +++ b/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", +] diff --git a/src/tabpfn/preprocessing/transform.py b/src/tabpfn/preprocessing/transform.py index 1bb4ba25a..261321172 100644 --- a/src/tabpfn/preprocessing/transform.py +++ b/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. """ diff --git a/src/tabpfn/regressor.py b/src/tabpfn/regressor.py index 9b245060c..42c114bba 100644 --- a/src/tabpfn/regressor.py +++ b/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, diff --git a/tests/test_preprocessing/test_target_transform.py b/tests/test_preprocessing/test_target_transform.py new file mode 100644 index 000000000..1eaa4ec3d --- /dev/null +++ b/tests/test_preprocessing/test_target_transform.py @@ -0,0 +1,151 @@ +# Copyright (c) Prior Labs GmbH 2026. + +"""Tests for the regression target transform pipelines.""" + +from __future__ import annotations + +import pickle + +import numpy as np +import pytest +from sklearn.pipeline import Pipeline + +from tabpfn.preprocessing.steps import ( + get_all_reshape_feature_distribution_preprocessors, +) +from tabpfn.preprocessing.target_transform import ( + UNSTANDARDIZE_STEP, + UnstandardizeTarget, + rebind_target_transform_statistics, + wrap_target_transform, +) + +TRANSFORM_NAMES = ["1_plus_log", "log", "safepower", "quantile_norm", "robust", "none"] + + +def _skewed_positive_target(n: int = 200) -> np.ndarray: + rng = np.random.default_rng(0) + return np.exp(rng.normal(5.0, 1.0, size=n)) + + +def _znorm(y: np.ndarray) -> tuple[np.ndarray, float, float]: + mean, std = float(np.mean(y)), float(np.std(y)) + 1e-20 + return (y - mean) / std, mean, std + + +def _get_transform(name: str, n: int): # noqa: ANN202 + return get_all_reshape_feature_distribution_preprocessors( + num_examples=n, random_state=0 + )[name] + + +@pytest.mark.parametrize("name", TRANSFORM_NAMES) +def test__wrap_target_transform__standardizes_the_transformed_raw_target( + name: str, +) -> None: + """The wrapped transform acts on the raw target and standardizes the result. + + This is the point of the wrapping: fed the z-normalised target, the + pipeline must produce exactly what applying the transform to the target in + its original units and standardizing afterwards would. + """ + y = _skewed_positive_target() + y_znorm, mean, std = _znorm(y) + + wrapped = wrap_target_transform( + _get_transform(name, len(y)), mean=mean, std=std + ).fit_transform(y_znorm.reshape(-1, 1)) + + unwrapped = ( + _get_transform(name, len(y)).fit_transform(y.reshape(-1, 1)).astype(float) + ) + expected = (unwrapped - unwrapped.mean()) / unwrapped.std() + + np.testing.assert_allclose(wrapped, expected, atol=1e-10) + + +@pytest.mark.parametrize("name", TRANSFORM_NAMES) +def test__wrap_target_transform__inverse_transform_returns_znormalized_target( + name: str, +) -> None: + """``inverse_transform`` must land back in the z-normalised space. + + The regressor maps the model's bar-distribution borders through it, and + those borders -- as well as the sanity limits applied to them -- live in + the z-normalised space. + """ + y = _skewed_positive_target() + y_znorm, mean, std = _znorm(y) + + pipeline = wrap_target_transform(_get_transform(name, len(y)), mean=mean, std=std) + transformed = pipeline.fit_transform(y_znorm.reshape(-1, 1)) + + np.testing.assert_allclose( + pipeline.inverse_transform(transformed).ravel(), y_znorm, atol=1e-8 + ) + + +def test__wrap_target_transform__none_transform_is_the_identity() -> None: + """Wrapping the identity transform leaves the z-normalised target alone. + + Keeps the ``None`` entry of ``REGRESSION_Y_PREPROCESS_TRANSFORMS`` and the + ``"none"`` preset equivalent, so the wrapping cannot silently rescale the + target of an estimator that is not supposed to transform it. + """ + y = _skewed_positive_target() + y_znorm, mean, std = _znorm(y) + + pipeline = wrap_target_transform(_get_transform("none", len(y)), mean=mean, std=std) + + np.testing.assert_allclose( + pipeline.fit_transform(y_znorm.reshape(-1, 1)).ravel(), y_znorm, atol=1e-10 + ) + borders = np.linspace(-5.0, 5.0, 21) + np.testing.assert_allclose( + pipeline.inverse_transform(borders.reshape(-1, 1)).ravel(), borders, atol=1e-10 + ) + + +def test__wrap_target_transform__is_picklable() -> None: + """Fitted configs are pickled for joblib workers and for `save_fit_state`.""" + y = _skewed_positive_target() + y_znorm, mean, std = _znorm(y) + + pipeline = wrap_target_transform( + _get_transform("1_plus_log", len(y)), mean=mean, std=std + ) + pipeline.fit(y_znorm.reshape(-1, 1)) + restored = pickle.loads(pickle.dumps(pipeline)) # noqa: S301 + + np.testing.assert_allclose( + restored.transform(y_znorm.reshape(-1, 1)), + pipeline.transform(y_znorm.reshape(-1, 1)), + ) + + +def test__unstandardize_target__transform_and_inverse_are_inverses() -> None: + step = UnstandardizeTarget(mean=3.0, std=2.0) + x = np.array([[-1.0], [0.0], [2.5]]) + + np.testing.assert_allclose(step.transform(x), x * 2.0 + 3.0) + np.testing.assert_allclose(step.inverse_transform(step.transform(x)), x) + + +def test__rebind_target_transform_statistics__updates_wrapped_transforms() -> None: + """The fine-tuning pipeline re-normalises per split and rebinds the stats.""" + y = _skewed_positive_target() + pipeline = wrap_target_transform(_get_transform("log", len(y)), mean=0.0, std=1.0) + + rebind_target_transform_statistics([None, pipeline], mean=7.0, std=3.0) + + step = pipeline.named_steps[UNSTANDARDIZE_STEP] + assert (step.mean, step.std) == (7.0, 3.0) + + +def test__rebind_target_transform_statistics__ignores_other_pipelines() -> None: + """Pipelines that do not come from `wrap_target_transform` are left alone.""" + other = Pipeline(steps=[("some_step", _get_transform("safepower", 100))]) + + rebind_target_transform_statistics([other], mean=7.0, std=3.0) + + assert UNSTANDARDIZE_STEP not in other.named_steps diff --git a/tests/test_regressor_interface.py b/tests/test_regressor_interface.py index 42a96609e..06daf8ccf 100644 --- a/tests/test_regressor_interface.py +++ b/tests/test_regressor_interface.py @@ -1365,6 +1365,104 @@ def test__predict_batched__uses_fitted_target_transforms() -> None: np.testing.assert_allclose(batched[i], ref.predict(X_tests[i]), atol=1e-4) +def _mk_skewed_reg_dataset( + seed: int, n: int = 200, f: int = 5 +) -> tuple[np.ndarray, np.ndarray]: + """A strictly positive, right-skewed target -- the log-transform case.""" + r = np.random.RandomState(seed) + X = r.randn(n, f) + y = np.exp(1.5 + X @ np.array([1.0, -0.5, 0.3, 0.0, 0.2]) + 0.3 * r.randn(n)) + return X, y + + +def test__fit__target_transform__applied_to_the_unnormalized_target() -> None: + """A target transform must see the target in its original units. + + The estimator z-normalises the target before the ensemble preprocessing, so + a transform such as ``1_plus_log`` would otherwise be applied to + standardized values, where it means something entirely different (and is + undefined below -1). + """ + X, y = _mk_skewed_reg_dataset(0) + reg = TabPFNRegressor( + n_estimators=1, + device="cpu", + random_state=42, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": ("1_plus_log",)}, + ) + reg.fit(X, y) + + expected = np.log1p(y) + expected = (expected - expected.mean()) / expected.std() + member_y = np.asarray(reg.executor_.ensemble_members[0].y_train) + + np.testing.assert_allclose(member_y, expected, atol=1e-8) + + +def test__fit__no_target_transform__leaves_the_target_znormalized() -> None: + """Without a target transform the model still sees the z-normalised target.""" + X, y = _mk_skewed_reg_dataset(0) + reg = TabPFNRegressor( + n_estimators=1, + device="cpu", + random_state=42, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": (None,)}, + ) + reg.fit(X, y) + + member_y = np.asarray(reg.executor_.ensemble_members[0].y_train) + + np.testing.assert_allclose(member_y, (y - y.mean()) / y.std(), atol=1e-8) + + +def test__fit__none_target_transform__is_treated_as_no_transform() -> None: + """The ``"none"`` preset is the identity, so it must not be composed. + + Composing it would leave the target unchanged but for floating-point noise + from the round trip through the z-normalisation, needlessly perturbing the + predictions of every config that uses it. + """ + X, y = _mk_skewed_reg_dataset(0) + reg = TabPFNRegressor( + n_estimators=1, + device="cpu", + random_state=42, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": ("none",)}, + ) + reg.fit(X, y) + + member = reg.executor_.ensemble_members[0] + assert member.config.target_transform is None + np.testing.assert_allclose( + np.asarray(member.y_train), (y - y.mean()) / y.std(), atol=1e-8 + ) + + +@pytest.mark.parametrize("transform", ["1_plus_log", "log", "safepower"]) +def test__predict__target_transform__predictions_are_finite_and_accurate( + transform: str, +) -> None: + """Transforming a skewed target must yield usable predictions. + + Guards the whole round trip: the transform is fitted on the unnormalized + target, while the model's borders are mapped back through it into the + z-normalised space the ensemble is aggregated in. + """ + X, y = _mk_skewed_reg_dataset(0) + X_train, X_test, y_train, y_test = X[:150], X[150:], y[:150], y[150:] + reg = TabPFNRegressor( + n_estimators=2, + device="cpu", + random_state=42, + inference_config={"REGRESSION_Y_PREPROCESS_TRANSFORMS": (transform,)}, + ) + reg.fit(X_train, y_train) + pred = reg.predict(X_test) + + assert np.isfinite(pred).all() + assert r2_score(y_test, pred) > 0.7 + + @pytest.mark.parametrize("device", devices) def test__predict_batched__matches_per_dataset_dataframe(device: str) -> None: """Batched prediction matches per-dataset on non-numeric DataFrame inputs.