Beyond fixed-compute, fixed-memory transformers — models that route compute where it's needed and remember across time.
MoE · MoD · early exit — different tokens, different compute
KV cache · SSMs · external memory banks
An end-to-end implementation in PyTorch
A standard decoder-only transformer is uniform and amnesiac. Both turn out to matter at scale.
Every token gets the same number of FLOPs, no matter how easy or hard:
"hello," (probably " world") takes the same compute as predicting the next token in a complex math proof.The model has no state outside its current context window:
The two are orthogonal — you can have one, the other, or both. Most frontier systems combine them.
A learned router decides which sub-network fires for each token at each layer. Most parameters stay dormant for any given token.
Replace each FFN with N experts. Router picks top-k per token.
In the wild: Mixtral 8×7B (47B total / 13B active), GPT-4 (rumored), DeepSeek-V3, Switch Transformer
Per-token, per-block routing. Router decides which tokens enter each transformer block; the rest skip via the residual.
In the wild: Google's "Mixture-of-Depths" (Raposo et al. 2024); related: PonderNet, Adaptive Computation Time
Stop early when confident. Each layer outputs an "exit signal." When confidence crosses a threshold, generation halts.
In the wild: CALM (Schuster et al. 2022), DeeBERT, Adaptive Computation Time (Graves 2016)
class Router(nn.Module): def __init__(self, d_model, n_choices): super().__init__() self.gate = nn.Linear(d_model, n_choices, bias=False) def forward(self, x, top_k=1): scores = self.gate(x) # (B, T, n_choices) weights, idx = scores.topk(top_k, dim=-1) # pick top-k weights = F.softmax(weights, dim=-1) # normalize return weights, idx # downstream uses these
Same routing primitive used by all three patterns above. The only difference is what the choices are: experts (MoE), entry-or-skip (MoD), continue-or-exit (ACT).
Replace the dense FFN with N specialized FFNs. A router picks top-k of them per token, weighted by softmax of the gate scores.
With 8 experts and top-k=2, only ¼ of FFN weights are touched per token. Mixtral 8×7B has 47B total params but only ~13B fire per token.
import torch import torch.nn as nn import torch.nn.functional as F class Expert(nn.Module): def __init__(self, d): super().__init__() self.fc1 = nn.Linear(d, 4*d) self.fc2 = nn.Linear(4*d, d) def forward(self, x): return self.fc2(F.gelu(self.fc1(x))) class MoEFFN(nn.Module): def __init__(self, d, n_experts=8, top_k=2): super().__init__() self.top_k = top_k self.router = nn.Linear(d, n_experts, bias=False) self.experts = nn.ModuleList([ Expert(d) for _ in range(n_experts) ]) def forward(self, x): # x: (B, T, d) scores = self.router(x) # (B, T, N) # Pick top-k experts per token weights, idx = scores.topk(self.top_k, dim=-1) weights = F.softmax(weights, dim=-1) # Sparse dispatch: send each token only # to its chosen experts out = torch.zeros_like(x) for e in range(len(self.experts)): for k in range(self.top_k): mask = (idx[..., k] == e) if mask.any(): out[mask] += ( weights[..., k:k+1][mask] * self.experts[e](x[mask]) ) return out # Drop-in replacement for the FFN sublayer # in a standard transformer block: # self.ffn = MoEFFN(d=128, n_experts=8, top_k=2)
Real production MoE adds load-balancing loss (so all experts get used) and capacity factor (token dropping when an expert overflows). Omitted for clarity.
Three flavors of accumulating state — from microsecond working memory to cross-session persistent knowledge.
Working memory within one generation.
Universal — every modern LLM serving framework uses it.
Compressed running summary. Replace attention with a fixed-size state vector updated each step.
In the wild: Mamba (Gu & Dao 2023), RWKV, RetNet, Jamba (hybrid).
Persistent across sessions. Learnable key-value store the model writes to and retrieves from.
In the wild: MemGPT, Memorizing Transformers, RAG (a degenerate form), Letta.
Replace attention with a tiny recurrent state updated each step. Inference becomes linear-time and constant-memory.
h_t = A · h_{t-1} + B · x_t # update state
y_t = C · h_t # emit output
A classical recurrence. The whole sequence is summarized into h_t — a fixed-size vector. Compute per step is constant.
| Attention | SSM | |
|---|---|---|
| Compute / token | O(N) | O(1) |
| Total compute | O(N²) | O(N) |
| State / token | grows with N | fixed |
| Random access | perfect | lossy compression |
| Trainable in parallel | yes | yes (with tricks) |
Mamba's clever bit: A, B, C are made input-dependent, giving it attention-like selectivity while keeping linear complexity.
class MinimalSSM(nn.Module): """Linear recurrent layer (simplified Mamba-style).""" def __init__(self, d_model, d_state=16): super().__init__() self.d_state = d_state # Learnable state-transition matrices self.A_log = nn.Parameter(torch.randn(d_state)) self.B = nn.Linear(d_model, d_state, bias=False) self.C = nn.Linear(d_state, d_model, bias=False) def forward(self, x): # x: (B, T, d_model) B, T, _ = x.shape A = -torch.exp(self.A_log) # stable, negative h = torch.zeros(B, self.d_state, device=x.device) outputs = [] for t in range(T): # Recurrence — same shape every step h = torch.exp(A) * h + self.B(x[:, t]) y = self.C(h) outputs.append(y) return torch.stack(outputs, dim=1) # Drop-in replacement for self-attention. # Notice: no QK^T matrix → no O(N²) cost. # 'h' carries all information about the past in d_state floats.
Real Mamba: parallel scan for fast training, hardware-aware kernel, selective parameters that gate input flow. This minimal version is the conceptual core.
A learnable key-value store that lives outside the context window — written to during conversation, read from across sessions.
class MemoryBank(nn.Module): """Persistent KV bank — read by attention, write by gate.""" def __init__(self, d_model, capacity=256): super().__init__() self.capacity = capacity self.d_model = d_model # Slots stored as buffers (saved with state_dict) self.register_buffer("keys", torch.zeros(capacity, d_model)) self.register_buffer("vals", torch.zeros(capacity, d_model)) self.register_buffer("size", torch.tensor(0)) self.write_gate = nn.Linear(d_model, 1) def read(self, q): # q: (B, T, d) if self.size == 0: return torch.zeros_like(q) K = self.keys[:self.size] # (M, d) V = self.vals[:self.size] scores = q @ K.T / (self.d_model ** 0.5) # (B, T, M) return F.softmax(scores, dim=-1) @ V def write(self, x): # x: (N, d) flattened importance = torch.sigmoid(self.write_gate(x)) keep = (importance > 0.5).squeeze(-1) new = x[keep] n = min(new.size(0), self.capacity - self.size) if n == 0: return self.keys[self.size:self.size+n] = new[:n] self.vals[self.size:self.size+n] = new[:n] self.size += n def save(self, path): torch.save({"keys": self.keys, "vals": self.vals, "size": self.size}, path)
The bank can be saved/loaded across sessions — that's what makes it persistent. Real systems add eviction policies, sharding, learned compression.
A drop-in transformer block that's adaptive (MoE FFN) AND has accumulating memory (external bank). ~80 lines of PyTorch.
Both adaptations are drop-in replacements. You can stack AnimaBlocks exactly like vanilla transformer blocks.
class AnimaBlock(nn.Module): """ Transformer block with: - adaptive compute (MoE FFN) - accumulating memory (external bank) """ def __init__(self, d_model, n_heads=4, n_experts=8, top_k=2, mem_capacity=256): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = nn.MultiheadAttention( d_model, n_heads, batch_first=True) self.ln2 = nn.LayerNorm(d_model) self.ffn = MoEFFN(d_model, n_experts, top_k) self.memory = MemoryBank(d_model, mem_capacity) def forward(self, x): # 1. Recall: pull relevant past memory recall = self.memory.read(x) h = self.ln1(x + recall) # 2. Attention sublayer attn_out, _ = self.attn(h, h, h, need_weights=False) x = x + attn_out # 3. Adaptive FFN sublayer x = x + self.ffn(self.ln2(x)) # 4. Persist what mattered with torch.no_grad(): self.memory.write(x.view(-1, x.size(-1))) return x # Stack them like normal transformer blocks: blocks = nn.ModuleList([ AnimaBlock(d_model=128) for _ in range(6) ]) # Forward pass: for block in blocks: x = block(x)
Memory writes use no_grad + detach in production — you don't want gradients flowing through saved memory across batches.
A non-exhaustive map of what's actually shipping today. Notice how often "frontier model" = "vanilla transformer + at least one of these tricks".
| System | Adaptive compute | Memory accumulation | Notable |
|---|---|---|---|
| Mixtral 8×7B (Mistral) | MoE — 8 experts, top-2 | Standard KV cache | 47B total params, ~13B active. Open-weights. |
| DeepSeek-V3 | MoE — 256 experts, top-8 + 1 shared | Multi-head Latent Attention (compressed KV) | 671B total / 37B active. Open weights. |
| GPT-4 / GPT-4o (rumored) | MoE | Long-context attention | Architecture not public; widely believed to be sparse MoE. |
| Mamba / Mamba-2 | — | Selective State-Space (recurrent) | Linear-time inference; competitive with transformers up to 7B. |
| Jamba (AI21) | MoE | Hybrid: SSM + attention layers | Best of both worlds — long context + adaptive compute. |
| RWKV | — | Linear-attention recurrent state | RNN-flavored transformer; very efficient inference. |
| MemGPT / Letta | — | External memory bank with self-managed read/write | Persistent personality / facts across sessions. |
| Memorizing Transformers | — | kNN over external memory cache (last layer only) | Showed external memory works at scale (Wu et al. 2022). |
| Mixture of Depths (Google) | Per-token block routing | Standard KV cache | Variable depth per token. ~50% FLOPs reduction at iso-quality. |
Pure attention transformers are still dominant for sub-30B models, but as scale grows, every frontier system adds at least one of: MoE for cheap inference, SSMs for long context, or external memory for persistence. The "vanilla" GPT-2 architecture from 2019 is increasingly the exception, not the rule.
MoEFFN in ~30 linesMinimalSSM as a recurrence primitiveMemoryBank with gated writes & attention readsAnimaBlock wiring it all togetherQuestions?