The Encoder-Decoder Architecture

Suppose you are translating Der Hund, den die Katze jagte, bellte into English: The dog that the cat chased barked. Notice what has to happen. German puts jagte (chased) at the end of the relative clause and bellte (barked) at the very end of the sentence; English puts both verbs in different places entirely. To emit the first English word you already need to know that Der Hund is the subject of a verb you have not yet reached.

So the two sides of this problem have opposite requirements. Reading the source, you want every word to see every other word, in both directions, all at once — there is no reason to hide the end of a sentence you already have. Writing the target, you must not see the future, because at inference time the future does not exist yet: you are generating it one token at a time.

The transformer's answer is two stacks with different rules. The encoder reads with unrestricted bidirectional attention. The decoder writes with causal attention plus a third mechanism that reaches back into the encoder's output. Getting the interaction between them right — especially the difference between how it trains and how it runs — is where most implementations go wrong.

The shape of the whole thing

          SOURCE                                TARGET (shifted right)             |                                          |      token embeddings                           token embeddings      + positional enc.                          + positional enc.             |                                          |   +---------v---------+                     +----------v-----------+   |  ENCODER LAYER 1  |                     |   DECODER LAYER 1    |   |  self-attn (full) |                     |  self-attn (causal)  |   |  feed-forward     |        +----------> |  cross-attn (K,V)    |   +---------+---------+        |            |  feed-forward        |             |                  |            +----------+-----------+            ...  (x N)          |                      ... (x N)             |                  |                       |   +---------v---------+        |            +----------v-----------+   |  ENCODER LAYER N  |--------+            |   DECODER LAYER N    |   +-------------------+   memory            +----------+-----------+                                                        |                                              linear -> softmax                                                        |                                             next-token probabilities

The base configuration from the original transformer paper: $N = 6$ layers each side, $d_{model} = 512$, $h = 8$ heads, $d_{ff} = 2048$, roughly 65 million parameters. Every encoder layer has the same shape as every other; the same for decoder layers. Only the learned weights differ.

One detail that catches people: every encoder layer output does not go to the matching decoder layer. The final encoder output — often called the memory — goes to all decoder layers. There are $N$ separate cross-attention modules but they all read the same memory tensor.

The encoder stack

What it is for

The encoder turns a sequence of token embeddings into a sequence of contextual representations, one per input position, same length in and out. Position $i$ of the output is a vector representing "the token at position $i$, as understood in the context of this whole sentence."

Bidirectionality is the point. Consider bank in these two sentences:

Both disambiguators come after the ambiguous word. A left-to-right model must commit to a representation of bank before seeing either. The encoder has no such constraint, so its representation of bank is built with loan or river already in view.

Sub-layer 1: multi-head self-attention

Every position attends to every position, itself included. For a source of length $T_{src}$ this produces a $T_{src} \times T_{src}$ attention matrix per head. The only masking is padding — blocking the [PAD] columns of short sequences in a batch — and never a causal mask.

A concrete size check. Batch 32, $T_{src} = 128$, 8 heads: the score tensor is $32 \times 8 \times 128 \times 128 = 4{,}194{,}304$ entries, 8 MiB in fp16. Comfortable. At $T_{src} = 1024$ the same tensor is 512 MiB per layer, which is why long source documents get expensive fast.

Sub-layer 2: the position-wise feed-forward network

Two linear layers with a nonlinearity between them, applied independently and identically to each position. "Position-wise" means exactly that: it is the same small network run at every position with no mixing between them. If you implemented it with a 1-D convolution of kernel size 1, you would get the same thing.

The dimensions matter. $W_1$ maps $512 \to 2048$ and $W_2$ maps $2048 \to 512$ — it expands four-fold and comes back. Count the parameters:

The feed-forward network holds twice as many parameters as the attention block. Two thirds of a transformer's parameters are in the FFNs. This surprises people who think of transformers as "attention models".

What is it doing? Attention mixes information across positions but is, per position, a linear combination of value vectors. Without the FFN the whole stack would collapse into something close to a single linear map plus softmax mixing. The FFN supplies the per-position nonlinear processing — the place where a mixed representation gets transformed into something new. A useful mental model: attention decides what to look at, the FFN decides what to think about it.

