Source-linked AI summary
Feature Space Adaptation for Effortless Gaussian Process Flows
Thomas Cowperthwaite, Louis Sharrock, Lachlan Astfalck, Henry Moss
TL;DR
Nonlinear GP conditioning remains challenging, and existing GP flows assume known kernel hyperparameters or rely on manual tuning. FeatureFlow moves the flow to finite-feature weights and uses guidance samples for hyperparameter optimization, achieving one-to-two-order speed-ups while scaling to a 69,081-point sea-level field.
Problem
Nonlinear GP conditioning is challenging, and existing GP-flow implementations assume known kernel hyperparameters or rely on manual tuning, with no equivalent optimization procedure for these tasks.
Source: §1 · p. 1, §3.4 · p. 5Method
Results
FeatureFlow gives one to two orders of magnitude speed-up over function-space GP flows and scales to a 69,081-point sea-level field inferred from over six million radar observations.
Source: §5 · p. 8Takeaways & Limitations
FeatureFlow makes hyperparameter optimization practical for GP flows and decouples flow cost from the resolution used to evaluate function samples.
Source: §1 · p. 1, §5 · p. 8Online parameter updates introduce a path-dependent approximation, so the conditional flow must be rerun with the fitted parameter to obtain samples from the fitted model.
Source: §3.4 · p. 6
Abstract
from arXiv · showhide
Outside the linear-Gaussian regime, conditional sampling from Gaussian processes (GPs) is challenging. Recent methods such as FlowGP (Moss et al., (2026)) can condition on arbitrary non-linear and non-Gaussian statements, but at considerable cost: an expensive iterative and high-dimensional diffusion that requires hand-specified kernel hyperparameters. In this paper, we alleviate two significant drawbacks of FlowGP by (1) introducing kernel approximations that enable scaling to high-resolution domains and (2) proposing a way to obtain the marginal likelihood by measuring the work needed to steer the diffusion towards conditioning statements. We enable, for the first time, hyperparameter optimisation within FlowGP and demonstrate our approach on probabilistic downscaling from areal summary statistics, PDE solution inference on irregular domains, and recovery of sea level anomaly fields from non-Gaussian satellite observations.
1 INTRODUCTION
Gaussian processes handle linear-Gaussian conditioning tractably, but scientific constraints often require nonlinear or non-Gaussian conditioning. FeatureFlow targets the computational cost and hyperparameter-tuning barriers of GP flows with finite-feature representations.
1 INTRODUCTION
- Scientific constraints include nonlinear differential equations, integral measurements, inequalities, and boundary conditions, for which existing approaches use approximations or specialized inference.FlowGP offers a general framework for nonlinear and non-Gaussian statements without bespoke inference schemes.Source: §1 · p. 1
- GP flows are costly on large grids because their computation grows with evaluation-point count, and repeated runs make kernel-hyperparameter optimization prohibitively expensive.The stated costs include an O(N^3) factorization and an N-dimensional ODE solve.Source: §1 · p. 1, §1.1 · p. 2
- FeatureFlow runs the flow over finite-feature weights rather than function values, separating flow dimension from evaluation resolution and making hyperparameter optimization practical.The feature representation also provides access to analytic derivatives.Source: §1 · p. 1
- The paper evaluates FeatureFlow on probabilistic downscaling, PDE inference on irregular domains, and a nonlinear oceanographic inverse problem.Figure 1 illustrates PDE conditioning with noisy observations and homogeneous Dirichlet boundary conditions.Source: §1 · p. 1, Figure 1 · p. 2
2 GAUSSIAN PROCESSES AND SAMPLING
Finite-feature GP approximations replace dense function-value computations with operations on a smaller set of weights. GP flows extend conditional sampling beyond linear-Gaussian information by using likelihood guidance along a diffusion path.
2 GAUSSIAN PROCESSES AND SAMPLING
- Dense GP sampling requires a covariance factorization costing O(N^3) time and O(N^2) memory, followed by O(N^2) sampling.The factorization is based on the covariance matrix over the N evaluation points.Source: §2 · p. 3
- Finite-feature approximations replace the full kernel matrix with a rank-R representation, yielding O(NR) sampling whose approximation quality depends on the tunable feature count.Random Fourier features and Laplacian eigenfunctions are examples for stationary covariance functions.Source: §2 · p. 3
- A feature basis can encode problem-specific structure exactly, while function samples can be evaluated at arbitrary locations from the lower-dimensional weights.This representation separates the number of weights R from the number of evaluation points N.Source: §2 · p. 3
2.1 Non-linear and Non-Gaussian Conditioning
- GP flows condition by adding likelihood-dependent guidance to a reverse-time diffusion, allowing nonlinear and non-Gaussian information to shape samples.The guidance term uses samples from the closed-form conditional distribution of the initial function given its noisy intermediate state.Source: §2.1 · p. 3
- The paper motivates scalable kernel approximations as a way to address the cost of repeated sampling and the derivative errors introduced by finite differences in PDE problems.The stated goals include making kernel tuning more feasible and enabling analytic derivatives.Source: §2.1 · p. 3
3 FEATUREFLOW
FeatureFlow runs conditional diffusion in the weights of a finite-feature GP, evaluates nonlinear guidance through a decoder, and incorporates conjugate observations analytically. It also estimates evidence gradients to adapt kernel hyperparameters, with an online option that requires a final conditional-flow rerun.
3 FEATUREFLOW
- FeatureFlow derives the generating diffusion in finite-feature coefficient space, separating sampling dimension from evaluation-grid resolution and supporting scalable kernel-hyperparameter estimation.The paper also describes joint estimation of the function and its hyperparameters.Source: §3 · p. 4
3.1 Conditional Flow in Weight Space
- For fixed hyperparameters, the weight-space flow has a standard-normal reference and a conditional score comprising the simple term −ξ_t and likelihood guidance.The chosen finite-feature GP defines the posterior recovered by exact conditional scoring and integration.Source: §3.1 · p. 4
3.2 Computing Weight-space Guidance
- The guidance is an intractable R-dimensional likelihood integral, so FeatureFlow estimates it with Monte Carlo bridge samples and differentiates through the condition and decoder.Conditions may be evaluated directly in weight space or at a smaller set of collocation points; cost depends on features and observations rather than output-grid resolution.Source: §3.2 · pp. 4–5
3.3 Exact Conjugate Updates by Matheron
- Conjugate linear-Gaussian observations are absorbed into an affine decoder using Matheron’s rule, leaving the flow to handle only the non-conjugate information.The latent is enlarged by one dimension per observation, and the resulting decoder yields samples from the finite-feature GP posterior given those observations.Source: §3.3 · p. 5
3.4 Hyperparameter Optimisation
- FeatureFlow estimates evidence gradients from bridge samples and uses them to optimize kernel hyperparameters, including through updates interleaved with a conditional solve.The evidence objective balances condition satisfaction against the relative-entropy cost of deforming the GP law, equivalent up to a boundary term to penalizing guidance-drift energy. Online updates add a path-dependent approximation, so the flow should be rerun at the fitted parameter to obtain samples from that fitted conditional model.Source: §3.4 · pp. 5–6
4 EMPIRICAL EVALUATION
The experiments test FeatureFlow on probabilistic downscaling, Allen–Cahn PDE inference, and sea-level anomaly recovery. They report comparable downscaling samples at higher resolution, improved PDE-condition performance, and more accurate joint sea-level predictions than the deterministic baseline.
4.1 Probabilistic Downscaling
- At a fixed compute constraint, FeatureFlow produces comparable downscaling samples at more than 10× as many evaluation points and achieves over two orders of magnitude speed-up at high resolutions.At the largest resolution considered, the GP flow requires an estimated ∼34 GB for its Gram matrix, versus ∼500 MB for FeatureFlow’s basis design matrix.Source: §4.1 · p. 6, Figure 3 · p. 7
4.2 Allen-Cahn Solution Fields
- Laplacian eigenfunction features impose homogeneous Dirichlet boundary conditions by construction, while analytic sample derivatives allow Allen–Cahn residual evaluation in O(RN) time.The GP flow instead imposes the boundary condition using artificial data.Source: §4.2 · p. 6
- FeatureFlow yields higher-quality Allen–Cahn samples by condition residuals, adheres more closely to boundary conditions, and gains more than an order of magnitude in speed at the highest resolutions.Hyperparameter estimation is more stable across resolutions than with the vanilla GP flow; the authors hypothesise improved PDE consistency may reflect removal of finite-difference error.Source: §4.2 · p. 7
4.3 Sea Level Anomaly Estimation
- The sea-level experiment uses non-Gaussian satellite observations to infer a field over 69,081 points, a scale described as computationally infeasible for a GP flow.The setup uses R = 512 random Fourier features and 100 flow-time steps, with the length scale learned from the non-Gaussian condition.Source: §4.3 · p. 8
- Joint predictions across the 69,081-point sea-level field provide an advantage in pointwise accuracy over the deterministic MLE3 baseline.FeatureFlow estimates were also substantially less biased than MLE3; MLE3 provides pointwise estimates and cannot be compared across the full 2-D field.Source: §4.3 · p. 8, Figure 4 · p. 8
5 CONCLUSION
FeatureFlow runs GP flows in the weight space of a finite-feature approximation, enabling evidence-gradient-based kernel hyperparameter optimisation and scaling to large fields. The conclusion reports one-to-two-order speed-ups, notes limitations from prior and adaptation bias, and describes the paper’s use of generative AI tools.
5 CONCLUSION
- FeatureFlow runs GP flows in finite-feature weight space, decoupling flow cost from evaluation resolution and absorbing conjugate data exactly via Matheron’s rule.The hyperparameter-independent bridge provides evidence gradients for kernel hyperparameter optimisation; the method scales to a 69,081-point sea-level field inferred from over six million radar observations.Source: §5 · p. 8
- The stated limitations are reliance on a largely stationary prior with an accurate finite basis and bias in the online adaptation scheme.Future work includes jointly estimating altimetry nuisance fields and characterising that bias.Source: §5 · p. 8
AI Use Statement
- Generative AI tools were used for feedback, implementation, synthetic data, figures, code, brainstorming, and information sourcing; the authors report reviewing AI-assisted work.They state that LLM-generated code was verified and tested for correctness.Source: AI Use Statement · p. 9
Supplementary Materials
The supplementary material describes two finite-feature GP approximations: random Fourier features for stationary kernels and Laplacian eigenfunctions that can incorporate domain geometry and boundary conditions.
A.1 Random Fourier features
- Random Fourier features approximate stationary kernels using sampled frequencies and phases to form a finite-dimensional feature representation.Fixing the standard-normal frequency draws while varying kernel hyperparameters gives differentiable features; spatial derivatives are also available analytically.Source: §A.1 · p. 12
A.2 Laplacian eigenfunction basis
- Laplacian eigenfunctions provide a deterministic low-rank approximation that can impose boundary conditions on bounded domains.For isotropic kernels, the covariance is approximated using the first R eigenfunctions weighted by the kernel spectral density.Source: §A.2 · p. 12
- On simple domains, eigenpairs can be analytic; on irregular domains, they can be computed by discretising the Laplacian and solving a sparse eigenvalue problem.Restricting the discretisation to points inside the domain makes approximate draws vanish on its boundary; changing the eigenvalue problem can incorporate other conditions, including homogeneous Neumann conditions.Source: §A.2 · p. 13
- Because the eigenbasis depends on the domain rather than kernel hyperparameters, it can remain fixed as those parameters vary, leaving the spectral weights to change.This parameterisation is differentiable with respect to kernel hyperparameters and replaces dense kernel operations with computations involving R basis functions.Source: §A.2 · p. 13
B CONDITIONING ON CONJUGATE DATA BY MATHERON’S RULE
This material develops feature-space conditioning and evidence-based hyperparameter optimisation for GP flows. It relates evidence to the cost of satisfying conditions and illustrates the approach in probabilistic downscaling.
B.1 Matheron’s rule targets the posterior
- Under the finite-feature prior, Matheron’s rule gives a latent target whose decoded functions have the exact finite-feature GP posterior given Gaussian observations.The construction appends observation-noise variables to the feature weights and conditions the decoder on the observations.Source: §B.1 · p. 13, B CONDITIONING ON CONJUGATE DATA BY… · p. 13
B.2 Diffusing under conjugacy is stiff
- Treating Gaussian observations as diffusion guidance produces a stiff score whose near-zero-time contraction restricts explicit solver step sizes.For an explicit solver such as Euler–Maruyama, stability requires Δt < 2/κ; the near-zero-time bound depends on the observation-direction eigenvalues and noise level.Source: §B.2 · p. 14
C.1 Evidence Decomposition with Conjugate Information
- When conditioning includes a linear-Gaussian component, its marginal likelihood can be evaluated exactly and differentiated with respect to kernel hyperparameters without additional Monte Carlo error.Including this term matters because the non-conjugate conditional term alone can admit degenerate solutions, such as reducing output variance to shrink a residual without explaining observations.Source: §C.1 · pp. 14–15
C.2 Algorithm
- The optimisation algorithm reuses fixed bridge draws, combines Monte Carlo gradients for non-conjugate conditions with an exact conjugate contribution, and refreshes parameter-dependent quantities after updates.The flow state and bridge samples are held fixed during differentiation; the conditional flow is rerun with the fitted hyperparameters afterward.Source: §C.2 · p. 15
C.3 A Static Interpretation
- For hard constraints with positive probability, evidence maximisation minimises the KL deformation from the unconditioned GP law to the constrained law.For soft conditions, evidence balances condition satisfaction against the KL deformation cost.Source: §C.3 · pp. 16–17
C.4 A Dynamic Interpretation
- Dynamically, evidence balances condition satisfaction against the guidance energy required to deform the reference process; for hard constraints in the complete-noising limit, maximisation reduces to minimising guidance energy.For soft constraints, the objective also includes expected residual loss.Source: §C.4 · p. 17
D.1 Probabilistic Downscaling
- The downscaling experiment compares function-space GP flow and FeatureFlow under shared kernels, observations, nonlinear constraints, diffusion schedule, and sampling budget.It conditions on regional maxima and minima; its reported resolution sweep uses 50 samples per method and resolution, with FeatureFlow also evaluated at higher grid sizes.Source: §D.1 · p. 18, §D.1.2 · p. 18, §D.1.3 · pp. 18–19
D.1.3 Sampling and resolution sweep
- FeatureFlow trades a small decrease in sample quality for a large computational speed-up as the number of features varies.Figure 5 compares sample quality and total sampling wall time against the number of features R.Source: Figure 5 · p. 19
D.1.4 Timing and Evaluation
This section reports timing and evaluation procedures, then details an irregular-domain PDE inference experiment with online length-scale adaptation.
D.1.4 Timing and Evaluation
- At a 0.5-second setup budget, the highest measured resolutions are 40×40 for function space and 144×144 for FeatureFlow.The plotted O(N^3) and O(N) slopes are fixed reference exponents, not fitted slopes.Source: §D.1.4 · p. 19
D.2 PDE Solutions
- The PDE experiment uses a stationary Allen–Cahn equation on the interior of a discrete star-shaped mask embedded in the unit square.Its reference field is generated by a numerical solve on a 64×64 grid, initialized with the eighth smallest Dirichlet eigenmode to obtain a high-frequency solution.Source: §D.2 · pp. 19–20
D.2.1 Observations and GP posterior
- For the PDE experiment, the GP prior is conditioned analytically on observations drawn from the numerical reference field, with observation-noise variance fixed at 10^-4.During adaptation, the length scale is the only optimized kernel parameter; the output scale remains fixed.Source: §D.2.1 · p. 20
D.2.2 Function-space and Feature-space Representations
- The dense GP flow decodes an N-dimensional whitened state using the posterior covariance eigendecomposition, while FeatureFlow uses a feature-space representation.The dense decoder is recomputed after each length-scale update.Source: §D.2.2 · p. 20, FeatureFlow · p. 20
FeatureFlow.
- FeatureFlow uses 256 Dirichlet Laplacian eigenfunctions plus 16 observation-noise features, for a total diffusion dimension of 272.The basis is constructed once on a fine star mask; spectral amplitudes and the conditioning correction are refreshed after adaptive updates.Source: FeatureFlow · p. 20
D.2.3 PDE and Boundary Likelihoods
- The PDE likelihood combines PDE and boundary residuals, with σPDE = 10^-3 and σ∂Ω = 10^-1; increasing the boundary sharpness parameter for GP flow gives very poor PDE residuals.The two methods use different PDE operators: nine-point finite differences for GP flow and an analytic operator for FeatureFlow.Source: §D.2.3 · p. 21
D.2.4 Guided Sampling
- Guided sampling uses 1,000 stochastic Euler–Maruyama steps, with five terminal Monte Carlo draws per particle and step in the PDE experiment.Fixed-prior ensembles use 50 particles, while adaptive runs use five.Source: §D.2.4 · p. 21
D.2.5 Online Hyperparameter Adaptation
- Adaptive runs update the length scale every 10 integration steps, and the reported optimisation dynamics appear more stable for FeatureFlow than for vanilla GP flow.The runs start at ℓ = 0.3 and use Adam with learning rate 0.05; Table 1 reports the realised final values, with output scale 1.0 in each case.Source: §D.2.5 · p. 21
D.2.6 Resolution Ablation, Metrics, and Timing
This section compares FeatureFlow and dense GP flow across spatial resolutions, defines the reported PDE and boundary errors, and specifies how fixed-run and adapted-run wall times are measured.
D.2.6 Resolution Ablation, Metrics, and Timing
- FeatureFlow runs at side lengths 16–128, while dense GP flow is restricted to H ≤64; both use the same 16 observations at each resolution.The resolution set is H ∈ {16, 22, 32, 45, 63, 90, 128}.Source: §D.2.6 · p. 21
- PDE error is the mean squared unscaled physical residual across 50 samples and all PDE evaluation points; boundary error is mean f^2 on the exterior ring.Source: §D.2.6 · p. 21
- Figure 6 shows PDE hyperparameter adaptation histories in function space on the left and weight space on the right.Source: Figure 6 · p. 22
- Table 1 reports final fitted lengthscales, with dashes marking GP-flow resolutions omitted because of dense covariance cost.Source: Table 1 · p. 22
- Adapted-result wall time includes the fresh final fixed run and the complete adaptive run, including setup and all 99 updates.Fixed-run timing is setup plus sampling; the reference timing curves fix exponents at three for GP flow and one for FeatureFlow.Source: §D.2.6 · p. 22
D.3.1 Data and Synthetic Observations
The sea-level anomaly experiment uses a reference field and synthetic satellite waveforms, then conditions an RFF spatial prior with guided sampling. Its setup fixes nuisance fields and sampler settings while adapting the length scale in a separate experiment.
D.3.1 Data and Synthetic Observations
- The reference is the 105 × 80 IMOS sea-level anomaly field for 29 August 2026; 16 satellite passes yield 63,458 simulated shots across about 18,276 km of track.The field is treated as a high-quality reference because it assimilates observational sources unavailable to the inversion.Source: §D.3.1 · p. 22
- Each shot produces a 104-gate waveform with independent Gamma speckle; the inversion observes all waveforms and treats four nuisance quantities as known.The exact joint Gamma log-likelihood is evaluated over 6,599,632 gates, with chunks of 256 used only to limit memory.Source: §D.3.1 · pp. 22–23
D.3.2 Spatial Prior and Basis Representation
- Longitude and latitude are mapped to local planar coordinates for an RBF prior, whose field mean and length scale are fitted to the complete reference field before waveform conditioning.The adaptive experiment retains the oracle mean and variance and treats only the length scale as unknown; the field uses 512 RFFs.Source: §D.3.2 · p. 23
- A shared 512-dimensional RFF weight vector is evaluated on both the display grid and shot coordinates, while land cells are masked for field metrics.Future work could investigate physics-informed boundary conditions, especially near coastlines, to inform samples beyond the non-Gaussian measurements.Source: §D.3.2 · p. 23
D.3.3 Guided Sampling
- Guided sampling uses 100 steps with equal log-SNR spacing, 50 particles per fixed-prior ensemble, one terminal guidance draw, a 1000 norm limit, and likelihood chunks of 256 shots.Source: §D.3.3 · p. 23
D.3.4 Adaptive Length Scale Experiment
The adaptive length-scale experiment starts from a prior set to three times the oracle length scale and updates only the length scale. It ends with a fitted fixed prior used for the final posterior ensemble.
D.3.4 Adaptive Length Scale Experiment
- The initial misspecified length scale is 226.25 km, or three times the oracle value, and the separate oracle fixed run uses 75.42 km.Source: §D.3.4 · pp. 23–24
- During adaptation, output variance and mean are fixed while the isotropic RBF length scale alone is trainable, with one Adam update every five integration steps.The samples from the misspecified fixed prior, adaptive path, and fitted fixed prior are shown in Figure 8.Source: §D.3.4 · p. 23, Figure 8 · p. 24
- The adaptive run fits a final length scale of 81.55 km, then freezes it for a fresh 50-draw ensemble used as the final posterior result.Source: §D.3.4 · p. 24
D.3.5 Baselines and evaluation
The evaluation compares FeatureFlow with MLE3 on satellite observations and reports reconstruction errors. The forward model maps sea-level anomaly and nuisance fields to synthetic radar waveforms with a non-Gaussian observation model.
D.3.5 Baselines and evaluation
- MLE3 is applied independently to each waveform, and its on-track error is reported because it has no two-dimensional field estimate; Table 2 reports observed reconstruction errors.The adaptive-path row is the single changing-prior trajectory used for fitting, not the final ensemble.Source: §D.3.5 · p. 24, Table 2 · p. 25
- FeatureFlow exhibits smaller RMS error than even post hoc debiased MLE3, a truth-corrected diagnostic that is not a deployable estimator.Source: §D.3.5 · p. 24, Figure 9 · p. 25
D.3.6 Forward Model
- The forward model deterministically maps the SLA and nuisance fields to mean waveforms, then draws observed powers from a Gamma speckle model with 90 looks.Source: §D.3.6 · pp. 25–26
- At each shot, sea-surface height determines two-way delay relative to the tracker reference; range-gate times are spaced by 3.125 ns around gate index 32.5 for zero SLA.Source: §D.3.6 · p. 25
- Significant wave height sets scattering-surface width, wind speed determines backscatter, and skewness and reflectivity tilt enter through Hermite coefficients in the expected gate power.Source: §D.3.6 · pp. 25–26