Beyond Transformers: GNNs, State-Space Models and Hybrids

Why look beyond the transformer

Transformers dominate, but they make two assumptions that do not always fit. They treat input as a sequence (or a set) of tokens, while molecules, road networks, supply chains and social graphs are graphs with explicit connections. And their attention cost grows with the square of sequence length while the KV cache grows linearly, which makes very long contexts and high-concurrency serving expensive. This lesson covers the main architectural answers: graph neural networks for relational data, state-space models and linear attention for linear-time sequence modelling, the hybrid models that mix them with attention, mixture-of-experts as a way to scale parameters without scaling compute, and transformers used as diffusion denoisers. The goal is to know what each one buys you and when to reach for it.

Intuition: match the architecture to the shape of the data and the cost you can pay

A CNN assumes a grid, an RNN assumes a chain, and a transformer assumes every token may need to talk to every other token. A graph neural network assumes information flows along given edges — like gossip spreading only between friends. A state-space model is a return to the RNN's running summary, rebuilt so that it trains in parallel and chooses what to remember. A hybrid model is a pragmatic compromise: most layers keep a cheap fixed-size memory, and a few layers keep full attention for the moments when the model must look something up exactly.

Part 1: Graph neural networks

A graph $G = (\mathcal{V}, \mathcal{E})$ has $N$ nodes with feature vectors $x_v \in \mathbb{R}^{d}$, and edges, encoded in an adjacency matrix $A \in \{0,1\}^{N \times N}$ with $A_{uv} = 1$ if $u$ and $v$ are connected. Tasks are node-level (classify each account as fraud or not), edge-level (predict a missing link: will this user buy this item?) or graph-level (predict a molecule's toxicity).

Message passing

Almost every GNN fits the message-passing template (formalized by Gilmer et al., 2017). In layer $l$, each node gathers messages from its neighbours $\mathcal{N}(v)$, aggregates them with a permutation-invariant function, and updates its state:

$\bigoplus$ is sum, mean or max (it must not depend on neighbour order, since graphs have none), $\psi$ and $\phi$ are small neural networks, and $e_{uv}$ are optional edge features. After $L$ layers a node's representation depends on its $L$-hop neighbourhood — the graph analogue of a receptive field. For graph-level tasks, a readout (sum or mean over all nodes) produces one vector per graph.

Graph convolutional network (GCN)

The GCN of Kipf and Welling (2016) is the simplest widely used instance. Add self-loops, $\tilde{A} = A + I$, let $\tilde{D}$ be the diagonal degree matrix of $\tilde{A}$, and set

$H^{(l)} \in \mathbb{R}^{N \times d_l}$ stacks node states, $W^{(l)}$ is a learned weight matrix shared by all nodes (weight sharing again), and $\sigma$ is a nonlinearity. The normalized matrix $\hat{A} = \tilde{D}^{-1/2}\tilde{A}\tilde{D}^{-1/2}$ has entries $\hat{A}_{uv} = 1/\sqrt{\tilde{d}_u \tilde{d}_v}$ for connected $u, v$: each node averages itself and its neighbours, down-weighting high-degree hubs so they do not dominate.

Worked example. Path graph 1 — 2 — 3 with scalar features $x = (1, 2, 3)$ and $W = 1$. With self-loops the degrees are $\tilde{d} = (2, 3, 2)$. So $\hat{A}_{11} = 1/2$, $\hat{A}_{12} = 1/\sqrt{6} = 0.408$, $\hat{A}_{22} = 1/3$, $\hat{A}_{23} = 0.408$, $\hat{A}_{33} = 1/2$.

• Node 1: $0.5 \cdot 1 + 0.408 \cdot 2 = 1.316$
• Node 2: $0.408 \cdot 1 + 0.333 \cdot 2 + 0.408 \cdot 3 = 2.300$
• Node 3: $0.408 \cdot 2 + 0.5 \cdot 3 = 2.316$

Each node's new value blends its own feature with its neighbours'. Repeat this many times and all nodes converge toward similar values — over-smoothing, the reason most GNNs use only 2–4 layers.

import torch, torch.nn as nn

class GCNLayer(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.lin = nn.Linear(d_in, d_out, bias=False)
    def forward(self, X, A):                         # X: (N, d_in); A: (N, N) 0/1
        A_t = A + torch.eye(A.shape[0])              # self-loops
        d_inv_sqrt = A_t.sum(1).pow(-0.5)
        A_hat = d_inv_sqrt[:, None] * A_t * d_inv_sqrt[None, :]   # D^-1/2 A D^-1/2
        return A_hat @ self.lin(X)

A = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=torch.float32)
X = torch.tensor([[1.0], [2.0], [3.0]])
layer = GCNLayer(1, 1)
with torch.no_grad():
    layer.lin.weight.fill_(1.0)
    print(layer(X, A).squeeze())                     # tensor([1.3165, 2.2997, 2.3165])

class GCN(nn.Module):                                # 2-layer node classifier
    def __init__(self, d_in, d_hidden, n_classes):
        super().__init__()
        self.l1, self.l2 = GCNLayer(d_in, d_hidden), GCNLayer(d_hidden, n_classes)
    def forward(self, X, A):
        return self.l2(torch.relu(self.l1(X, A)), A)  # per-node logits

Dense adjacency is fine for teaching; real graphs are sparse, so libraries such as PyTorch Geometric and DGL use edge lists and scatter operations. Other important variants: GraphSAGE samples neighbours to scale to huge graphs; graph attention networks (GAT) learn per-edge weights with attention; graph transformers let all nodes attend to each other with structural position encodings. A transformer is in fact a GNN on a fully connected graph. GNNs are used for molecular property prediction, recommendation over user–item graphs, fraud rings in transaction graphs, traffic forecasting and chip design.

Part 2: State-space models

The continuous model and its two faces

A linear state-space model (SSM) maps an input signal $x(t)$ to an output $y(t)$ through a hidden state $h(t) \in \mathbb{R}^N$:

with $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$, $C \in \mathbb{R}^{1 \times N}$. To apply it to tokens, discretize with a step size $\Delta$. The zero-order-hold rule gives

This is a linear RNN: constant work and memory per step at inference. Because $\bar{A}, \bar{B}, C$ are the same at every step (linear time-invariant), unrolling gives $y_t = \sum_{k \ge 0} C \bar{A}^k \bar{B}\, x_{t-k}$, which is a convolution with kernel $\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, \dots)$. Training can therefore run as one long parallel convolution (via FFT), and inference as a recurrence.

