Writing the underlying pieces as ordinary code exposes their intermediate data and operations. This tutorial builds three of them from scratch in TypeScript: a byte-pair encoding (BPE) tokenizer, cosine similarity vector search, and scaled dot-product attention.
Why Build LLM Primitives in TypeScript?
Products like ChatGPT and Claude can make LLM primitives look like a black box. Writing the underlying pieces as ordinary code exposes their intermediate data and operations. This tutorial builds three of them from scratch in TypeScript: a byte-pair encoding (BPE) tokenizer, cosine similarity vector search, and scaled dot-product attention. Each exposes a different mechanism: BPE converts text into token IDs, vector similarity compares numerical representations, and attention computes a weighted combination of value vectors from query-key scores.
These are three inspectable primitives used in LLM-related workflows. They are not an LLM. The Transformer described in "Attention Is All You Need" also relies on learned projections, positional information, feed-forward sublayers, residual connections, and normalization. Implementing attention does not reproduce the behavior of any production model.
These are three inspectable primitives used in LLM-related workflows. They are not an LLM.
TypeScript is useful here because interfaces document data flow. The contract is simple: the tokenizer emits number[], a toy embedding lookup maps each ID to a fixed-dimension Float64Array, and attention consumes an ordered sequence of those vectors. Type annotations do not verify matrix dimensions at runtime, however, so every function below checks shapes explicitly.
The result is a CPU-based learning sandbox. It is not a compatible tokenizer, a trained embedding model, a RAG application, or a performant LLM.
Prerequisites and Project Setup
This article targets Node 24.12+ and TypeScript 7.x. Verify the release and native-TypeScript requirements against the official Node.js and TypeScript documentation before relying on them. The version details that follow reflect the time of writing, so confirm them against those official sources. TypeScript 7.0 shipped in July 2026. Its native compiler does not yet expose the older programmatic compiler API. A plain tsc workflow is unaffected, but tools that depend on that API need a compatibility check. Node 24.21.0 is the current LTS. Node 22.18.0 and later enable native type stripping by default, but the feature became stable on the 24.x line at 24.12.0. Readers staying on Node 22 need at least 22.18.
Prerequisites: Node.js 24.12+ with npm available. After creating the shown files, install the development dependencies before running the scripts.
"Zero-dependency" has a precise meaning here: no external ML/math libraries and no runtime dependencies. The project installs TypeScript and Node type declarations only as development dependencies for type-checking. Tests use Node's built-in runner instead of Vitest, so you don't need to install a test framework. The test script passes an explicit glob for tests/**/*.test.ts rather than relying on default discovery, which can report zero tests and still exit successfully. Confirm the glob and missing-file behavior on your target Node version.
One distinction matters throughout. Native Node strips types without type-checking them. A strict tsconfig.json does not make node src/demo.ts perform strict checks. The typecheck script must run separately.
{
"name": "llm-primitives-ts",
"version": "0.1.0",
"private": true,
"type": "module",
"engines": { "node": ">=24.12.0" },
"scripts": {
"typecheck": "tsc",
"test": "node --test \"tests/**/*.test.ts\"",
"demo": "node src/demo.ts"
},
"devDependencies": {
"@types/node": "24.0.0",
"typescript": "7.0.0"
}
}The example pins the development dependency versions exactly so that every reader resolves the same toolchain. Before first use, confirm that the pinned TypeScript release exists and ships a tsc binary (npm view typescript@7.0.0 bin), commit the generated lockfile, and update the pins deliberately when you revalidate the article. By default, npm only warns when the installed Node version does not satisfy the engines field. Add an .npmrc next to package.json to make the check enforce:
engine-strict=true```
{
"compilerOptions": {
"target": "ES2022",
"module": "nodenext",
"strict": true,
"noEmit": true,
"rewriteRelativeImportExtensions": true,
"erasableSyntaxOnly": true,
"verbatimModuleSyntax": true,
"types": ["node"]
},
"include": ["src", "tests"]
}
llm-primitives-ts/ ├── .npmrc ├── package.json ├── tsconfig.json ├── src/ │ ├── tokenizer/bpe.ts │ ├── vectors/index.ts │ ├── attention/attention.ts │ └── demo.ts └── tests/ ├── bpe.test.ts ├── vectors.test.ts ├── attention.test.ts └── demo.test.ts
The pipeline section below lists every test file in full. `npm test` only verifies the properties those files assert.
Node's own recommended configuration uses `target: "esnext"`. The `ES2022` target here is a deliberate compatibility choice. It does not cause native Node to transpile syntax that Node cannot run.
Every example executes `.ts` files directly. For that to work, relative runtime imports must use the `.ts` extension, and type-only imports must use `import type`. Mixing this approach with a compile-to-`.js` import strategy silently breaks reproducibility. Native Node also ignores `tsconfig.json` path aliases. It does not support runtime `enum`, parameter properties, or `.tsx` under ordinary type stripping, so the code below avoids all three. As with the version details above, confirm these runtime behaviors against the official Node.js TypeScript documentation for your target release.
## Part 1: Building a Byte-Pair Encoding (BPE) Tokenizer
### What Tokenization Actually Solves
A model operates on integers, not strings. Word-level vocabularies explode in size and still miss unseen words. Character-level sequences are compact in vocabulary but long, which makes every downstream sequence operation more expensive. Subword tokenization sits between the two: frequent fragments get their own IDs, and rare strings decompose into smaller pieces. The trade-off is vocabulary size versus sequence length.
Character-level BPE has a subtle failure mode. If the training corpus omits a character that later appears at encoding time, the tokenizer cannot necessarily encode it. **Byte-level BPE** avoids this. It initializes the vocabulary with all 256 byte values, so any valid UTF-8 text is representable, emoji and unseen scripts included. OpenAI's educational `tiktoken` implementation uses this design (check its source for the exact details), and this article builds the same design.
### The BPE Algorithm, Step by Step
1. Split text into pieces with a regex before anything else. Merges happen only *within* pieces, never across them. Whether the split retains separators matters a great deal: a whitespace split that discards spaces makes`decode(encode(text)) === text` impossible. The implementation below therefore checks that the pieces concatenate back to the input and throws if they do not. OpenAI's educational implementation applies the same regex pre-tokenization during training and encoding, which makes the pre-tokenizer part of the tokenizer's definition.
2. Next, encode each piece as UTF-8 bytes and count every adjacent pair of symbols across all pieces.
3. The core loop selects the most frequent pair, assigns it a new ID, replaces its occurrences, recounts, and repeats. It stops after `numMerges` merges or when no adjacent pair remains. Tie-breaking must be deterministic; here, ties go to the smallest left ID, then the smallest right ID.
4. Every merge receives an increasing rank. Starting from 256 base IDs, the vocabulary grows by exactly one token per *successful* merge, so`numMerges` caps growth rather than fixing the final vocabulary size. The tokenizer defines no special tokens; the educational`tiktoken` implementation omits them too.
Encoding differs from training in one important way. It does **not** pick the most frequent pair in the new input. It repeatedly applies the eligible pair with the **lowest learned rank**, replaying merges in the order they were learned.
### Implementing BPE in TypeScript
The `Vocabulary`, `MergeRule`, and `Tokenizer` types are article-defined, not standard-library APIs. Each token stores its exact byte sequence. A display string would be insufficient because a merged token can be a fragment of a multi-byte character.
// src/tokenizer/bpe.ts /** One base token per possible byte value. */ export const BYTE_VOCAB_SIZE = 256; /**
- Simplified GPT-2-style pre-tokenization. Every character matches exactly one
- branch (letters, numbers, other symbols, whitespace), so the pieces always
concatenate back to the original text: separators are retained, not dropped.
*/ export const DEFAULT_PATTERN = String.raw
?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+; export interface MergeRule { readonly left: number; // ID of the left symbol readonly right: number; // ID of the right symbol readonly rank: number; // lower rank = learned earlier = applied first readonly id: number; // ID of the merged token (BYTE_VOCAB_SIZE + rank) } export interface Vocabulary { /** idToBytes[id] is the exact byte sequence represented by that token ID. */ readonly idToBytes: readonly Uint8Array[]; } export interface Tokenizer { readonly vocab: Vocabulary; readonly merges: readonly MergeRule[]; /** Persist this with the merges: changing it changes the tokenizer. */ readonly pattern: string; } export interface PairCount { readonly left: number; readonly right: number; readonly count: number; } export function pairKey(left: number, right: number): string { return${left},${right}; } export function preTokenize(text: string, pattern: string = DEFAULT_PATTERN): string[] { const pieces = Array.from(text.matchAll(new RegExp(pattern, 'gu')), (match) => match[0]); if (pieces.join('') !== text) { throw new Error('Pre-tokenizer pattern does not cover the input; round trip would be lossy'); } return pieces; } /**- Counts adjacent symbol pairs inside each piece (never across pieces).
- Overlapping occurrences are counted ("aaa" counts (a,a) twice), while
mergePair replaces non-overlapping ones, so the count can exceed the merge yield.
*/ export function getPairFrequencies( pieces: readonly (readonly number[])[], ): Map<string, PairCount> { const counts = new Map<string, PairCount>(); for (const piece of pieces) { for (let i = 0; i < piece.length - 1; i++) { const left = piece[i]; const right = piece[i + 1]; const key = pairKey(left, right); const previous = counts.get(key)?.count ?? 0; counts.set(key, { left, right, count: previous + 1 }); } } return counts; } /** Highest count wins; ties go to the smallest left ID, then smallest right ID. */ function pickMostFrequent(counts: ReadonlyMap<string, PairCount>): PairCount | undefined { let best: PairCount | undefined; for (const candidate of counts.values()) { if ( best === undefined || candidate.count > best.count || (candidate.count === best.count && (candidate.left < best.left || (candidate.left === best.left && candidate.right < best.right))) ) { best = candidate; } } return best; } /** Replaces non-overlapping occurrences of (left, right), scanning left to right. */ export function mergePair( symbols: readonly number[], left: number, right: number, newId: number, ): number[] { const out: number[] = []; let i = 0; while (i < symbols.length) { if (i < symbols.length - 1 && symbols[i] === left && symbols[i + 1] === right) { out.push(newId); i += 2; } else { out.push(symbols[i]); i += 1; } } return out; } export function mergeMostFrequentPair( pieces: readonly (readonly number[])[], newId: number, ): { pieces: number[][]; merged: PairCount } | undefined { const best = pickMostFrequent(getPairFrequencies(pieces)); if (best === undefined) return undefined; // no adjacent pair left anywhere const { left, right } = best; return { pieces: pieces.map((p) => mergePair(p, left, right, newId)), merged: best }; } function concatBytes(a: Uint8Array, b: Uint8Array): Uint8Array { const out = new Uint8Array(a.length + b.length); out.set(a, 0); out.set(b, a.length); return out; } export function trainBPE(corpus: string, numMerges: number): Tokenizer { if (!Number.isInteger(numMerges) || numMerges < 0) { throw new RangeError(
numMerges must be a non-negative integer, got ${numMerges}); } // TextEncoder would silently replace unpaired surrogates with U+FFFD. if (/\p{Cs}/u.test(corpus)) throw new RangeError('Corpus contains unpaired surrogates'); const encoder = new TextEncoder(); let pieces: number[][] = preTokenize(corpus).map((p) => Array.from(encoder.encode(p))); const idToBytes: Uint8Array[] = []; for (let byte = 0; byte < BYTE_VOCAB_SIZE; byte++) idToBytes.push(Uint8Array.of(byte)); const merges: MergeRule[] = []; for (let rank = 0; rank < numMerges; rank++) { const id = BYTE_VOCAB_SIZE + rank; const step = mergeMostFrequentPair(pieces, id); if (step === undefined) break; // stop early: vocabulary may be < 256 + numMerges pieces = step.pieces; const { left, right } = step.merged; merges.push({ left, right, rank, id }); idToBytes.push(concatBytes(idToBytes[left], idToBytes[right])); } return { vocab: { idToBytes }, merges, pattern: DEFAULT_PATTERN }; }
### Encoding and Decoding Text
The outline's `encode(text)` becomes `encode(tokenizer, text)`. Passing the tokenizer explicitly keeps its state as plain, persistable data. `getLookup` builds the merge lookup table once per tokenizer object and caches it, which assumes callers treat tokenizers as immutable, as their `readonly` types declare. Append the following to the same file:
// src/tokenizer/bpe.ts (continued)
// Assumes Tokenizer objects are treated as immutable (as their readonly types declare).
const lookupCache = new WeakMap<Tokenizer, Map<string, MergeRule>>();
function getLookup(tokenizer: Tokenizer): Map<string, MergeRule> {
let lookup = lookupCache.get(tokenizer);
if (lookup === undefined) {
lookup = new Map();
for (const rule of tokenizer.merges) {
const key = pairKey(rule.left, rule.right);
const existing = lookup.get(key);
if (existing === undefined || rule.rank < existing.rank) lookup.set(key, rule);
}
lookupCache.set(tokenizer, lookup);
}
return lookup;
}
export function encode(tokenizer: Tokenizer, text: string): number[] {
// TextEncoder would silently replace unpaired surrogates with U+FFFD.
if (/\p{Cs}/u.test(text)) throw new RangeError('Input contains unpaired surrogates');
const lookup = getLookup(tokenizer);
const encoder = new TextEncoder();
const ids: number[] = [];
for (const piece of preTokenize(text, tokenizer.pattern)) {
let symbols = Array.from(encoder.encode(piece)); // start from raw bytes 0..255
while (symbols.length > 1) {
// Apply the LOWEST-ranked learned merge present, not the most frequent pair.
let best: MergeRule | undefined;
for (let i = 0; i < symbols.length - 1; i++) {
const rule = lookup.get(pairKey(symbols[i], symbols[i + 1]));
if (rule !== undefined && (best === undefined || rule.rank < best.rank)) best = rule;
}
if (best === undefined) break;
symbols = mergePair(symbols, best.left, best.right, best.id);
}
for (const s of symbols) ids.push(s); // no spread: safe for arbitrarily long pieces
}
return ids;
}
export function decode(tokenizer: Tokenizer, tokenIds: readonly number[]): string {
const table = tokenizer.vocab.idToBytes;
let total = 0;
for (let i = 0; i < tokenIds.length; i++) {
const id = tokenIds[i];
if (!Number.isInteger(id) || id < 0 || id >= table.length) {
throw new RangeError(Unknown token ID: ${id} at index ${i});
}
total += table[id].length;
}
// Concatenate ALL bytes first: a single token may hold part of a multi-byte character.
const bytes = new Uint8Array(total);
let offset = 0;
for (const id of tokenIds) {
bytes.set(table[id], offset);
offset += table[id].length;
}
// fatal: true throws on malformed UTF-8 instead of inserting U+FFFD.
// ignoreBOM: true keeps a leading U+FEFF instead of silently stripping it.
return new TextDecoder('utf-8', { fatal: true, ignoreBOM: true }).decode(bytes);
}
Training on `"low lower lowest"` with three merges proceeds as follows. Pre-tokenization yields `"low"`, `" lower"`, `" lowest"`. The pairs `l+o` and `o+w` tie at three occurrences, and the tie-break selects `l+o` (ID 256). Next comes `lo+w` (257). Finally `" "+low` ties with `low+e` at two occurrences and wins on the smaller left ID (258). Running `encode` and `decode` on a few inputs produces:
"lowest low" -> [257, 101, 115, 116, 258] -> "lowest low" "héllo 👋" -> [104, 195, 169, 108, 256, 32, 240, 159, 145, 139] -> "héllo 👋" vocabulary size: 259 (256 bytes + 3 merges)
The second line shows why this tutorial uses byte-level BPE. `é` and `👋` never appear in the corpus, so they fall back to their raw UTF-8 bytes and still round-trip.
The error policy is deliberate. Unknown IDs throw, and the message reports their position. Malformed UTF-8 also throws, and this can happen with arbitrary ID sequences, such as unconstrained model-generated IDs that cut a character in half. `encode` rejects input containing unpaired surrogates up front, because UTF-8 cannot represent it without loss. `decode` preserves a leading byte-order mark (U+FEFF) rather than stripping it. Applications that stream partial output need a different policy.
A test suite should cover repeated pairs, tied frequencies, empty input, tabs and newlines, leading and trailing spaces, a leading U+FEFF, very long single pieces, emoji, and Unicode absent from training. Any lossy normalization or change to the pre-tokenizer voids the round-trip guarantee, which is why the tokenizer saves `pattern` alongside `merges`. A persisted pattern that fails to cover some input now throws instead of silently dropping text. Because `preTokenize` compiles the persisted pattern into a `RegExp`, treat tokenizer files as trusted input only; a pattern from untrusted storage can cause catastrophic backtracking or unintended splitting. Encoding cost also grows quadratically with the length of a single piece in the worst case, so set a maximum input length before encoding untrusted text. This article supplies no universal limit; choose one from the worst-case encoding time and memory use you measure in your deployment. Byte-level BPE alone does not make this GPT-2 or `tiktoken` compatible. Compatibility would also require matching pre-tokenization, ranks, and special-token handling.
## Part 2: Representing Meaning as Vectors
### From Tokens to Embeddings (Conceptually)
In a model, an embedding table maps each token ID to a learned vector. This tutorial substitutes deterministic mock vectors. That substitution has a direct consequence. Cosine similarity ranks **existing** vectors by direction; it does not create meaning. Random or hand-assigned vectors can exercise the calculation and the index, but their search results say nothing about semantic retrieval. Embedding-based retrieval, as OpenAI's embeddings guide describes it, depends on vectors produced by an embedding model.
Cosine similarity ranks **existing** vectors by direction; it does not create meaning.
The examples store vectors in `Float64Array` values, which use 8 bytes per element. The choice is about transparent, fixed-width arithmetic. It is not a claim about throughput. Nothing in the typed-array specification establishes a performance advantage, and a real system would need to benchmark precision, memory, and deployment trade-offs.
### Measuring Similarity: The Cosine Similarity Function
The dot product sums element-wise products. Magnitude is the square root of a vector's dot product with itself. Cosine similarity divides the dot product by the product of the two magnitudes, giving a value from -1 (opposite) through 0 (orthogonal) to 1 (same direction).
Squaring components directly can overflow to `Infinity` for very large values or underflow to 0 for very small ones, which would produce `NaN` scores or wrongly reject valid vectors. The implementation therefore scales each vector by its largest absolute component before squaring, normalizes each vector before taking the dot product, and clamps the result to [-1, 1] to absorb rounding drift.
The formula is undefined when either magnitude is zero. Returning `NaN` or silently scoring zero would hide bugs, so the implementation throws. pgvector likewise excludes zero vectors from cosine-distance indexes (confirm this in the pgvector documentation for your version).
// src/vectors/index.ts
function assertFinite(v: Float64Array, name: string): void {
for (let i = 0; i < v.length; i++) {
if (!Number.isFinite(v[i])) throw new RangeError(${name}[${i}] is not finite);
}
}
function assertSameDimension(a: Float64Array, b: Float64Array): void {
if (a.length !== b.length) {
throw new RangeError(Dimension mismatch: ${a.length} vs ${b.length});
}
}
export function dotProduct(a: Float64Array, b: Float64Array): number {
assertSameDimension(a, b);
let sum = 0;
for (let i = 0; i < a.length; i++) sum += a[i] * b[i];
return sum;
}
/** Scaled Euclidean norm: avoids overflow and underflow when squaring components. */
export function magnitude(v: Float64Array): number {
let scale = 0;
for (let i = 0; i < v.length; i++) scale = Math.max(scale, Math.abs(v[i]));
if (scale === 0 || !Number.isFinite(scale)) return scale;
let sum = 0;
for (let i = 0; i < v.length; i++) {
const x = v[i] / scale;
sum += x * x;
}
return scale * Math.sqrt(sum);
}
export function cosineSimilarity(a: Float64Array, b: Float64Array): number {
assertSameDimension(a, b);
assertFinite(a, 'a');
assertFinite(b, 'b');
const magA = magnitude(a);
const magB = magnitude(b);
if (magA === 0 || magB === 0) {
throw new RangeError('Cosine similarity is undefined for a zero-magnitude vector');
}
if (!Number.isFinite(magA) || !Number.isFinite(magB)) {
throw new RangeError('Vector magnitude overflows float64');
}
let dot = 0;
for (let i = 0; i < a.length; i++) dot += (a[i] / magA) * (b[i] / magB);
return Math.min(1, Math.max(-1, dot)); // clamp rounding drift
}
### Building a Minimal Vector Index
The index validates dimensions and `topK` and returns each result's ID and score. It breaks score ties by ID in ascending UTF-16 code-unit order, so numeric-string IDs do not sort numerically (`"10"` sorts before `"9"`). `add()` validates the input once and stores a unit-normalized copy, so a caller who later mutates the original cannot corrupt previously indexed data. `search()` validates and normalizes the query once, even when the index is empty, and then scores each entry with a single dot product. The class uses private `#` fields because they are plain JavaScript and survive type stripping. Parameter properties would not.
// src/vectors/index.ts (continued)
export interface SearchResult {
readonly id: string;
readonly score: number;
}
function toUnit(v: Float64Array, name: string): Float64Array {
assertFinite(v, name);
const mag = magnitude(v);
if (mag === 0) throw new RangeError(Zero vector rejected: ${name});
if (!Number.isFinite(mag)) throw new RangeError(Magnitude overflows float64: ${name});
return Float64Array.from(v, (x) => x / mag); // also serves as the defensive copy
}
export class VectorIndex {
readonly dimension: number;
readonly #entries: { id: string; vector: Float64Array }[] = [];
readonly #ids = new Set<string>();
constructor(dimension: number) {
if (!Number.isInteger(dimension) || dimension < 1) {
throw new RangeError(dimension must be a positive integer, got ${dimension});
}
this.dimension = dimension;
}
get size(): number {
return this.#entries.length;
}
add(id: string, vector: Float64Array): void {
if (vector.length !== this.dimension) {
throw new RangeError(Expected dimension ${this.dimension}, got ${vector.length});
}
if (this.#ids.has(id)) throw new Error(Duplicate id: ${id});
const unit = toUnit(vector, id);
this.#ids.add(id);
this.#entries.push({ id, vector: unit });
}
/** Linear scan: O(count x dimension) scoring plus O(count log count) sort per query. */
search(query: Float64Array, topK: number): SearchResult[] {
if (!Number.isInteger(topK) || topK < 1) {
throw new RangeError(topK must be a positive integer, got ${topK});
}
if (query.length !== this.dimension) {
throw new RangeError(Expected dimension ${this.dimension}, got ${query.length});
}
const q = toUnit(query, 'query'); // validated even when the index is empty
const scored = this.#entries.map((e) => ({
id: e.id,
score: Math.min(1, Math.max(-1, dotProduct(q, e.vector))),
}));
scored.sort((x, y) => y.score - x.score || (x.id < y.id ? -1 : x.id > y.id ? 1 : 0));
return scored.slice(0, topK); // empty index -> []; topK > count -> all results
}
}
This scan scores every vector and returns the first k under cosine score followed by the documented ID tie-break. When several vectors tie at the k-th score, more than one mathematically valid top-k set exists; the ascending-ID rule deterministically selects one of them. Scoring costs O(Nd); this implementation then sorts all N scores, adding O(N log N). Production approximate nearest-neighbor indexes such as HNSW and IVFFlat give up some recall and spend extra memory and build effort so that a query does not have to score every stored vector. Benchmark query latency on your own workload before choosing one. pgvector's documentation notes that it performs exact search by default until you add an approximate index (check the documentation for your pgvector version). This article supplies no universal collection-size cutoff. A linear scan is a reasonable choice when measured query latency at your expected N and d fits your budget, or wherever exactness matters more than latency.
The index relies on one optimization, and it works only because `add()` and `search()` normalize the vectors. When vectors have unit length, a dot product produces the same ranking as cosine similarity, which is why both methods normalize before scoring. OpenAI documents this for its own normalized embedding outputs; confirm the exact models covered in OpenAI's embeddings guide. The shortcut does not hold for arbitrary toy vectors that nobody has normalized first.
This index is related to RAG, but it implements nearest-vector ranking only. It does not chunk documents, generate embeddings, filter on metadata, or assemble prompts. Tests should cover orthogonal, opposite, and identical vectors; mismatched lengths; zero and non-finite vectors; very large and very small magnitudes; empty indexes; `topK > count`; and deterministic ties.
## Part 3: Self-Attention: Computing Context-Dependent Token Representations
The outline's heading framed attention as "the mechanism behind understanding." The equation supports a narrower claim: it computes weighted combinations of value vectors. The heading above reflects that narrower claim.
The equation supports a narrower claim: it computes weighted combinations of value vectors.
### Why Attention Was a Breakthrough
Attention lets every position in a sequence score its relevance to every other position, then build its output as a weighted average of their values. Each token's representation therefore depends on its context.
### The Scaled Dot-Product Attention Formula
The operation from "Attention Is All You Need" is `softmax(QKᵀ / √dₖ)V`. The algorithm applies softmax across **each query's row** of key scores. The shapes are:
- `Q` is`L×dₖ` : one query row per output position.
- `K` is`S×dₖ` : one key row per attended position.
- `V` is`S×dᵥ` : one value row per attended position.
The output is `L×dᵥ`. In self-attention, `L = S` is common.
Under the usual simplifying assumptions that query and key components are independent, zero-mean, and similarly scaled, dot-product variance grows with `dₖ`; dividing by `√dₖ` controls that growth. Without the `1/√dₖ` factor, score magnitudes keep growing with `dₖ` and push softmax toward saturation, which shrinks its gradients. Softmax then converts each row of scores into non-negative weights that sum to one.
In a real Transformer, Q, K, and V come from **learned projections** of the input. Setting `Q = K = V = embeddings` is a deliberate simplification for this demonstration.
### Implementing Matrix Operations from Scratch
The outline specifies `number[][]`. That conflicts with the `Float64Array` constraint, so matrices here are `readonly Float64Array[]` (one typed array per row). These are proposed article interfaces, not standard APIs. The matrix helpers reject ragged matrices and incompatible shapes.
Softmax uses the numerically stable form `exp(xᵢ − max) / Σ exp(xⱼ − max)`. A row that is entirely `-Infinity` (fully masked) throws. PyTorch issue #103749 documents how that case otherwise turns into `NaN`.
// src/attention/attention.ts
export type Matrix = readonly Float64Array[];
/** Returns [rows, cols]; rejects empty and ragged matrices. */
export function shape(m: Matrix): [number, number] {
if (m.length === 0 || m[0].length === 0) throw new RangeError('Matrix must be non-empty');
const cols = m[0].length;
for (const row of m) {
if (row.length !== cols) throw new RangeError('Ragged matrix: rows differ in length');
}
return [m.length, cols];
}
export function matmul(a: Matrix, b: Matrix): Float64Array[] {
const [n, k] = shape(a);
const [kb, m] = shape(b);
if (k !== kb) throw new RangeError(matmul: ${n}x${k} cannot multiply ${kb}x${m});
const out = Array.from({ length: n }, () => new Float64Array(m));
for (let i = 0; i < n; i++) {
for (let p = 0; p < k; p++) {
const aip = a[i][p];
for (let j = 0; j < m; j++) out[i][j] += aip * b[p][j];
}
}
return out;
}
export function transpose(matrix: Matrix): Float64Array[] {
const [rows, cols] = shape(matrix);
const out = Array.from({ length: cols }, () => new Float64Array(rows));
for (let i = 0; i < rows; i++) {
for (let j = 0; j < cols; j++) out[j][i] = matrix[i][j];
}
return out;
}
/** Stable softmax. -Infinity entries (masked) get weight 0; NaN/+Infinity are rejected. */
export function softmax(row: Float64Array): Float64Array {
if (row.length === 0) throw new RangeError('softmax: empty row');
let max = -Infinity;
for (const x of row) {
if (Number.isNaN(x) || x === Infinity) throw new RangeError('softmax: invalid score');
if (x > max) max = x;
}
if (max === -Infinity) throw new RangeError('softmax: row is fully masked');
const out = new Float64Array(row.length);
let sum = 0;
for (let i = 0; i < row.length; i++) {
out[i] = Math.exp(row[i] - max); // subtracting max prevents overflow
sum += out[i];
}
for (let i = 0; i < out.length; i++) out[i] /= sum;
return out;
}
### Putting It Together: The Attention Function
The function returns the attention **weights** alongside the output. Looking only at the output hides the weighted-average step, which is the central part of the mechanism. An optional causal mask sets future-key scores to `-Infinity` *before* softmax. The mask never touches the diagonal, so no row can end up fully masked.
// src/attention/attention.ts (continued)
export interface AttentionOptions {
/** Mask future positions (decoder-style). Requires Q and K to have the same length. */
readonly causal?: boolean;
}
export interface AttentionResult {
readonly weights: Float64Array[]; // L x S, each row sums to ~1
readonly output: Float64Array[]; // L x dV
}
export function scaledDotProductAttention(
Q: Matrix,
K: Matrix,
V: Matrix,
options: AttentionOptions = {},
): AttentionResult {
const [L, dK] = shape(Q); // Q: L x dK
const [S, dKk] = shape(K); // K: S x dK
const [Sv] = shape(V); // V: S x dV
if (dK !== dKk) throw new RangeError(Q and K key dims differ: ${dK} vs ${dKk});
if (S !== Sv) throw new RangeError(K and V lengths differ: ${S} vs ${Sv});
if (options.causal === true && L !== S) {
throw new RangeError('Causal mask requires L === S');
}
// QK^T: raw query-key scores, shape L x S
const scores = matmul(Q, transpose(K));
// Divide by sqrt(dK) so score variance does not grow with dK; mask before softmax.
const scale = 1 / Math.sqrt(dK);
for (let i = 0; i < L; i++) {
for (let j = 0; j < S; j++) {
scores[i][j] = options.causal === true && j > i ? -Infinity : scores[i][j] * scale;
}
}
// softmax(...) applied independently to each query's row
const weights = scores.map((row) => softmax(row));
// (...)V: each output row is a weighted average of V's rows, shape L x dV
const output = matmul(weights, V);
return { weights, output };
}
**Worked example.** Take three tokens with `Q = K = V = I₃`, the 3×3 identity matrix, and no mask. `QKᵀ` is the identity. After scaling, each diagonal score is `1/√3 ≈ 0.5774` and each off-diagonal score is `0`. The softmax of `[0.5774, 0, 0]` is approximately `[0.4711, 0.2645, 0.2645]`, so the weights are:
weights (unmasked) output = weights x I [0.4711, 0.2645, 0.2645] [0.4711, 0.2645, 0.2645] [0.2645, 0.4711, 0.2645] [0.2645, 0.4711, 0.2645] [0.2645, 0.2645, 0.4711] [0.2645, 0.2645, 0.4711] weights (causal: true) [1.0000, 0.0000, 0.0000] [0.3595, 0.6405, 0.0000] [0.2645, 0.2645, 0.4711]
Because `V` is the identity, each output row equals its weight row. The example gives concrete expected values without implying trained "understanding."
The masked version highlights an important constraint. Unmasked attention is bidirectional: token 0 draws on tokens 1 and 2. Decoder-style next-token prediction requires the causal mask, so the unmasked version is not an LLM's next-token attention. Attention alone also carries no notion of token order. The original Transformer adds positional information to supply it.
Cost is the other constraint. The score matrix grows quadratically in self-attention. Keep tutorial inputs bounded, and before accepting untrusted inputs, set a sequence-length limit based on the score-matrix memory and latency you measure in your deployment. The same caution applies to the other primitives: BPE training recounts every pair after each merge, and the vector index scores and sorts every stored entry per query. Production frameworks such as PyTorch can select optimized attention kernels instead of materializing this sequence of ordinary matrix operations.
## Wiring the Primitives Into a Mini Pipeline
After embedding lookup, the pipeline splits into **two branches**. Vector search ranks stored vectors against a query. Self-attention contextualizes an ordered sequence. Neither feeds the other, so drawing them as a single chain would misrepresent their roles.
The lookup table covers every vocabulary ID, including learned merges. A table sized only for the 256 base bytes would fail as soon as `encode` emits ID 256 or higher. The demo generates vectors deterministically; fully reproducible runs additionally require a pinned toolchain and dependencies.
// src/demo.ts import { trainBPE, encode, decode } from './tokenizer/bpe.ts'; import { VectorIndex } from './vectors/index.ts'; import { scaledDotProductAttention } from './attention/attention.ts'; const DIM = 4; const tokenizer = trainBPE('low lower lowest', 3); const ids = encode(tokenizer, 'lowest low'); if (ids.length === 0) throw new Error('encode produced no tokens; nothing to search or attend over'); // Deterministic mock embeddings for EVERY ID: 256 bytes + learned merges. const table = tokenizer.vocab.idToBytes.map((_, id) => Float64Array.from({ length: DIM }, (_, d) => Math.sin((id + 1) * (d + 1)))); const sequence = ids.map((id) => table[id]); // Branch A: exact cosine neighbours over toy token vectors (not semantic retrieval). const index = new VectorIndex(DIM); table.forEach((vector, id) => index.add(String(id), vector)); console.log('exact toy cosine neighbours:', index.search(sequence[0], 3)); // Branch B: untrained, unprojected, unmasked self-attention over the ordered sequence. const { weights, output } = scaledDotProductAttention(sequence, sequence, sequence); console.log('token ids:', ids, '| round trip:', decode(tokenizer, ids)); console.log('attention weights:', weights); console.log('untrained attention outputs:', output);
Branch A indexes the **token**, not a document chunk. The top neighbor is the query token itself, with a score of 1 (up to floating-point rounding). The output labels make the difference between the two branches explicit.
`npm test` verifies only what the test files assert. The following files cover the tokenizer, vector, and attention properties discussed above, plus one end-to-end run of the demo.
// tests/bpe.test.ts import { test } from 'node:test'; import assert from 'node:assert/strict'; import { trainBPE, encode, decode } from '../src/tokenizer/bpe.ts'; const tok = trainBPE('low lower lowest', 3); test('matches the documented trace', () => { assert.deepEqual(encode(tok, 'lowest low'), [257, 101, 115, 116, 258]); assert.equal(tok.vocab.idToBytes.length, 259); }); test('round-trips a leading BOM, whitespace, and unseen Unicode', () => { for (const s of ['\uFEFFabc', '', ' \t low ', 'héllo 👋', '日本語']) { assert.equal(decode(tok, encode(tok, s)), s); } }); test('long single piece does not overflow the stack', () => { assert.equal(encode(tok, 'x'.repeat(500_000)).length, 500_000); }); test('rejects lone surrogates and unknown IDs', () => { assert.throws(() => encode(tok, 'a\uD800b'), RangeError); assert.throws(() => decode(tok, [9999]), RangeError); });
// tests/vectors.test.ts import { test } from 'node:test'; import assert from 'node:assert/strict'; import { cosineSimilarity, VectorIndex } from '../src/vectors/index.ts'; test('extreme magnitudes stay finite and correct', () => { assert.equal(cosineSimilarity(Float64Array.of(1e200), Float64Array.of(1e200)), 1); const idx = new VectorIndex(1); assert.doesNotThrow(() => idx.add('tiny', Float64Array.of(1e-200))); }); test('deterministic ties and zero query rejected on empty index', () => { const idx = new VectorIndex(2); idx.add('b', Float64Array.of(1, 0)); idx.add('a', Float64Array.of(2, 0)); assert.deepEqual(idx.search(Float64Array.of(1, 0), 2).map((r) => r.id), ['a', 'b']); assert.throws(() => new VectorIndex(2).search(Float64Array.of(0, 0), 1), RangeError); });
// tests/attention.test.ts import { test } from 'node:test'; import assert from 'node:assert/strict'; import { scaledDotProductAttention } from '../src/attention/attention.ts'; const I3 = [Float64Array.of(1, 0, 0), Float64Array.of(0, 1, 0), Float64Array.of(0, 0, 1)]; test('identity example and causal row values', () => { const { weights } = scaledDotProductAttention(I3, I3, I3); assert.ok(Math.abs(weights[0][0] - 0.4711) < 1e-4); for (const row of weights) assert.ok(Math.abs(row.reduce((s, x) => s + x, 0) - 1) < 1e-12); const causal = scaledDotProductAttention(I3, I3, I3, { causal: true }).weights; assert.equal(causal[0][1], 0); assert.ok(Math.abs(causal[1][1] - 0.6405) < 1e-4); });
// tests/demo.test.ts (integration) import { test } from 'node:test'; import assert from 'node:assert/strict'; import { spawnSync } from 'node:child_process'; test('demo runs end-to-end', () => { const r = spawnSync(process.execPath, ['src/demo.ts'], { encoding: 'utf8' }); assert.equal(r.status, 0, r.stderr); assert.match(r.stdout, /round trip: lowest low/); assert.match(r.stdout, /id: '257'/); // query token is its own top neighbour });
With the test files in place, install the pinned dependencies, then run type checking, tests, and the demo:
node --version # should satisfy the engines field (>=24.12.0) npm --version npm install # first run: generates package-lock.json; commit it npm ci # subsequent clean installs from the committed lockfile npm run typecheck npm test npm run demo
`npm test` runs `node --test` over `tests/**/*.test.ts`, and `npm run typecheck` runs `tsc` separately (with `noEmit` set in `tsconfig.json`), because running the tests does not type-check them. As a quick sanity check of the BOM round trip, run:
`node --input-type=module -e "import {trainBPE,encode,decode} from './src/tokenizer/bpe.ts'; const t=trainBPE('low lower lowest',3); console.log(decode(t,encode(t,'\uFEFFlow'))==='\uFEFFlow')"`
Expected output: `true`.
## What This Does (and Doesn't) Teach You About Real LLMs
This project implements single-head attention without learned embeddings or learned Q/K/V projections. The mask is optional and off in the demo. The code also leaves out training entirely (backpropagation and the training loop), along with positional information, multi-head attention, the rest of the Transformer block (feed-forward sublayers, residual connections, normalization), and GPU parallelism. PyTorch's attention API exposes masking, dropout, batch and head dimensions, and kernel selection, none of which appear here.
A clean round trip, or attention weights whose rows sum to one, respect the causal mask, and match the worked example, shows that the implementation behaves correctly. It does not demonstrate model quality, retrieval quality, or compatibility with any particular model.
A reasonable progression from here is: causal masking by default, then learned projections and positional information, then multi-head attention, then a full Transformer block.
## Key Takeaways
- The toy byte-level BPE tokenizer round-trips well-formed text it encoded itself (including a leading U+FEFF), provided its persisted pre-tokenizer stays unchanged, and it persists merge ranks. It rejects unpaired surrogates and patterns that do not cover the input, and arbitrary token-ID sequences can still fail UTF-8 decoding.
- Vector search here means an exhaustive, in-memory cosine-similarity scan over supplied vectors. It guards against overflow and underflow and breaks ties deterministically by ID.
- For attention, you get a single-head scaled dot-product implementation with shape checks, a stable softmax, and optional causal masking.
The project uses no external ML/math libraries and has no runtime dependencies. Each primitive can be tested on its own, and none of them has trained language ability.