Attention & Transformers

A language model must assign probabilities to sequences of tokens. A translation model must map one sequence to another. Both tasks require capturing dependencies between distant tokens: in "The animal didn't cross the street because it was too tired", resolving "it" requires attending back to "animal", not just the adjacent words.

RNNs encode this history into a fixed-size hidden state $h_t \in \mathbb{R}^d$ via: $$h_t = f(W_h h_{t-1} + W_x x_t)$$ The problem is the vanishing gradient: gradients of the loss with respect to early hidden states shrink exponentially with sequence length because they flow through repeated matrix multiplications. LSTMs and GRUs mitigate this with gating, but do not eliminate it. Even with LSTMs, encoding a 512-token sequence into a single context vector is a severe information bottleneck.

A second problem is sequential computation: $h_t$ depends on $h_{t-1}$, so you cannot compute the hidden states in parallel. On modern hardware (GPUs, TPUs), this sequential bottleneck is catastrophically slow compared to what could be achieved with full parallelism.

The Transformer solves both problems at once by abandoning recurrence. The price is quadratic memory in sequence length — a trade-off that was correct for the scales of 2017 and remains acceptable today for most tasks.

# Illustrating the information bottleneck in a seq2seq RNN encoder
import torch
import torch.nn as nn

encoder = nn.LSTM(input_size=256, hidden_size=512, batch_first=True)
x = torch.randn(1, 128, 256)   # batch=1, seq_len=128, embed=256
outputs, (h_n, c_n) = encoder(x)
# h_n shape: (1, 1, 512) — 128 tokens compressed into 512 floats
# Everything the decoder sees about the input lives here
print(h_n.shape)  # torch.Size([1, 1, 512])

Before the Transformer, Bahdanau et al. (2015) introduced attention for seq2seq as an augmentation to RNNs: at each decoder step, instead of relying solely on the final encoder hidden state, let the decoder compute a weighted average over all encoder hidden states.

The key intuition: think of attention as a database lookup, but soft and differentiable:

