What a language model does

Story

When you type a message, your phone suggests the next word. A language modelLanguage model A program that reads text and gives every possible next piece of text a probability. Open the glossary does the same thing, only much better. You give it some text, and it answers with every piece of text that could come next, each with a chance.

Numbers

Here is what that answer could look like after “Once upon a”. These numbers are made up to show the idea.

The chances, called probabilitiesProbability A number from 0 to 1 that says how likely something is. The probabilities of all options add up to 1. Open the glossary, always add up to 1, because something has to come next.

Formula

The model assigns a probability to each possible next token, given all the tokens before it.

For each position the probabilities over the whole vocabulary sum to 1.

Code

const trace = forward(model, tokenIds);
const lastRow = rowView(trace.logits, tokenIds.length - 1);
const chances = probabilities(lastRow, 1); // one number per token in the vocabulary

Story

To write more than one word, the model picks one of the likely pieces, adds it to the text, and guesses again. Repeat this a few hundred times and you have a paragraph. Chatbots like ChatGPT write their answers this way, one piece at a time.

Numbers

Starting from “Once upon a”, three rounds could go like this.

1. “Once upon a” gives “time” the highest chance. The text becomes “Once upon a time”.
2. “Once upon a time” makes “,” likely. The text becomes “Once upon a time,”.
3. “Once upon a time,” makes “there” likely. The text becomes “Once upon a time, there”.

Formula

The probability of a whole sequence is the product of the next-token probabilities, one position at a time.

Generating text means choosing each $t_i$ from that distribution and feeding it back in.

Code

const result = generate(model, promptIds, { maxNewTokens: 20, temperature: 0.8, seed: 1 });

Story

The model that makes these guesses is called a transformerTransformer The design behind most modern language models. It passes text through layers of attention and feed-forward steps. Open the glossary (Vaswani et al., 2017). It turns the text into numbers, passes those numbers through a stack of layers that let each word gather information from the words before it, and ends with a score for every possible next piece. The next chapters open each of those steps.

Numbers

The small model on this site has a vocabulary of 4,096 pieces. It turns each piece into a list of 128 numbers, runs them through 4 layers, and has about 1.3 million numbers it learned during training. Large chat models follow the same plan with tens of thousands of times more numbers, or more.

Formula

A transformer computes scores for every token in the vocabulary, called logits, from the tokens so far. Softmax turns the logits into probabilities.

Here $\theta$ stands for all of the model’s learned parameters.

Code

const model = loadModel(
  await (await fetch(modelUrl)).arrayBuffer(), // modelUrl is where the site keeps the file
);
console.log(model.config); // { vocabSize: 4096, dModel: 128, nLayers: 4, ... }