Flow-Transformed Implicit Processes for Function-Space Variational Inference

Implicit-process priors define distributions over functions through flexible generative mechanisms, making them attractive for Bayesian function-space modelling. The gist: Flow-Transformed Implicit Processes (FTIP) is a variational inference method that uses a normalizing flow to define a richer variational distribution over surrogate variables, enabling it to capture asymmetric and multimodal posterior structure in function space while preserving tractable optimization.

Background: Function-Space Variational Inference for Implicit Processes

Function-space inference targets the posterior distribution over functions rather than over parameters, which aligns inference with prediction. For flexible priors like implicit stochastic processes, the induced function-space prior is rarely available through a tractable density. Practical methods rely on approximations to represent the prior and parameterize the posterior. Standard Variational Implicit Processes (VIP) approximate this by using a finite-dimensional surrogate constructed from prior samples, defining a finite-rank Gaussian process whose empirical mean and covariance match those of the sampled implicit prior on the span of sampled features.

Flow-Transformed Implicit Processes (FTIP)

FTIP is designed to retain the scalability and sample-forward tractability of VIP while substantially enlarging the flexibility of the predictive posteriors it can represent, e.g., non-Gaussian, asymmetric, and multimodal predictive distributions over functions. FTIP starts from a finite collection of prior function samples and places a normalizing-flow variational distribution over the posterior. Posterior functions are obtained by sampling base noise, transforming it through an invertible flow, and mapping the transformed variables through the sampled function surrogate. This construction separates two sources of modelling capacity: the implicit process defines a flexible prior over functions, while the normalizing flow increases the expressiveness of the variational posterior used for inference.

Mechanism and Tractability

The FTIP objective retains the coefficient-space ELBO structure of VIP but replaces the elliptically contoured Gaussian posterior with a more flexible pushforward distribution. The variational objective is given by:

(13)

LFTIP = Eqψ(a) [log p(y1:N F(·; a)) − KLqψ(a)∥ p(a).

This objective is optimized by reparameterizing the coefficients as a flow of base noise, where the induced density is tractable through the change-of-variables formula. The flow transformation, defined by the trainable parameters in Equation (19), is constructed using an initial affine map followed by L = 2 rational-quadratic spline coupling layers interleaved with 1 × 1 LU mixing layers. This architecture preserves exact sampling, inversion, and density evaluation while allowing the posterior to represent non-Gaussian geometry.

Training Objective and Optimization

FTIP can be trained using either the standard evidence lower bound or a Black-Box α objective. The Black-Box α objective modifies the likelihood aggregation across posterior samples while retaining the same coefficient-space KL regularization, with Equation (18) defining it:

(18)

Lα = Ndata Nbatch Xn∈B 1/α log " 1/K X Kk=1 exp α log p(yn F(k)n) -KLqψ(a)∥ p(a).

The parameter α controls how the likelihoods of different posterior samples are combined. As α → 0, the objective reduces to the standard Monte Carlo ELBO, which tends to favor a mode-seeking approximation when the posterior predictive distribution is multimodal. For larger α, it emphasizes samples that explain each observation well, encouraging more inclusive or mass-covering behavior.

Experimental Results and Significance

Experiments on synthetic diagnostics show that FTIP captures separated predictive branches in the bimodal setting and asymmetric uncertainty in the skewed setting. On UCI regression benchmarks, FTIP improves over FBNN, MFVI, and TFSVI across all reported datasets. Specifically, on Energy and Protein data when α = 1.0, FTIP substantially improves NLL and CQM over VIP by assigning probability mass more adaptively near high-density regions. On large-scale regression tasks like YearPredictionMSD, FTIP obtains the best NLL, CRPS, and CQM while retaining the sample-forward scalability of VIP without requiring neural network Jacobians. In classification tasks (FashionMNIST and CIFAR-10), FTIP improves calibration metrics like NLL by correcting the severe overconfidence exhibited by VIP, demonstrating that posterior expressiveness is beneficial for improving approximation quality in these settings.

Limitations and Future Work

The additional posterior expressiveness of FTIP introduces a more complex optimization problem, as the normalizing flow increases Monte Carlo variability. Furthermore, it is noted that posterior expressiveness is only useful when the task requires it, as FTIP often matches VIP when the predictive distribution is close to unimodal or Gaussian. The method remains limited by its reliance on a finite prior-sample surrogate: if this basis does not contain the relevant functional structure, the flow posterior can only partially compensate.

Here are the specific improvements to AI systems that can be made by implementing Flow-Transformed Implicit Processes (FTIP), based on this paper:

1. The system will gain the ability to model and predict outcomes with complex, non-Gaussian uncertainty structures, specifically capturing:

2. Non-Gaussian predictive distributions in continuous regression tasks (e.g., asymmetric skewness and heavy tails). The improved AI can assign appropriate probability mass to distinct modes of the target distribution rather than collapsing them into a single Gaussian approximation.

3. Better calibration in classification tasks, particularly when dealing with overconfident predictions from standard variational methods (like VIP). The improved AI will produce better-calibrated posterior predictive distributions, as evidenced by significant reductions in Negative Log-Likelihood (NLL) and improvement in the Centered Quantile Metric (CQM).

4. More robust performance under prior misspecification. The system can partially compensate for a restrictive or misspecified Gaussian process prior by utilizing the flow-transformed posterior to represent non-Gaussian functional explanations, allowing it to capture separated functional modes that a standard GP posterior would average out.

5. Scalable function-space inference with flexible priors. The system retains the sample-forward scalability of Variational Implicit Processes (VIP) while introducing a more expressive variational family over the surrogate coefficients, making it suitable for implicit stochastic processes where closed-form density approximations are unavailable.

6. Enhanced predictive performance across diverse benchmarks (UCI regression, large-scale regression, pedestrian path-length forecasting). The improved system will achieve superior likelihood-based and calibration-sensitive metrics (NLL and CQM) compared to existing methods (like FBNN or TFSVI) in settings where the true conditional distribution is multimodal or skewed.

7. Improved uncertainty representation in classification. Even when the likelihood function itself is non-Gaussian (like Bernoulli), FTIP improves the quality of the posterior approximation over a Gaussian surrogate, leading to better calibration and accuracy improvements on image classification tasks (FashionMNIST, CIFAR-10).

In summary, FTIP allows AI systems to move beyond simple unimodal or Gaussian uncertainty assumptions in function space inference, enabling them to handle real-world data distributions characterized by multimodality and skewness with higher fidelity.

Implicit-process priors define distributions over functions through flexible generative mechanisms, making them attractive for Bayesian function-space modelling. However, performing posterior inference with such priors is challenging because their induced function-space distributions are typically not available in closed form. One practical strategy is to approximate the prior using a finite collection of sampled functions, and then represent posterior functions as learned combinations of these samples. Existing approaches commonly place a Gaussian variational distribution over the combination weights. While tractable, this choice limits the shapes of posterior uncertainty that can be represented, especially when the true posterior is asymmetric, heavy-tailed, or multimodal. We propose Flow-Transformed Implicit Processes (FTIP), a variational inference method that makes this finite-dimensional function-space approximation more expressive. Instead of using a Gaussian distribution over the combination weights, FTIP uses a normalizing flow to define a richer variational distribution. This induces a flexible posterior distribution over functions while preserving tractable optimization. We train the model using a Black-Box α objective, allowing us to compare mass-covering and mode-seeking variational behaviour. Experiments show that FTIP captures asymmetric and multimodal posterior structure in function space that Gaussian coefficient approximations tend to smooth or collapse.