Worked example. Scalar $A = -1$, $B = C = 1$, $\Delta = 0.5$. Then $\bar{A} = e^{-0.5} = 0.6065$ and $\bar{B} = \frac{1}{-0.5}(0.6065 - 1)(0.5) = 0.3935$. Inputs $x = (1, 0, 0, 2)$:

• Recurrence: $h_1 = 0.3935$, $h_2 = 0.6065 \cdot 0.3935 = 0.2387$, $h_3 = 0.1447$, $h_4 = 0.6065 \cdot 0.1447 + 0.3935 \cdot 2 = 0.0878 + 0.7869 = 0.8747$.
• Convolution: $\bar{K} = (0.3935, 0.2387, 0.1447, 0.0878)$, so $y_4 = \bar{K}_0 x_4 + \bar{K}_3 x_1 = 0.7869 + 0.0878 = 0.8747$. Same answer.

The state is a decaying memory: the impulse at step 1 fades by a factor of 0.6065 per step.

S4: making SSMs work for long sequences

Gu, Goel and Ré's S4 (2021) made deep SSMs practical. It initializes $A$ with a structure derived from the HiPPO theory of compressing history onto orthogonal polynomials, and parameterizes $A$ (normal plus low-rank) so the convolution kernel can be computed efficiently and stably. S4 reached state-of-the-art results on every task of the Long Range Arena benchmark, including the 16k-length Path-X task that earlier models had failed. But on language it still trailed transformers.

Mamba: selective state spaces

