Self-attention is O(n²)

Background

Self-attention scores every token against every other, so it materialises an $n \times n$ attention matrix per head. That's the quadratic cost that makes long context expensive:

$$\text{entries} = h \cdot n^2$$

for h heads and sequence length n.

Problem statement

Implement attention_matrix_size(n, num_heads) returning the number of entries in the attention matrices.

Input

• n — int, sequence length.
• num_heads — int, number of attention heads.

Output

Returns an int: num_heads * n * n.

Examples

Example 1

Input:  n = 4, num_heads = 1
Output: 16

Example 2 — 8 heads

Input:  n = 100, num_heads = 8
Output: 80000

Constraints

• Return num_heads * n * n.
• Return a plain int.

Notes

• Doubling the sequence length quadruples the attention memory — the pressure behind sparse and linear-attention research.