NextStep-1: continuous images in a causal language model
An interactive field guide
A shared causal Transformer can predict discrete text with one head and continuous image patches with another.
Before this article: MAR · Attention and autoregression · Diffusion and flow matching
What was already possible
Section titled “What was already possible”MAR shows how a diffusion head can model continuous next-token distributions. Text language models already provide causal sequence backbones. NextStep-1 combines these interfaces at scale, keeping strict causal ordering rather than requiring random-order image generation.
What changes, and why
Section titled “What changes, and why”Initialize the backbone from Qwen2.5-14B. Text positions use categorical cross-entropy; image positions use a patch-wise flow-matching MLP. The next image token is sampled through an inner flow process conditioned on the backbone state. The shared sequence can contain both text IDs and projected continuous image vectors.
Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.
Standardize and pack
Use the tokenizer-specific channel statistics.
z = (vae.encode(image)-channel_mean)/channel_std
image_tokens = space_to_depth(z,2)Shared causal sequence
The original paper specifies 1D rotary positions.
inputs = merge(text_embeddings,project(image_tokens))
h = causal_transformer(inputs,position_ids)Two output distributions
Use CE for text and flow matching for image targets.
loss_text = ce(text_head(h[:-1]),text_targets[1:])
loss_image = flow_loss(image_head,h[:-1],image_targets[1:])Work through the mechanism
Section titled “Work through the mechanism”A fine-tuned Flux VAE produces a 32 × 32 × 16 latent for 256² images. Group each 2 × 2 region into channels: the sequence becomes 256 tokens of width 64. The scalar count stays 16384. The packing demo changes group size and checks this conservation; shorter sequences alone are not stronger information compression.
The main hidden state at position i predicts the value at i+1. Shift both target values and image/text masks together. If the image value being predicted is already included in the same context position, a low training loss may reflect target leakage rather than learning.
Change a parameter, inspect the computed values, then compare this small example with the model described above.
Analytical teaching model · controls change the example, not a trained research model.
Packing is a reversible rearrangement. The scalar count remains H × H × 16.
| Quantity | Value |
|---|---|
| Packed positions | 256 |
| Values per position | 64 |
| Total scalar values | 16384 |
Training is not sampling
Section titled “Training is not sampling”The paper uses standard 1D RoPE. Strict causal history permits KV caching, while image tokens still require sequential outer generation. The image head has 12 layers and about 157M parameters. Decoder noise augmentation and per-channel latent standardization are part of the tokenizer interface, not details that can be silently omitted.
Reported evidence, not a local reproduction
Section titled “Reported evidence, not a local reproduction”| Setting | Metric | Value | Interpretation |
|---|---|---|---|
| ImageNet-256; NextStep tokenizer | PSNR ↑ / SSIM ↑ | 30.60 / 0.89 | Table 8, reconstruction |
| Same table; Flux tokenizer | PSNR ↑ / SSIM ↑ | 31.64 / 0.91 | Generation suitability cannot be inferred from this ranking alone |
What this does not establish
Section titled “What this does not establish”The current repository also includes NextStep-1.1. A 2D-index helper existing in the code does not prove that the original model calls it. The paper links strong-CFG artifacts to token statistics drift and leaves some spatial artifacts unresolved; this is not proof that 1D positions alone cause all grid defects. Its latency table is a theoretical estimate, not our measurement.
Check your understanding
Section titled “Check your understanding”Does 2 × 2 packing discard latent scalar values?
No. It rearranges four positions into one wider vector. Any information loss comes from other steps, such as the tokenizer, not this reversible reshape.
Sources and next steps
Section titled “Sources and next steps”Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.