Anisotropic Trust Region Warping for Prior-Data Fitted Networks in Bayesian Optimisation

AI summary

Overview

Research area: Bayesian optimisation (BO), surrogate modelling, and amortised Bayesian inference with Prior-Data Fitted Networks (PFNs).

Technical level: Intermediate to Advanced. The paper assumes familiarity with Gaussian processes, kernels, lengthscales, acquisition functions and trust-region methods, but its central argument can be followed by anyone who knows what a surrogate model does.

Scope: The paper diagnoses a specific failure mode of PFN surrogates — degradation when the objective's lengthscale lies far from the mode of the PFN's meta-training prior — and proposes a training-free input-space correction, Anisotropic Trust Region (ATR) warping, evaluated on Matérn-3/2 synthetic functions in 2, 5 and 10 dimensions.

What This Paper Is About

PFNs replace the expensive hyperparameter refitting of Gaussian process surrogates in Bayesian optimisation with a single forward pass, but their prior is frozen at meta-training time. The authors show empirically that a PFN's predictive log-likelihood falls progressively behind a refitted GP as the objective's lengthscale grows beyond the mode of that frozen prior, and that the deficit carries over into optimisation performance. Their goal is to fix this mismatch without retraining the network, by rescaling the input space so the objective looks like the lengthscale the PFN expects.

Key Contributions

1. An empirical characterisation of prior mismatch. A stress test across 17 lengthscales in [0.021, 2.0] and 10 dimensions shows the PFN's predictive log-likelihood falling increasingly behind a refitted GP as lengthscale grows, with the crossover point moving to larger lengthscales as dimension increases.
2. Anisotropic Trust Region (ATR) warping. A method that estimates per-dimension lengthscales by kernel target alignment (KTA), then applies an affine rescaling of the input space that maps the objective onto the PFN's target lengthscale, with a trust region around the current anchor determining which observations are retained.
3. Preservation of PFN scalability. The correction runs in O(d·N²), matching the PFN's inference cost, and requires no retraining of the network.
4. A public implementation at https://github.com/youngers2006/Trust-Region-Warping-For-PFNs, plus ablations (supplying the true lengthscale instead of the KTA estimate) and an analysis of estimator accuracy.

Main Findings

• Monotone degradation above the prior mode. In one dimension, the GP and PFN are indistinguishable up to approximately l = 0.1, after which the PFN degrades steadily, reaching a deficit of roughly 2.5 nats at l = 2.0. The onset coincides with the PFN's nominal prior mode, l = 0.1385, even though the tested lengthscales lie within the bulk of the prior (mean 0.80) — so degradation tracks distance from the mode, not departure from support.
• The onset shifts with dimension. At d = 10 the same shape appears with the onset moved to approximately l = 0.5 and a smaller deficit of roughly 0.4 nats. The authors attribute this to sampling density: under a budget of 20·d observations, observations per axis scale as (20·d)^(1/d), so data per correlation length drops sharply with dimensionality.
• Downstream optimisation is affected, but not uniformly. In 2D the gap in normalised simple regret between GP and PFN peaks at intermediate lengthscales and narrows at both extremes. At d = 10 only the GP makes progress, and only at the longer lengthscales.
• ATR warping helps at both extremes. On Matérn-3/2 draws, ATR-PFN improves on the base PFN at both ends of the lengthscale range, though the improvement is more modest at the lower end. The same pattern holds at 5D and 10D with smaller margins.
• Behaviour near the prior mode is benign. For in-distribution lengthscales the ratio l_j / l_target is close to one, the trust region covers most of the domain, warping is near identity, and ATR-PFN behaves much like the PFN; residual gains there come from the trust region and restarts.
• Longer budgets widen the advantage at small lengthscales. Repeating the 2D small-lengthscale configurations (7 lengthscales in [0.021, 0.1]) with 150 iterations instead of 50 widened the advantage substantially. The authors do not separate the contributions of warping and the trust region in that regime.
• Estimator variance is not the driver. An ablation supplying the true lengthscale in place of the KTA estimate shows the improvement is not driven by estimator noise.
• Above the prior mode, gains are attributable to warping alone. There L_j = 1, so the region spans the domain, restarts cannot move it, and neither localisation mechanism is active.
• The KTA estimator undercorrects. Estimates are monotone in the true lengthscale but shrunk toward the centre of the range, crossing the identity line near l = 0.1; the warp therefore partially corrects at both extremes, and never over-warps anywhere in the tested range.

Methodology in Plain English

The authors first run a controlled comparison. They draw synthetic functions from a Matérn-3/2 kernel with a known lengthscale, sweep that lengthscale from very rough (0.021) to very smooth (2.0), and ask both a GP and a PFN to predict held-out values given 20·d observations. The GP is allowed to refit its lengthscale by marginal likelihood; the PFN cannot, because its prior is baked in. Because absolute log-likelihood rises with lengthscale for both models, they report the difference between the two, which isolates the effect of the fixed prior.

