Memory Efficient Neural Processes via Constant Memory Attention Block
This addresses memory constraints for meta-learning applications in low-resource settings, offering an incremental improvement over existing NP methods.
The paper tackled the memory inefficiency of Neural Processes (NPs) due to expensive attention mechanisms, proposing Constant Memory Attentive Neural Processes (CMANPs) that achieve state-of-the-art results on benchmarks while using constant memory.
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.