Embeddings and Attention

Table of Contents

1. Why Words Need to Become Numbers
2. A Word as a Point in Space
3. Hand-Made Word Vectors
4. Dot Product and Norm
5. Cosine Similarity: How Close Are Two Words?
6. Nearest Neighbours
7. Real Embeddings Are Learned
8. Searching by Meaning
9. The Problem: One Vector per Word Is Not Enough
10. Attention: Words Looking at Words
11. Attention in Code
12. Reading the Attention Table
13. Context Changes the Meaning
14. The Real Transformer, Honestly
15. Recipe: The embeddings and attention Modules
16. Try It Yourself
17. Quick Review: Common Questions

1. Why Words Need to Become Numbers

A computer cannot do arithmetic on the word "coffee". In Tokens and Tokenizers you turned text into token numbers, and in How a Model Is Trained a word became a row number in a table. But a row number is just a label. Word number 12 is not "close" to word number 13 in any useful way, and the model has no idea that coffee and tea are alike while coffee and dog are not.

We want something better: a representation where similar words get similar numbers. Then whatever the model learns about coffee automatically helps with tea, because their numbers are near each other. That representation is called an embedding: a word (or token) turned into a vector, which is just a list of numbers.

This idea is one of the most useful in all of modern AI. It is how a model begins to get at meaning, and it also powers search by meaning, which you will use in later series.

2. A Word as a Point in Space

You already know a vector with two numbers: a point on a map, (x, y). Two numbers place a point on a flat page. Three numbers place a point in a room. Four, ten, or a thousand numbers place a point in a space we cannot picture, but the arithmetic works the same way.

Imagine a map where similar things are placed near each other:

           animal
             ^
   cat  dog  |
    puppy    |
             |
             +------------> drink
                  tea
                coffee  milk

cat, dog and puppy cluster together on one side. coffee, tea and milk cluster on another. The direction and distance between points carry meaning. A real embedding does this with hundreds or thousands of dimensions, but here we will use four, so you can read every number.

3. Hand-Made Word Vectors

In a real model nobody writes the numbers by hand. For learning, we will, so you can see why they work. Each word gets four numbers. We decide the four dimensions mean [drink, animal, tech, person]: how much the word is about each topic, from 0.0 to 1.0.

Create src/embeddings.js:

// src/embeddings.js
// Hand-made 4-number vectors. Dimensions: [drink, animal, tech, person].
// Real models LEARN their vectors, with hundreds or thousands of numbers.
export const vectors = {
  coffee:  [0.9, 0.0, 0.1, 0.0],
  tea:     [0.8, 0.0, 0.0, 0.1],
  milk:    [0.7, 0.2, 0.0, 0.0],
  cat:     [0.1, 0.9, 0.0, 0.1],
  dog:     [0.0, 0.9, 0.0, 0.2],
  puppy:   [0.0, 0.8, 0.0, 0.3],
  phone:   [0.0, 0.0, 0.9, 0.1],
  laptop:  [0.0, 0.0, 0.9, 0.2],
  app:     [0.0, 0.0, 0.8, 0.3],
  student: [0.1, 0.0, 0.2, 0.9],
  teacher: [0.0, 0.0, 0.2, 0.9],
  sokha:   [0.2, 0.0, 0.1, 0.9],
};

Twelve words, four numbers each. coffee is mostly "drink" with a hint of "tech" (laptops and coffee shops, a bit of cheating for fun). puppy is mostly "animal" with a little "person". Words about the same topic get similar rows of numbers.

In a real model the dimensions have no names. Nobody decides that dimension 3 means "tech". The numbers are learned, and the meaning is spread over all dimensions together. We label ours only so you can follow the idea.

4. Dot Product and Norm

To measure how related two vectors are, we need two small tools.

The dot product multiplies two vectors position by position and adds up the results:

dot([1, 2, 3], [4, 5, 6]) = 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32

Vectors that are big in the same positions give a large dot product. Vectors that are big in different positions give a small one. So the dot product is a rough score of "how much do these two agree".

The norm is the length of a vector, from Pythagoras: square every number, add them, take the square root. For [3, 4] it is sqrt(9 + 16) = 5, the same as the diagonal of a 3-4-5 triangle.

Add both functions to src/embeddings.js:

export function dot(a, b) {
  return a.reduce((sum, x, i) => sum + x * b[i], 0);
}

export function norm(a) {
  return Math.sqrt(dot(a, a));
}

(norm is the square root of the dot product of a vector with itself: 3*3 + 4*4 = 25, and sqrt(25) = 5.) Check them:

