On Difficulties of Attention Factorization through Shared Memory

TL;DR

The paper addresses the quadratic $O(L^2)$ complexity of Transformer attention for long sequences and evaluates external-memory–based architectures that attempt to linearize computation. It uncovers memory degradation, where memory vectors collapse during training and reduce effective memory utilization. It proposes ConvLuna with input filtering via convolution or pooling on the packing keys/values and a learnable softmax temperature to improve information transfer between input and memory. Empirically, ConvLuna achieves improvements on Long Range Arena tasks, including cases where a single memory cell surpasses the vanilla Transformer, indicating that pre-filtering inputs before memory interaction is a promising direction for scalable long-sequence models.

Abstract

Transformers have revolutionized deep learning in numerous fields, including natural language processing, computer vision, and audio processing. Their strength lies in their attention mechanism, which allows for the discovering of complex input relationships. However, this mechanism's quadratic time and memory complexity pose challenges for larger inputs. Researchers are now investigating models like Linear Unified Nested Attention (Luna) or Memory Augmented Transformer, which leverage external learnable memory to either reduce the attention computation complexity down to linear, or to propagate information between chunks in chunk-wise processing. Our findings challenge the conventional thinking on these models, revealing that interfacing with the memory directly through an attention operation is suboptimal, and that the performance may be considerably improved by filtering the input signal before communicating with memory.

Figures (3)

• Figure 1: Memory degradation illustrated. The horizontal axis is the feature dimension, the vertical one is the memory vector index, and the color indicates element values. Despite that memory matrices (learned parameters fed to the first model block as a memory input $P$ on the Figure \ref{['fig:modules']}) were randomly initialized, during training they converge to a small number of unique vectors, see the vertical stripes on a) and b) and similar blocks on c) and d).
• Figure 2: Achieved significances (p-values) of the Friedman test across all considered memory sizes with the H0 = "expected accuracies are equal". Color codes and numbers correspond to the p-values, corrected using the Holm metod garcia. We could not reject the null hypothesis on the 5% level of significance for any setup.
• Figure 3: Comparison of the ordinary Transformer, Luna and the proposed ConvLuna encoder blocks (we refer to them as modules). "Packing" and "unpacking" attention layers in Luna are analogous to the attentions in the vanilla Transformer, while the Rescaled Attention additionally multiplies attention logits with a learnable parameter. The "unpacking" attention of ConvLuna is also denoted as rescalable; however, in our experiments we keep the normalization fixed and identical to the vanilla Transformer.