Skip to main content Home Creators adu2021 skillxiv nextstep-1-autoregressive-images
nextstep-1-autoregressive-images Train a unified autoregressive model to generate images and text by directly handling continuous image tokens with flow matching, eliminating the need for quantization or separate diffusion models.
Jump to install Skills Marketplace Discover and explore AI skills built by the community.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
Show prompt details A direct command skips the review prompt. Inspect the source before running it.
npx skills add https://github.com/ADu2021/skillXiv --skill nextstep-1-autoregressive-imagesThe command stays on one line. Scroll horizontally to inspect it before copying.
Prefer a local copy? Download the files currently available to SkillsMP.
Download Zip Downloading... More from this repository
Related occupations SOC
Based on SOC occupation classification
name nextstep-1-autoregressive-images title NextStep-1: Autoregressive Image Generation with Continuous Tokens version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2508.10711 keywords ["image-generation","autoregressive-modeling","continuous-tokens","flow-matching","text-to-image"] description Train a unified autoregressive model to generate images and text by directly handling continuous image tokens with flow matching, eliminating the need for quantization or separate diffusion models.
NextStep-1: Autoregressive Image Generation with Continuous Tokens
Core Concept
Most text-to-image models use separate pipelines: text encoders, diffusion models, and discrete image tokenizers. NextStep-1 unifies image and text generation in a single autoregressive architecture that directly processes continuous image tokens using flow matching instead of diffusion.
The key innovation is combining a large autoregressive language model (14B parameters) with a smaller flow matching head (157M parameters) to predict both discrete text tokens and continuous image tokens in a unified next-token prediction framework.
Architecture Overview
Unified Autoregressive Backbone : Single 14B-parameter transformer predicts both text and image tokens sequentiallyFlow Matching Head : 157M-parameter decoder with flow matching objectives for continuous image token predictionMixed Token Types : Seamlessly handles discrete tokens (text) and continuous vectors (image features) in one sequenceNext-Token Prediction : Standard language model training objective applied across both modalitiesEfficient Inference : Direct autoregressive generation without iterative diffusion sampling loops
Implementation Steps
1. Encode Images to Continuous Tokens
Convert images into continuous token representations using a learned encoder. These are embedded directly into the autoregressive sequence without quantization.
import torch
import torch.nn as nn
from transformers import AutoModel
class ContinuousImageTokenizer (nn.Module):
"""
Encodes images to continuous token embeddings
Similar to VQ-VAE but without quantization
"""
def __init__ (self, image_size=256 , token_dim=768 , num_tokens=1024 ):
super ().__init__()
self .image_size = image_size
self .token_dim = token_dim
self .num_tokens = num_tokens
self
3
64
4
2
1
64
128
4
2
1
128
256
4
2
1
self
256
32
32
def
encode
self, images
"""
Args:
images: [batch, 3, H, W] image tensors
Returns:
tokens: [batch, num_tokens, token_dim] continuous embeddings
"""
self
0
1
self
self
self
return
def
decode
self, tokens
"""
Args:
tokens: [batch, num_tokens, token_dim]
Returns:
images: [batch, 3, H, W] reconstructed images
"""
0
1
self
256
32
32
256
128
4
2
1
128
64
4
2
1
64
3
4
2
1
return
2. Create Mixed Token Sequences Combine text tokens and continuous image tokens into single sequences that the autoregressive model can process.
def create_mixed_sequence (text_ids, image_tokens, tokenizer, image_token_start=50257 ):
"""
Create interleaved sequence of text and image tokens.
Format: [text_tokens] [IMAGE_START] [image_tokens] [IMAGE_END]
Args:
text_ids: [seq_len] text token IDs
image_tokens: [num_image_tokens, token_dim] continuous embeddings
tokenizer: text tokenizer
image_token_start: special token ID marking image start
Returns:
input_ids: [total_seq_len] mixed token sequence (text IDs as scalars)
embeddings: [total_seq_len, embedding_dim] actual embeddings
"""
text_embedding_dim = 768
image_token_dim = 768
text_embeddings = tokenizer.transformer.wte(text_ids)
input_ids = torch.cat([
text_ids,
torch.tensor([image_token_start], device=text_ids.device),
torch.arange(50257 + 1 , 50257 + image_tokens.shape[0 ] + 1 ,
device=text_ids.device),
torch.tensor([image_token_start + 1 ], device=text_ids.device)
])
image_start_embed = text_embeddings[-1 :] * 0.5
embeddings = torch.cat([
text_embeddings,
image_start_embed,
image_tokens,
image_start_embed
], dim=0 )
return input_ids, embeddings
3. Define the Unified Autoregressive Model Build a transformer that predicts next tokens, treating continuous image tokens like any other embedding.
class UnifiedAutoregressive (nn.Module):
"""
14B parameter autoregressive model for text + continuous image tokens
"""
def __init__ (self, vocab_size=50257 , hidden_size=1024 , num_layers=24 ,
num_image_tokens=1024 , image_token_dim=768 ):
super ().__init__()
self .vocab_size = vocab_size
self .hidden_size = hidden_size
self .num_image_tokens = num_image_tokens
self .embed_tokens = nn.Embedding(vocab_size, hidden_size)
self .transformer = nn.ModuleList([
nn.TransformerDecoderLayer(
d_model=hidden_size, nhead=16 ,
dim_feedforward=4096 , batch_first=True
) for _ in range (num_layers)
])
self .text_head = nn.Linear(hidden_size, vocab_size)
self .image_token_head = nn.Linear(hidden_size, image_token_dim)
def forward (self, input_ids, embeddings ):
"""
Args:
input_ids: [batch, seq_len] (for masking/identifying token types)
embeddings: [batch, seq_len, hidden_size] mixed embeddings
Returns:
text_logits: [batch, seq_len, vocab_size] for text token prediction
image_logits: [batch, seq_len, image_token_dim] for image token prediction
"""
hidden = embeddings
for layer in self .transformer:
hidden = layer(hidden, hidden)
text_logits = self .text_head(hidden)
image_logits = self .image_token_head(hidden)
return text_logits, image_logits
4. Implement Flow Matching Loss for Image Tokens Use flow matching instead of diffusion to train the continuous image token predictions.
class FlowMatchingHead (nn.Module):
"""
Flow matching objective for continuous image tokens.
Simpler than diffusion, more stable than MSE loss.
"""
def __init__ (self, token_dim=768 , hidden_size=1024 ):
super ().__init__()
self .mlp = nn.Sequential(
nn.Linear(token_dim + hidden_size, 512 ),
nn.GELU(),
nn.Linear(512 , token_dim)
)
def forward (self, target_tokens, context_hidden, t=None ):
"""
Flow matching loss: learn to map noise to data
target: continuous image tokens [batch, num_tokens, token_dim]
context: model hidden states [batch, seq_len, hidden_size]
"""
batch_size, num_tokens, token_dim = target_tokens.shape
noise = torch.randn_like(target_tokens)
t = torch.rand(batch_size, 1 , 1 , device=target_tokens.device)
z_t = (1 - t) * noise + t * target_tokens
context_repeated = context_hidden[:, -num_tokens:, :].unsqueeze(1 )
velocity_pred = self .mlp(torch.cat([z_t, context_repeated.expand_as(z_t)], dim=-1 ))
target_velocity = target_tokens - noise
flow_loss = ((velocity_pred - target_velocity) ** 2 ).mean()
return flow_loss
5. Training Loop with Mixed Objectives Combine text prediction loss (cross-entropy) with flow matching loss for image tokens.
def train_nextstep (model, flow_head, image_tokenizer, dataloader, num_epochs=10 ):
"""
Train unified autoregressive model with mixed objectives
"""
optimizer = torch.optim.AdamW(
list (model.parameters()) + list (flow_head.parameters()) +
list (image_tokenizer.parameters()),
lr=1e-4
)
for epoch in range (num_epochs):
for batch_idx, (images, captions) in enumerate (dataloader):
image_tokens = image_tokenizer.encode(images)
text_ids = tokenizer.encode(captions)
input_ids, embeddings = create_mixed_sequence(
text_ids, image_tokens, tokenizer
)
text_logits, image_logits = model(input_ids, embeddings)
text_loss = F.cross_entropy(
text_logits.view(-1 , model.vocab_size),
text_ids.view(-1 )
)
image_context = embeddings[:, :text_ids.shape[1 ], :]
flow_loss = flow_head(image_tokens, image_context)
total_loss = text_loss + 0.5 * flow_loss
optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0 )
optimizer.step()
if batch_idx % 100 == 0 :
print (f"Epoch {epoch} , Step {batch_idx} " )
print (f" Text Loss: {text_loss:.4 f} , Flow Loss: {flow_loss:.4 f} " )
6. Inference and Image Generation Use the trained model to generate images autoregressively from text prompts.
@torch.no_grad()
def generate_image (model, flow_head, image_tokenizer, prompt, max_new_tokens=1024 ):
"""
Generate image autoregressively from text prompt
"""
input_ids = tokenizer.encode(prompt)
input_ids = torch.cat([
input_ids,
torch.tensor([50257 ], device=input_ids.device)
])
generated_tokens = []
hidden_state = None
for step in range (max_new_tokens):
embeddings = model.embed_tokens(input_ids.unsqueeze(0 ))
text_logits, image_logits = model(input_ids.unsqueeze(0 ), embeddings)
next_token_logits = image_logits[0 , -1 , :]
next_token = torch.randn_like(next_token_logits) * 0.1 + next_token_logits
generated_tokens.append(next_token)
if step % 100 == 0 :
print (f"Generated {step} /{max_new_tokens} image tokens" )
image_tokens = torch.stack(generated_tokens)
image = image_tokenizer.decode(image_tokens.unsqueeze(0 ))
return image
Practical Guidance
Hyperparameters & Configuration
Token Dimension : 768 (align with language model embedding size)Number of Image Tokens : 1024 (depends on desired spatial resolution)Flow Matching Weight : 0.5 relative to text loss (balance both objectives)Learning Rate : 1e-4 (conservative due to mixed objectives)Batch Size : 64-256 (computational dependent)Gradient Clipping : max_norm=1.0 (prevent divergence)
When to Use NextStep-1 Approach
You want unified text-to-image generation in one model
You prefer autoregressive generation over iterative sampling
Computational efficiency (no diffusion loops) is important
You need both text and image generation capabilities
You want to avoid quantization artifacts from VQ-VAE
When NOT to Use NextStep-1
Image quality is paramount (diffusion models still produce better results)
You need fine-grained control over image generation (e.g., inpainting)
Inference latency must be absolute minimum (autoregressive is slower than diffusion)
You only need image-to-text or don't need text+image in same model
Your computational resources are very limited
Common Pitfalls
Imbalanced Loss Weights : If text loss dominates, image quality suffers. Use 1:0.5 ratio and monitor both.Insufficient Image Tokens : Too few tokens (< 256) loses spatial detail. 1024 is reasonable baseline.Poor Continuous Token Learning : Without good flow matching, image tokens become noisy. Use proper velocity field training.Ignoring Text-Image Alignment : Text and image should be processed with awareness of each other, not independently.Slow Inference : Autoregressive generation is 100x slower than diffusion for same image. Use strategies like speculative decoding.
Reference Unified autoregressive model for text and continuous image tokens with flow matching, eliminating quantization loss and diffusion overhead while enabling efficient image generation and editing.
