Constant Memory Attention Block

Published: 20 Jun 2023, Last Modified: 16 Jul 2023ES-FoMO 2023 PosterEveryoneRevisionsBibTeX
Keywords: Attention, Efficient Computation, Neural Processes, Meta-Learning, Uncertainty Estimation, Temporal Point Processes
TL;DR: We propose a novel attention mechanism, Constant Memory Attention Block (CMABs), showing on two settings models based on CMABs achieve results competitive with state-of-the-art and only requiring constant memory.
Abstract: Modern foundation model architectures rely on attention mechanisms to effectively capture context. However, these methods require linear or quadratic memory in terms of the number of inputs/datapoints, limiting their applicability in low-compute domains. In this work, we propose Constant Memory Attention Block (CMAB), a novel general-purpose attention block that computes its output in constant memory and performs updates in constant computation. Highlighting CMABs efficacy, we introduce methods for Neural Processes and Temporal Point Processes. Empirically, we show our proposed methods achieve results competitive with state-of-the-art while being significantly more memory efficient.
Submission Number: 4