Attention-Based Machine Translation

A transformer trained to translate English into Pig Latin, an artificial mapping with a clean, rule-based alignment between input and output. That regularity makes it a good testbed: the attention weights should recover the mapping, and they sharpen toward that alignment as the model trains.

Attention works as a weighted lookup. Each output position forms a query and compares it against a key for every input position; the match scores become weights on the corresponding values. Scaled dot-product attention computes all positions at once:

where $Q$, $K$, $V$ are the query, key, and value matrices and $d_k$ is the key dimension. The $\sqrt{d_k}$ divisor keeps the dot products from growing large enough to push softmax into flat regions where gradients vanish. The softmax row for an output token is a distribution over input tokens — exactly the alignment being learned.

Attention grew out of a fixed-vector bottleneck. Earlier sequence-to-sequence translators encoded the whole source sentence into a single fixed-length vector and decoded from that alone. The vector was a bottleneck: a long sentence had to be squeezed into the same width as a short one, and translation quality fell off as length grew. Attention removes the bottleneck by letting every decoder step read a weighted sum of all encoder states, choosing what to look at per output token instead of leaning on one summary of the source.

The encoder self-attends over the source sentence; the decoder attends over its own prefix and, through cross-attention, over the encoder output. Because the decoder generates left to right, its self-attention is causally masked so position $i$ never sees positions after it. Multiple attention heads run in parallel, each with its own learned query, key, and value projections, so one head can follow subject-verb agreement while another tracks adjacency or word order; their outputs are concatenated and mixed by a final linear map. Since attention alone is order-agnostic, positional encodings are added to the token embeddings to inject sequence order.

The alignment sharpens over training. The cross-attention weights, drawn as a heatmap of output positions against input positions, start diffuse, with every output token attending everywhere, and concentrate along the true correspondence as training proceeds. For Pig Latin, that correspondence is nearly diagonal with a shift at word boundaries, and the weights converge toward it, a direct and legible record of what the model has learned.

You can browse the project repository.