• Query ($Q$): what am I looking for? (decoder's current hidden state)
• Keys ($K$): what does each memory slot advertise? (encoder hidden states)
• Values ($V$): what information is actually stored? (also encoder hidden states, or a transformation of them)

The lookup works like this:

1. Compute a compatibility score between the query and each key: $e_i = \text{score}(q, k_i)$
2. Normalize with softmax: $\alpha_i = \frac{\exp(e_i)}{\sum_j \exp(e_j)}$ — these are the attention weights
3. Return the weighted sum of values: $\text{output} = \sum_i \alpha_i v_i$

The result is a differentiable operation: gradients flow through the softmax and back to whichever keys/values had high attention. The model learns which tokens to attend to.

In a hard lookup, you'd retrieve exactly one entry. In soft attention, you retrieve a convex combination of all entries — more like a fuzzy, probabilistic recall. This smoothness is what makes it trainable with gradient descent.

import torch
import torch.nn.functional as F

def soft_attention(query, keys, values):
    """query: (d,), keys: (n, d), values: (n, dv) -> output: (dv,)"""
    # Step 1: dot-product scores
    scores = keys @ query          # (n,)
    # Step 2: normalize
    weights = F.softmax(scores, dim=0)  # (n,)
    # Step 3: weighted sum
    output = weights @ values      # (dv,)
    return output, weights

# Example: 5 memory slots, d=4, dv=4
keys   = torch.randn(5, 4)
values = torch.randn(5, 4)
query  = torch.randn(4)
out, w = soft_attention(query, keys, values)
print(w)  # attention weights sum to 1

The Transformer uses a specific, elegant form of attention called scaled dot-product attention, where queries, keys, and values are all derived from linear projections of the input:

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

Here $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{m \times d_k}$, $V \in \mathbb{R}^{m \times d_v}$, where $n$ is the number of queries, $m$ is the number of key-value pairs, $d_k$ is the key/query dimension.

Why the $\sqrt{d_k}$ scaling? Consider $d_k$ independent unit-variance random variables summed: the dot product $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ has variance $d_k$, so standard deviation $\sqrt{d_k}$. Without scaling, for large $d_k$ the dot products grow large in magnitude, pushing the softmax into saturation regions where gradients are near zero — effectively vanishing gradients through the attention operation itself. Dividing by $\sqrt{d_k}$ brings the inputs to softmax back to unit variance, keeping gradients healthy.

Complexity: computing $QK^T$ is $O(n^2 d_k)$ time and $O(n^2)$ memory. For sequence length $n=1024$, $d_k=64$, the attention matrix is $1024 \times 1024 = 1M$ entries — feasible. For $n=100000$, it becomes the bottleneck that motivates sparse and linear attention research.

Masking: for autoregressive (decoder) attention, we add a causal mask: set attention scores to $-\infty$ for positions $j > i$ before the softmax, so token $i$ can only attend to tokens $\leq i$. After softmax, $-\infty$ becomes exactly 0.

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (batch, heads, seq_q, d_k)
    K: (batch, heads, seq_k, d_k)
    V: (batch, heads, seq_k, d_v)
    mask: (batch, 1, seq_q, seq_k) boolean, True = mask out
    """
    d_k = Q.size(-1)
    # Compute scores: (batch, heads, seq_q, seq_k)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask, float('-inf'))
    
    weights = F.softmax(scores, dim=-1)   # attend over seq_k
    output  = torch.matmul(weights, V)    # (batch, heads, seq_q, d_v)
    return output, weights

# Causal mask for autoregressive decoding (seq_len=5)
seq_len = 5
causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
print(causal_mask)
# tensor([[False,  True,  True,  True,  True],
#         [False, False,  True,  True,  True], ...])

In the original Bahdanau attention, the query came from the decoder and the keys/values from the encoder — cross-attention between two sequences. The Transformer's key innovation is self-attention: the queries, keys, and values all come from the same sequence.

Given an input sequence $X \in \mathbb{R}^{n \times d_{\text{model}}}$, we learn three projection matrices: $$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$ where $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$.

Now every token queries every other token's key. Token $i$'s output is a weighted sum of all tokens' values, where the weights reflect how relevant each token is to token $i$'s query.

What can self-attention capture that RNNs struggle with?

• Long-range dependencies: the path length between any two tokens is O(1) in self-attention (one attention operation), versus O(n) in an RNN. Gradients flow directly between any pair of positions.
• Syntactic structure: heads learn to attend to subject-verb agreement, pronoun antecedents, etc.
• Contextual representations: the same word "bank" gets different representations depending on context (river vs. financial), because the values it aggregates differ.

The output of self-attention for token $i$ can be thought of as: "given what token $i$ is looking for (its query), what information should it pull from the entire context?" This is fundamentally different from an RNN, which only has access to what was passed forward sequentially.

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.scale = math.sqrt(d_k)

    def forward(self, X):  # X: (batch, seq, d_model)
        Q = self.W_Q(X)    # (batch, seq, d_k)
        K = self.W_K(X)
        V = self.W_V(X)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale  # (b, seq, seq)
        weights = torch.softmax(scores, dim=-1)
        return torch.bmm(weights, V)  # (batch, seq, d_k)

# Each token's output is a blend of all tokens' values
attn = SelfAttention(d_model=64, d_k=32)
X = torch.randn(2, 10, 64)  # batch=2, seq_len=10
out = attn(X)
print(out.shape)  # (2, 10, 32)

A single attention head computes one weighted average over values. But natural language has multiple, simultaneous relationship types: syntactic dependencies, semantic similarity, coreference, positional proximity. A single head must compromise.

Multi-head attention runs $h$ attention heads in parallel, each with its own learned projections $(W_Q^i, W_K^i, W_V^i)$, then concatenates their outputs and projects back:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) , W^O$$ $$\text{head}_i = \text{Attention}(Q W_Q^i, K W_K^i, V W_V^i)$$

Typically $d_k = d_v = d_{\text{model}} / h$, so the total computation is similar to a single head at full dimension. In the original paper: $d_{\text{model}} = 512$, $h = 8$, $d_k = d_v = 64$.

Why does this work? Each head can specialize. Empirically, different heads in trained Transformers learn interpretable patterns:

• Some heads track syntactic dependencies (verb → subject)
• Some heads track coreference ("it" → "animal")
• Some heads attend to adjacent positions (local context)
• Some heads attend to positions offset by a fixed amount (capturing n-gram patterns)

The $W^O$ output projection mixes the information from all heads, allowing them to interact and compose. Multi-head attention is one of the key reasons Transformers are so expressive: they can simultaneously represent multiple relationship types at each layer.

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.h = num_heads
        self.d_k = d_model // num_heads
        # Single projection for all heads (then split)
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def split_heads(self, x, batch):
        # (batch, seq, d_model) -> (batch, h, seq, d_k)
        x = x.view(batch, -1, self.h, self.d_k)
        return x.transpose(1, 2)

    def forward(self, Q_in, K_in, V_in, mask=None):
        B = Q_in.size(0)
        Q = self.split_heads(self.W_Q(Q_in), B)  # (B, h, seq, d_k)
        K = self.split_heads(self.W_K(K_in), B)
        V = self.split_heads(self.W_V(V_in), B)
        
        scale = math.sqrt(self.d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / scale
        if mask is not None:
            scores = scores.masked_fill(mask, float('-inf'))
        weights = torch.softmax(scores, dim=-1)
        
        # Merge heads
        out = torch.matmul(weights, V)            # (B, h, seq, d_k)
        out = out.transpose(1, 2).contiguous()    # (B, seq, h*d_k)
        out = out.view(B, -1, self.h * self.d_k)  # (B, seq, d_model)
        return self.W_O(out)

mha = MultiHeadAttention(d_model=512, num_heads=8)
X = torch.randn(2, 20, 512)
print(mha(X, X, X).shape)  # (2, 20, 512)

Self-attention is permutation-equivariant: if you shuffle the input tokens, the output is shuffled the same way. There is no notion of position baked in. An RNN has position implicitly: token 5 is always processed after token 4. Attention has no such ordering.

To give the model positional information, the Transformer adds a positional encoding to each token's embedding before the first layer: $$\text{input}_i = \text{embed}(x_i) + \text{PE}(i)$$

The original paper uses sinusoidal encodings: $$\text{PE}(i, 2j) = \sin!\left(\frac{i}{10000^{2j/d_{\text{model}}}}\right)$$ $$\text{PE}(i, 2j+1) = \cos!\left(\frac{i}{10000^{2j/d_{\text{model}}}}\right)$$

Why sinusoids? Three properties:

1. Uniqueness: each position $i$ gets a unique encoding vector
2. Relative position via linear transforms: for any fixed offset $k$, $\text{PE}(i+k)$ can be expressed as a linear function of $\text{PE}(i)$ — the model can learn to attend to "the token $k$ positions to my left" via a single learned linear transformation
3. Extrapolation: since the encoding is a function (not a lookup table), it can generalize to sequence lengths longer than those seen in training

Modern Transformers often use learned positional embeddings (a trainable matrix $E_{\text{pos}} \in \mathbb{R}^{T_{\max} \times d_{\text{model}}}$, indexed by position), or more recent innovations like Rotary Position Embeddings (RoPE) which encode relative position directly in the attention computation rather than adding it to embeddings. RoPE is used in LLaMA, GPT-NeoX, and many current LLMs.

import torch
import math

def sinusoidal_encoding(seq_len, d_model):
    """Returns PE of shape (seq_len, d_model)"""
    PE = torch.zeros(seq_len, d_model)
    pos = torch.arange(seq_len).unsqueeze(1).float()  # (seq_len, 1)
    # Frequencies: 1/10000^(2j/d_model) for j = 0..d_model/2-1
    div = torch.exp(
        torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
    )
    PE[:, 0::2] = torch.sin(pos * div)   # even indices
    PE[:, 1::2] = torch.cos(pos * div)   # odd indices
    return PE

pe = sinusoidal_encoding(seq_len=50, d_model=512)
print(pe.shape)       # (50, 512)
print(pe[0, :4])      # first position, first 4 dims
print(pe[49, :4])     # last position — different encoding
# The encoding at position 0 differs from position 49
# The model learns to use these differences to infer order

The Transformer encoder takes an input sequence of token embeddings (plus positional encodings) and transforms them into a sequence of contextual representations — each token's embedding now encodes information about the entire context.

Each encoder layer has two sub-layers:

1. Multi-Head Self-Attention: as described above, $Q = K = V = $ layer input
2. Position-wise Feed-Forward Network (FFN): applied independently to each position: $$\text{FFN}(x) = \max(0,, x W_1 + b_1) W_2 + b_2$$ where $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$, $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$, and $d_{ff} = 2048$ (4$\times$ wider than $d_{\text{model}} = 512$ in the original paper).

Both sub-layers use residual connections and Layer Normalization: $$\text{output} = \text{LayerNorm}(x + \text{sublayer}(x))$$

Why residual connections? They give gradients a direct path back to early layers, enabling training of deep networks (6 layers in the original paper, up to 96+ in GPT-3). This is the same trick from ResNets.

Why LayerNorm? Unlike BatchNorm (which normalizes over the batch dimension), LayerNorm normalizes over the feature dimension for each token independently — it works well at sequence tasks where batch size can be small and sequence lengths vary.

Why the FFN? Self-attention is essentially a linear operation followed by a weighted average. The FFN adds nonlinearity and per-position computation. Research suggests the FFN layers act as a key-value memory storing factual associations, while attention heads route information between positions. The interplay of attention (routing) and FFN (storage) is central to how Transformers store and retrieve knowledge.

import torch
import torch.nn as nn

class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads)  # from above
        self.ff   = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop  = nn.Dropout(dropout)

    def forward(self, x, mask=None):       # x: (batch, seq, d_model)
        # Sub-layer 1: self-attention + residual
        attn_out = self.drop(self.attn(x, x, x, mask))
        x = self.norm1(x + attn_out)      # residual + LayerNorm
        # Sub-layer 2: FFN + residual
        ff_out = self.drop(self.ff(x))
        x = self.norm2(x + ff_out)        # residual + LayerNorm
        return x

# Stack N=6 encoder layers
class Encoder(nn.Module):
    def __init__(self, N=6, **kwargs):
        super().__init__()
        self.layers = nn.ModuleList([EncoderLayer(**kwargs) for _ in range(N)])
    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return x

For tasks like machine translation, the Transformer uses an encoder-decoder architecture. The encoder processes the source sequence into contextual representations. The decoder generates the target sequence autoregressively — one token at a time, conditioned on all previously generated tokens and the encoder output.

Each decoder layer has three sub-layers:

1. Masked Multi-Head Self-Attention: the decoder attends to its own previously generated tokens. The causal mask ensures token $i$ cannot see tokens $j > i$ — otherwise the decoder would cheat by looking at the answer.
2. Cross-Attention (Encoder-Decoder Attention): queries come from the decoder's current representation; keys and values come from the encoder output. This is the bridge between source and target: each decoder position asks "which parts of the source are relevant to generating the next target token?"
3. Position-wise FFN: same as in the encoder.

All three sub-layers use residual + LayerNorm.

During training, the decoder processes the entire target sequence in parallel using teacher forcing (feeding ground-truth tokens), with the causal mask preventing future token access. During inference, it generates autoregressively: at each step, it runs the full decoder on the tokens generated so far (or uses a KV-cache to avoid redundant computation).

For encoder-only models (BERT): use only the encoder, trained with masked language modeling. For decoder-only models (GPT): use only the decoder (with causal masking), trained with next-token prediction. The encoder-decoder architecture (T5, BART) is used for tasks with explicit input-output pairs (translation, summarization).

class DecoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn  = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.ff   = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.drop  = nn.Dropout(dropout)

    def forward(self, tgt, memory, tgt_mask=None, mem_mask=None):
        # 1. Masked self-attention (causal)
        x = self.norm1(tgt + self.drop(
            self.self_attn(tgt, tgt, tgt, tgt_mask)))
        # 2. Cross-attention: queries from decoder, keys/values from encoder
        x = self.norm2(x + self.drop(
            self.cross_attn(x, memory, memory, mem_mask)))
        # 3. FFN
        x = self.norm3(x + self.drop(self.ff(x)))
        return x

A complete Transformer for sequence-to-sequence tasks consists of:

1. Source embedding + positional encoding → encoder layers → encoder output
2. Target embedding + positional encoding → decoder layers (with cross-attention to encoder) → decoder output
3. Linear projection + softmax over vocabulary → token probabilities

The training objective is cross-entropy loss over the target sequence (teacher forcing): $$\mathcal{L} = -\sum_{t} \log P(y_t \mid y_{<t}, x)$$

Embedding weight tying: the embedding matrix (vocab $\times$ $d_{\text{model}}$) is often shared between input embedding, output embedding, and the final projection layer — a regularization trick that reduces parameters significantly.

Learning rate schedule: the original paper uses a warmup-then-decay schedule: $$\text{lr}(\text{step}) = d_{\text{model}}^{-0.5} \cdot \min(\text{step}^{-0.5},; \text{step} \cdot \text{warmup}^{-1.5})$$ This warms up the learning rate for the first ~4000 steps (important: attention weights start random and diverge easily with a large lr), then decays as inverse square root of step.

Dropout: applied to attention weights, FFN outputs, and embeddings during training. Rates of 0.1 (base) to 0.3 (large) are used.

Label smoothing: instead of one-hot targets, use $\epsilon_{LS}=0.1$, distributing 0.1 probability uniformly over the vocabulary. This prevents overconfident predictions and acts as regularization.

At inference time, decoder-only models use KV-cache: since the keys and values of previously generated tokens don't change, they are cached and reused, reducing per-step computation from $O(n^2)$ to $O(n)$ at step $n$.

# Complete minimal Transformer (encoder-decoder) skeleton
import torch
import torch.nn as nn

class Transformer(nn.Module):
    def __init__(self, vocab_src, vocab_tgt, d_model=512,
                 num_heads=8, N=6, d_ff=2048, dropout=0.1):
        super().__init__()
        self.src_embed = nn.Embedding(vocab_src, d_model)
        self.tgt_embed = nn.Embedding(vocab_tgt, d_model)
        # positional encoding (sinusoidal, not learned)
        self.encoder = Encoder(N, d_model=d_model,
                               num_heads=num_heads, d_ff=d_ff)
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(N)
        ])
        self.proj = nn.Linear(d_model, vocab_tgt)

    def encode(self, src, src_mask=None):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, tgt, memory, tgt_mask=None, mem_mask=None):
        x = self.tgt_embed(tgt)
        for layer in self.decoder_layers:
            x = layer(x, memory, tgt_mask, mem_mask)
        return x

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        memory = self.encode(src, src_mask)
        dec_out = self.decode(tgt, memory, tgt_mask, src_mask)
        return self.proj(dec_out)  # (batch, tgt_len, vocab_tgt)

Why do Transformers outperform RNNs and CNNs as scale (data, parameters, compute) increases?

Inductive bias is the set of assumptions a model class bakes in. RNNs assume locality: recent context matters more. CNNs assume translation invariance and local connectivity. These biases help when data is scarce, because they reduce the hypothesis space the model must search. But they hurt when data is plentiful, because they prevent the model from learning relationships the architecture cannot easily represent.

Transformers have weaker inductive biases: they can in principle represent any pairwise relationship between any positions, with no prior preference for local vs. distant, or sequential vs. hierarchical. This makes them harder to train (they need more data to learn basic structure like locality that RNNs get for free). But given enough data, they can learn more — they are not constrained by the architecture's prior.

The scaling laws (Kaplan et al. 2020, Chinchilla 2022) show that loss decreases predictably as a power law in data, model size, and compute — with no sign of a hard ceiling. This empirical finding made the case for investing more compute into larger Transformers, leading to GPT-3, GPT-4, and beyond.

The hardware fit is also crucial: self-attention and FFN layers reduce to matrix multiplications (GEMMs), which are maximally efficient on GPU/TPU tensor cores. An RNN's sequential dependency prevents batching over time steps. A Transformer trains all positions simultaneously — a perfect match for massively parallel hardware.

Finally, Transformers compose naturally: stack more layers, add more heads, widen the FFN. Every component communicates through the same $d_{\text{model}}$-dimensional interface. This architectural regularity makes scaling straightforward to engineer.

The original Transformer architecture has spawned a rich family of variants. Understanding the design space is essential for practitioners:

BERT (encoder-only, 2018): pre-trained with masked language modeling (predict randomly masked tokens) and next-sentence prediction. Produces bidirectional contextual representations. Used for classification, NER, QA. Fine-tuned on downstream tasks.

GPT (decoder-only, 2018–): pre-trained with causal language modeling (predict next token). Autoregressive generation. GPT-2, GPT-3, GPT-4 follow the same pattern at increasing scale.

T5 (encoder-decoder, 2020): frames every NLP task as text-to-text. Systematic study of pre-training objectives.

Flash Attention (Dao et al. 2022): rewrites attention computation to be IO-aware — tiles the computation to fit in GPU SRAM, reducing memory reads/writes from $O(n^2)$ to $O(n)$ by never materializing the full attention matrix. 2-4$\times$ speedup, same outputs. Essential for long contexts.

Rotary Position Embeddings (RoPE): encodes position by rotating query/key vectors in 2D planes in the complex plane. $q_i \cdot k_j$ depends only on $i - j$, giving clean relative position awareness. Used in LLaMA, Mistral, GPT-NeoX.

Grouped Query Attention (GQA): share key/value heads across multiple query heads to reduce KV-cache memory. Used in LLaMA-2, Mistral, Gemma.

Mixture of Experts (MoE): replace the FFN with a sparse MoE layer — a router selects $k$ of $E$ expert FFNs per token. Achieves very large model capacity with sublinear compute. Used in Mixtral, GPT-4 (reportedly), and Switch Transformer.