The Full Transformer
Encoder-decoder with cross-attention
25.1 What You Will Learn
Chapter 24 built one transformer block and stopped there, which was the right place to stop because a block is the repeating unit and everything else is arrangement. This chapter does the arranging. The architecture from Vaswani and colleagues in 2017 has two stacks rather than one, an encoder that reads the input and a decoder that writes the output, and a mechanism connecting them that Chapter 19 already introduced under a different name. Getting from a single block to that architecture requires two additions and one deletion, and all three are small.
The first addition is a mask. A decoder writes its output one token at a time and must never see a token it has not yet produced, so its self-attention needs a way to blank out every position after the current one. The mask does that with a single comparison and a large negative number, and we will see shortly how few lines it takes. The second addition is cross-attention, which is the same scaled dot-product mechanism from Chapter 20 with the queries taken from one sequence and the keys and values from another, so that each decoder position can read whichever parts of the input are relevant to the word it is currently producing. The deletion is that the encoder gets no mask at all, since it has the whole input available from the start and no reason to hide any of it.
One point of vocabulary before starting, since the word decoder is doing double duty in this book. Chapter 18 used it for the recurrent half of a sequence to sequence model, and here it means a stack of transformer blocks with a mask and a cross-attention sublayer. The two share a job rather than a mechanism, and the transformer version has no hidden state carried between steps at all.
Assembling those pieces gives the architecture that started everything. It is not the architecture most current language models use, and Chapter 26 explains why the field mostly dropped the encoder, but the encoder-decoder arrangement is still what you reach for when the input and the output are genuinely different sequences, which covers translation, summarization and a large amount of what production systems do. Building it here also makes Chapter 26 a subtraction rather than a fresh start.
25.2 The Full Architecture
The original transformer contains three attention operations that look identical in code and differ entirely in what they are allowed to see. Encoder self-attention lets every input position attend to every other input position in both directions, which is what Chapter 24 built and what makes an encoder good at understanding and useless at generating. Decoder self-attention lets each output position attend only to positions at or before itself, which is the causal or masked variety, and the restriction exists because at generation time the later positions genuinely do not exist yet. Cross-attention lets each decoder position attend to every encoder position, with queries coming from the decoder and keys and values from the encoder, and it is the only channel through which the input reaches the output.
Figure 25-1 marks all three on the architecture. They are the same operation in code, and the numbers say what each one is allowed to look at. The two stacks around them are what Chapters 22 through 24 already built, which is why the only new material in this chapter is the mask and the wiring that carries encoder output across to the decoder.
The stacks are built from those. An encoder block is what Chapter 24 produced, self-attention followed by a feedforward network, each wrapped in a normalization and a residual addition. A decoder block inserts a third sublayer between the two, so it runs causal self-attention, then cross-attention, then the feedforward network, with the same wrapping around all three. That extra sublayer is the entire structural difference between the two stacks, and it costs a decoder block roughly half again as many attention parameters as an encoder block.
The asymmetry between the two stacks is larger than it looks on a diagram. An encoder runs once over a fixed input and every position is computed in parallel, so its cost is one forward pass regardless of how the output is produced. A decoder runs once per output token during generation, and each run must redo its causal self-attention over everything produced so far. For a task where the input is long and the output is short, summarization for instance, that split is efficient and the encoder earns its parameters. For a task where the output is as long as the input, or where there is no separate input at all, the encoder is carrying weight for no return, which is the beginning of the argument Chapter 26 finishes.
Notice what the ordering inside a decoder block implies. Causal self-attention comes first, so a decoder position gathers context from the output produced so far before it goes looking at the input. Cross-attention comes second, so the query it sends to the encoder is informed by that gathered context rather than by the current token alone. A decoder generating the fourth word of a translation therefore asks the source sentence a question shaped by the three words it has already committed to, which is a different and better question than the one the raw fourth embedding would have asked.
25.3 The Causal Mask
The mask is the smallest piece of machinery in the chapter and the one with the largest consequence. Before the softmax runs, every score belonging to a position later than the current one is replaced with a very large negative number, which the exponential inside softmax then maps to essentially zero, leaving those positions with no weight at all. Implementations write that large negative number rather than an actual negative infinity, and the choice is practical. An exact infinity produces a NaN the moment it meets a zero or another infinity anywhere in the arithmetic, and a value like negative 1e9 exponentiates to zero in single precision without any of that risk. The program below runs the same attention twice on the same input, once without the mask and once with it, so that the mask is the only thing that differs between the two tables.
/* 131_Causal_Mask.c */
#include <stdio.h>
#include <math.h>
#include <float.h>
#define SEQ 5
static void softmax(float *x, int n)
{
float mx = -FLT_MAX, s = 0;
int i;
for (i = 0; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
int main(void)
{
/* Raw attention scores (before masking) */
float scores[SEQ][SEQ];
int i, j;
/* Fill with uniform scores */
for (i = 0; i < SEQ; i++)
for (j = 0; j < SEQ; j++)
scores[i][j] = 1.0f;
printf("Without causal mask (bidirectional):\n\n");
for (i = 0; i < SEQ; i++) {
float row[SEQ];
for (j = 0; j < SEQ; j++) row[j] = scores[i][j];
softmax(row, SEQ);
printf(" pos %d attends to: ", i);
for (j = 0; j < SEQ; j++)
printf("%.2f ", row[j]);
printf("\n");
}
printf("\nWith causal mask (decoder-style):\n\n");
for (i = 0; i < SEQ; i++) {
float row[SEQ];
for (j = 0; j < SEQ; j++) {
if (j > i)
row[j] = -1e9f;
/* mask future positions */
else
row[j] = scores[i][j];
}
softmax(row, SEQ);
printf(" pos %d attends to: ", i);
for (j = 0; j < SEQ; j++)
printf("%.2f ", row[j]);
printf("\n");
}
printf("\nPosition 0 can only see itself.\n");
printf("Position 4 can see all 5 positions.\n");
printf("This is autoregressive, each position\n");
printf("generating from what came before it.\n");
return 0;
}
Figure 25-2 runs the same attention twice, once masked and once not. The unmasked table is uniform, with every position giving 0.20 to every other position including itself, and that uniformity is deliberate rather than a failure of the example. The inputs were chosen so that all the raw scores come out equal, which means softmax has nothing to distinguish them and the only structure that can appear in the second table is structure the mask put there. Isolating one mechanism by flattening everything else makes for a clean demonstration, and recognizing the difference matters because Chapter 19 and Chapter 20 both contained listings where uniform weights were an accident rather than a choice.
The masked table is the triangle. Position 0 gives 1.00 to itself and nothing to anyone else, because it is the first token and there is nothing before it, so softmax over a single surviving score must return 1. Position 1 splits 0.50 and 0.50 across the two positions available to it. Position 2 gives 0.33 three ways, position 3 gives 0.25 four ways, and position 4, which can see the whole sequence, returns to the same 0.20 across five positions that the unmasked run produced. Reading down the diagonal shows the pattern clearly, since every row is a uniform distribution over exactly the prefix that position is permitted to see.
That last row is worth dwelling on because it explains something about how transformers train. The final position of a masked sequence sees everything, so its computation is identical to the unmasked case, and the restriction only bites for earlier positions. This is why a decoder can be trained on a whole sequence in one pass rather than one token at a time. Feed the entire target sentence in, apply the mask, and every position simultaneously computes the prediction it would have made had it only seen its own prefix. Position 3 predicts token 4 while position 7 predicts token 8, in the same forward pass, with no leakage between them. Remove the mask and that trick collapses immediately, because every position would see the answer it was supposed to predict and the model would learn to copy rather than to continue. The cost shows up at generation time instead, where the parallelism is unavailable regardless. Producing a hundred token output requires a hundred forward passes, each one a token longer than the last, because token 50 cannot be computed until token 49 exists. Training is parallel over positions and inference is not, and that asymmetry is behind most of the engineering effort that goes into serving language models.
25.4 Cross-Attention
Cross-attention is the operation from Chapter 20 with one change in the wiring. There the queries, keys and values were all projected from the same sequence, which is what self-attention means. Here the queries come from the decoder and the keys and values come from the encoder, so a decoder position is asking a question of the input rather than of its own output. Nothing in the arithmetic changes, no new formula is needed, and the same scaled dot-product and softmax and weighted sum run exactly as before.
The projections below were chosen rather than trained, for the reason given in Chapters 20 and 21. Random weights produce near uniform attention that demonstrates the plumbing and hides the point, and the point here is that different target words read different parts of the source.
/* 132_Cross_Attention.c */
#include <stdio.h>
#include <math.h>
#define ENC_LEN 5
#define DEC_LEN 3
#define DM 6
#define DK 4
static void softmax(float *x, int n)
{
float mx = x[0], s = 0;
int i;
for (i = 1; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
int main(void)
{
/* Encoder outputs, one per source word */
float enc[ENC_LEN][DM] = {
{ 0.5f, 0.1f, -0.3f, 0.8f, 0.2f, -0.1f },
{ 0.3f, 0.7f, 0.1f, -0.2f, 0.5f, 0.4f },
{ -0.1f, 0.4f, 0.6f, 0.3f, -0.4f, 0.2f },
{ 0.2f, -0.3f, 0.5f, 0.6f, 0.1f, -0.5f },
{ 0.8f, 0.2f, -0.1f, 0.4f, 0.3f, 0.6f },
};
const char *src[ENC_LEN] = {
"I", "love", "big", "fat", "cats"
};
/* Decoder hidden states, one per target word */
float dec[DEC_LEN][DM] = {
{ 0.4f, 0.6f, -0.2f, 0.3f, 0.1f, 0.5f },
{ 0.1f, 0.3f, 0.5f, -0.1f, 0.4f, -0.3f },
{ 0.6f, -0.1f, 0.2f, 0.7f, -0.2f, 0.3f },
};
const char *tgt[DEC_LEN] = { "J'", "aime", "les" };
/* Chosen rather than trained, so the alignment
is legible. Training would find its own. */
float W_Q[DK][DM] = {
{ +0.65f, +3.69f, -1.87f, -0.15f,
+1.01f, +2.19f },
{ +1.19f, +2.48f, +3.21f, -0.19f,
+2.68f, -1.47f },
{ +0.53f, -0.21f, +0.76f, +0.54f,
+0.05f, -0.20f },
{ +2.14f, -1.92f, +2.03f, +2.77f,
-0.98f, -0.06f },
};
float W_K[DK][DM] = {
{ -0.32f, +0.95f, -0.99f, +0.91f,
-0.08f, -0.61f },
{ -0.08f, +0.94f, +0.26f, -0.26f,
+0.94f, -0.46f },
{ -0.03f, +0.24f, +1.05f, +0.46f,
-0.33f, -0.01f },
{ +1.11f, -1.27f, +1.03f, -0.12f,
+0.27f, +0.73f },
};
/* W_V keeps the first four encoder components */
float W_V[DK][DM] = {
{ 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f },
};
float scale = 1.0f / sqrtf((float)DK);
float w[ENC_LEN], ctx[DK];
int i, j, k, d, best;
printf("Cross-attention, decoder queries against "
"encoder keys\n\n");
for (i = 0; i < DEC_LEN; i++) {
float q[DK];
for (k = 0; k < DK; k++) {
q[k] = 0;
for (d = 0; d < DM; d++)
q[k] += W_Q[k][d] * dec[i][d];
}
for (j = 0; j < ENC_LEN; j++) {
float key[DK];
float sc = 0;
for (k = 0; k < DK; k++) {
key[k] = 0;
for (d = 0; d < DM; d++)
key[k] += W_K[k][d] * enc[j][d];
sc += q[k] * key[k];
}
w[j] = sc * scale;
}
softmax(w, ENC_LEN);
for (k = 0; k < DK; k++) {
ctx[k] = 0;
for (j = 0; j < ENC_LEN; j++) {
float v = 0;
for (d = 0; d < DM; d++)
v += W_V[k][d] * enc[j][d];
ctx[k] += w[j] * v;
}
}
best = 0;
for (j = 1; j < ENC_LEN; j++)
if (w[j] > w[best]) best = j;
printf(" \"%s\" reads the source\n", tgt[i]);
printf(" weights: ");
for (j = 0; j < ENC_LEN; j++)
printf("%s=%.2f ", src[j], w[j]);
printf("\n strongest: \"%s\"\n", src[best]);
printf(" context: [%+.3f, %+.3f, %+.3f, "
"%+.3f]\n\n",
ctx[0], ctx[1], ctx[2], ctx[3]);
}
printf("Each target word pulls a "
"different mix of\n");
printf("source words, and gets a different "
"context\n");
printf("vector as a result. Chapter 18 gave the\n");
printf("decoder one vector for the whole "
"sentence.\n");
return 0;
}
Figure 25-3 has each target word reading a different part of the source. Three target words, three quite different readings of the same five word source. The French “J’” puts 0.62 on “I”, which is the pronoun it translates, and spreads the remaining 0.38 thinly across the other four. “aime” puts 0.63 on “love”, the verb it corresponds to. “les” is the interesting row, splitting 0.45 on “cats” and 0.33 on “fat” with very little elsewhere, which is roughly what a French determiner should do given that it agrees with a noun phrase rather than a single word.
The context vectors underneath differ accordingly, and they are what actually leaves the sublayer. “J’” receives [+0.425, +0.175, −0.091, +0.582], dominated by the encoder state for “I”, while “aime” receives [+0.309, +0.477, +0.139, +0.067] with the second component much larger because “love” contributes most of it. Three positions in the same decoder, three different summaries of the same source sentence, each assembled on demand.
Set this against Chapter 18 and the improvement is the entire reason transformers displaced what came before. There the encoder compressed the input into one fixed vector and handed it to the decoder once, and we measured what that cost back then, with the first token of a twenty token input having exactly zero measurable influence on the vector the decoder received. Cross-attention removes the compression completely. Every encoder state remains available in full, at every decoder step, and the decoder decides afresh each time which ones matter. The information is not summarized better, it is simply never discarded. What that costs is memory and arithmetic that scale with the input length rather than staying fixed, since the decoder holds every encoder state for the whole of generation and scores against all of them at every step. Chapter 18′s single vector was cheap and lossy, and this is the opposite trade.
The alignment here is learned rather than supplied, which is the last thing to register about it. Nobody tells the model that “aime” corresponds to “love”, and no dictionary is consulted. Training adjusts W_Q and W_K until decoder queries land near the right encoder keys, and the alignments that emerge often match what a linguist would draw, which is why early attention papers spent so much space on printed alignment matrices.
25.5 The Decoder Block
A decoder block runs three sublayers where an encoder block runs two, and the extra one is cross-attention sitting between the self-attention and the feedforward network. Each of the three gets its own normalization on the way in and its own residual addition on the way out, so the pattern established in Chapter 24 repeats three times rather than twice.
/* 133_Decoder_Block.c */
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <stdlib.h>
static float randf(void)
{
return(float)rand()/RAND_MAX;
}
static float gelu(float x)
{
float c = 0.7978846f;
float u = x + 0.044715f * x * x * x;
return 0.5f * x * (1.0f + tanhf(c * u));
}
static float dot(const float *a, const float *b,
int n)
{
float s = 0;
int i;
for (i = 0; i < n; i++) s += a[i] * b[i];
return s;
}
static void softmax(float *x, int n)
{
float mx = -FLT_MAX, s = 0;
int i;
for (i = 0; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
static void rmsnorm(const float *x, float *out,
int n)
{
int i;
double ss = 0;
float eps = 1e-5f;
for (i = 0; i < n; i++) ss += x[i] * x[i];
float ri = 1.0f / sqrtf((float)(ss/n) + eps);
for (i = 0; i < n; i++) out[i] = x[i] * ri;
}
#define EL 4 /* encoder length */
#define DL 3 /* decoder length */
#define DM 6 /* model dimension */
#define DK 3 /* head dimension */
#define DF 12 /* FFN intermediate */
/* Simplified single-head attention */
static void attention(const float Q[][DM], int q_len,
const float K[][DM], int k_len,
const float V[][DM],
const float WQ[DK][DM],
const float WK[DK][DM],
const float WV[DK][DM],
const float WO[DM][DK],
int causal, float out[][DM])
{
float scale = 1.0f / sqrtf((float)DK);
int i, j, k, d;
for (i = 0; i < q_len; i++) {
float q[DK];
for (k = 0; k < DK; k++) {
q[k] = 0;
for (d = 0; d < DM; d++)
q[k] += WQ[k][d] * Q[i][d];
}
float scores[8]; /* max seq len */
for (j = 0; j < k_len; j++) {
float key[DK];
for (k = 0; k < DK; k++) {
key[k] = 0;
for (d = 0; d < DM; d++)
key[k] += WK[k][d] * K[j][d];
}
scores[j] = dot(q, key, DK) * scale;
if (causal && j > i) scores[j] = -1e9f;
}
softmax(scores, k_len);
/* Weighted sum of values, then project */
float head[DK] = {0};
for (j = 0; j < k_len; j++) {
for (k = 0; k < DK; k++) {
float val = 0;
for (d = 0; d < DM; d++)
val += WV[k][d] * V[j][d];
head[k] += scores[j] * val;
}
}
for (d = 0; d < DM; d++) {
out[i][d] = 0;
for (k = 0; k < DK; k++)
out[i][d] += WO[d][k] * head[k];
}
}
}
int main(void)
{
/* Encoder output, as if from the stack */
float enc[EL][DM] = {
{0.5f, 0.1f, -0.3f, 0.8f, 0.2f, -0.1f},
{0.3f, 0.7f, 0.1f, -0.2f, 0.5f, 0.4f},
{-0.1f, 0.4f, 0.6f, 0.3f, -0.4f, 0.2f},
{0.8f, 0.2f, -0.1f, 0.4f, 0.3f, 0.6f},
};
/* Decoder input */
float dec[DL][DM] = {
{0.4f, 0.6f, -0.2f, 0.3f, 0.1f, 0.5f},
{0.1f, 0.3f, 0.5f, -0.1f, 0.4f, -0.3f},
{0.6f, -0.1f, 0.2f, 0.7f, -0.2f, 0.3f},
};
const char *dec_words[] = { "J'", "aime", "les" };
/* Weights for three attention layers + FFN */
float WQ1[DK][DM], WK1[DK][DM];
float WV1[DK][DM], WO1[DM][DK];
float WQ2[DK][DM], WK2[DK][DM];
float WV2[DK][DM], WO2[DM][DK];
float W1[DF][DM], W2[DM][DF];
int i, j;
srand(42);
float *all_w[] = {
(float*)WQ1, (float*)WK1, (float*)WV1,
(float*)WO1,
(float*)WQ2, (float*)WK2, (float*)WV2,
(float*)WO2,
(float*)W1, (float*)W2
};
int sizes[] = { DK*DM, DK*DM, DK*DM, DM*DK,
DK*DM, DK*DM, DK*DM, DM*DK,
DF*DM, DM*DF };
for (i = 0; i < 10; i++)
for (j = 0; j < sizes[i]; j++)
all_w[i][j] = (randf()*2-1)*0.15f;
printf("Decoder Block: 3 sublayers\n\n");
/* === Sublayer 1: Causal self-attention === */
float norm_dec[DL][DM], sub[DL][DM];
for (i = 0; i < DL; i++)
rmsnorm(dec[i], norm_dec[i], DM);
attention(norm_dec, DL, norm_dec, DL, norm_dec,
WQ1, WK1, WV1, WO1, 1 /* causal */, sub);
for (i = 0; i < DL; i++)
for (j = 0; j < DM; j++) dec[i][j] += sub[i][j];
printf(" 1. Causal self-attention, the decoder\n");
printf(" attending to itself\n");
/* === Sublayer 2: Cross-attention === */
for (i = 0; i < DL; i++)
rmsnorm(dec[i], norm_dec[i], DM);
attention(norm_dec, DL, enc, EL, enc,
WQ2, WK2, WV2, WO2, 0 /* no mask */, sub);
for (i = 0; i < DL; i++)
for (j = 0; j < DM; j++) dec[i][j] += sub[i][j];
printf(" 2. Cross-attention, the decoder\n");
printf(" attending to the encoder\n");
/* === Sublayer 3: FFN === */
for (i = 0; i < DL; i++)
rmsnorm(dec[i], norm_dec[i], DM);
for (i = 0; i < DL; i++) {
float hid[DF], fo[DM];
int f;
for (f = 0; f < DF; f++) {
float z = 0;
for (j = 0; j < DM; j++)
z += W1[f][j] * norm_dec[i][j];
hid[f] = gelu(z);
}
for (j = 0; j < DM; j++) {
float z = 0;
for (f = 0; f < DF; f++)
z += W2[j][f] * hid[f];
dec[i][j] += z;
}
}
printf(" 3. Feed-forward network\n\n");
printf("Output:\n");
for (i = 0; i < DL; i++) {
printf(" %-5s [", dec_words[i]);
for (j = 0; j < DM; j++)
printf("%+.3f%s", dec[i][j],
j<DM-1?",":"");
printf("]\n");
}
printf("\nA decoder block has 3 sublayers where\n");
printf("an encoder block has 2.\n");
printf("Cross-attention is the bridge across.\n");
return 0;
}
Figure 25-4 runs a decoder block with three sublayers instead of two. The program prints the three sublayers in order and then the output, and the output is three vectors of six components that differ from each other, which is the minimum evidence that the block is doing something. “J’” leaves as [+0.328, +0.614, −0.260, +0.337, +0.038, +0.418] and “aime” as [+0.072, +0.325, +0.509, −0.057, +0.363, −0.328], and those two are not close to each other in any component, nor is either close to the vector it entered as.
The ordering of the three sublayers is fixed and the reasons are worth separating. Causal self-attention must come before cross-attention for the reason given earlier, since the decoder should consult the source with a query informed by what it has already written. The feedforward network must come last because it is the only sublayer that cannot move information between positions, so putting it earlier would waste it on representations that had not yet gathered anything. Any other ordering runs, produces numbers, and learns worse.
Parameter arithmetic follows from the count in Chapter 24. An encoder block spends 4 times d_model squared on its single attention and 2 times d_model times d_ff on its feedforward network. A decoder block spends twice the attention figure, because it has two attention sublayers, plus the same feedforward cost, plus one extra normalization. At the original paper’s dimensions of 512 for the model and 2048 for the feedforward hidden layer, an encoder block comes to 1.05M for attention plus 2.10M for the feedforward network, totalling 3.15M. A decoder block comes to 2.10M for its two attentions plus the same 2.10M, totalling 4.19M, so it is a third larger. Across a six and six arrangement that is 18.9M of encoder against 25.2M of decoder, and the encoder is carrying a third of the total parameter budget purely to read an input the decoder could in principle have read itself. Exercise 5 asks you to total a full six and six arrangement and compare it against a twelve layer decoder-only stack, and the answer explains part of why the field moved.
25.6 The Complete Encoder-Decoder
Everything now goes into one program that runs the whole path from token indices to predicted tokens. Source tokens are embedded and given positional encodings, the encoder stack processes them, target tokens are embedded and encoded the same way, the decoder stack runs its three sublayers with the encoder output available to its cross-attention, and an output projection turns the final decoder states into a score for every entry in the vocabulary.
/* 134_Full_Transformer.c */
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <stdlib.h>
static float randf(void)
{
return(float)rand()/RAND_MAX;
}
static float gelu(float x)
{
float c = 0.7978846f;
float u = x + 0.044715f * x * x * x;
return 0.5f * x * (1.0f + tanhf(c * u));
}
static float dot(const float *a, const float *b,
int n)
{
float s = 0;
int i;
for (i = 0; i < n; i++) s += a[i] * b[i];
return s;
}
static void softmax(float *x, int n)
{
float mx = -FLT_MAX, s = 0;
int i;
for (i = 0; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
static void rmsnorm(const float *x, float *out,
int n)
{
int i;
double ss = 0;
float eps = 1e-5f;
for (i = 0; i < n; i++) ss += x[i] * x[i];
float ri = 1.0f / sqrtf((float)(ss/n) + eps);
for (i = 0; i < n; i++) out[i] = x[i] * ri;
}
#define DM 8
#define DK 4
#define DF 16
#define VOCAB 10
#define MAX_LEN 6
/* Simple single-head attention */
static void attn(const float *Q, int ql,
const float *K, int kl,
const float *V,
const float *WQ, const float *WK,
const float *WV, const float *WO,
int causal, float *out)
{
float scale = 1.0f / sqrtf((float)DK);
int i, j, k, d;
for (i = 0; i < ql; i++) {
float q[DK];
for (k = 0; k < DK; k++) {
q[k] = 0;
for (d = 0; d < DM; d++)
q[k] += WQ[k*DM + d] * Q[i*DM + d];
}
float scores[MAX_LEN];
for (j = 0; j < kl; j++) {
float key[DK];
for (k = 0; k < DK; k++) {
key[k] = 0;
for (d = 0; d < DM; d++)
key[k] += WK[k*DM + d]
* K[j*DM + d];
}
scores[j] = dot(q, key, DK) * scale;
/* the mask, and the whole of it */
if (causal && j > i) scores[j] = -1e9f;
}
softmax(scores, kl);
float head[DK] = {0};
for (j = 0; j < kl; j++)
for (k = 0; k < DK; k++) {
float v = 0;
for (d = 0; d < DM; d++)
v += WV[k*DM + d] * V[j*DM + d];
head[k] += scores[j] * v;
}
for (d = 0; d < DM; d++) {
out[i*DM + d] = 0;
for (k = 0; k < DK; k++)
out[i*DM + d] += WO[d*DK + k] * head[k];
}
}
}
int main(void)
{
int enc_len = 4, dec_len = 3;
/* Token IDs */
/* source sentence */
int enc_tokens[] = { 1, 3, 5, 7 };
/* target, shifted right */
int dec_tokens[] = { 0, 2, 4 };
/* Embedding table */
float embed[VOCAB][DM];
srand(42);
int i, j;
for (i = 0; i < VOCAB; i++)
for (j = 0; j < DM; j++)
embed[i][j] = (randf()*2-1)*0.3f;
/* Sinusoidal positional encoding */
float pe[MAX_LEN][DM];
for (i = 0; i < MAX_LEN; i++)
for (j = 0; j < DM; j++) {
float e = (float)(j/2*2) / DM;
float angle = i / powf(10000.0f, e);
pe[i][j] = (j%2==0)
? sinf(angle)
: cosf(angle);
}
/* === EMBED + POSITION === */
float enc_x[MAX_LEN*DM], dec_x[MAX_LEN*DM];
for (i = 0; i < enc_len; i++)
for (j = 0; j < DM; j++)
enc_x[i*DM + j] =
embed[enc_tokens[i]][j] + pe[i][j];
for (i = 0; i < dec_len; i++)
for (j = 0; j < DM; j++)
dec_x[i*DM + j] =
embed[dec_tokens[i]][j] + pe[i][j];
/* Weights (all random for demonstration) */
float enc_WQ[DK*DM], enc_WK[DK*DM];
float enc_WV[DK*DM], enc_WO[DM*DK];
float enc_W1[DF*DM], enc_W2[DM*DF];
float dec_WQ1[DK*DM], dec_WK1[DK*DM];
float dec_WV1[DK*DM], dec_WO1[DM*DK];
float dec_WQ2[DK*DM], dec_WK2[DK*DM];
float dec_WV2[DK*DM], dec_WO2[DM*DK];
float dec_W1[DF*DM], dec_W2[DM*DF];
float out_W[VOCAB*DM];
float *all[] = {
enc_WQ, enc_WK, enc_WV, enc_WO,
enc_W1, enc_W2,
dec_WQ1, dec_WK1, dec_WV1, dec_WO1,
dec_WQ2, dec_WK2, dec_WV2, dec_WO2,
dec_W1, dec_W2, out_W};
int sz[] = {DK*DM, DK*DM, DK*DM, DM*DK,
DF*DM, DM*DF,
DK*DM, DK*DM, DK*DM, DM*DK,
DK*DM, DK*DM, DK*DM, DM*DK,
DF*DM, DM*DF, VOCAB*DM};
for (i = 0; i < 17; i++)
for (j = 0; j < sz[i]; j++)
all[i][j] = (randf()*2-1)*0.1f;
printf("Full Encoder-Decoder Transformer\n");
printf(" Encoder %d tokens, decoder %d tokens, "
"d_model=%d\n\n", enc_len, dec_len, DM);
/* === ENCODER (1 block) === */
printf("1. Encoder self-attention + FFN\n");
float norm[MAX_LEN*DM], sub[MAX_LEN*DM];
/* Self-attention */
for (i = 0; i < enc_len; i++)
rmsnorm(&enc_x[i*DM], &norm[i*DM], DM);
attn(norm, enc_len, norm, enc_len, norm,
enc_WQ, enc_WK, enc_WV, enc_WO, 0, sub);
for(i = 0;i<enc_len*DM;i++) enc_x[i] += sub[i];
/* FFN */
for (i = 0; i < enc_len; i++)
rmsnorm(&enc_x[i*DM], &norm[i*DM], DM);
for(i = 0;i<enc_len;i++) {
float hid[DF];
int f;
for (f = 0; f < DF; f++) {
float z = 0;
for (j = 0; j < DM; j++)
z += enc_W1[f*DM + j] * norm[i*DM + j];
hid[f] = gelu(z);
}
for (j = 0; j < DM; j++) {
float z = 0;
for (int f = 0; f < DF; f++)
z += enc_W2[j*DF + f] * hid[f];
enc_x[i*DM + j] += z;
}
}
/* === DECODER (1 block) === */
printf("2. Decoder causal self-attention\n");
for (i = 0; i < dec_len; i++)
rmsnorm(&dec_x[i*DM], &norm[i*DM], DM);
/* causal flag on, this is self-attention */
attn(norm, dec_len, norm, dec_len, norm,
dec_WQ1, dec_WK1, dec_WV1, dec_WO1, 1, sub);
for(i = 0;i<dec_len*DM;i++) dec_x[i] += sub[i];
printf("3. Decoder cross-attention to encoder\n");
for (i = 0; i < dec_len; i++)
rmsnorm(&dec_x[i*DM], &norm[i*DM], DM);
/* queries from dec, keys and values from enc */
attn(norm, dec_len, enc_x, enc_len, enc_x,
dec_WQ2, dec_WK2, dec_WV2, dec_WO2, 0, sub);
for(i = 0;i<dec_len*DM;i++) dec_x[i] += sub[i];
printf("4. Decoder FFN\n");
for (i = 0; i < dec_len; i++)
rmsnorm(&dec_x[i*DM], &norm[i*DM], DM);
for(i = 0;i<dec_len;i++) {
float hid[DF];
int f;
for (f = 0; f < DF; f++) {
float z = 0;
for (j = 0; j < DM; j++)
z += dec_W1[f*DM + j] * norm[i*DM + j];
hid[f] = gelu(z);
}
for (j = 0; j < DM; j++) {
float z = 0;
for (int f = 0; f < DF; f++)
z += dec_W2[j*DF + f] * hid[f];
dec_x[i*DM + j] += z;
}
}
/* === OUTPUT HEAD === */
printf("5. Output projection -> vocabulary "
"logits\n\n");
for (i = 0; i < dec_len; i++) {
float logits[VOCAB];
for (j = 0; j < VOCAB; j++) {
logits[j] = 0;
for (int k = 0; k < DM; k++)
logits[j] += out_W[j*DM+k]
* dec_x[i*DM+k];
}
softmax(logits, VOCAB);
int pred = 0;
for (j = 1; j < VOCAB; j++)
if (logits[j] > logits[pred]) pred = j;
printf(" Dec pos %d: token %d at prob %.3f\n",
i, pred, logits[pred]);
}
printf("\nComplete pipeline:\n");
printf(" source -> embed+PE -> encoder\n");
printf(" target -> embed+PE -> decoder\n");
printf(" -> output head -> softmax -> tokens\n");
return 0;
}
Figure 25-5 follows the whole path from tokens to predictions. The five numbered lines trace the path and the three prediction lines report the result, which is that position 0 predicts token 8 at probability 0.125, position 1 predicts token 1 at 0.111 and position 2 predicts token 1 at 0.113. The vocabulary has ten entries, so a model that had learned nothing at all would spread 0.100 across every token, and the winning probabilities here sit between one and two and a half percentage points above that floor. Two of the three positions also agree on the same token despite having different inputs, which is what a model with no preferences looks like when a tiny random asymmetry in the output projection decides every row. The predictions are noise and the confidence figures say so.
That is exactly what an untrained transformer should produce, and saying so plainly beats letting a reader wonder whether something broke. Every weight in this program came from a random number generator scaled small, following the argument from the last chapter that a freshly initialized block should sit close to the identity. A network built from near identity blocks and finished with a random output projection has no reason to prefer one token over another, and if it did prefer one strongly that would indicate a bug rather than intelligence. What the program demonstrates is that the shapes line up end to end, that the encoder output reaches the decoder through cross-attention, that the causal flag is set on one attention call and clear on the other two, and that the whole thing produces a well formed probability distribution per position.
The causal flag is the detail to look for in the code. The same attn function is called three times with different arguments, and the only difference between encoder self-attention, decoder self-attention and cross-attention is which sequences are passed for Q, K and V and whether the causal parameter is 1 or 0. Encoder self-attention passes the encoder states three times with the flag clear. Decoder self-attention passes the decoder states three times with the flag set. Cross-attention passes decoder states for Q and encoder states for K and V, with the flag clear because the decoder is allowed to see the whole input. Three attention types, one implementation, and the entire difference is in the call site.
25.7 Key Takeaways
The original transformer runs two stacks. The encoder reads the input with bidirectional self-attention and the decoder writes the output with causal self-attention plus cross-attention, and the two meet only at the cross-attention sublayer.
The causal mask sets scores for future positions to a large negative number before softmax, which the exponential drives to zero. The mask produced the expected triangle, with position 0 giving 1.00 to itself and position 4 spreading 0.20 across all five.
The last position of a masked sequence sees everything, so its computation matches the unmasked case. That is what lets a decoder train on a whole target sequence in one forward pass, with every position predicting its own next token and no leakage between them.
Cross-attention is the mechanism from Chapter 20 with the queries taken from the decoder and the keys and values from the encoder. No new arithmetic is involved and the same scaled dot-product runs unchanged.
Cross-attention measured real alignment, with a target word putting 0.62 on the source pronoun it translates and another putting 0.63 on the matching verb, each producing a different context vector as a result.
Cross-attention replaces the fixed context vector of Chapter 18 entirely. There the first token of a twenty token input had zero measurable influence on what the decoder received, and here every encoder state stays available in full at every decoder step.
A decoder block runs three sublayers against an encoder block’s two, with cross-attention inserted between the self-attention and the feedforward network. Each still gets its own normalization and residual.
The ordering is not arbitrary. Self-attention first lets the decoder query the source with a question shaped by what it has already written, and the feedforward network last means it operates on representations that have already gathered from both sequences.
A decoder block costs about a third more than an encoder block, because it carries two attention sublayers instead of one while the feedforward cost is unchanged.
The assembled model produced winning probabilities between 0.111 and 0.125 against a uniform baseline of 0.100 over a ten token vocabulary, which is the correct output for an untrained network and confirms only that the shapes connect end to end.
One attention implementation serves all three uses. Which sequences are passed for Q, K and V, and whether the causal flag is set, is the entire difference between encoder self-attention, decoder self-attention and cross-attention.
25.8 Exercises
Print the causal attention weight matrix for a 6-position sequence. Verify it is lower-triangular (zeros above the diagonal).
In cross-attention, what happens if the encoder output is all zeros? What does the decoder receive?
Stack 2 encoder blocks and 2 decoder blocks. Does the decoder output change with depth?
The original transformer paper uses post-norm placement. Modify 134_Full_Transformer.c to use post-norm and compare. Which produces more stable norms after 4 layers?
Count the total parameters for a 6-layer encoder + 6-layer decoder transformer with d_model=512, d_ff=2048, and 8 heads. How does it compare to a 12-layer decoder-only transformer with the same d_model?
Implement beam search for the decoder, where at each step you keep the top-k most likely partial sequences and expand each one. Compare the output to greedy decoding.