Linear attention scales O(n) in sequence length — making it practical for image feature maps without the O(n^2) cost of full attention

Standard multi-head self-attention has O(n^2) time and memory complexity in sequence length n (where n = H×W for image features). For large feature maps this is prohibitive. Linear attention reformulates attention using a kernel trick that computes a context matrix in O(n) time by reordering the einsum operations. DDPM uses two attention variants in its U-Net: full multi-head self-attention at the bottleneck (where spatial resolution is small) and linear attention at other levels. The tradeoff is approximation quality vs. compute. Both are implemented as nn.Conv2d-based attention operating on spatial feature maps.

Examples

In the U-Net: downs stages use LinearAttention; the mid_attn uses full Attention. Switching all to full attention would work for small images but become infeasible for high-resolution generation.

Assessment

Explain why O(n^2) attention is problematic for image feature maps. Describe the key difference between how linear attention and full attention compute their output.