Efficiency Frontiers › Efficient Architectures

The original transformer architecture was designed for correctness and flexibility, not for computational efficiency. Its self-attention mechanism computes relationships between every pair of tokens, which scales quadratically with sequence length. For a 1,000-token sequence, this means one million attention computations per layer. For a 100,000-token sequence, it means ten billion. The architecture also allocates compute uniformly across all positions and all layers, even when most positions are not informative and most layers are doing similar work.

Efficient architecture research attacks these inefficiencies from several directions.

Standard attention computes a full $n \times n$ attention matrix for a sequence of $n$ tokens. The quadratic cost arises because every query token attends to every key token:

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

where $Q$, $K$, $V \in \mathbb{R}^{n \times d_k}$ are the query, key, and value matrices, and $d_k$ is the head dimension. This is $O(n^2 d_k)$ in time and $O(n^2)$ in memory just to store the attention matrix.

Sparse attention replaces the full matrix with a structured pattern. Instead of every token attending to every other token, each token attends to a limited neighborhood. The two most influential patterns are:

• Local window attention: each token attends only to the $w$ tokens to its left and right. This reduces complexity to $O(n \cdot w)$, which is linear in sequence length for fixed window size.
• Strided attention: tokens attend locally plus to tokens at fixed stride intervals, allowing global information to propagate over multiple layers.

Longformer and BigBird combine local attention for most tokens with global attention for special tokens like [CLS]. The global tokens can see all positions, giving the model a way to aggregate information across the entire sequence without paying the full $O(n^2)$ cost everywhere.

Linear attention takes a different approach. Recall that the softmax in standard attention creates a non-linear function that cannot be factored:

$$\text{softmax}(QK^T) \cdot V \neq Q \cdot (K^T V)$$

If we replace softmax with a kernel function $\phi$ such that $\phi(q)^T \phi(k)$ approximates the original attention score, we can rewrite attention as:

$$\text{Attention}(q, K, V) = \frac{\phi(q)^T \sum_i \phi(k_i) v_i^T}{\phi(q)^T \sum_i \phi(k_i)}$$

where:

• $\phi(q)$ is the feature map applied to the query vector
• $\phi(k_i)$ is the feature map applied to the $i$-th key
• $v_i^T$ is the $i$-th value row

The key insight is that $\sum_i \phi(k_i) v_i^T$ can be computed once and reused for every query, reducing the total cost to $O(n \cdot d^2)$, which is linear in sequence length. Performers and Linear Transformers use this idea with random Fourier features as the kernel approximation.

The tradeoff is that linear attention loses some expressiveness. The softmax non-linearity helps attention sharply focus on a few relevant tokens, while kernel approximations tend to produce softer distributions. In practice, linear attention works well on tasks that do not require sharp token selection, but lags on tasks requiring precise retrieval.

A standard transformer layer applies the same feed-forward network to every input token. This means the model uses the same parameters regardless of whether it is processing a Python variable name or a French sentence. Mixture-of-Experts (MoE) architectures replace the single feed-forward network with many expert sub-networks and route each token to only a few of them.

The gating mechanism uses a learned linear transformation followed by softmax:

$$g(x) = \text{softmax}(\text{TopK}(x W_g + \epsilon, k))$$

where:

• $x \in \mathbb{R}^d$ is the token embedding
• $W_g \in \mathbb{R}^{d \times E}$ is the gating weight matrix
• $\text{TopK}(\cdot, k)$ keeps only the top-$k$ logits and sets the rest to $-\infty$
• $\epsilon$ is optional noise for load balancing during training

The output of the MoE layer is:

$$y = \sum_{i \in \text{TopK}(g(x))} g_i(x) \cdot \text{Expert}_i(x)$$

The efficiency gain is striking. A model with 8 experts and top-2 routing has the same computational cost per token as a model with one feed-forward network, but four times the total parameter count. Mixtral 8x7B, for example, has 46.7B total parameters but only uses about 12.9B active parameters per token, giving it performance comparable to 70B dense models at a fraction of the inference cost.

