Transformers and Attention from Scratch

01Text becomes a sequence of vectors

A language transformer operates on numbers. A token ID selects a row from a learned embedding table. A token need not be a word: it can represent a character, a word fragment, or a byte sequence. The tokenizer defines the conversion between text and IDs.

BPE is one way to learn a subword vocabulary. In Sennrich et al.'s character-based version, each word starts as characters plus a word-end marker. Count adjacent pairs across the training vocabulary with word frequencies, merge a most-frequent pair, and repeat. The learned merge sequence becomes part of the tokenizer.^[1]

The toy corpus below represents frequency by repeated words. Each Step learns one merge. Pair counts are recomputed after merging, so the next winning pair can change. Ties use a deterministic lexical rule for this demo. The training process is shown; ordinary tokenization later applies the learned ranked merges to new text rather than relearning the vocabulary from each prompt.

The integer ID itself carries no geometry. Token 80 is not "twice as meaningful" as token 40. Learned embeddings supply continuous coordinates, and later layers contextualize them. The same token can therefore acquire different hidden representations in different sequences.

Tokenization is part of the model contract. Using a different vocabulary or merge table changes which embedding rows are read. Sequence length is measured in tokens, so the same character count can occupy different amounts of context depending on language and content. A character-count budget does not establish a token-count budget.

02Queries and keys choose how values mix

A dense layer mixes features within a vector. Attention mixes information across sequence positions, using weights computed from the input. The hidden representation at each position produces a query, a key, and a value through learned projections. A query-key dot product becomes a score; a softmax turns a row of scores into mixing weights.^[2]

X has sequence length N and model width D. For one head, Q and K are N×d_k, V is N×d_v, A is N×N, and Y is N×d_v. Softmax normalizes each query row over keys.

The letters identify roles, not meanings guaranteed by training. Keys do not have to encode topic labels, and attention weights are not probabilities that a statement is true. They are coefficients in a weighted sum of values. A large score makes one value contribute more strongly relative to the other available values.

Dividing by √dₖ controls the scale of a dot product. Under independent, zero-mean, unit-variance component assumptions, the sum has variance dₖ. That argument motivates the factor; real learned vectors need not satisfy those assumptions exactly. The result still depends on vector lengths as well as angles.^[2]

Multi-head attention runs several such mixtures. Concatenate the head outputs and apply an output projection. Heads can learn different relations, but assigning a fixed human role such as "the grammar head" is not an architectural guarantee. The total representation width can stay fixed while each head uses a smaller width.^[2]

03Hide future positions before softmax

A next-token decoder must not read the answer from a later input position. Causal attention permits the current position and its prefix, then excludes later positions. The exclusion belongs in the score calculation before normalization. Masked scores are effectively −∞, so their exponentials are zero.^[2]

Replacing a masked score by zero is wrong: exp(0) is one, so the masked value still receives probability mass. Zeroing the attention weights after softmax without renormalizing is also different: the remaining weights no longer sum to one. A row with no permitted keys needs an explicit implementation policy; softmax over an all-−∞ row is undefined in the usual subtraction-of-maximum implementation.

For a toy input sequence A,B,C, next-token targets are B,C,D. Row one can read A and B when predicting C. Including the current input position is therefore correct; the target is shifted one position forward. Input padding introduces a separate validity mask. Combine it with the causal rule so padded keys receive no probability mass.

During training, all query rows can be evaluated together under this mask. During generation, the next unknown token still depends on the previously generated one. Parallel computation across known training positions does not remove the sequential dependency in autoregressive decoding.^[2]

04Position changes the relation between vectors

Unmasked self-attention without position information is equivariant to reordering the input rows: the output rows reorder with them. The original transformer adds positional vectors to embeddings. Other architectures encode position inside the attention calculation. The causal mask itself also introduces an ordering constraint; it should not be confused with an explicit position representation.^[2]

Rotary position embedding, or RoPE, rotates coordinate pairs in queries and keys. For one pair, write R(θ) for a two-dimensional rotation matrix. The position-scaled rotations preserve length and make the score depend on relative position.^[3]

m and n are token positions; θ is the rotation in radians per token position for this coordinate pair. A full head uses multiple pairs and frequencies.

The equality follows from the rotation identity R(mθ)ᵀR(nθ)=R((n−m)θ). Shifting both positions by the same amount leaves the dot product unchanged for the same underlying vectors. Changing their gap changes the rotation difference. The widget computes both sides independently.

Rotary coordinates do not give a trained model unlimited context. Position formulas may accept a large index while the model behaves poorly outside its training distribution. Cache positions must also use the same indexing convention as the original forward pass. Restarting positions at zero for a cached query changes its relationship to the stored keys.

