Mamba Flow Matching Neural Processes: Linear-Time Inference for Irregularly Observed Spatial Fields

Spatial interpolation from sparse irregular observations is a recurring task across the environmental and health sciences. Neural processes amortise this inference into a single forward pass at test time. Joint-predictive variants capture correlations across query locations but typically rely on self-attention over context and query points, incurring $O(K^2)$ compute in the combined point count $K$ and capping practical use at a few thousand points. State-space models offer a linear-time alternative, but existing state-space NPs compress the context into a single hidden state before reading queries, reducing the information available to the decoder. We propose MambaFlowNP, a joint-predictive NP with linear compute trained with a conditional flow-matching objective. A multi-directional Mamba-2 scan reads context and query tokens in a single shared sequence, so each query's velocity is read from its own scan position rather than from a compressed context summary. MambaFlowNP is competitive with the strongest attention-based joint-predictive baselines on synthetic Gaussian-process benchmarks, and across regional and continental-scale NOAA temperature interpolation tasks it attains the lowest RMSE on every tier and the lowest CRPS at the California and continental US tiers, while running $170$–$760\times$ faster than the strongest joint-predictive baseline. A single forward pass scales to $10^6$ points on one GPU in under a second. This scalability unlocks amortised joint-predictive inference for continental-scale, irregularly sampled sensor networks, spanning climate monitoring to disease vector surveillance.