The weakness of time-invariant SSMs is that they treat every token the same; they cannot decide, based on content, to remember this name and ignore that filler word. Mamba (Gu and Dao, 2023) makes the SSM selective: the step size $\Delta_t$ and the matrices $B_t$ and $C_t$ become functions of the current input $x_t$ (via small linear projections), while $A$ stays a learned diagonal matrix with negative entries:

(Mamba uses this simplified $\bar{B} \approx \Delta B$.) A large $\Delta_t$ drives $e^{\Delta_t A}$ toward 0 — forget the past and focus on this token; a small $\Delta_t$ keeps the state and nearly ignores the input. This is an input-dependent gate, much like an LSTM's forget gate, but on a linear recurrence.

The price is that the model is no longer time-invariant, so the convolution trick is gone. Mamba instead uses a hardware-aware parallel scan: the linear recurrence is associative, so it can be computed in $O(\log T)$ parallel depth, and the kernel keeps the expanded state in fast SRAM rather than writing it to HBM — the same IO-awareness idea as FlashAttention. The paper reports training cost linear in sequence length, inference throughput about 5× that of transformers of similar size, and improving results on real data up to million-token sequences. Mamba-2 (Dao and Gu, 2024) restricted $A$ further to a scalar per head, showed a formal duality between this SSM and a form of masked linear attention, and used that duality for a faster, matrix-multiply-based training algorithm.

import torch, torch.nn as nn, torch.nn.functional as F

class NaiveSelectiveSSM(nn.Module):
    """Mamba-style selective SSM as a slow sequential loop, for understanding only."""
    def __init__(self, d, N=16):
        super().__init__()
        self.A_log = nn.Parameter(torch.log(torch.arange(1, N + 1).float()).repeat(d, 1))
        self.to_B, self.to_C, self.to_delta = nn.Linear(d, N), nn.Linear(d, N), nn.Linear(d, d)
    def forward(self, x):                                   # x: (B, T, d)
        Bsz, T, d = x.shape
        A = -torch.exp(self.A_log)                          # (d, N), negative = decaying
        delta = F.softplus(self.to_delta(x))                # (B, T, d), input-dependent
        Bt, Ct = self.to_B(x), self.to_C(x)                 # (B, T, N), input-dependent
        h = x.new_zeros(Bsz, d, A.shape[1])                 # fixed-size state
        ys = []
        for t in range(T):
            dA = torch.exp(delta[:, t, :, None] * A)        # (B, d, N)
            dB = delta[:, t, :, None] * Bt[:, t, None, :]   # (B, d, N)
            h = dA * h + dB * x[:, t, :, None]
            ys.append((h * Ct[:, t, None, :]).sum(-1))      # (B, d)
        return torch.stack(ys, dim=1)                       # (B, T, d)

print(NaiveSelectiveSSM(d=8)(torch.randn(2, 20, 8)).shape)  # torch.Size([2, 20, 8])

Each of the $d$ channels runs its own $N$-dimensional state. A real Mamba block wraps this in an input projection, a short causal convolution, a gating branch and an output projection, and replaces the Python loop with the fused scan kernel.

Linear attention and the delta rule

A parallel line of work reaches the same place from the attention side. Replace $\operatorname{softmax}(q_t^\top k_s)$ with a kernel $\phi(q_t)^\top \phi(k_s)$; then causal attention becomes

— a recurrence with a fixed-size matrix state $S_t \in \mathbb{R}^{d_k \times d_v}$. Plain additive updates never forget and blur memories together, so newer variants add gates (decay of $S$) and the delta rule, which overwrites the value stored under a key instead of adding to it: in a common form, $S_t = S_{t-1} + \beta_t\, k_t (v_t - S_{t-1}^\top k_t)^\top$ with a learned write strength $\beta_t$. Gated DeltaNet (Yang et al., 2024) combines gating and the delta rule, and Kimi Delta Attention is a refinement of it.

The fundamental trade-off