In a full model, hidden representations already depend on their context. The identity holds for rotating fixed vectors; it does not assert that moving words around leaves those contextual vectors unchanged.

05Attention is one part of a transformer block

A transformer block also applies a position-wise feed-forward network and residual additions. The feed-forward weights are reused at every sequence position, but the network operates on each position's features separately. Attention supplies cross-position mixing; the feed-forward sublayer transforms the mixed features.^[2]

Layer normalization measures mean and variance across a chosen feature dimension. In the token-wise case here, each token gets its own statistics. Unlike batch normalization during training, it does not borrow statistics from other examples in the batch.^[4]

The variance averages squared deviations over the feature count. ε avoids division by zero; the learned affine parameters need not preserve zero mean or unit variance.

Adding a constant to every feature cancels during centering. Multiplying all features by a positive constant mostly cancels during normalization when the variance dominates ε. The small stabilizing ε means scale invariance is not mathematically exact. Learned gain and bias then change the distribution again.

The order of normalization and residual addition matters. The original transformer uses post-LN, LayerNorm(x+F(x)). A pre-LN block uses x+F(LayerNorm(x)); its residual stream bypasses the normalization inside the branch. These define different functions and different gradient paths. Xiong et al. study this architectural distinction; neither formula can be substituted for the other when loading an existing model.^[5]

06Reuse keys and values during generation

For a causal decoder with fixed position rules at inference, appending a token does not change earlier tokens' hidden states. Their keys and values can be stored in a KV cache. Compute the new token's query, key, and value, append its key and value, then attend over the stored prefix. Model parameters are fixed, and inference-time dropout must be disabled for this reuse to match the reference calculation.^[6]

The new query still reads earlier keys and values. Caching avoids recomputing old projections and hidden states; it does not make full-context decoding independent of context length. Prompt processing, often called prefill, can process many known positions together. Later decoding commonly processes one new position per sequence.

For an ordinary dense cache, payload bytes are 2×layers×batch×KV-heads×tokens×head-width×bytes-per-entry. The factor two accounts for keys and values. This excludes allocator overhead and any separate metadata. Grouped-query attention changes the KV-head count; it does not merely reinterpret the cache of an arbitrarily trained multi-head model.^[7]

Queries in a cached chunk have absolute offsets. Local row zero in a one-token chunk may be absolute position 100. Its causal mask must allow the stored keys through position 100. An unadjusted small triangular mask can silently allow only the first cached key. The code below makes the offset explicit.^[6]

Dense attention has quadratic score-pair arithmetic in a full N-token prefill. Storing an N×N score matrix is an implementation choice. FlashAttention computes exact attention in tiles and avoids materializing that full matrix in device memory, reducing memory transfers. This changes storage and execution, not the attention equation or its dense pair count.^[8]

07A causal attention kernel with a cache offset

These complete programs implement one head with two query/key coordinates and two value coordinates. Legal keys are selected before normalization. Subtracting the largest legal score prevents exponentials of large positive numbers. Each program checks a one-key row, cached/full agreement, and immunity of an earlier row to a changed future value.

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <vector>
using Vector = std::array<double,2>;
using Sequence = std::vector<Vector>;
Sequence attention(const Sequence& queries, const Sequence& keys,
                   const Sequence& values, int queryOffset = 0) {
    assert(keys.size() == values.size() && !keys.empty());
    Sequence output(queries.size(), Vector{0,0});
    for (int row = 0; row < static_cast<int>(queries.size()); ++row) {
        // queryOffset is the number of cached tokens before this query chunk.
        const int lastKey = row + queryOffset;
        assert(lastKey >= 0 && lastKey < static_cast<int>(keys.size()));
        std::vector<double> scores(lastKey+1);
        for (int column = 0; column <= lastKey; ++column)
            scores[column] = (queries[row][0]*keys[column][0]
                            + queries[row][1]*keys[column][1])/std::sqrt(2.0);
        const double maximum = *std::max_element(scores.begin(), scores.end());
        double denominator = 0;
        for (double& score : scores) {
            score = std::exp(score-maximum); // Prevent large positive exponentials.
            denominator += score;
        }
        for (int column = 0; column <= lastKey; ++column)
        for (int feature = 0; feature < 2; ++feature)
            output[row][feature] += scores[column]/denominator * values[column][feature];
    }
    return output;
}
int main() {
    const Sequence queries{{1,0},{0,1},{1,1}};
    Sequence values{{2,0},{0,3},{-1,1}};
    const auto full = attention(queries,queries,values);
    assert(std::abs(full[0][0]-2.0) < 1e-12); // First token has one legal key.
    const auto cached = attention(Sequence{queries.back()},queries,values,2);
    for (int feature = 0; feature < 2; ++feature)
        assert(std::abs(full.back()[feature]-cached[0][feature]) < 1e-12);
    values.back() = {100,100};
    const auto changed = attention(queries,queries,values);
    assert(full[1] == changed[1]); // Future values cannot influence row one.
}

