All writing

How a language model chooses its next token

Published · Updated

How a language model chooses its next token

Trace next-token generation from transformer concepts to a small word model that runs on the browser CPU.

I built Latent Space Observatory because the phrase "next-token prediction" hides most of the process. A large language model (LLM) must turn its current text into a probability distribution. It selects one token, adds that token to the text, and starts again.

The Observatory now runs a small word model on the browser CPU. You can select a token and change the sampling temperature while the next token and continuation update. A graphics processing unit (GPU) renders the moving field behind these controls.

The word model is real, but it is not a transformer or an LLM. It learns transition counts from four built-in sentences. The attention matrix remains a teaching function. Keeping that boundary visible matters more than making the demonstration look sophisticated.

Follow the generation loop

A decoder language model uses the following sequence for each new token:

  1. Split the available text into tokens.
  2. Convert each token and its position into a vector representation.
  3. Transform the representations with attention and feed-forward layers.
  4. Calculate one score, called a logit, for each token in the vocabulary.
  5. Convert the logits into probabilities and select one token.
  6. Add the selected token to the context and repeat the sequence.

The model does not write a complete paragraph before it shows the first word. It makes a new selection after each generated token. An early selection can change all later selections.

Tokens are the working units

An LLM does not usually process complete words. A tokenizer maps text to token identifiers. A token can contain a word, part of a word, punctuation, whitespace, or a byte fragment.

The Observatory uses nine readable parts from the following sentence:

The city wakes before the sun remembers its name.

The parts look like words because this format makes the interaction easy to read. A production tokenizer can split the same sentence differently.

Tokenization defines one generation step. For example, a tokenizer can split a rare name into four tokens. The model must then make four related selections to produce that name. Token count also determines how much text fits in the context window.

Embeddings create representations

A token identifier is only an index. The model uses the index to get a learned vector, called a token embedding. The model combines this vector with position information.

The term "latent space" can suggest one fixed map of all concepts. A transformer does not have one such map. It maintains many representations, and each layer changes them. The surrounding tokens also affect each representation.

Consider the token bank in the following sentences:

We sat on the river bank.
The bank approved the loan.

Both occurrences start with the same token embedding. Later layers use the surrounding tokens to create different representations. For example, river and loan provide different context.

The moving field in the Observatory is a visual metaphor for these changing representations. The field does not show a projection of model data.

Attention moves information between positions

Self-attention lets each token position use information from other permitted positions. A transformer creates three vectors from each input representation:

Scaled dot-product attention uses the following operation:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V

The query-key dot products produce compatibility scores. The scale factor prevents large vector dimensions from producing excessively large scores. Softmax converts the scores into weights. The weighted values produce a new representation for each position.

A decoder-only model also applies a causal mask. The mask lets a position use earlier tokens and itself. The mask blocks tokens that the model has not generated.

Production models repeat this calculation across many attention heads and layers. One heatmap cannot explain the complete calculation. A bright cell in one head does not fully explain a generated token.

The Observatory uses a simpler 9 by 9 matrix. The matrix combines token distance, selected-token distance, grammatical-role matching, and a causal bias:

const distance = Math.abs(rowIndex - columnIndex);
const focusPull = 1 / (1 + Math.abs(columnIndex - focusIndex) * 0.72);
const localPull = 1 / (1 + distance * 0.64);
const roleMatch = rowToken.role === columnToken.role ? 0.12 : 0;
const causalBias = columnIndex <= rowIndex ? 1 : 0.22;

const weight = (localPull * 0.4 + focusPull * 0.48 + roleMatch) * causalBias;

The function normalizes each row by its largest value. When you select a token, the matrix changes immediately. The interaction teaches one point: relevance depends on the current position and its context. The displayed numbers are not learned attention weights.

Logits become probabilities

The final transformer layer produces a representation for the current position. The model maps this representation to one logit for each vocabulary token. A logit is an unnormalized score, not a probability.