A fixed-size state is what makes these models cheap, and it is also their limit: a state of a few megabytes cannot hold an exact copy of an arbitrarily long context. Pure recurrent models do well on language modelling perplexity but tend to struggle on tasks that require exact retrieval or copying of specific earlier content, where attention, with its growing cache, can simply look back. That observation drives the hybrid designs below.

Part 3: Hybrid architectures

A hybrid model interleaves a minority of full-attention layers with a majority of SSM or linear-attention layers. The attention layers provide precise lookup; the others provide cheap, constant-memory sequence mixing. As of November 2026, publicly documented examples include:

The memory effect is easy to compute with the KV-cache formula $2 \times \text{layers} \times \text{KV heads} \times \text{head dim} \times \text{bytes}$. A 32-layer transformer with 8 KV heads of dimension 128 in 16-bit stores $2 \times 32 \times 8 \times 128 \times 2 = 131{,}072$ bytes (128 KiB) per token, which is 32 GiB at 262,144 tokens. Keep attention in only 4 of 32 layers (a 1:7 ratio) and the cache becomes 16 KiB per token, or 4 GiB — an 8× reduction, matching Jamba's reported ratio — plus a fixed recurrent state for the other layers that does not grow with context. Faster decoding follows, because decode speed is largely bounded by bytes read per token (see tokens-per-second).

The trade-offs: fewer mature kernels and runtimes support SSM layers (IBM ships a pure-transformer Granite 4.0 Micro for platforms without hybrid support), quantization and serving tools lag behind, and the right ratio and placement of attention layers are still active research.

Part 4: Mixture-of-experts, recapped

Mixture-of-experts (MoE) is orthogonal to the attention-versus-recurrence question: it changes the MLP. A router $g(x) = W_r x$ scores $E$ expert MLPs; the token is sent to the top $k$ and the outputs are combined with normalized router weights:

Compute per token scales with $k$, while total parameters scale with $E$, so MoE buys capacity without proportional FLOPs. An auxiliary load-balancing loss keeps the router from sending everything to a few experts, and many designs add an always-on shared expert. Note from the table that hybrids and MoE combine naturally (Jamba, Granite 4.0-H Tiny/Small, Qwen3-Next). The memory catch: all experts must be resident, so memory follows total parameters while speed follows active parameters. Details are in mixture-of-experts.

Part 5: Diffusion transformers

Diffusion image generators were long built on U-Net convolutional denoisers. Peebles and Xie's Diffusion Transformer (DiT) (2022) showed that a plain transformer works as the denoiser: the image's latent representation (from a pretrained autoencoder) is cut into patches exactly as in a ViT, the timestep and class label condition each block through adaptive layer norm, and the output patches predict the noise. DiT's quality improved consistently with forward-pass compute (Gflops), from deeper or wider models or more tokens, and DiT-XL/2 reached an FID of 2.27 on class-conditional ImageNet 256×256, state of the art at the time. Transformer denoisers have since become common in large image and video generators; see diffusion-models.

Where each architecture wins

Common pitfalls

• Treating graph data as a table: flattening relational data throws away the structure a GNN would exploit — and conversely, forcing a GNN onto data whose edges carry no signal.
• Deep GNNs: stacking many GCN layers over-smooths node features; prefer 2–4 layers plus residuals or jumping-knowledge connections.
• Leakage in graph splits: in link prediction, test edges must be removed from the message-passing graph.
• Judging SSMs by perplexity alone: evaluate retrieval and long-context tasks too, where fixed-state models may lag.
• Assuming a hybrid runs everywhere: check that your inference runtime, quantization format and hardware support its layer types before committing.
• Confusing MoE active and total parameters when sizing memory.

In practice

• Use PyTorch Geometric or DGL for GNNs, start with a 2-layer GCN or GraphSAGE baseline, and compare against a gradient-boosted model on hand-built graph features.
• For long-context or high-concurrency LLM serving, compare candidate models on cost per token at your real context lengths; hybrid and linear-attention models can win by large factors there while matching transformers at short context.
• Verify claims in model cards against your own evaluations, especially long-context recall.
• Keep transformers as the default when tooling, ecosystem and exact retrieval matter more than raw efficiency.