type Vector = [f64;2];
fn attention(queries: &[Vector], keys: &[Vector], values: &[Vector],
             query_offset: usize) -> Vec<Vector> {
    assert!(keys.len() == values.len() && !keys.is_empty());
    let mut output = vec![[0.0;2];queries.len()];
    for row in 0..queries.len() {
        // query_offset is the number of cached tokens before this query chunk.
        let last_key = row + query_offset;
        assert!(last_key < keys.len());
        let mut scores = Vec::with_capacity(last_key+1);
        for key in &keys[..=last_key] {
            scores.push((queries[row][0]*key[0]+queries[row][1]*key[1])/2.0_f64.sqrt());
        }
        let maximum = scores.iter().copied().fold(f64::NEG_INFINITY,f64::max);
        let mut denominator = 0.0;
        for score in &mut scores {
            *score = (*score-maximum).exp(); // Prevent large positive exponentials.
            denominator += *score;
        }
        for column in 0..=last_key { for feature in 0..2 {
            output[row][feature] += scores[column]/denominator * values[column][feature];
        }}
    }
    output
}
fn main() {
    let queries = [[1.0,0.0],[0.0,1.0],[1.0,1.0]];
    let mut values = [[2.0,0.0],[0.0,3.0],[-1.0,1.0]];
    let full = attention(&queries,&queries,&values,0);
    assert!((full[0][0]-2.0).abs() < 1e-12); // First token has one legal key.
    let cached = attention(&[queries[2]],&queries,&values,2);
    for feature in 0..2 {
        assert!((full[2][feature]-cached[0][feature]).abs() < 1e-12);
    }
    values[2] = [100.0,100.0];
    let changed = attention(&queries,&queries,&values,0);
    assert_eq!(full[1],changed[1]); // Future values cannot influence row one.
}

This reference kernel omits learned projection matrices, multiple heads, positional transforms, padding masks, batching, and optimized matrix multiplication. It defines the mixing operation rather than a full language model. Its caller must supply finite vectors whose dot-product scores remain finite, and at least one legal key per row. Subtracting the maximum cannot repair a dot product that has already overflowed. A complete decoder also needs a token embedding, a block stack, and a vocabulary-output projection trained with shifted targets.

For the block stack, preserve tensor shapes and normalization order exactly. Apply the chosen positional transform to new queries and keys using absolute positions, store keys in the same representation the reference pass uses, and compare the entire decoded prefix against an uncached implementation. A cached implementation can return plausible numbers while using the wrong tokens.

08Check semantics as well as tensor shapes

Use small asymmetric vectors when verifying a kernel. Equal keys produce uniform probabilities and can hide transposes or mask errors. Change a forbidden future value and assert that earlier outputs stay fixed. Compare cached and uncached results for multiple prompt lengths, multi-token chunks, and position offsets.

Keep mask conventions explicit when calling a library: one API can interpret true as "allowed" while another interprets it as "masked." A correct boolean array with reversed semantics produces a different model. Check the actual library contract and one hand-computed case.

Model quality needs a separate evaluation. These widgets compute exact attention on small arrays; their values and kernels have not been trained for language. Building the layer establishes numerical behavior. Useful predictions require appropriate data, a training objective, and evaluation on the intended tasks.

09What's next

Reinforcement Learning changes the source of the learning signal. An agent chooses an action, observes a reward, and must account for outcomes that arrive later. The neural network becomes one possible function approximator inside that learning process.

10Sources

The cited equations define the algorithms. Widget readouts report the displayed toy arrays and settings; they are not hardware benchmarks.

1. Rico Sennrich, Barry Haddow, Alexandra Birch, 2016. Neural Machine Translation of Rare Words with Subword Units. Frequency-based subword merges and a word-end marker.
2. Ashish Vaswani et al., 2017. Attention Is All You Need. Scaled dot-product attention, multiple heads, causal masks, and the original transformer block.
3. Jianlin Su et al., 2021. RoFormer: Enhanced Transformer with Rotary Position Embedding. Query/key rotations and relative-position dot products.
4. Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton, 2016. Layer Normalization. Feature normalization and learned gain and bias.
5. Ruibin Xiong et al., 2020. On Layer Normalization in the Transformer Architecture. The distinction between pre-LN and post-LN blocks.
6. Hugging Face Transformers contributors, cache documentation. How caching works. Key/value reuse, cache dimensions, and attention-mask alignment.
7. Joshua Ainslie et al., 2023. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Groups of query heads sharing key/value heads.
8. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Tiled exact attention without materializing a full score matrix.