// demo-dot.js
import { dot, norm } from "./src/embeddings.js";
console.log(dot([1, 2, 3], [4, 5, 6]));
console.log(norm([3, 4]));
// 32
// 5

Both lines are exact and stable.

5. Cosine Similarity: How Close Are Two Words?

The raw dot product has a flaw: a long vector gives a large dot product just because it is long. We care about direction, "do these point the same way?", not length. So we divide by both lengths. The result is cosine similarity:

// 1 = same direction, 0 = unrelated, -1 = opposite.
export function cosine(a, b) {
  return dot(a, b) / (norm(a) * norm(b));
}

Read the output as a similarity score. 1 means pointing exactly the same way, 0 means nothing in common, and (with vectors that can go negative) -1 means opposite. Our hand-made vectors have no negative numbers, so ours will stay between 0 and 1.

// demo-cosine.js
import { vectors, cosine } from "./src/embeddings.js";
console.log(cosine(vectors.coffee, vectors.tea).toFixed(3));
console.log(cosine(vectors.coffee, vectors.dog).toFixed(3));
console.log(cosine(vectors.cat, vectors.dog).toFixed(3));
console.log(cosine(vectors.phone, vectors.laptop).toFixed(3));
// 0.986
// 0.000
// 0.988
// 0.994

coffee and tea score 0.986: nearly the same direction. coffee and dog score exactly 0.000: nothing in common. cat with dog and phone with laptop also score above 0.98. The numbers say what you would say: coffee is like tea, a cat is like a dog, a phone is like a laptop, and coffee is not like a dog.

All four lines are stable, since there is no randomness. The values are our hand-made ones, so they prove nothing about real language, only that the method works when the vectors are good.

6. Nearest Neighbours

Given any vector, we can find the closest words by scoring every word and sorting. That is nearest:

// The k words whose vectors are most similar to vec.
export function nearest(vec, k = 3, table = vectors) {
  return Object.entries(table)
    .map(([word, v]) => ({ word, score: cosine(vec, v) }))
    .sort((a, b) => b.score - a.score)
    .slice(0, k);
}

It returns k objects like { word: "tea", score: 0.986 }. The closest word to a word's own vector is always the word itself (score 1), so to see its neighbours we ask for k + 1 and skip the first:

// demo-nearest.js
import { vectors, nearest } from "./src/embeddings.js";

const fmt = (r) => `${r.word} ${r.score.toFixed(2)}`;
for (const w of ["coffee", "dog", "laptop", "teacher"]) {
  console.log(w, "->", nearest(vectors[w], 4).slice(1).map(fmt).join(", "));
}

// A made-up query: "something to drink", all weight on dimension 0.
console.log("query [1,0,0,0] ->", nearest([1, 0, 0, 0], 3).map(fmt).join(", "));
// coffee -> tea 0.99, milk 0.96, sokha 0.23
// dog -> puppy 0.99, cat 0.99, milk 0.27
// laptop -> phone 0.99, app 0.99, teacher 0.42
// teacher -> student 0.99, sokha 0.97, app 0.55
// query [1,0,0,0] -> coffee 0.99, tea 0.99, milk 0.96

The nearest neighbours of coffee are tea and milk; of dog, puppy and cat; of laptop, phone and app. After the first two or three, the scores drop sharply: the third neighbour of coffee is sokha at 0.23, which is basically "nothing in common".

The last line is the interesting one. We did not look up a word. We made up a direction ("drink") and asked: what is close to it? The answer is the three drinks. This "find what is close to this vector" operation is the heart of vector search.

7. Real Embeddings Are Learned

Where do real vectors come from? From training, using the same loop you saw in the last chapter. The vectors are parameters, a table with one row per token, and they start as random numbers. As the model trains on text, gradient descent nudges them. Words used in similar contexts ("a cup of ___", "I drank some ___") receive similar nudges and drift toward each other. After enough training, coffee and tea end up close because the model needed them close to predict text well.

Three honest points:

• Nobody labels the dimensions. There is no "drink" axis. The structure emerges from data.
• They carry the data's habits. If a corpus uses some words together in a biased way, the vectors reflect that. This is one way the biases from the training chapter reach the model.
• There are different kinds. Some embedding models are built only to turn text into vectors for search. The vectors inside a chat model are built to help predict the next token. Both are learned, with different goals.

Real vectors usually have hundreds to thousands of numbers, not four, and real models store one for every token in their vocabulary. Nothing about our four-number toy is the real structure, but the operations (dot, cosine, nearest) are exactly what is used.