Encoder in code

1import torch2import torch.nn as nn34class FeedForward(nn.Module):5    def __init__(self, d_model, d_ff, dropout=0.1):6        super().__init__()7        self.net = nn.Sequential(8            nn.Linear(d_model, d_ff),9            nn.ReLU(),10            nn.Dropout(dropout),11            nn.Linear(d_ff, d_model),12        )1314    def forward(self, x):15        return self.net(x)1617class EncoderLayer(nn.Module):18    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):19        super().__init__()20        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)21        self.ff = FeedForward(d_model, d_ff, dropout)22        self.norm1 = nn.LayerNorm(d_model)23        self.norm2 = nn.LayerNorm(d_model)24        self.drop = nn.Dropout(dropout)2526    def forward(self, x, src_mask=None):27        # residual + norm around each sub-layer28        a, attn = self.self_attn(x, x, x, mask=src_mask)29        x = self.norm1(x + self.drop(a))30        f = self.ff(x)31        x = self.norm2(x + self.drop(f))32        return x, attn3334class Encoder(nn.Module):35    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1):36        super().__init__()37        self.layers = nn.ModuleList(38            EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)39        )4041    def forward(self, x, src_mask=None):42        maps = []43        for layer in self.layers:44            x, attn = layer(x, src_mask)45            maps.append(attn)46        return x, maps          # x is the "memory"

Note that self_attn(x, x, x) — the same tensor three times — is what makes it self-attention.

The decoder stack

Three sub-layers, not two

Sub-layer 1: causal self-attention

Position $i$ may attend only to positions $\le i$. This is enforced by adding $-\infty$ to the scores of future positions before softmax.

Work it through for a four-token target at position 3 (0-indexed position 2). Scaled scores $[2.0,\ 1.0,\ 3.0,\ 5.0]$, with position 4 in the future:

Without the mask, 83% of the attention at this position goes to the token the model is being asked to predict. Training loss would drop near zero and generation would produce noise, because at inference that column is empty. Note also that the masked weights renormalise to sum to 1 over the legal positions — which is why you must mask the scores rather than zeroing the weights afterwards.

A causal mask is not a regulariser or a heuristic. It is the constraint that makes training conditions match inference conditions; break it and the two diverge completely.

Sub-layer 2: cross-attention

This is where the two stacks meet, and its asymmetry is the thing to internalise:

The attention matrix is $T_{tgt} \times T_{src}$ — rectangular, not square. Row $i$ says: while producing target token $i$, how much did I draw on each source token?

Take the German example. Source tokens: Der Hund , den die Katze jagte , bellte. When the decoder emits barked, a trained model's cross-attention row typically looks like:

Sum: $0.04+0.11+0.01+0.01+0.01+0.02+0.03+0.01+0.76 = 1.00$. The mass sits on bellte, with a secondary lobe on Hund — the subject it agrees with. Nobody supplied word alignments; this is learned from sentence pairs alone. This is also why cross-attention maps are the single most useful diagnostic when a translation model misbehaves.

No causal mask here. The whole source exists from the start, so hiding parts of it would only remove information. The only mask is source padding.

Decoder in code

1class DecoderLayer(nn.Module):2    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):3        super().__init__()4        self.self_attn  = MultiHeadAttention(d_model, num_heads, dropout)5        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)6        self.ff = FeedForward(d_model, d_ff, dropout)7        self.norm1 = nn.LayerNorm(d_model)8        self.norm2 = nn.LayerNorm(d_model)9        self.norm3 = nn.LayerNorm(d_model)10        self.drop  = nn.Dropout(dropout)1112    def forward(self, x, memory, tgt_mask=None, src_mask=None):13        a, self_map = self.self_attn(x, x, x, mask=tgt_mask)14        x = self.norm1(x + self.drop(a))1516        # Q from the decoder, K and V from the encoder memory17        c, cross_map = self.cross_attn(x, memory, memory, mask=src_mask)18        x = self.norm2(x + self.drop(c))1920        f = self.ff(x)21        x = self.norm3(x + self.drop(f))22        return x, self_map, cross_map

