UniTok: split a code, share the representation

An interactive field guide

UniTok gives reconstruction and image-text alignment access to the same quantized representation.

Before this article: Semantic representations · VQ-VAE

What was already possible

A single small codebook can restrict how much detail and semantic information survive quantization. Reconstruction and contrastive objectives also put different demands on features. UniTok tests a joint training recipe with more flexible quantization, rather than assuming these demands are always incompatible.

What changes, and why

Multi-codebook quantization splits the channel vector into M sub-vectors. Each selects an entry from its own codebook; concatenate the selected vectors. This differs from residual quantization, which repeatedly quantizes the remaining error. UniTok uses a hybrid CNN/Transformer backbone and sends the quantized features to both reconstruction and image-text branches.

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.

1. Hybrid visual encoderimage → [B,256,D]2. Channel-wise codebooks[B,256,64] → [B,256,8] IDs3. Shared readoutsquantized features → pixels and alignment

Hybrid visual encoder

Project features to the quantization bottleneck.

z = projection(visual_encoder(image))

Channel-wise codebooks

Quantize eight independent 8D slices in the main configuration.

parts = split_channels(z,8)
quantized = concat([vq_m(part) for m,part in enumerate(parts)])

Shared readouts

Both branches use the representation after quantization.

image_hat = decoder(quantized)
image_embedding = pool(project_back(quantized))
loss = reconstruction(image_hat,image)+contrastive(image_embedding,text)

Inspect the pinned source

$$|\mathcal{C}_{\mathrm{combined}}|=K^M,\qquad B_{\mathrm{position}}=M\log_2K$$

Work through the mechanism

For M codebooks of K entries, store M×K vectors but represent K^M combinations. The demo computes both quantities. Combinations are not an explicitly stored giant vocabulary. Each position requires M log₂K nominal bits when all sub-codes use fixed length.

The main configuration has 256 spatial positions, eight 4096-entry sub-codebooks, and eight dimensions per sub-vector. That is 2048 IDs, 64 latent channels, and 24576 nominal bits per image, or 0.375 bpp at 256². Calling these 256 visual tokens describes the spatial interface, not 256 individual codes.

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.

Independent sub-codebooks M

2

Entries per sub-codebook K

4

Reset

Independent sub-codes concatenate. This is not repeated residual quantization of the same vector.

Computed values for the current example

Quantity	Value
Stored entries	8
Combinations	16
Nominal bits / position	4.000

Training is not sampling

Train reconstruction and image-text contrastive objectives jointly. For the multimodal model, merge sub-code embeddings at each spatial position into a main-sequence token. A depth Transformer predicts sub-codes inside that position. Tokenizer training and LLM adaptation remain separate stages. The main tokenizer data is DataComp-1B, not an ImageNet-only budget.

Reported evidence, not a local reproduction

Setting	Metric	Value	Interpretation
UniTok, from scratch, ImageNet-256 eval	rFID ↓ / zero-shot top-1 ↑	0.41 / 70.8%	Main table, jointly trained tokenizer
UniTok†, pretrained initialization	rFID ↓ / zero-shot top-1 ↑	0.38 / 78.6%	Different initialization, not a matched pretraining cost

What this does not establish

The pretrained-initialization model has better zero-shot classification but does not lead every downstream VQA result. Keep initialization, data, and task protocol attached to each score. The current encoder includes convolutional spatial structure; an absent explicit position table would not make it position-free.

Open the derivation

Each of M independent subspaces chooses one of K entries. The product rule gives K^M possible combinations; fixed-length coding needs log₂ of this count. The storage count for entries is only M×K.

$$\log_2(K^M)=M\log_2K,\qquad 256\cdot8\cdot\log_2 4096=24576\ \mathrm{bits}$$

Check your understanding

Are eight sub-codebooks equivalent to eight residual quantization stages?

No. MCQ quantizes different channel slices. RQ repeatedly quantizes the residual of the same vector.

Sources and next steps

• UniTok · 2502.20321v3
• Method / pinned implementation source
• VILA-U
• AToken

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.