8. Searching by Meaning

A keyword search finds documents that contain the same letters as your query. Search for "puppy" and a document that only says "dog" is missed. Embeddings let you search by meaning: turn both the query and every document into vectors, then find the nearest ones.

For a quick sentence vector we use the simplest trick: average the word vectors.

// demo-search.js
import { vectors, cosine } from "./src/embeddings.js";

// A sentence vector: the average of its word vectors (words we know).
function embed(text) {
  const vecs = text.split(" ").filter((w) => vectors[w]).map((w) => vectors[w]);
  return vecs[0].map((_, d) => vecs.reduce((s, v) => s + v[d], 0) / vecs.length);
}

const docs = [
  "coffee tea milk",
  "cat dog puppy",
  "phone laptop app",
  "student teacher sokha",
];

function search(query) {
  const q = embed(query);
  return docs
    .map((doc) => ({ doc, score: cosine(q, embed(doc)) }))
    .sort((a, b) => b.score - a.score);
}

for (const query of ["puppy", "tea", "laptop app", "teacher"]) {
  const best = search(query)[0];
  console.log(query.padEnd(11), "->", best.doc, best.score.toFixed(2));
}
// puppy       -> cat dog puppy 0.99
// tea         -> coffee tea milk 0.99
// laptop app  -> phone laptop app 1.00
// teacher     -> student teacher sokha 0.99

The query puppy finds the document cat dog puppy, and teacher finds the people document, even though the query word is not the whole document. Real systems do the same with a learned embedding model and far more documents.

This one idea, embed everything, then find the nearest vectors, is how many AI applications let a model look things up in your own files. It is the basis of a technique called retrieval-augmented generation (RAG), which a later series covers. For now, remember that it is plain cosine similarity, exactly what you just wrote.

9. The Problem: One Vector per Word Is Not Enough

Embeddings are a great start, but they have a limit. Each word has one vector, no matter where it appears. Think of the word mouse. In "the mouse ran away from the cat" it is an animal. In "click the mouse to open the laptop" it is a device. A fixed vector cannot be both. We could give mouse a mixed vector, half animal and half tech, like [0.0, 0.5, 0.5, 0.0], but then it is only halfway right everywhere.

Language is full of this: "bank" (river or money), "light" (not heavy or not dark), "match" (a game or a fire). What we need is a way for a word's vector to change depending on the words around it. The mouse near cat should lean toward animal, the mouse near laptop toward tech.

That is exactly what attention does.

10. Attention: Words Looking at Words

The idea of attention is simple to say in plain words:

For each word, look at all the other words in the sentence, decide how much each one matters to this word, and blend their information in using those amounts.

Three steps, all of which you already know:

1. Score. For one word (the one doing the looking), compute the dot product with every word in the sentence, including itself. A high score means "these two are related".
2. Weights. Pass the scores through softmax (from Probabilities, Temperature and Sampling) so they become percentages that add up to 1. These are the attention weights.
3. Blend. Make a new vector for the word: a weighted average of all the word vectors, using those weights. A word with weight 0.8 contributes 80% of the new vector.

The new vector is the same word, in context. The word has looked at its neighbours and absorbed something from them. A real Transformer does this for every word at once, in every layer.

11. Attention in Code

Create src/attention.js. It imports dot from the file you just wrote and softmax from the sampling chapter:

// src/attention.js
import { dot } from "./embeddings.js";
import { softmax } from "./sampling.js";

// For each word, score every word (dot product), turn scores into
// weights with softmax, then blend the vectors using those weights.
export function attention(sentence, table, scale = 1) {
  const vecs = sentence.map((w) => table[w]);
  const weights = [];
  const blended = [];
  for (const q of vecs) {
    const scores = vecs.map((k) => dot(q, k) * scale);
    const w = softmax(scores);
    weights.push(w);
    blended.push(
      q.map((_, d) => w.reduce((sum, wi, i) => sum + wi * vecs[i][d], 0)),
    );
  }
  return { weights, blended };
}

export function printTable(sentence, weights) {
  const pad = Math.max(...sentence.map((w) => w.length)) + 1;
  console.log("".padEnd(pad) + sentence.map((w) => w.padStart(pad)).join(""));
  sentence.forEach((w, i) => {
    console.log(
      w.padEnd(pad) + weights[i].map((x) => x.toFixed(2).padStart(pad)).join(""),
    );
  });
}

attention(sentence, table, scale) takes an array of words, a lookup from word to vector, and a scale. For each word q in the sentence it:

• scores every word k with dot(q, k), multiplied by scale;
• turns the scores into w with softmax;
• builds the blended vector: for each dimension d, the sum of weight * that word's number.

It returns two things: weights (one row of percentages per word) and blended (the new, context-aware vectors).

The scale is a sharpness knob: a bigger value makes the softmax pick its favourites more strongly. Real Transformers have a fixed scaling factor built in (they divide by something related to the vector length) so you can think of it as part of the recipe.

12. Reading the Attention Table

Try it on a three-word sentence. We add a vector for drinks, because it is not in our word list:

// demo-attention.js
import { vectors } from "./src/embeddings.js";
import { attention, printTable } from "./src/attention.js";

const table = { ...vectors, drinks: [0.8, 0.0, 0.0, 0.3] };
const sentence = ["sokha", "drinks", "tea"];

console.log("scale 1:");
printTable(sentence, attention(sentence, table, 1).weights);
console.log("scale 4:");
const { weights, blended } = attention(sentence, table, 4);
printTable(sentence, weights);
console.log("tea before:", table.tea.join(" "));
console.log("tea after: ", blended[2].map((x) => x.toFixed(2)).join(" "));
// scale 1:
//          sokha drinks    tea
// sokha     0.46   0.30   0.25
// drinks    0.28   0.37   0.35
// tea       0.25   0.38   0.37
// scale 4:
//          sokha drinks    tea
// sokha     0.79   0.14   0.07
// drinks    0.14   0.48   0.38
// tea       0.09   0.47   0.44
// tea before: 0.8 0 0 0.1
// tea after:  0.75 0.00 0.01 0.27

Each row is one word doing the looking. Each column is a word being looked at. A row adds up to 1 (to rounding). Read the scale 4 table:

• Row tea: when processing "tea", 47% of the attention goes to drinks, 44% to itself, and only 9% to sokha. The word tea looks mostly at the verb drinks and itself, because those are both "drink" words with similar numbers.
• Row sokha: the person word mostly looks at itself (79%), because nothing else in the sentence is about a person.

At scale 1 the table is nearly flat (everything gets about a third), because our vectors are small numbers and the scores do not differ much. A bigger scale sharpens the picture. Real models also learn how to look, which we skip.

Now the effect. tea started as 0.8 0 0 0.1. After attention it is 0.75 0.00 0.01 0.27. It gained some "person" (from drinks, which has a person-like component, and sokha): it now carries a trace of "this tea is being drunk by someone". The word was changed by its context.

You can check one row by hand. For tea, the three scores are the dot products with sokha, drinks and tea:

// demo-hand.js
import { vectors, dot } from "./src/embeddings.js";
import { softmax } from "./src/sampling.js";

const tea = vectors.tea;
const sentence = ["sokha", "drinks", "tea"];
const table = { ...vectors, drinks: [0.8, 0.0, 0.0, 0.3] };
const scores = sentence.map((w) => dot(tea, table[w]));
console.log("scores :", scores.map((s) => s.toFixed(2)).join(" "));
console.log("scaled :", scores.map((s) => (s * 4).toFixed(2)).join(" "));
console.log("weights:", softmax(scores.map((s) => s * 4)).map((p) => p.toFixed(2)).join(" "));
// scores : 0.25 0.67 0.65
// scaled : 1.00 2.68 2.60
// weights: 0.09 0.47 0.44

These are the same 0.09 0.47 0.44 as the tea row above. No magic: dot product, softmax, weighted average. Every number in these outputs is stable across runs.

13. Context Changes the Meaning

Now the payoff, with the ambiguous word mouse from section 9. Give it the mixed vector, half animal and half tech, then put it in two different two-word contexts:

// demo-mouse.js
import { vectors, nearest } from "./src/embeddings.js";
import { attention, printTable } from "./src/attention.js";

const table = { ...vectors, mouse: [0.0, 0.5, 0.5, 0.0] };
const fmt = (x) => `${x.word} ${x.score.toFixed(2)}`;

console.log("mouse alone:", nearest(table.mouse, 3).map(fmt).join(", "));

for (const other of ["cat", "laptop"]) {
  const s = ["mouse", other];
  const r = attention(s, table, 4);
  console.log(`\n"${s.join(" ")}"`);
  printTable(s, r.weights);
  console.log("mouse now closest to:", nearest(r.blended[0], 2).map(fmt).join(", "));
}
// mouse alone: phone 0.70, cat 0.70, dog 0.69
//
// "mouse cat"
//        mouse   cat
// mouse   0.55  0.45
// cat     0.18  0.82
// mouse now closest to: cat 0.93, dog 0.91
//
// "mouse laptop"
//          mouse laptop
// mouse     0.55   0.45
// laptop    0.17   0.83
// mouse now closest to: phone 0.93, laptop 0.92