The argument order (x, memory, memory) in the cross-attention call is the whole architecture in one line. Swap it to (memory, x, x) and you get a shape error if the lengths differ — or, if source and target happen to be the same length, no error at all and a model that quietly learns the wrong thing.

Training and inference are different programs

Teacher forcing

During training the correct target is known, so the whole target sequence is fed in at once and every position is predicted in parallel. The input is the target shifted right by one, with a start token prepended.

Position 3 receives dog as input, can attend to positions 1–3 only, and must predict barked. Because the causal mask prevents it from seeing positions 4 and 5, all five predictions can be computed in a single forward pass without any of them cheating. This is the entire reason transformers train faster than recurrent models: teacher forcing plus causal masking turns a sequential problem into a parallel one.

1decoder_input = target[:, :-1]     # drop the final token2labels        = target[:, 1:]      # drop the start token34logits = model(source, decoder_input)                     # (B, T-1, vocab)5loss = F.cross_entropy(6    logits.reshape(-1, logits.size(-1)),7    labels.reshape(-1),8    ignore_index=PAD_ID,9)

The off-by-one here is the single most common bug in sequence-to-sequence code. If you forget the shift and feed the full target, position $i$ sees token $i$ and is asked to predict token $i$ — the model learns the identity function, training loss drops to near zero within an epoch, and generation emits the start token forever.

Autoregressive generation

At inference there is no target. You generate one token, append it, and run again:

1@torch.no_grad()2def greedy_decode(model, source, src_mask, max_len, sos_id, eos_id):3    memory, _ = model.encode(source, src_mask)        # encoder runs ONCE4    ys = torch.full((source.size(0), 1), sos_id, device=source.device)56    for _ in range(max_len - 1):7        tgt_mask = causal_mask(ys.size(1), ys.device)8        out = model.decode(ys, memory, tgt_mask, src_mask)9        logits = model.generator(out[:, -1])          # only the last position matters10        next_tok = logits.argmax(-1, keepdim=True)11        ys = torch.cat([ys, next_tok], dim=1)12        if (next_tok == eos_id).all():13            break14    return ys

That last row is exposure bias: the model is trained on perfect prefixes and deployed on its own imperfect ones. A model with 95% per-token accuracy under teacher forcing has, naively, a $0.95^{20} \approx 0.36$ chance of producing a 20-token sequence with no errors, and each error shifts the input distribution further from anything it saw in training. Beam search mitigates this by keeping several hypotheses alive rather than committing to the argmax at every step.

The $T$ decoder passes are also where inference cost lives. Each pass recomputes keys and values for every previous position, which is $O(T^2)$ work over the whole generation. A KV cache stores them instead: at step $i$ you compute K and V for the new token only and concatenate onto the cache, reducing the total to $O(T)$ recomputation. Every production inference stack does this.

Encoder and decoder side by side

You do not always need both stacks

Decoder-only models are worth a second look, because they have become dominant and the reason is not obvious. If you delete cross-attention from a decoder layer you get exactly an encoder layer with a causal mask. So a decoder-only model handles source-to-target tasks by concatenating them into one sequence:

Translate to English: Der Hund bellte => The dog barked

The causal mask means the source tokens can only see earlier source tokens — strictly less context than a bidirectional encoder would give them. That is a genuine disadvantage on understanding tasks. What decoder-only models buy in exchange is uniformity: one stack, one objective (next-token prediction), any text corpus as training data, no need for paired input-output examples. At scale that trade has generally been worth it.

What this means when you build something

Pick the architecture from the shape of the task before anything else. If your output is a label, an encoder-only model is smaller, faster, and bidirectional — do not reach for a generative model to do classification. If your output is free text conditioned on a specific input document, an encoder-decoder gives the input full bidirectional context and keeps the two sequences cleanly separated. If your task is open-ended continuation, decoder-only.

Then check these four things, in this order, because each has a characteristic and misleading symptom:

And keep the parameter arithmetic in mind when you size a model. Doubling $d_{ff}$ from 2048 to 4096 adds about 2.1 M parameters per layer — more than doubling the head count would, which adds none at all. If you need capacity, the feed-forward width is usually where it should go.