Direct Flow Neural Processes: Efficient Sampling via Flow Step Amortization

Neural Processes (NPs) aim to combine the flexible uncertainty modeling of stochastic processes with the scalability of deep learning, providing efficient amortized inference for distributions of arbitrary conditionals. While recent approaches based on autoregressive modeling or flow matching have achieved state-of-the-art posterior approximation, they fundamentally compromise the efficiency of NPs by relying on slow iterative sampling methods. This computational bottleneck can make these models impractical for tasks involving sequential decision-making loops for which NPs are best suited. In this work, we propose Direct Flow Neural Processes (D-FlowNP), a method that enables rapid inference within highly expressive flow-based NPs for sequential tasks. By amortizing the computation of the flow inference steps, our model can generate consistent posterior samples across variable computational budgets. In particular, the number of inference steps can be substantially reduced with only negligible quality degradation. We validate our approach on 1D Gaussian processes, 2D images, and scientific data including Navier Stokes flows and black hole imaging. Experiments on Bayesian optimization demonstrate that our method unlocks downstream capabilities previously infeasible for flow based NPs.