Memory Efficient Neural Processes via Constant Memory Attention Block

TL;DR

The paper tackles the memory burden of attention-based Neural Processes by introducing Constant Memory Attentive Neural Processes (CMANPs). It develops Constant Memory Attention Blocks (CMABs) that ensure constant memory with respect to the context set and provide an exact Cross Attention update to efficiently incorporate new data, enabling scalable conditioning, querying, and updating. The CMANP framework, including an Autoregressive Not-Diagonal extension (CMANP-AND), achieves competitive or state-of-the-art performance on NP benchmarks while using substantially less memory than prior methods. These contributions enable reliable uncertainty estimation in meta-learning under tight resource constraints and offer a pathway to extend constant-memory attention to other modalities and tasks.

Abstract

Neural Processes (NPs) are popular meta-learning methods for efficiently modelling predictive uncertainty. Recent state-of-the-art methods, however, leverage expensive attention mechanisms, limiting their applications, particularly in low-resource settings. In this work, we propose Constant Memory Attentive Neural Processes (CMANPs), an NP variant that only requires constant memory. To do so, we first propose an efficient update operation for Cross Attention. Leveraging the update operation, we propose Constant Memory Attention Block (CMAB), a novel attention block that (i) is permutation invariant, (ii) computes its output in constant memory, and (iii) performs constant computation updates. Finally, building on CMAB, we detail Constant Memory Attentive Neural Processes. Empirically, we show CMANPs achieve state-of-the-art results on popular NP benchmarks while being significantly more memory efficient than prior methods.

Figures (11)

• Figure 1: Constant Memory Attention Block (CMAB). CMABs are stackable attention blocks that (i) are permutation invariant in the input data, (ii) compute their output in constant memory, and (iii) compute updates in constant computation per new data point. Notably, $L_B$ is a learned set of latents unique to the block, allowing CMABs to take advantage of the efficient update property of Cross Attention.
• Figure 2: Constant Memory Attentive Neural Processes (CMANPs). CMANPs perform the conditioning, querying, and updating phases in constant memory with respect to the size of the context dataset $|\mathcal{D}_C|$.
• Figure 3: Computational Complexity Comparison Plots. (Left) Empirical memory usage comparison of CMANP-AND with state-of-the-art Not-Diagonal NPs. (Middle) Empirical memory usage comparison of CMANPs with state-of-the-art NPs. (Right) Empirical runtime comparison of CMANPs with state-of-the-art NPs.
• Figure 4: Analyses Plots. (Left) CMANP's runtime relative to the predictive block size ($b_Q$). (Middle) CMANP's performance relative to the predictive block size ($b_Q$). (Right) CMANP's performance relative to the number of latent vectors ($|L_I| = |L_B|$).
• Figure 5: Constant Memory Hawkes Processes