Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Autoregressive Token World Models › Tokenizing the World

Michael BrenndoerferJune 27, 202619 min read

World Models HandbookAll chapters

Part 5 of 12World-Model Architectures

Chapter 27Autoregressive Token World Models › Tokenizing the World

Reading section 3 of 11 · Tokenizing the World · 19 minSection 3 of 11 · 19 min

Section 3 of 11 · 19 min read

Autoregressive Token World Models › Tokenizing the WorldOne section · Sections in this chapter3 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

A token world model is only as good as its tokenizer. This section covers three common tokenization design problems: turning images into tokens, turning actions into tokens, and fusing many modalities into one vocabulary.

Tokenization is not necessarily learned or lossy: our toy uses an exact hand-designed mapping. The visual tokenizers below do compress observations. If two complete retained token histories are identical but require different correct predictions or decisions, no downstream model can always distinguish them. A variable missing from one encoded frame may still be inferable from other frames or actions. Temporal stride sets the latent grid rate, not by itself the finest event timing the codes can preserve or the decoder can reconstruct.

Visual Tokenizers: From Pixels to SymbolsLink Copied

A $256\times256$ RGB image contains 196,608 scalar channel values. iGPT applied autoregression to small images represented by clustered pixel colors, with VQ-VAE compression for larger inputs. At one color token per pixel, a $64\times64$ frame requires 4,096 ordinary sequential decoding steps. Each step evaluates a decoder, although KV caching avoids recomputing the full prefix. Learned visual compression reduces the number of positions by encoding an image or clip as a smaller code grid.

The workhorse is the vector-quantized variational autoencoder (VQ-VAE), introduced by van den Oord, Vinyals, and Kavukcuoglu (2017). It works in three stages:

An encoder $E$ maps an image $o$ to a downsampled feature map $z = E(o)$ . A $256 \times 256$ frame with 8× downsampling becomes a $32 \times 32$ grid of feature vectors, one per spatial cell.
Each feature vector $z_i$ is quantized: we replace it with the nearest entry in a learned codebook $\{e_1, \dots, e_K\}$ :

\begin{aligned} k^\star &= \arg\min_k \lVert z_i - e_k \rVert_2 && \text{(find the nearest codebook entry)} \\ \tilde{z}_i &= e_{k^\star} && \text{(replace } z_i \text{ with it)} \end{aligned}

where:

$z_i$ : the feature vector at spatial cell $i$ of the encoder output
$e_k$ : the $k$ -th entry of the learned codebook, one of $K$ entries in total
$\lVert \cdot \rVert_2$ : the Euclidean norm, so $k^\star$ is the index of the codebook entry closest to $z_i$
$\tilde{z}_i$ : the quantized replacement for $z_i$ , namely the selected codebook entry

