Sequence Modeling

The Challenge of Sequences

A sequence model processes ordered inputs such as text tokens or audio samples, and an output may depend on inputs far earlier in the sequence. We compare how attention and state-space models use earlier inputs, and how their memory and computation grow with sequence length.

This page covers three designs: attention, which compares every pair of positions; state-space models, which carry a fixed-size state from step to step; and Mamba, which lets that state depend on the input. They trade off parallel training, memory during generation, and access to distant tokens. Chapter 12 of Bishop & Bishop (2024) gives a textbook treatment of attention and transformers.

Self-Attention

In a Transformer (Vaswani et al., 2017), each token's output is a weighted average of value vectors from the tokens it can see, with weights set by comparing its query vector with their key vectors. No recurrence is needed. The weight from token i to token j is:

Here d_k is the number of entries in each query and key vector. All positions are processed in parallel during training, but every pair of positions gets a score, so the work grows as O(n²) in sequence length n. Memory does too if the n × n matrix of weights is stored; FlashAttention avoids storing it, though the O(n²) work remains (Dao et al., 2022).

RoPE: Rotary Position Embeddings

Attention by itself ignores word order. RoPE (Su et al., 2024) encodes position by rotating queries and keys. Each vector is split into pairs of coordinates, and pair i of a vector at position m is turned by the angle mθ_i, with frequencies θ_i = 10000^−2i/d that run from fast to slow:

Because rotations compose, the score between a query at position m and a key at position n depends on position only through the offset n − m. This relative form does not by itself let a model use contexts longer than those it was trained on: past the training length the slow pairs reach angles never seen in training, and quality degrades (Press et al., 2022; Chen et al., 2023). Position interpolation (Chen et al., 2023) divides every position by the extension factor, so a longer input stays within the trained range. YaRN (Peng et al., 2024) instead rescales the rotation frequencies by band: it leaves the fast pairs unchanged, interpolates the slow ones, blends those in between, and adds a temperature to the attention logits. Both papers fine-tune briefly and extend the usable context to some multiple of the training length.

KV Cache

When a model generates text one token at a time, the keys and values of earlier tokens do not change, so they can be stored and reused. Only the new token's query, key and value need computing:

The attention work for each new token is then O(t·d) per layer, where d is the model width, instead of recomputing the whole prefix. The price is memory: 2·t·d numbers per layer for keys and values. With long contexts and large batches the cache can take more memory than the model's weights (Pope et al., 2023).

Grouped Query Attention

In multi-head attention (MHA) every query head has its own key and value heads. Multi-query attention (MQA; Shazeer, 2019) shares a single key head and value head across all query heads. Grouped-query attention (GQA; Ainslie et al., 2023) sits in between: each key–value head serves a group of query heads.

The factor 2 counts keys and values. Ainslie et al. report that GQA reaches quality close to MHA at a decoding speed close to MQA.

SwiGLU Activation

The feed-forward layer of many recent transformers uses a gated activation. SwiGLU multiplies a Swish-activated projection of x, elementwise, by a second linear projection (Shazeer, 2020):

The second projection acts as a gate that scales each hidden unit up or down depending on the input. Shazeer found that GLU variants such as SwiGLU gave better quality than ReLU or GELU feed-forward layers at equal parameter count and computation, which he kept equal by shrinking the hidden width to two thirds.

Mixture of Experts

A mixture-of-experts (MoE) layer replaces one feed-forward network with many "expert" networks and a router that sends each token to only a few of them (Shazeer et al., 2017). The parameter count grows with the number of experts, while the computation per token grows only with the number each token uses:

How many experts a token uses is a design choice. Switch Transformers send each token to one expert (Fedus et al., 2022), and Mixtral to two of eight (Jiang et al., 2024). DeepSeekMoE splits experts into many smaller ones, routes each token to several of them, and adds shared experts that every token uses (Dai et al., 2024). In each case the model holds far more parameters than any one token passes through, at the cost of routing and, when experts sit on different devices, communication.

The Full Transformer

Many recent language models, LLaMA among them (Touvron et al., 2023), stack copies of a pre-norm decoder block. Each sublayer reads a normalized copy of the residual stream x and adds its output back:

• 1. h = x + Attention(Norm(x)): causal multi-head self-attention, with RoPE
• 2. output = h + FFN(Norm(h)): a SwiGLU feed-forward network

Stack N of these blocks, put a token embedding before them and a final normalization and unembedding after them, and you have the layout of LLaMA, which uses RMSNorm for Norm. The figure shows the original encoder–decoder Transformer (Vaswani et al., 2017), which normalized after each residual addition instead (post-norm). Xiong et al. (2020) traced post-norm's need for a learning-rate warmup to large gradients near the output at initialization; pre-norm avoids them.

LayerNorm vs RMSNorm

Normalization keeps the scale of activations in check during training. LayerNorm (Ba et al., 2016) normalizes each token's vector across its features, using the mean μ and standard deviation σ of its entries:

Here γ and β are learned per-feature scale and shift, and RMS(x) is the root mean square of the entries. RMSNorm drops the mean subtraction. The form shown here also has no shift β; that is an implementation choice, since the general form of Zhang & Sennrich (2019) keeps a learned bias. They report quality comparable to LayerNorm with less computation, and LLaMA uses RMSNorm (Touvron et al., 2023).

Sliding Window Attention

With full attention, each new token attends to every earlier one, so both its cost and the KV cache grow with sequence length. Sliding-window attention lets each token attend only to a fixed window of recent tokens (Beltagy et al., 2020):

Stacking layers widens the reach: after L layers, information can travel up to L × w positions back. For Mistral 7B, with w = 4096 and 32 layers, that is about 131K tokens (Jiang et al., 2023). This is an upper bound on how far information can move, not a guarantee that the model uses it.

State Space Models

A state-space model carries a fixed-size hidden state h through the sequence and updates it linearly with each input (Gu et al., 2022). In discrete time:

When A, B and C are the same at every step, unrolling the recurrence makes y a convolution of x with the kernel (CB, CAB, CA²B, …), plus the Dx_t term. Training can then process the whole sequence in parallel with fast Fourier transforms in O(n log n) time, while generation runs the recurrence and keeps only h: constant time and memory per new token, however long the sequence so far.

Mamba: Selective State Spaces

In a fixed SSM every input is written into the state in the same way, so the model cannot choose what to keep based on content. Mamba (Gu & Dao, 2024) makes B, C and the step size Δ functions of the current input. A stays a learned matrix that does not depend on the input; it is discretized with the input-dependent step:

A large Δ_t lets the current input overwrite the state; a small one keeps the state and largely ignores the input. Because the parameters now change from step to step, the convolution form no longer applies. Mamba computes the recurrence with a hardware-aware parallel scan instead, which is linear in sequence length.

Tradeoffs

Costs are in sequence length n, with model width held fixed.