cs.LGFeb 21, 2026
SaveIncremental Transformer Neural Processes
Organizations: University of Cambridge · Helmholtz AI · Technical University of Munich · Munich Center for Machine Learning ·
Abstract
Neural Processes (NPs), and specifically Transformer Neural Processes (TNPs), have demonstrated remarkable performance across tasks ranging from spatiotemporal forecasting to tabular data modelling. However, many of these applications are inherently sequential, involving continuous data streams such as real-time sensor readings or database updates. In such settings, models should support cheap, incremental updates rather than recomputing internal representations from scratch for every new observation -- a capability existing TNP variants lack. Drawing inspiration from Large Language Models, we introduce the Incremental TNP (incTNP). By leveraging causal masking, Key-Value (KV) caching, and a data-efficient autoregressive training strategy, incTNP matches the predictive performance of standard TNPs while reducing the computational cost of updates from quadratic to linear time complexity. We empirically evaluate our model on a range of synthetic and real-world tasks, including tabular regression and temperature prediction. Our results show that, surprisingly, incTNP delivers performance comparable to -- or better than -- non-causal TNPs while unlocking orders-of-magnitude speedups for sequential inference. Finally, we assess the consistency of the model's updates -- by adapting a metric of "implicit Bayesianness", we show that under a one-at-a-time streaming protocol, incTNP retains a prediction rule as implicitly Bayesian as standard non-causal TNPs, demonstrating that incTNP achieves the computational benefits of causal masking without sacrificing the consistency required for streaming inference.
Figures & tables
Figure 1: Computational complexity analysis in the streaming setting (see theoretical analysis in Appendix F ). Top : As the data stream grows (steps ), the model updates its context and predicts targets . In AR mode, the model must perform forward passes (FP) at every step, amplifying the computational burden. Bottom Right : A comparison of the FP mechanics. Standard TNP must re-encode the entire context history at every step (cost ), whereas incTNP leverages KV caching to process only the incremental update (cost ). Bottom Left : The cumulative cost over a stream of length . Because TNP recomputes the full attention matrix repeatedly, its total cost scales cubically ( ). incTNP reduces this to quadratic scaling ( ), making it the only viable option for long streams, particularly when combined with the expensive AR decoding loop.
Table 1: Average Test Log-Likelihoods ( ) for (a) Factorised and (b) AR deployment. We report the absolute performance of the reference non-incremental TNP-D. For all the other methods we report the difference ( ) relative to the reference performance. Held-out task are shown in purple, while transfer scenarios, such as Sim-to-Real (Tabular) and Forecasting (Temperature) are coloured in orange. Best results (including ties within one SEM) are bolded, with standard errors provided in Table 2 .
Figure 2: Test log-likelihoods ( ) on synthetic GP (RBF kernel, up to ) and Tabular datasets across multiple training configurations. incTNP-Seq exhibits reduced variance compared to TNP-D and incTNP , demonstrating the robustness offered by its training strategy.
Figure 3: Performance (Joint NLL gap relative to optimal) versus implicit Bayesianness (KL Gap). incTNP-Seq achieves a KL gap similar to the non-causal TNP-D , whereas the streaming GP baseline ( WISKI ) performs worse in both metrics. Diamond markers ( ) denote the exchangeable versions of the models.
Figure 4: Test log-likelihood ( ) on Tabular (real-world datasets). Performance generally increases with stream length for all models. While incTNP and TNP-D remain competitive on smaller datasets, incTNP-Seq demonstrates greater robustness on Protein.
Figure 5: Test log-likelihood ( ) on the Tabular (Synthetic) dataset. (Left) Performance improves with more data, with incTNP-Seq consistently outperforming TNP-D and incTNP —even when evaluating beyond the training limit of points. (Right) While achieving similar accuracy, the incremental variants demonstrate significantly better scaling efficiency as context size grows.
Figure 6: Test log-likelihood ( ) on the Temperature Prediction task (streaming AR mode).
Appendix figures & tables29 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Architecture of the incremental Transformer Neural Process ( incTNP ), which uses interwoven M-MHSA and MHCA layers.
Figure 8: Schematic of incTNP-Seq during training. All M-MHSA and M-MHCA layers use causal masking.
Figure 9: Distribution of the stations across the restricted latitude and longitude used for our station temperature prediction task, using patches.
Table 2: Average Test Log-Likelihoods ( ) for (a) Factorised and (b) AR deployment. We report the absolute performance of the reference non-incremental TNP-D. For all the other methods we report the difference ( ) relative to the reference performance. Held-out task are shown in purple, while transfer scenarios, such as Sim-to-Real (Tabular) and Forecasting (Temperature) are coloured in orange. The values indicate absolute or relative performance the standard error of the mean (SEM) of each method.
Figure 10: Mean Log-Likelihood boxplots for RBF and Mixed kernels and for different context sizes ( ).
Figure 11: Mean Log-Likelihood boxplots for tabular data prediction with up to .
Figure 12: Test log-likelihoods ( ) on synthetic GP (RBF kernel, up to ) and Tabular datasets across multiple training configurations without any outlier removal. The highest learning rate we used led to failure of some models ( TNP-D and incTNP in the GP case, and incTNP in the tabular case) to converge, resulting in outliers.
Figure 13: The gradient noise-to-signal ratio (GNSR) for incTNP and incTNP-Seq as training progresses on the RBF kernel with . For each checkpoint, we compute gradients over a fixed collection of mini-batches, estimate the signal as the squared norm of the mean gradient, and estimate the noise as the variance of mini-batch gradients around this mean. We estimate this using k samples with batch size .
Table 3: Average Test Log-Likelihoods ( ) for (a) Factorised and (b) AR deployment. We report the absolute performance of the reference non-incremental TNP-D. For all the other methods we report the difference ( ) relative to the reference performance. All tasks are considered out-of-distribution and coloured in orange. The values indicate absolute or relative performance the standard error of the mean (SEM) of each method.
Figure 14: Cost of each AR mode deployment step as weather observations arrive sequentially on the HadISD task. At each step, models condition on all streamed data and perform AR prediction on a fixed target set of using sample unrolls. The dotted orange line represents the cost of updating incTNP ’s context representation. Both the original unoptimised TNP-A model from Nguyen and Grover (2022) , TNP-A (Original) , and our optimised version with KV caching implemented, TNP-A (Cached) , are compared. We only run TNP-A (Original) over context points and AR TNP-D over context points due to their expensive scaling, but fit quadratic functions to them (with dotted black lines) to show their expected scaling. We also fit cost curves to TNP-A (Original) and incTNP with dotted black lines to highlight their respective quadratic and linear per-step AR scaling. At , TNP-A (Cached) is slower than incTNP despite having approximately half as many parameters.
Figure 15: Performance (Joint NLL gap relative to optimal) versus implicit Bayesianness (KL Gap). Similarly to Figure 3 , incTNP-Seq achieves a KL gap similar to the non-causal TNP-D and a better performance. In contrast, incTNP ’s performance is poorer, and the KL gap is slightly larger, but within errors. Diamond markers ( ) denote the permutation-invariant versions of the models.
Figure 16: Test log-likelihood ( ) on the Tabular (Synthetic) dataset including baselines ( LBANP and WISKI ). (Left) The TNP variants show significantly better performance than the baselines. (Right) Although the TNP-based models show better performance, the TNP-D is more costly than the baselines at large streaming lengths. In contrast, the incTNP variants have both the best performance and lowest computational cost.
Figure 17: Test log-likelihood ( ) on Tabular (real-world datasets). The TNP-based variants show better performance on the real-world datasets too.
Figure 18: Test log-likelihood ( ) on the temperature prediction task using AR predictions, reporting streamed performance for both spatiotemporal interpolation (left) and forecasting (right). Results are averaged over 2,000 samples, with the shaded areas representing the standard error of the mean (SEM).
Figure 19: Standard deviation of the log-likelihood over different target orderings (following Figure 9 in Bruinsma et al. (2023) ). Empty markers indicate small, fixed-size contexts (1 or 10 tabular datapoints), while filled markers represent larger, variable contexts sampled uniformly per the legend. We observe similar behaviour between TNP-D and incTNP-Seq , while incTNP tends to show slightly higher sensitivity to target ordering.
| Querying Cost (Step) | ||||
|---|---|---|---|---|
| Model | Update Cost (Step) | Factorised | Autoregressive | Training Cost |
| incTNP | ||||
| incTNP-Seq | ||||
| TNP-D | ||||
Table 4: Summary of per-step streamed conditioning and querying cost for NP models, using both factorised and AR predictions. The training step cost for context points and targets is also included.
Figure 20: Cumulative runtime cost (measured in seconds) on the HadISD dataset of deploying in AR mode as the context set grows ( ). Runtime measured as station observations arrive in a streamed fashion, resulting in updated predictions at stations using sample unrolls across 32 batches. We fit quadratic and cubic cumulative cost functions to incTNP and TNP-D (plotted with dotted black lines), to highlight the cumulative AR scaling cost of both approaches.
| Querying Cost (Cumulative) | ||||
|---|---|---|---|---|
| Model | Update Cost (Cumulative) | Factorised | Autoregressive | |
| incTNP ( -Seq ) | ||||
| TNP-D | ||||
Table 5: Summary of cumulative streamed conditioning and querying cost for NP models, using both factorised and AR predictions. Both incTNP and incTNP-Seq share the same streamed cumulative scaling and are presented together.
| Querying Cost | ||||
|---|---|---|---|---|
| Model | Update Cost | Factorised | Autoregressive | Attention |
| incTNP | Mem-Eff | |||
| TNP-D | — | Mem-Eff | ||
| incTNP | Standard | |||
| TNP-D | — | Standard | ||
Table 6: Summary of per-step streamed conditioning and querying transient memory costs for NP models, using both factorised and AR predictions. Cost is shown for the memory-efficient kernel ( Mem-Eff ) and standard attention kernel ( Standard ).
| Cost | |||
|---|---|---|---|
| Model | Factorised | Autoregressive | Attention |
| incTNP | Mem-Eff | ||
| TNP-D | Mem-Eff | ||
| incTNP | Standard | ||
| TNP-D | Standard | ||
Table 7: Summary of per-step streamed maximum memory costs for NP models, using both factorised and AR predictions. Cost is shown for the memory-efficient kernels ( Mem-Eff ) and standard attention kernels ( Standard ).
Figure 21: Measured memory scaling of TNP-D and incTNP models on the HadISD temperature prediction task as the streamed context set grows for factorised predictions on locations. Peak memory use is reported when using FlashAttention (left - ) and when using standard quadratic memory attention (right - ). A dotted orange line is used to indicate the persistent memory costs of incTNP (i.e. the memory cost of the KV cache and model weights), which contribute to the peak memory. Memory measured on a single NVIDIA A100 80GB GPU, using 16 bit floats and PyTorch’s torch.nn.attention.SDPBackend FlashAttention and Math backends. Dotted black lines are fit to peak memory plots to highlight linear or quadratic memory scaling.
Figure 22: Performance (NLL gap relative to optimal for factorised predictions) versus context permutation sensitivity (context KL gap) for context points and target points. Diamond markers ( ) denote context permutation invariant model constructions. TNP-D is context permutation invariant (up to numerical noise) and thus is represented as a dotted line. We plot the standard error of the mean (SEM) of all models. WISKI is pretrained on the first 64 context points, and updated in a streaming manner for the remaining 64 context points. All NP models are trained with .
Figure 23: Performance (NLL gap relative to optimal for factorised predictions) versus context permutation sensitivity (context KL gap) for context points and target points. Diamond markers ( ) denote context permutation invariant model constructions. TNP-D is context permutation invariant (up to numerical noise) and thus is represented as a dotted line. We plot the standard error of the mean (SEM) of all models. We plot the test time characteristics of TNP-A when unrolling over and autoregressive sample unrolls. All NP models are trained with .
Figure 24: Distribution of test set log-likelihoods across random context permutations for TNP-D , incTNP and incTNP-Seq . We consider the distribution for both (left) and (right) RBF GP regression tasks. We approximate over test set permutations per model (processing billion samples per model) for both tasks. TNP-D is context permutation invariant and thus represented with a dotted line. For incremental NPs, we plot distribution density with a violin plot, using internal dashed lines to signify the distribution quartiles.
Figure 25: Distribution of test set log-likelihoods across random context permutations for TNP-D , incTNP and incTNP-Seq on synthetic tabular data. We approximate over test set permutations per model. TNP-D is context permutation invariant and thus represented with a dotted line. For incremental NPs, we plot distribution density with a violin plot, using internal dashed lines to signify the distribution quartiles.
Figure 26: Distribution of test set log-likelihoods across random context permutations for an incTNP-Seq variant trained to condition on the null context set that we term incTNP-Prior . We measure performance over test set permutations for 1D GP regression using the RBF kernel and . Our test set consists of samples. incTNP-Seq ’s distribution density is shown using a violin plot, with internal dashed lines signifying distribution quartiles. We also benchmark incTNP-Prior ’s modelling accuracy when using variance-based greedy contextual ordering as outlined in Algorithm 2 . We consider all selection strategies , corresponding to Var (Highest), Var (Median) and Var (Lowest) respectively.
| Reference | Ours ( ( ) relative to Ref.) | ||||||
|---|---|---|---|---|---|---|---|
| Dataset | Model | TNP-D | Gain ( ) | ||||
| 1D GP (64) | incTNP | ||||||
| incTNP-Seq | |||||||
| 1D GP (512) | incTNP | ||||||
| incTNP-Seq | |||||||
| Tabular | incTNP | ||||||
Table 8: Average Test Log-Likelihoods ( ) for factorised deployment. We report the absolute performance of the non incremental TNP-D . For incTNP and incTNP-Seq we report the difference ( ) relative to the reference performance when averaging model predictions over permutations of the context set. We also report the performance gain for context permutations over a single permutation . Values indicate absolute or relative performance the standard error of the mean (SEM) of each method. We measure performance on 1D GP regression with an RBF kernel for both and , as well as on the synthetic tabular dataset.
Figure 27: Test log-likelihood ( ) on the distribution shift experiment (transitioning from kernel to ) using models trained on the mixed kernel GP task. The red line indicates the crossover point , where both kernels and are mixed equally. The shaded purple region represents the transition from to . Shaded region starts at 5% mixing interval (distribution is drawn mostly from ) and ends at 95% mixing interval (distribution is drawn mostly from ), using . The dotted grey line indicates the ground truth log-likelihood (GT LL). Results are averaged over 4,096 tasks.
Figure 28: Test log-likelihood ( ) on the distribution shift task for the RBF kernel as the transition rate between kernels changes. Smaller values result in a much more abrupt distribution shift, whilst larger values lead to a smoother transition. is used throughout. Results are averaged over 4,096 tasks.