The assignment geometry is easiest to see in two dimensions. The next cell simulates a small feature space: five clusters of encoder outputs (the analogue of one spatial cell's feature vector across many images), a codebook of 12 entries, and the nearest-neighbor rule from the equation above. Everything is deterministic given the seed.

In[5]:

Code

import numpy as np

rng = np.random.default_rng(11)

## Five clusters of encoder outputs in 2-D, standing in for the feature map of
## one spatial cell across many images.
cluster_centers = np.array(
    [[0.0, 0.0], [2.4, 1.1], [-1.6, 2.1], [1.7, -1.7], [-2.1, -1.0]]
)
enc_outputs = np.vstack(
    [c + 0.35 * rng.standard_normal((60, 2)) for c in cluster_centers]
)

## A codebook of 12 entries.
VQ_K = 12
vq_codes = rng.uniform(-3.0, 3.0, (VQ_K, 2))

## Nearest-code assignment: argmin over squared Euclidean distance.
vq_assign = (
    ((enc_outputs[:, None, :] - vq_codes[None, :, :]) ** 2).sum(-1).argmin(1)
)
vq_error = np.linalg.norm(enc_outputs - vq_codes[vq_assign], axis=1)
print(f"mean quantization error per assignment: {vq_error.mean():.3f}")

## Dense grid for the Voronoi (nearest-code) partition background.
gx, gy = np.meshgrid(np.linspace(-3.6, 3.6, 220), np.linspace(-3.2, 3.6, 220))
grid_pts = np.stack([gx.ravel(), gy.ravel()], axis=1)
grid_assign = (
    ((grid_pts[:, None, :] - vq_codes[None, :, :]) ** 2).sum(-1).argmin(1)
)
grid_assign = grid_assign.reshape(gx.shape)

Out[5]:

Console

mean quantization error per assignment: 0.755

Out[6]:

Visualization

Scatter of encoder outputs colored by nearest code with stars marking codebook entries. — Each encoder output (dot) is replaced by its nearest codebook entry (star). Thin connectors show the quantization error each assignment introduces, and the shaded background marks each code's Voronoi cell.

A decoder $D$ reconstructs the image from the quantized grid: $\hat{o} = D(\tilde{z})$ .

Tokenization emits code indices, one per grid cell. With 8,192 entries, each index needs 13 bits in fixed-width storage and carries at most 13 bits of entropy. For one $32\times32$ grid $Z_t$ , $I(O_t;Z_t)\le H(Z_t)\le32\times32\times13=13{,}312$ bits. This bounds information about that image in that grid, not all information available across history. Arbitrary 8-bit RGB values for the original $256\times256$ frame require 1,572,864 bits uncompressed. Whether the smaller representation retains task-relevant distinctions depends on the data, architecture, objective, and learned parameters.

The nearest index is locally constant away from assignment boundaries, so ordinary differentiation through the index does not give a useful encoder gradient. This does not block the derivative of a decoder loss with respect to the selected code vector itself. Standard VQ-VAE deliberately separates these paths: a straight-through estimator routes reconstruction gradients to the encoder, while a codebook loss or EMA update trains dictionary entries.

Straight-Through Estimator

In the backward pass, pretend the quantizer was the identity function. The gradient of the loss with respect to the encoder output $z_i$ is approximated by the gradient with respect to the quantized value $\tilde{z}_i$ :

\frac{\partial \mathcal{L}}{\partial z_i} \approx \frac{\partial \mathcal{L}}{\partial \tilde{z}_i},

where:

$\mathcal{L}$ : the training loss being differentiated, here the reconstruction loss
$z_i$ : the encoder output at spatial cell $i$
$\tilde{z}_i$ : its quantized replacement from the codebook

The codebook entries themselves are updated by a separate loss (or by exponential moving averages of the encoder outputs assigned to them).

The straight-through estimator is biased: it differentiates as if a small change of the encoder output changed the quantized value continuously, although hard assignments are piecewise constant. Encouraging encoder outputs to stay near their assigned codes can help training, but closeness alone does not prove that the surrogate is a useful descent direction. The estimator does not guarantee convergence to the optimum of the discrete reconstruction problem.

A VQ-VAE-style objective can combine a chosen reconstruction loss with codebook and commitment terms:

\mathcal{L} = \underbrace{\lVert o - D(\tilde{z}) \rVert_1}_{\text{reconstruction}} + \underbrace{\lVert \mathrm{sg}(z) - \tilde{z} \rVert_2^2}_{\text{codebook loss}} + \underbrace{\beta \, \lVert z - \mathrm{sg}(\tilde{z}) \rVert_2^2}_{\text{commitment loss}},

where:

$\mathcal{L}$ : the total VQ-VAE training objective
$o$ : the input image
$z$ : the encoder's continuous feature grid
$\tilde{z}$ : the quantized version of that grid
$D$ : the decoder that reconstructs the image from $\tilde{z}$
$\lVert \cdot \rVert_1$ and $\lVert \cdot \rVert_2$ : the L1 and Euclidean norms
$\mathrm{sg}(\cdot)$ : "stop-gradient", which treats its argument as a constant during backpropagation
$\beta\ge0$ : the coefficient weighting commitment relative to the other loss terms

The displayed reconstruction term is L1. The original VQ-VAE image experiments used squared error within a more general reconstruction-likelihood formulation; the VQGAN implementation uses L1. Squared error selects a conditional mean; L1 selects coordinatewise conditional medians. Neither universally guarantees sharp or semantically faithful reconstruction. In the standard gradient routing, the codebook term moves assigned entries toward encoder outputs and the commitment term moves encoder outputs toward fixed entries. Without commitment, encoder scale can drift; without a codebook loss or a replacement such as EMA, that dictionary-update path is absent. These are alignment mechanisms, not the posterior-to-prior matching performed by a VAE KL term.

Reconstruction-only point decoders can blur when an encoded representation leaves incompatible images plausible. Squared error returns their conditional mean; coordinatewise L1 medians may or may not form a coherent joint image. The mean need not lie between modes for every multimodal distribution. The asymmetric mixture below is a specific example where the squared-error optimum does lie in a low-density region.

The next cell constructs a continuous scalar Gaussian mixture, not a categorical token distribution. Its three components have equal mixture weights. The plotted mean is the squared-error-optimal point prediction; a manually chosen component center provides a visual contrast, not a learned or sampled adversarial output.

In[7]:

Code

import numpy as np

## Next-value distribution for a continuous scalar when three futures are
## equally plausible: a trimodal Gaussian mixture with the modes placed
## asymmetrically, so the mixture mean lands in a low-density region
## between them rather than on the middle future.
xs_blend = np.linspace(-4.0, 4.0, 800)
BLEND_MODES = np.array([-3.0, -1.5, 3.0])
BLEND_SIGMA = 0.45


def mixture_pdf(x):
    comp = np.exp(
        -0.5 * ((x[:, None] - BLEND_MODES[None, :]) / BLEND_SIGMA) ** 2
    )
    return comp.sum(axis=1) / (
        len(BLEND_MODES) * BLEND_SIGMA * np.sqrt(2 * np.pi)
    )


blend_density = mixture_pdf(xs_blend)

## The L2-optimal single prediction is the mixture mean, which sits in a
## low-density region between the modes and matches none of them.
blend_prediction = float(BLEND_MODES.mean())
chosen_mode_center = float(BLEND_MODES[0])

Out[8]:

Visualization

Trimodal density curve with the regression optimum marked in the low-density region between the modes. — For this equal-weight mixture, the squared-error-optimal point prediction is the mean. The dashed line is a manually chosen component center, not a sampled prediction or the result of adversarial training.

Averaging incompatible fine-detail arrangements can blur a reconstruction. VQGAN adds perceptual and adversarial objectives and reports improved perceptual quality; these losses encourage realism without guaranteeing the correct conditional mode or turning a deterministic decoder into a sampler. Learned codebooks can also become persistently underused. An entry with no assignments receives no new assignment contribution in that batch, but momentum, weight decay, EMA statistics, smoothing, or resets may still change it. One finite batch's zero counts do not establish permanent collapse.

We can inspect initial utilization in the same two-dimensional setting. The next cell assigns points from five clusters to 48 randomly initialized entries. It measures which entries are unused in this sample before any training; persistent collapse would require following utilization during training.

In[9]:

Code

import numpy as np

rng = np.random.default_rng(5)

## Synthetic encoder outputs concentrated near five chosen centers.
## This illustrative sample makes no claim about tokenized datasets generally.
collapse_centers = rng.uniform(-2.0, 2.0, (5, 2))
collapse_data = np.vstack(
    [c + 0.3 * rng.standard_normal((80, 2)) for c in collapse_centers]
)

## A codebook of 48 entries initialized uniformly over the box.
CB_K = 48
cb_codes = rng.uniform(-3.0, 3.0, (CB_K, 2))

## One round of nearest-neighbor assignment, before any training.
cb_assign = (
    ((collapse_data[:, None, :] - cb_codes[None, :, :]) ** 2).sum(-1).argmin(1)
)
cb_counts = np.bincount(cb_assign, minlength=CB_K)
usage_sorted = np.sort(cb_counts)[::-1]
n_dead = int((cb_counts == 0).sum())
print(f"unused codes after one assignment round: {n_dead} of {CB_K}")
print(
    f"mean quantization error on live codes: "
    f"{np.linalg.norm(collapse_data - cb_codes[cb_assign], axis=1).mean():.3f}"
)

Out[9]:

Console

unused codes after one assignment round: 25 of 48
mean quantization error on live codes: 0.461

Out[10]:

Visualization

Sorted bar chart of assignment counts with 25 unused codes in this batch. — Nearest-neighbor assignment counts in one synthetic batch. Zero counts indicate unused codes in this batch, not persistent unused codes. No codebook training is performed.

EMA is an alternative centroid update, not a guarantee of reviving unused entries; dead-code resets are another possible intervention. Finite Scalar Quantization uses fixed per-coordinate levels, and MAGVIT-v2 uses lookup-free quantization. Without learned centroids, there are no dead learned code vectors, but an encoder can still underuse the implicit code space. The fixed grid constrains quantizer coordinates while learned surrounding networks shape the effective representation. These are alternatives to learned vector codebooks, not evidence of their universal displacement. Residual VQ instead approximates a feature, subtracts that approximation, and quantizes the residual. With $R$ stages of $K$ entries there are $K^R$ index tuples at a cost of $R$ tokens per cell; collisions or unreachable tuples can yield fewer distinct reconstructions. Additional well-trained stages can improve precision, but stage count alone does not guarantee smaller error.

The next plot compares the number of possible index tuples with token count for 512 entries per residual stage. It does not measure distinct reconstructed vectors or achieved reconstruction quality.

In[11]:

Code

import numpy as np

RV_K = 512  # entries per codebook
rv_stages = np.arange(
    1, 7, dtype=np.int64
)  # stacked codebooks per spatial cell

## R stacked codebooks each with K entries can represent K^R combinations.
rv_combinations = RV_K**rv_stages
## while costing only R tokens per cell.
rv_tokens = rv_stages.astype(float)
print(f"combinations at R=6: {float(rv_combinations[-1]):.3e}")

Out[11]:

Console

combinations at R=6: 1.801e+16

Out[12]:

Visualization

Log-scale chart: possible index tuples rise exponentially while tokens per cell rise linearly. — R codebooks of 512 entries admit 512^R index tuples using R tokens per cell. Reachable distinct reconstructions can be fewer. Enumerating every tuple at R=6 in one table would require about 1.80e16 entries.

Video tokenizers can encode time as well as space. MAGVIT uses 3D operations; a compressed code can summarize multiple frames, for example a $2\times16\times16$ spatiotemporal block. Temporal stride spaces latent positions farther apart, but a code may retain within-block timing and a decoder may output every frame. Two-frame encoding is therefore not equivalent to dropping every other frame. Loss of timing depends on capacity and the learned representation. A noncausal encoder that waits for a second unseen frame adds latency; causal or overlapping schemes can have different online interfaces. Actual downsampling aliases unresolved temporal frequencies unless filtering or prior constraints make them recoverable. These distinctions connect token design to the observation-design questions in Part III.

Continuous-token designs omit codebook rounding but may retain other projection or compression losses. A continuous head can be deterministic, Gaussian, or an iterative generative module. MAR is an autoregressive image model with continuous tokens and a diffusion loss. Iterative sampling adds work, but continuous outputs are not inherently more expensive than categorical decoding across many positions. The relevant comparison includes evaluations per output, width, parallelism, representation quality, and the task.

Action TokenizersLink Copied

Actions enter the token stream too. Many environments have naturally discrete actions; many do not.

Natively discrete actions need an agreed symbol mapping and timing convention, but no continuous-value quantization. Atari's action space has fewer than 20 actions; this makes token mapping straightforward. IRIS evaluates on Atari 100k, but this convenience alone does not explain its historical benchmark choice.
Continuous actions must be quantized only for a fully categorical interface. RT-1 uses 256 bins per scalar action variable; its arm variables describe position, orientation, and gripper opening rather than seven joint angles. Any seven-scalar command takes seven tokens under one-token-per-variable encoding. Gato bins scalar observations and actions into 1,024 levels, while images use continuous patch embeddings. Hybrid observation-token models can instead accept or emit continuous actions.

Binning creates several distinct costs. Boundary jitter occurs when a varying command crosses a quantizer threshold repeatedly: small smooth changes become switches between adjacent output levels. This can add higher-frequency components to the command. It is not necessarily bang-bang control, and whether the physical trajectory oscillates depends on the plant and controller.

The next cell illustrates command quantization without a learned policy or a physical plant. A sinusoid crosses an interior bin edge in an eight-bin partition of $[-0.5,0.5]$ ; the figure shows command discontinuities, not a simulated trajectory.

In[13]:

Code

import numpy as np

## A smooth controller output regulating a setpoint that lands exactly on a
## bin edge, with a small illustrative sinusoidal variation. The
## slow crossings of that edge make the binned action flip back and forth
## between the two adjacent bins: boundary jitter.
t_cmd = np.linspace(0.0, 6.0, 600)
smooth_command = 0.25 + 0.03 * np.sin(3.0 * t_cmd)  # in [-0.5, 0.5]

N_BINS = 8
bin_edges = np.linspace(-0.5, 0.5, N_BINS + 1)
bin_index = np.clip(np.digitize(smooth_command, bin_edges) - 1, 0, N_BINS - 1)
bin_midpoints = 0.5 * (bin_edges[:-1] + bin_edges[1:])
binned_action = bin_midpoints[bin_index]

## How often consecutive samples flip between two adjacent bins.
n_flips = int((np.abs(np.diff(bin_index)) == 1).sum())

Out[14]:

Visualization

Step plot of a binned action zigzagging near bin edges against a smooth command curve. — A smooth command near a bin edge switches between adjacent reconstructed bin centers. This illustrates command discontinuity; no physical plant or closed-loop response is simulated.

Precision depends on dynamics as well as bin width. With midpoint decoding of an in-range scalar command, quantization error is bounded by half a bin width. That is not a bound on position error or a proof of unreachable targets: target alignment, feedback, integration, and execution horizon matter. Latency arises when action variables are decoded sequentially. Seven categorical variables require seven ordinary decoding steps in that design, but parallel heads or externally supplied candidate actions need not. Candidate rollouts also cost computation, although batching can share hardware throughput; Oasis's video-generation speed is not evidence about seven-token action decoding.

Continuous action heads can share an observation-token backbone. Coarse-to-fine tokens add precision with additional codes, but task-adaptive allocation requires an extra mechanism. Genie's latent action model instead uses prior frames and the next frame to learn compact transition codes without labeled actions. Such a code can provide a useful controllable interface, as evaluated in Genie, but observational frame changes can also reflect camera motion or exogenous events. Compactness alone does not identify a causal intervention. We return to this distinction in Part VI.

Multimodal Tokenizers and Unified VocabulariesLink Copied

A fully unified categorical design can represent observations, actions, rewards, and termination flags in one stream. Decision Transformer instead conditions a policy on embedded return-to-go scalars; this is not the same as predicting categorical reward symbols. An integrated reward predictor can support planning if it predicts the task's relevant reward accurately. Termination is supervised only when included as a target, and an integrated head does not remove the need to evaluate planning decisions.

A fully discrete implementation can use disjoint token ranges and optional modality-boundary markers. Chameleon demonstrates early-fusion text and image modeling. Gato uses a shared policy transformer over a mix of discrete symbols and continuous image-patch embeddings, without predicting image observations. Modality-specific processing can also be fused later at an intermediate or output representation. These designs share some sequence machinery without being identical categorical world models.

Why go through the trouble? Three reasons:

One objective, shared processing. Supervised discrete output positions can share cross-entropy. Tokenizers may require reconstruction or other objectives, and input-only modalities need no output loss. Adding a sensor can also require synchronization, missing-data conventions, masks, positional encodings, loss balancing, and evaluation. A common optimization loop is reusable, but equal loss forms do not imply equal scaling curves.
One context. Shared attention provides a path for cross-modal conditioning without a dedicated pairwise fusion module. A past command can affect a frame prediction when it remains accessible under the context and mask. Learning that dependency still requires suitable data and optimization. Controlled tasks and modality ablations help distinguish genuine conditioning from shortcuts.
Reusable infrastructure. Causal-transformer optimization, data parallelism, and cache techniques are reusable. Context-parallel kernels must support the chosen masks; quantization needs calibration and accuracy checks; speculative decoding needs compatible draft probabilities and action-conditioning semantics. Reuse is a practical advantage, not a promise that every component transfers unchanged.

Modality token counts can differ substantially: a 1,024-token image may sit beside one scalar token. Equal logits would give the image keys 1,024 times the scalar's total attention mass, but learned logits can overcome this count difference. Comparable information density is not required. Modality embeddings identify type and positions identify order; whether either helps an imbalance-related learning problem needs measurement. Token budgets, masks, and training mixtures are design choices rather than an intrinsic one-dimensional bottleneck.