MoE architectures introduce two key challenges. The first is load balancing: if the gating network routes all tokens to the same few experts, the unused experts contribute nothing. Training typically includes an auxiliary load-balancing loss that encourages uniform expert utilization. The second is communication overhead: in distributed training, different experts may live on different devices. Routing tokens to remote experts requires network communication, which can become a bottleneck at scale.

A different approach to the quadratic attention problem is to abandon the attention mechanism entirely and replace it with architectures that process sequences with linear or sub-linear complexity.

State Space Models (SSMs) model the sequence transformation as a linear dynamical system. The input sequence $x_1, x_2, \ldots, x_n$ is processed through a recurrence:

$$\begin{aligned} h_t &= \bar{A} h_{t-1} + \bar{B} x_t \\ y_t &= C h_t + D x_t \end{aligned}$$

where:

• $h_t \in \mathbb{R}^N$ is the hidden state
• $\bar{A}$, $\bar{B}$, $C$, $D$ are learned matrices
• The bars indicate that $A$ and $B$ are discretized from continuous-time parameters

The recurrent form processes each token in $O(N)$ time with $O(N)$ memory (constant in sequence length). SSMs can also be unrolled into a convolutional form for efficient parallel training:

$$y = K \ast x$$

where $K$ is a learned convolution kernel derived from the SSM parameters. The convolutional form enables parallel training on GPUs while the recurrent form enables efficient autoregressive inference.

Mamba is the most influential SSM architecture for language modeling. Its key innovation is selective state spaces: the transition matrices $\bar{A}$ and $\bar{B}$ depend on the input, making the model input-dependent rather than linear and time-invariant. This selective mechanism allows Mamba to focus its hidden state on the most relevant parts of the input history, similar to how attention focuses on relevant tokens. Mamba achieves competitive performance with transformers on many language tasks at $O(n)$ complexity, though its performance on tasks requiring precise token retrieval still lags behind full attention.

RetNet takes a hybrid approach, showing that transformers with decay-based attention (rather than softmax) can be computed in three equivalent ways: recurrently for inference (linear complexity), in parallel for training (quadratic complexity), and in chunks for a middle ground. This "train in parallel, infer recurrently" duality was an important insight because it separated the training computation graph from the inference computation graph.

Transformers typically apply the same number of layers to every input token. A trivial question, like "What is 2 + 2?", receives the same number of forward passes as a complex reasoning problem. This is wasteful.

Early exit architectures add classifier heads after intermediate layers. During inference, if the model's confidence at layer $l$ exceeds a threshold $\tau$, it exits early without computing the remaining layers:

$$\hat{y}_l = \text{ClassifierHead}(h_l) \quad \text{if } \max_c P(\hat{y}_l = c) > \tau$$

This reduces average inference depth on easy inputs while preserving full-depth computation for hard ones. DeeBERT and FastBERT demonstrated that classification tasks can often exit by layer 3-4 out of 12, reducing computation by 50-70% with minimal accuracy loss.

A related idea is mixture-of-depths: routing tokens to skip certain layers. If a token is already well-represented after 6 layers, it may not need to pass through layers 7-12. This allows the model to dynamically allocate depth to the tokens that need it.

Both early exit and mixture-of-depths challenge a core assumption of the standard transformer: that all computation should be homogeneous. In practice, language is heterogeneous. Some tokens are rare and semantically loaded; others are grammatical connectives that carry little information. Some sentences require careful multi-step reasoning; others are simple pattern completions. Depth efficiency exploits this heterogeneity by allocating computation adaptively. The catch is that training these models is harder: the model must learn useful representations and reliable policies for when to stop. Distillation from a full-depth teacher helps, as does using differentiable threshold mechanisms that allow the exit decision to be trained end-to-end.