Google

If the memory component $$\mathrm{Mem}$$ used in the attention operation is defined as a fixed-size variable, the computational cost of performing the attention function $$\mathrm{Att}(\mathbf{q}_i, \mathrm{Mem})$$ will be fixed. By representing keys and values using this fixed-size memory model, the cost remains constant regardless of the sequence length. This foundational concept opens up several alternative ways to design the memory $$\mathrm{Mem}$$.

Fixed-Size Memory for Constant Attention Cost

A simple and effective method for creating a fixed-size memory component, denoted as $$\mathrm{Mem}$$, in attention mechanisms is to use a sliding window. This approach, which is a form of local attention, considers only a limited, constant number of the most recent key and value pairs. By restricting attention to this local neighborhood, the memory size is capped and prevented from growing with the input sequence length.

Fixed-Size Window Memory as a Form of Local Attention

An alternative to using a sliding window for the memory component (`Mem`) is to define it as a pair of summary vectors. This approach creates a more compressed representation of the sequence's history, rather than storing a subset of the raw key-value pairs.

Summary Vectors for Memory Compression in Attention

The update process for a memory component in a memory-based attention mechanism can be described by a general recurrent function. At each time step `i`, the new memory state, `Mem`, is computed by a function `f`. This function takes the current key-value pair, $(\mathbf{k}_i, \mathbf{v}_i)$, and the previous memory state, $Mem_{pre}$, as its inputs. The formula is expressed as: $$Mem = f((\mathbf{k}_i, \mathbf{v}_i), Mem_{pre})$$ This general framework can be instantiated with specific models for the update function `f`, such as a recurrent neural network or a simple moving average.

General Recurrent Formula for Memory Update

Window-based and moving average-based caches offer different approaches to storing historical key-value pairs for attention mechanisms. A window-based cache directly stores a fixed number of recent pairs; for instance, a window of four pairs results in a memory size of 4x2. In contrast, a moving average-based cache compresses the same four pairs into a single summary pair by averaging the keys and values independently. This compression reduces the memory size to a constant 1x2, providing a more memory-efficient representation.

Comparison of Memory Storage in Window-based and Moving Average Caches

A hybrid cache is a memory management strategy that combines two types of memory to efficiently handle long sequences. As illustrated in the diagram, it consists of a 'Local Memory' and a 'Compressed Memory'. The Local Memory (e.g., size 4x2) stores a fixed number of the most recent key-value pairs in their original, uncompressed form. As new data arrives, the oldest key-value pairs are evicted from the Local Memory. These evicted pairs are then passed through a compression function and stored in the Compressed Memory (e.g., size 2x2). This two-level approach allows a model to maintain high-fidelity information about the recent past while retaining a summarized, space-efficient representation of the more distant past.

Hybrid Cache for Attention Mechanisms

An attention mechanism is designed to use a memory component that has a constant, fixed size, regardless of how long the input sequence becomes. What is the primary computational consequence of this design choice as the input sequence length increases significantly?

Consider two language models processing a very long sequence of text one token at a time. Model A uses an attention mechanism where the memory component it attends to has a constant, predetermined size. Model B uses a standard attention mechanism where the memory component grows to include every previous token. Compare how the computational cost of calculating attention for each *new* token changes as the sequence gets longer for Model A versus Model B. Explain the fundamental reason for this difference.

Computational Cost Scaling in Attention Mechanisms

An engineer is developing a model for a real-time task that processes a continuous stream of data. The model uses an attention mechanism where, for each new data point, it must relate it to the entire history of previously seen data points. The engineer observes that as the stream continues and the history grows, the time required to process each new data point increases proportionally, eventually making the system too slow for real-time use.

Learn Before

Related