Perceiver: Latent Cross-Attention

Table of Contents

• 1. Motivation — The Quadratic Bottleneck
• 2. Architecture Overview
• 3. The Latent Array
• 4. Cross-Attention: Input → Latents
  • 4.1 Formal Definition
  • 4.2 Worked Example
• 5. Self-Attention Among Latents
• 6. Stacking Blocks
• 7. Complexity Analysis
• 8. References

1. Motivation — The Quadratic Bottleneck 📐

Standard transformer self-attention on an input of length \(N\) computes all pairwise interactions between tokens, giving \(O(N^2)\) cost per layer. For long sequences — images flattened to pixels, audio waveforms, long documents — this is prohibitive.

The Perceiver (Jaegle et al., 2021) breaks this bottleneck with a simple architectural change: instead of letting the input attend to itself, a small fixed-size set of latent tokens attends to the input via cross-attention. The expensive self-attention then operates only on these latents, whose count \(M\) is chosen freely and held fixed regardless of input length \(N\).

2. Architecture Overview 🗺️

A single Perceiver block consists of two sub-layers applied in sequence:

flowchart LR
    I["Input array
N × d_in"] -->|keys, values| CA["Cross-Attention"]
    L["Latent array
M × d_lat"] -->|queries| CA
    CA --> L2["Updated latents
M × d_lat"]
    L2 --> SA["Self-Attention"]
    SA --> L3["Refined latents
M × d_lat"]

This block is repeated \(B\) times. Each repetition re-reads the original input array via cross-attention, allowing the latents to iteratively refine their representation.

3. The Latent Array 🔑

Definition (Latent Array). The latent array \(\mathbf{Z} \in \mathbb{R}^{M \times d_\text{lat}}\) is a matrix of \(M\) learned vectors, initialized randomly and updated by gradient descent during training. The integer \(M\) is a hyperparameter chosen independently of the input length \(N\).

The latent tokens are not tied to any input position. They carry no intrinsic semantic meaning at initialization — they learn, through training, to act as flexible query vectors that extract task-relevant summaries from the input.

4. Cross-Attention: Input → Latents 📐

4.1 Formal Definition

Let the input array be \(\mathbf{X} \in \mathbb{R}^{N \times d_\text{in}}\) and the latent array be \(\mathbf{Z} \in \mathbb{R}^{M \times d_\text{lat}}\).

Definition (Perceiver Cross-Attention). Queries are projected from the latents; keys and values are projected from the input:

\[ \mathbf{Q} = \mathbf{Z} \mathbf{W}_Q \in \mathbb{R}^{M \times d_k}, \qquad \mathbf{K} = \mathbf{X} \mathbf{W}_K \in \mathbb{R}^{N \times d_k}, \qquad \mathbf{V} = \mathbf{X} \mathbf{W}_V \in \mathbb{R}^{N \times d_v} \]

The cross-attention output is:

\[ \text{CrossAttn}(\mathbf{Z}, \mathbf{X}) = \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}}\right) \mathbf{V} \in \mathbb{R}^{M \times d_v} \]

The attention weight matrix \(\mathbf{A} = \text{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d_k}) \in \mathbb{R}^{M \times N}\) has one row per latent and one column per input token. Entry \(\mathbf{A}_{ij}\) is the weight latent \(i\) assigns to input token \(j\).

The output shape is \(M \times d_v\) — the same number of rows as latents, regardless of \(N\).

4.2 Worked Example

Sentence: "The cat sat on the mat" → \(N = 6\) token embeddings, \(d_\text{in} = 64\). Latent array: \(M = 4\) latents, \(d_\text{lat} = 64\).

After projecting to \(d_k = d_v = 32\):

\[ \mathbf{A} \approx \begin{pmatrix} 0.05 & \mathbf{0.60} & 0.10 & 0.05 & 0.05 & 0.15 \\ 0.10 & 0.05 & \mathbf{0.65} & 0.05 & 0.05 & 0.10 \\ \mathbf{0.50} & 0.10 & 0.05 & 0.10 & 0.20 & 0.05 \\ 0.05 & 0.05 & 0.05 & 0.10 & 0.05 & \mathbf{0.70} \end{pmatrix} \]

Each row sums to 1. In this illustration: latent 0 specializes on “cat”, latent 1 on “sat”, latent 2 on “The/the”, latent 3 on “mat”. The updated latent 0 is approximately \(0.60 \cdot \mathbf{v}_\text{cat} + 0.15 \cdot \mathbf{v}_\text{mat} + \cdots\), a learned weighted aggregate of input vectors.

5. Self-Attention Among Latents 🔄

After cross-attention, the \(M\) updated latent vectors are passed through standard multi-head self-attention:

\[ \mathbf{Z}' = \text{SelfAttn}(\mathbf{Z}_\text{updated}) \in \mathbb{R}^{M \times d_\text{lat}} \]

This costs \(O(M^2 \cdot d_\text{lat})\). Since \(M \ll N\), this is the cheap step. Its role is to let latents communicate with each other — so that the “cat” latent and the “sat” latent can form a joint representation of the subject-verb pair, for instance.

6. Stacking Blocks 🏗️

A full Perceiver stacks \(B\) blocks:

\[ \mathbf{Z}^{(b)} = \text{SelfAttn}\!\left(\text{CrossAttn}\!\left(\mathbf{Z}^{(b-1)},\ \mathbf{X}\right)\right), \quad b = 1, \ldots, B \]

where \(\mathbf{Z}^{(0)}\) is the learned initial latent array. The same input \(\mathbf{X}\) is re-read at every block — the latents iteratively refine their summary of the input over \(B\) steps.

7. Complexity Analysis ⚡

Per block, with input length \(N\), latent count \(M\), and hidden dimension \(d\):

Total over \(B\) blocks: \(O(B \cdot M \cdot N \cdot d + B \cdot M^2 \cdot d)\).

Compare to a vanilla transformer: \(O(B \cdot N^2 \cdot d)\).

For \(M = 512\), \(N = 65536\) (e.g., a 256×256 image flattened):

\[ \frac{N^2}{M \cdot N} = \frac{N}{M} = \frac{65536}{512} = 128\times \text{ reduction in cross-attention cost} \]

Key conclusion: the Perceiver’s complexity is linear in \(N\) (for fixed \(M\)), not quadratic. This enables the same architecture to process images, audio, video, and text with no input-specific engineering.

8. References 📚