On its own, mouse is equally near phone and cat (about 0.70), so a fixed vector cannot choose. After attention in "mouse cat", the new mouse vector is closest to cat and dog. After attention in "mouse laptop", it is closest to phone and laptop. Same word, different surroundings, different vector. That is how a Transformer can handle "bank" in a river sentence and in a money sentence.

This is why attention is such a big deal. It makes meaning depend on context, using only the simple steps you just ran.

14. The Real Transformer, Honestly

What you built is a simplified picture of the idea at the heart of the Transformer, the design behind modern LLMs. Here is what is different in the real thing. The details are conceptual; do not take our toy as an implementation guide.

• Learned lookups. In our demo, each word is scored using its own vector. A real model first passes every vector through learned matrices to make three versions of it (usually called query, key and value): one for "what I am looking for", one for "what I offer", one for "what I will contribute". These matrices are parameters, trained like everything else.
• Many heads. Real models run many attention operations in parallel, called heads, each free to learn a different kind of relationship (one might track grammar, another which noun a pronoun refers to). We built one head, with no learning.
• Many layers. The result of attention goes through more learned arithmetic and then into another round of attention, stacked dozens of times. Early layers tend to capture simpler patterns and later ones more abstract ones, though how exactly is an active research area.
• Order matters. Our toy ignores word order. Real models add information about each token's position, or "dog bites man" would look the same as "man bites dog".
• No peeking ahead. When predicting the next token, a model is only allowed to attend to tokens before the current one, not future ones.
• Scale. Real vectors have hundreds or thousands of numbers, and the vocabulary has tens of thousands of tokens.

The pipeline from the series overview now has more shape: text is split into tokens, each token becomes an embedding, stacked attention layers blend context into every vector, and the last vector is turned by one more learned step into the scores for every possible next token. Those scores go through softmax and sampling, exactly as in chapter 4.

Be careful with the intuition, too. It is tempting to say "attention is how the model understands", but that is a story we tell. Researchers study what attention heads do, and the picture is partial: some heads do seem to do something easy to describe, many do not. Treat the table you just printed as a teaching illustration of a mechanism, not a window into a mind.

15. Recipe: The embeddings and attention Modules

You created both files as you went. Together:

• src/embeddings.js exports vectors, dot(a, b), norm(a), cosine(a, b) and nearest(vec, k = 3, table = vectors).
• src/attention.js exports attention(sentence, table, scale = 1) and printTable(sentence, weights). It needs softmax from src/sampling.js, from Probabilities, Temperature and Sampling.

Every demo in this chapter is a small file next to src/ (for example demo-nearest.js) that you run with node demo-nearest.js. Make sure "type": "module" is set in package.json, as in chapter 1, so import works.

16. Try It Yourself

1. Add a word. Add juice: [0.8, 0.0, 0.0, 0.0] and tablet (pick its numbers). Print nearest for each. Do they land where you expected?
2. Break the vectors. Give tea the numbers [0.0, 0.0, 0.9, 0.0] (making it look like tech) and print cosine(coffee, tea) again. What happens to search results for "tea"? What does this tell you about the importance of good embeddings?
3. Odd one out. Write oddOneOut(words) that returns the word with the lowest average cosine similarity to the others. Try ["coffee", "tea", "milk", "laptop"].
4. More documents. Add four more documents to the search demo, including one with a word that is not in vectors (the code skips unknown words). What breaks if a document contains only unknown words, and how would you guard against it?
5. Change the scale. Run the "sokha drinks tea" attention table with scale set to 0.1, 1, 10 and 50. Describe what happens to each row. At very large scale, what does each word attend to?
6. New ambiguous word. Invent bank with a vector halfway between [drink, animal, tech, person] dimensions of your choice, and test it next to a "money" word and a "river" word you also invent. Print the nearest words after attention, as in section 13.

17. Quick Review: Common Questions

Questions beginners often ask about this chapter. Try answering each one yourself first, then click the question to open the answer.

Where to go next. You can now explain how words become numbers and how words look at each other. A model can only look at a limited amount of text at once, though, and it forgets everything between requests unless you resend it. That limit, and the tricks used to live with it, are the subject of Context Windows and Memory.