To fix the mismatch without retraining, they borrow a trick from input warping: dividing the input domain by a factor s divides the function's apparent lengthscale by s. So if you can estimate the objective's lengthscale, you can rescale the domain to make the objective look like what the PFN expects. The estimate comes from kernel target alignment — fit a Matérn-3/2 kernel's per-dimension lengthscales by gradient descent so that its kernel matrix aligns as closely as possible with the matrix of centred observed outputs. This is cheap, scaling as O(d·N²).

The warp needs guarding in two directions. If the objective is rougher than the prior expects, warping the whole domain would push queries outside the unit cube the PFN was trained on, so the method restricts queries to a smaller trust region and warps that region up to the full domain — this is what makes warping admissible below the mode. If the objective is smoother, the region extends beyond the domain, all data is retained, and the real domain simply occupies part of the warped space. The region is centred on the current best point, translated back inside the domain rather than truncated so its width is preserved, and the anchor is restarted to a uniformly drawn point after 5 consecutive non-improving iterations. To stop aggressive shrinking from starving the PFN of context, a floor is placed on the region size using the 4 observations nearest the anchor in a lengthscale-scaled Chebyshev norm, guaranteeing at least that many points remain in context.

Optimisation experiments use 5·d initial points, 50 iterations, 15 lengthscales in [0.021, 1.5], dimensions 2, 5 and 10, expected improvement over 100000 Sobol candidates common to all surrogates, and 21 repeats per configuration, with regret normalised by the objective's range.

Why This Matters

Impact on research. The paper turns a vague concern about amortised surrogates — that they only work under their training prior — into a measured curve: a specific onset at the prior mode, a specific deficit (roughly 2.5 nats at l = 2.0 in 1D), and a specific dimensional trend. It also offers a correction that lives entirely in the input space, so it does not require touching the network and composes naturally with other prior corrections. This reframes prior mismatch as an estimation problem rather than a retraining problem.

Real-world applications (BO settings where the surrogate must handle objectives whose smoothness is unknown and possibly far from a pretrained default):

• Hyperparameter tuning of machine learning models, where response surfaces range from very rough to nearly flat across hyperparameters.
• Design and screening in materials or molecular discovery, where each evaluation is an expensive simulation or experiment and surrogate scalability matters.
• Engineering design optimisation with costly simulator calls, where GP refitting at O(N³) becomes the bottleneck.
• Tuning of controller or system parameters, where roughness varies sharply between parameters.

Industry relevance. PFNs promise to remove the cubic refitting cost that blocks BO at scale, and this paper identifies the practical condition under which that promise silently fails — objectives smoother than the meta-training prior, which is precisely the regime of expensive, low-noise, physically smooth design problems. A correction that keeps the inference cost at O(d·N²) and requires no retraining is directly deployable by anyone already using a frozen PFN checkpoint such as the HEBO+ prior from PFNs4BO.

Future Directions

• A better low-sample lengthscale estimator. The authors note KTA estimates are monotone but shrunk toward the centre of the range and have high variance on sparse data, which propagates into run-to-run optimisation variance. They also note that high alignment is only a sufficient, not necessary, sign of a good kernel.
• Testing under multiple simultaneous misspecifications. The paper isolates a single axis of prior mismatch and states that behaviour when several are misspecified at once remains to be tested on standard benchmarks.
• Composition with other corrections. Since ATR leaves the network untouched, it could be combined with corrections for other prior axes, such as HEBO's Kumaraswamy input warping.
• Resolving the context-versus-correction tension in high dimensions. Shrinking the region discards evaluations, and the k-nearest-neighbour floor keeps at least k points in context at the cost of capping the warp; in high dimensions, where the k nearest evaluations are far apart, the floor keeps the region large and the warp has little effect. Separating the contributions of warping and localisation at short lengthscales is also left open.

Target Audience

Researchers and practitioners working on Bayesian optimisation and surrogate modelling, especially those already using or evaluating PFN-based surrogates (for example the PFNs4BO checkpoint). It will also interest anyone studying amortised inference and the generalisation limits of meta-learned priors, and engineers deploying BO in settings where GP refitting cost is prohibitive but objective smoothness is unknown. Readers without a background in kernels and acquisition functions will find the diagnostic figures accessible but the method section demanding.

Authors’ abstract

Prior-Data Fitted Networks (PFNs) replace the hyperparameter refitting of Gaussian Process (GP) surrogates in Bayesian optimisation with a single forward pass, at the cost of a prior that is fixed at meta-training time. We show empirically that a PFN's predictive log-likelihood falls increasingly behind that of a refitted GP as the objective function's lengthscale grows beyond the mode of the PFN's meta-training prior, and that this deficit carries over to optimisation performance. To correct the mismatch without retraining the network, we introduce Anisotropic Trust Region (ATR) warping. Per-dimension lengthscales are estimated by kernel target alignment, and an affine rescaling of the input space maps the objective onto the lengthscale the PFN expects, with a trust region around the anchor determining which observations are retained, leaving the PFN's $\mathcal{O}(d \cdot N^2)$ inference cost unchanged. On functions drawn from a Matérn-3/2 kernel, ATR warping improves optimisation performance of the PFN at both extremes of the lengthscale range. Our implementation is publicly available at https://github.com/youngers2006/Trust-Region-Warping-For-PFNs.