Constant Memory Attention Block
Leo Feng Frederick Tung Hossein Hajimirsadeghi Yoshua Bengio Mohamed Osama Ahmed
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.