Add prediction_scaling to rescale predictions toward the training target distribution

[psinger-prior]

Summary

Adds a prediction_scaling argument to both estimators that rescales predictions toward the training target distribution. Every mode is one multiplicative correction of the predicted distribution and they differ only in where the factors come from: none, balanced (the existing balance_probabilities), sampler (undo the label shift introduced by majority_downsample row subsampling, for free), and holdout (fit the correction on held-out rows). The default auto resolves to sampler when the row sampler shifted the prior and to none otherwise, so behavior is unchanged for every configuration that does not use majority_downsample, and that mode now comes out calibrated by default. Follow-up to #1253.

What changed

 TabPFNClassifier.logits_to_probabilities
   temperature
   softmax / average
-  balance by training class counts        (if balance_probabilities)
+  multiply by prediction_scaling_weights_  (per-class vector, fitted at fit time)

 TabPFNRegressor
   _reduce_accumulated_logits
     temperature
+    + prediction_scaling_log_weights_      ("sampler": per-bucket reweighting)
   _rebuild_raw_space_bardist
     borders * std + mean
+    * scale + shift                        ("holdout": affine map of the target axis)

• New module tabpfn.prediction_scaling holds the mode enum, the resolver, and the pure factor functions: class weights for balanced/sampler, the holdout fixed-point fit for class weights, per-bucket log weights for the regression sampler mode, and the holdout affine fit.
• Classifier: one weight vector applied after averaging. balanced is the inverse training prior (identical to the old behavior), or the inverse context prior when the sampler shifted it. sampler is training prior over context prior. holdout iterates w_c <- w_c * freq_c / mean_pred_c on the held-out probabilities until the mean predicted probability per class matches the observed frequency.
• Regressor: sampler reweights the bar-distribution bucket holding the most frequent target value by its training share over its context share, and every other bucket by the ratio of the complements. holdout fits a scalar (or additive shift for signed targets) on held-out mean predictions and applies it to the raw-space borders, so mean, median, quantiles, and output_type="full" all move together. balanced is rejected.
• holdout reuses the split-fit-predict of the tuning machinery. With tuning_config set, it shares the holdout predictions with temperature calibration at no extra cost; without it, it runs one holdout pass with the tuning defaults. Tuning clones are forced to prediction_scaling="none" so the holdout rows are scored unscaled. Order is temperature, then scaling, then decision thresholds.
• balance_probabilities=True still works, emits a DeprecationWarning, and maps to prediction_scaling="balanced"; combining it with another mode raises.
• predict_proba_batched / predict_batched raise when a mode would actually apply per-dataset state, and keep working under auto when no majority_downsample sampler is configured.
• TabPFNEnsemblePreprocessor exposes the resolved row-sampling method and a sampler_shifted_prior flag.

Results

Same laptop-scale setups as #1253: 4 estimators, 10k rows of context per estimator, seed 0, MPS. Ranking metrics are unchanged by the classifier modes by construction, so only calibration moves there.

Credit card fraud (0.17% positive, 200k-row pool, 50k test rows, base rate 0.0015)

arm	ROC AUC	avg precision	log loss	Brier	mean pred
auto / none	0.970	0.691	0.00378	0.00068	0.0015
majority_downsample / none	0.982	0.773	0.00627	0.00114	0.0048
majority_downsample / sampler (new default)	0.982	0.773	0.00320	0.00056	0.0014
majority_downsample / holdout	0.982	0.773	0.00318	0.00057	0.0016

The downsampled context ranks better but predicts three times too many positives. Both corrections bring the mean prediction back to the base rate and make it the best-calibrated arm on every metric.

freMTPL2 loss cost (3.7% nonzero, 100k-row pool, 20k test rows, raw target with default transforms)

arm	gini	gini in exposure deciles	Tweedie skill	cost ratio
auto / none	0.281	0.203	-0.81	0.11
auto / holdout	0.281	0.203	0.032	0.89
majority_downsample / none	0.455	0.278	-0.09	2.98
majority_downsample / sampler (new default)	0.419	0.289	0.041	0.56
majority_downsample / holdout	0.455	0.278	0.098	0.82

For regression the sampler correction is the analytic label-shift fix and turns Tweedie skill positive, but it overshoots the level: the model adopts the context prior only partially, so removing the full shift lands below the true mean. It also changes the ranking slightly, since the bucket reweighting is nonlinear per row; raw Gini drops while in-decile Gini improves. holdout measures the actual level and is the stronger choice for regression when the extra fit is affordable. Both beat leaving the shift in place.

Review guide

• Start with src/tabpfn/prediction_scaling.py. It is short, pure, and carries the math; the estimator changes are plumbing around it.
• In the classifier, read _resolve_prediction_scaling and the two-line change in logits_to_probabilities. Note that the ensemble preprocessor is now constructed before the tuning step, because the free modes need the context prior and the tuning step applies whatever weights are known. The two are independent otherwise.
• In the regressor, read _rebuild_raw_space_bardist, _reduce_accumulated_logits, and the extended _maybe_calibrate_ensemble_temperature. The affine map has to be in place before the raw-space borders are built, which is why the holdout fit happens in the tuning step and the sampler weights after the preprocessor exists.
• Mechanical: constructor arguments and docstrings, the forced prediction_scaling="none" in both tuning clones, the batched-predict guards.
• Design choice: balanced under a shifted context divides by the context prior rather than the training prior. Balancing relative to a prior the model never saw would be wrong; this composes the sampler correction with plain balancing. Without subsampling it is bit-identical to the old balance_probabilities.
• Design choice: the holdout fit measures a first-moment correction only (per-class prior for classification, level for regression). It does not attempt full recalibration; temperature calibration remains the tool for that.
• Not confident about: whether sampler should stay the auto default for regression given the overshoot above, or whether auto should mean none for regressors and leave sampler/holdout opt-in. The classifier case is clear-cut.

Testing

• tests/test_prediction_scaling.py: 36 tests covering every factor function (including an exactness check that the sampler weights recover the training posterior under synthetic label shift, and that the legacy balancing is reproduced bit for bit), the resolver, and the estimator wiring for both tasks: auto resolution, ranking invariance, batch independence (predict(X[:1]) == predict(X)[:1]), the deprecation alias and its conflict error, holdout fitting without a tuning_config, rejection of balanced for regressors and of holdout under differentiable input, and the batched-predict guards.
• Existing suites pass locally: classifier and regressor interfaces, inference tuning, save/load, ensemble preprocessing, consistency, finetuning, batched shape safety. The only failures on this machine are checkpoint downloads gated on license acceptance (v2.5, v2.6, v3.5 variants) and are unrelated.
• Benchmark numbers come from ad hoc scripts outside the repo and are not part of CI.

Breaking changes / follow-ups

• balance_probabilities is deprecated in favor of prediction_scaling="balanced". No behavior change yet.
• Default behavior changes only for SAMPLE_SUBSAMPLING_METHOD="majority_downsample", which now applies the sampler correction. Pass prediction_scaling="none" to get the previous uncorrected output.
• Follow-up: a fair-budget comparison on a GPU with the full training book, and a second seed on the regression benchmark to firm up the auto default question raised above.

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit b8242ee. Configure here.}

[psinger-prior]

Will simplify in follow-up PR