Softmax converts the logits into a probability distribution. The following equation shows how temperature changes that distribution:

p_i = exp(z_i / T) / sum_j exp(z_j / T)

In this equation, z_i is the logit for token i. The value T is the temperature.

A low temperature increases the difference between the scaled logits. Softmax then gives more probability to the strongest candidates. A high temperature reduces the difference, so weaker candidates get more probability.

Temperature does not increase model understanding. It changes how the sampler uses the logits. Production systems can also use top-k sampling, top-p sampling, repetition penalties, or constrained decoding.

The Observatory applies the same operation to a much smaller vocabulary. At startup, it counts words and adjacent word pairs in four short training sentences. For each step, it scores candidate words with three signals:

return local * 0.42 + focused * 0.48 + prior * 0.1 + repetitionPenalty;

local measures how often the candidate follows the previous word. focused measures how often it follows the selected word. prior is the candidate's frequency in the complete training set. Additive smoothing keeps unseen transitions from receiving a probability of zero.

The model divides these scores by the temperature before softmax. It then uses a seeded sample, so the same token and temperature produce the same result after a reload. The result is reproducible, and the output still changes across the temperature range.

After the first generated word, that word becomes the new focus. The model repeats the calculation until it has produced seven words. It does not choose a complete continuation in advance.

The temperature control also sends a value to the shader. It changes the field distortion, brightness, and interference. The text and graphics respond to the same control, but they use separate code paths.

The browser runs real GPU code

I wanted the background to behave like a field instead of a video. The page first requests a high-performance WebGPU adapter. It then compiles a pipeline with WebGPU Shading Language (WGSL).

If WebGPU is unavailable, the page uses WebGL 2. The fallback renderer uses OpenGL Shading Language (GLSL). Both renderers receive the following eight floating-point values:

width, height, pointer x, pointer y,
time, scroll progress, temperature, motion

The vertex shader draws one large triangle that covers the viewport. The fragment shader calculates the color of each pixel. It creates eight moving nodes and measures the energy near each node. It also draws lines between selected node pairs. Pointer position and temperature change the final field.

React controls the selected token, the temperature, and the renderer label. React does not render each animation frame. The canvas code writes the current values to a GPU uniform buffer. The shader then runs outside the React render cycle.

This design prevents the canvas from causing 60 React renders each second. The page also limits the device pixel ratio. It stops rendering when the tab is hidden. It reduces the frame rate when the visitor requests reduced motion.

Separate facts from examples

The following table identifies which parts use real system behavior and which parts use examples:

Part of the experienceSource of the displayed behavior
Token inspectorA simplified token sequence written for the demonstration
Attention matrixRelationships calculated by a teaching function, not learned weights
Candidate probabilitiesOutput from the small word transition model
Temperature continuationsSeven word-by-word samples from the local model
Moving fieldA visual metaphor, not a hidden-state projection
WebGPU and WebGL 2 renderersReal browser GPU pipelines
Pointer, scroll, and temperature inputsReal values used by both shader implementations

These labels are necessary because model explanations can look more authoritative than their data supports. A visualization should identify model output and teaching data as separate things.

Know what the small model cannot show

The browser model makes sampling observable, but it cannot explain how a transformer learns contextual representations. It has no embedding vectors, feed-forward layers, learned attention heads, or subword tokenizer. Its transition table only knows the small corpus bundled with the page.

A transformer-backed version could return token identifiers, decoded tokens, and next-token logits. Attention data would need more controls because a model can contain many layers and heads. If the user interface combined those values, it would also need to explain the combination method.

The shader could use a two-dimensional projection of hidden states. That would make the field depend on model data, but the projection would add another transformation that the interface must explain.

I kept the model small so you can inspect the complete inference path in one TypeScript file. It is useful because its limits are obvious.

Open Latent Space Observatory. Select different tokens, then change the temperature. The source code contains the word model, both shader implementations, and the browser lifecycle.