Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Linear Attention › Memory Efficiency

Michael BrenndoerferUpdated June 28, 20254 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 123Linear Attention › Memory Efficiency

Reading section 7 of 15 · Memory Efficiency · 4 minSection 7 of 15 · 4 min

Section 7 of 15 · 4 min read

Linear Attention › Memory EfficiencyOne section · Sections in this chapter7 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Beyond compute savings, linear attention also dramatically reduces memory requirements. Standard attention must store the full $n \times n$ attention matrix during the forward pass (for use in backpropagation). Linear attention only stores the $d \times d$ key-value aggregate.

In[17]:

Code

def memory_standard_attention(n, d, dtype_bytes=4):
    """Memory for attention matrix in standard attention."""
    # Store the n x n attention weights
    return n * n * dtype_bytes


def memory_linear_attention(n, d, D=None, dtype_bytes=4):
    """Memory for key-value aggregate in linear attention."""
    if D is None:
        D = d
    # Store the D x d KV aggregate and D-dimensional key sum
    kv_memory = D * d * dtype_bytes
    norm_memory = D * dtype_bytes
    return kv_memory + norm_memory

In[18]:

Code

d = 64
memory_results = []
for n in [512, 1024, 2048, 4096, 8192, 16384]:
    std_mem = memory_standard_attention(n, d)
    lin_mem = memory_linear_attention(n, d)
    reduction = std_mem / lin_mem
    memory_results.append((n, std_mem, lin_mem, reduction))

Out[19]:

Console

Memory Comparison (fp32):

Seq Len    Standard        Linear             Reduction
-------------------------------------------------------
512        1.0 MB          16.2 KB                 63x
1024       4.2 MB          16.2 KB                252x
2048       16.8 MB         16.2 KB               1008x
4096       67.1 MB         16.2 KB               4033x
8192       268.4 MB        16.2 KB              16132x
16384      1073.7 MB       16.2 KB              64528x

At 16,384 tokens, standard attention requires over 1 GB just for the attention matrix, while linear attention uses only about 16 KB. This memory efficiency is what enables linear attention to scale to extremely long sequences that would be impossible with standard attention.

Out[20]:

Visualization

Log-log plot showing standard attention memory growing from kilobytes to gigabytes as sequence length increases, while linear attention remains flat at 16 KB. — Memory usage comparison between standard and linear attention across sequence lengths. Standard attention memory grows quadratically (slope 2 on log-log scale), while linear attention remains constant regardless of sequence length. The divergence becomes dramatic at long sequences.

The log-log visualization makes the scaling difference stark. Standard attention memory grows as a straight line with slope 2 (quadratic), while linear attention remains essentially flat at around 16 KB regardless of sequence length. At 65,536 tokens, standard attention would require over 16 GB just for the attention matrix, exceeding even high-end GPU memory, while linear attention uses the same 16 KB it needed at 128 tokens.