Transformers and Neural State-Space Architectures › Attention over Trajectories

This section develops self-attention as the core mechanism of transformer world models: how a trajectory becomes tokens, how attention computes a content-based weighted average over the past, and why that mechanism is so well suited to partially observable decision problems.

A transformer can retain multiple token representations instead of depending only on a single recurrent history summary. For each query it computes a new fixed-width readout from that memory. Tokenization, embeddings, value projections, and the readout can all discard information; this is not lossless storage. One possible tokenization, used below, embeds an observation-action pair. For observation width $d_o$, action width $d_a$, and embedding width $d$, define $o_t\in\mathbb R^{d_o}$, $a_t\in\mathbb R^{d_a}$, and $z_t,p_t\in\mathbb R^d$:

$$z_t = W_{\text{in}} \begin{bmatrix} o_t \\ a_t \end{bmatrix} + p_t,$$

where:

• $o_t$: the observation received at timestep $t$
• $a_t$: the action taken at timestep $t$
• $W_{\text{in}}\in\mathbb R^{d\times(d_o+d_a)}$: a learned embedding of the stacked observation-action pair
• $p_t$: an additive positional encoding for timestep $t$

The projection embeds the pair; it need not reduce dimension or information. Our code expands two inputs to sixteen features. Without positional signals or an order-dependent mask, self-attention is permutation-equivariant: permuting input rows permutes output rows, and the key-value collection has no intrinsic order. A causal mask changes that statement because different positions have different accessible prefixes; repeated content can then produce different outputs even without explicit positions. Positional encodings provide additional temporal information. The original Transformer uses sinusoidal or learned additive encodings. Rotary encodings instead rotate query and key coordinates inside attention, so they are not literally the additive $p_t$ in this equation. Transformer blocks combine attention, feed-forward layers, residual connections, and normalization. Residual paths allow incremental updates and can improve gradient propagation, but they do not guarantee stable gradients.

The heart of the mechanism is scaled dot-product attention. Each token emits three vectors through learned linear maps:

$$Q = ZW_Q, \qquad K = ZW_K, \qquad V = ZW_V,$$

where:

• $Z\in\mathbb R^{T\times d}$: the matrix with row $t$ equal to $z_t^\top$
• $W_Q,W_K\in\mathbb R^{d\times d_k}$ and $W_V\in\mathbb R^{d\times d_v}$: learned projections, giving $Q,K\in\mathbb R^{T\times d_k}$ and $V\in\mathbb R^{T\times d_v}$
• $q_i=Q_{i,:}^\top$, $k_j=K_{j,:}^\top$, and $v_j=V_{j,:}^\top$: column vectors for the query, key, and value; the attention output below is also a column

It is worth pausing on why the same tokens produce three different roles. The query encodes what the current position is looking for; the key encodes what each position advertises about itself, the label under which it can be found; and the value encodes the information each position contributes once selected. Separating these roles, rather than using the raw embedding for all three, is what lets the network learn asymmetric relationships: a token predicting a collision can learn to search for tokens advertising recent velocities, while advertising, under its own key, something unrelated such as its object identity. The attention output for token $i$ is a weighted average over all value vectors (in the causal models trained below, the sum runs only over $j \le i$, but we first present the unmasked form):

$$\text{Attn}(Q, K, V)_i = \sum_j \frac{\exp\!\big(q_i^\top k_j / \sqrt{d_k}\big)}{\sum_{j'} \exp\!\big(q_i^\top k_{j'} / \sqrt{d_k}\big)} \, v_j,$$

where:

• $j$: the index ranging over all tokens in the sequence
• $j'$: a dummy index used only inside the normalizing constant, so that the sum in the denominator is independent of the sum in the numerator
• $q_i^\top k_j$: the dot product between the query of token $i$ and the key of token $j$, measuring how relevant token $j$ is to token $i$
• $\sqrt{d_k}$: the square root of the key dimension $d_k$, by which the scores are scaled to keep the softmax gradients well behaved when $d_k$ is large
• $\sum_{j'} \exp\!\big(q_i^\top k_{j'} / \sqrt{d_k}\big)$: the normalizing constant that turns the scores into convex weights summing to one
• $v_j$: the value vector of token $j$, the content that gets retrieved

For finite logits and a nonempty allowed set, the softmax weights are nonnegative and sum to one. Thus the per-head output before its output projection is a convex combination of value vectors. The unmasked expression includes current and future positions; a causal mask restricts it to $j\le i$. This convexity does not describe the entire block after residual, projection, and feed-forward operations. If query and key components are independent, zero-mean, unit-variance variables, independent across coordinates, then $\operatorname{Var}(q^\top k)=d_k$. Dividing by $\sqrt{d_k}$ gives unit variance under those assumptions. It can reduce softmax saturation; learned projections can still produce saturated scores.

Two properties explain why this matters for world models:

• Content-based addressing. Scores depend on token content and any positional mechanism, rather than only on a fixed decay schedule. An old retained token can receive high weight if the learned query matches it. A Kalman filter instead weights the current innovation using predicted covariance, measurement covariance, and the observation model; its fixed-size state can still be sufficient for its specified task. Attention recomputes compatibility with retained learned keys, not necessarily raw evidence.
• Short gradient paths. Full attention provides a token-to-token path within a layer whose length does not grow with temporal separation. A recurrent path may contain many state-update Jacobians. Attention therefore shortens temporal dependency paths, but saturated scores, depth, and normalization can still cause optimization problems.

The price depends on which operation is measured. At fixed widths and layer count, dense full-sequence attention performs $O(T^2)$ pairwise arithmetic. Materializing its score matrix uses $O(T^2)$ memory, but this is not an intrinsic memory lower bound: FlashAttention computes exact attention with memory-efficient tiled evaluation. An uncompressed fixed-width KV cache uses $O(T)$ storage. With that cache, one new query performs $O(T)$ attention work; a fixed-width recurrent update instead has cost independent of elapsed history. These are asymptotic statements at fixed batch size, precision, and architecture, not wall-clock speed predictions.

Multi-head attention repeats the lookup with separate projections, then concatenates and linearly mixes the head outputs. Heads may learn different patterns, but assigning one head to each physical state factor requires evidence. The predictive representation is produced jointly by embeddings, attention, residual paths, and feed-forward layers; heads alone do not define a Bayesian belief. Inspecting their weights is one diagnostic, alongside state probes, interventions, ablations, and predictive tests. Attention is not Explanation gives examples where attention weights do not uniquely explain predictions.