The Transformer Architecture Explained for Developers

// RNN: processes tokens sequentially
// Token 1 → hidden state → Token 2 → hidden state → Token 3 → ...
// Information from Token 1 degrades by the time we reach Token 100
 
// Transformer: processes all tokens in parallel
// Every token attends to every other token simultaneously
// Token 100 has direct access to Token 1's information
 
interface SequentialModel {
  process(tokens: number[]): number[] {
    let hiddenState = initializeState();
    const outputs: number[] = [];
    // Must process one at a time — slow, forgets early tokens
    for (const token of tokens) {
      hiddenState = updateState(hiddenState, token);
      outputs.push(computeOutput(hiddenState));
    }
    return outputs;
  }
}
 
interface TransformerModel {
  process(tokens: number[]): number[] {
    // All tokens processed simultaneously — parallel, no forgetting
    const embeddings = embed(tokens);
    const attended = selfAttention(embeddings); // Every token sees all others
    return feedForward(attended);
  }
}

// Self-attention in pseudocode
// Input: sequence of token embeddings (each is a vector)
// Output: sequence of context-enriched embeddings
 
function selfAttention(
  embeddings: number[][],  // [seqLength x embeddingDim]
  Wq: number[][],         // Query weight matrix
  Wk: number[][],         // Key weight matrix
  Wv: number[][]          // Value weight matrix
): number[][] {
  const seqLength = embeddings.length;
  const dk = Wk[0].length; // Key dimension for scaling
 
  // Step 1: Project each embedding into Query, Key, Value spaces
  const queries = matMul(embeddings, Wq);  // What am I looking for?
  const keys = matMul(embeddings, Wk);     // What do I contain?
  const values = matMul(embeddings, Wv);   // What information do I carry?
 
  // Step 2: Compute attention scores (dot product of queries and keys)
  // scores[i][j] = how much should token i attend to token j?
  const scores = matMul(queries, transpose(keys));
 
  // Step 3: Scale scores to prevent softmax saturation
  const scaled = scaleMatrix(scores, 1 / Math.sqrt(dk));
 
  // Step 4: Softmax converts scores to probabilities (rows sum to 1)
  const weights = softmaxRows(scaled);
 
  // Step 5: Weighted sum of values — the output for each position
  // output[i] = sum of all value vectors, weighted by attention to each
  const output = matMul(weights, values);
 
  return output;
}

// Concrete example: "The cat sat on the mat"
// When processing "sat", the attention weights might look like:
const attentionForSat = {
  'The': 0.05,   // Low — "The" isn't very relevant to "sat"
  'cat': 0.45,   // High — "cat" is the subject of "sat"
  'sat': 0.15,   // Medium — self-reference
  'on':  0.10,   // Medium — grammatical relationship
  'the': 0.05,   // Low
  'mat': 0.20,   // Medium-high — "sat" relates to location
};
// The model learns these weights during training
// "cat" gets high attention because it's the most semantically relevant

// Sinusoidal positional encoding (from the original paper)
function positionalEncoding(
  seqLength: number,
  embeddingDim: number
): number[][] {
  const encoding: number[][] = [];
 
  for (let pos = 0; pos < seqLength; pos++) {
    const posVector: number[] = [];
 
    for (let i = 0; i < embeddingDim; i++) {
      const angle = pos / Math.pow(10000, (2 * Math.floor(i / 2)) / embeddingDim);
 
      if (i % 2 === 0) {
        posVector.push(Math.sin(angle));  // Even dimensions: sin
      } else {
        posVector.push(Math.cos(angle));  // Odd dimensions: cos
      }
    }
 
    encoding.push(posVector);
  }
 
  return encoding;
}
 
// The final input to the Transformer:
// input[i] = tokenEmbedding[i] + positionalEncoding[i]
// Now "cat" at position 1 is different from "cat" at position 5

// ❌ Single attention head — limited perspective
function singleHeadAttention(embeddings: number[][]): number[][] {
  return selfAttention(embeddings, Wq, Wk, Wv);
  // Can only capture one type of relationship at a time
}
 
// ✅ Multi-head attention — multiple perspectives combined
function multiHeadAttention(
  embeddings: number[][],
  numHeads: number,
  embeddingDim: number
): number[][] {
  const headDim = embeddingDim / numHeads;
  const headOutputs: number[][][] = [];
 
  for (let h = 0; h < numHeads; h++) {
    // Each head has its own Q, K, V projections
    const Wq_h = getWeightMatrix(h, 'query', headDim);
    const Wk_h = getWeightMatrix(h, 'key', headDim);
    const Wv_h = getWeightMatrix(h, 'value', headDim);
 
    const headOutput = selfAttention(embeddings, Wq_h, Wk_h, Wv_h);
    headOutputs.push(headOutput);
  }
 
  // Concatenate all head outputs and project back to embedding dimension
  const concatenated = concatHeads(headOutputs);
  const projected = matMul(concatenated, Wo);  // Output projection
 
  return projected;
}
 
// GPT-3 175B uses 96 attention heads
// GPT-4 likely uses even more
// Each head learns to focus on different linguistic relationships

// One Transformer block (decoder-only, like GPT)
function transformerBlock(
  input: number[][],           // [seqLength x embeddingDim]
  numHeads: number,
  ffDim: number                // Feed-forward hidden dimension (usually 4x embedding)
): number[][] {
  // Step 1: Multi-head self-attention with residual connection
  const attended = multiHeadAttention(input, numHeads, input[0].length);
  const afterAttn = layerNorm(addMatrices(input, attended));  // Residual + norm
 
  // Step 2: Feed-forward network with residual connection
  // This is where the model "thinks" — applies learned transformations
  const ffOutput = feedForwardNetwork(afterAttn, ffDim);
  const output = layerNorm(addMatrices(afterAttn, ffOutput));  // Residual + norm
 
  return output;
}
 
function feedForwardNetwork(
  input: number[][],
  hiddenDim: number
): number[][] {
  // Two linear transformations with GELU activation between
  const hidden = gelu(matMul(input, W1));   // Project up to hiddenDim
  const output = matMul(hidden, W2);         // Project back to embeddingDim
  return output;
}
 
// Full model: stack N blocks
function transformer(
  tokenIds: number[],
  numLayers: number,
  numHeads: number,
  embeddingDim: number,
  ffDim: number
): number[][] {
  // Embed tokens and add positional encoding
  let x = addMatrices(
    tokenEmbedding(tokenIds, embeddingDim),
    positionalEncoding(tokenIds.length, embeddingDim)
  );
 
  // Pass through N transformer blocks
  for (let layer = 0; layer < numLayers; layer++) {
    x = transformerBlock(x, numHeads, ffDim);
  }
 
  return x;
}
 
// GPT-3 175B: 96 layers, 96 heads, embedding dim 12288, ff dim 49152
// That's a LOT of matrix multiplications