FlowMo: a Transformer reconstructs through flow

An interactive field guide

FlowMo is a direct example of a Transformer flow decoder conditioned on compressed image codes.

Before this article: DiTo · DiT · TiTok

What was already possible

TiTok shows that global slots can replace a spatial bottleneck. DiTo shows that iterative denoising can train an image tokenizer. FlowMo combines a Transformer-based discrete bottleneck with a flow decoder, then improves its sampled reconstructions through post-training. These are useful comparisons, not a claim that the architectures are identical.

What changes, and why

A dual-stream Transformer exchanges information between image patches and latent queries. Lookup-free quantization, or LFQ, converts encoder outputs to binary codes. A second dual-stream Transformer predicts pixel-space velocity from the noisy image, the codes, and time. Reconstruction samples a conditional distribution instead of returning RGB in one call.

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. Dual-stream encoderimage patches + K queries → [B,K,d]2. Conditional flow decoderxt, c, t → velocity with image shape3. Iterative reconstructionnoise → x̂ conditioned on c

Dual-stream encoder

Queries gather image information; LFQ turns their values into bits.

z = encoder(image_patches,latent_queries)
c = lfq(z)

Conditional flow decoder

Image coordinates and slot identities use distinct position conventions.

v = decoder(noisy_image,c,t)
x_estimate = noisy_image - t*v

Iterative reconstruction

Every integration step calls the Transformer.

state = noise
for t,t_next in schedule:
    state += (t_next-t)*decoder(state,c,t)

Inspect the pinned source

$$\mathcal{L}_{\mathrm{flow}}=\mathbb{E}\|v_\theta(x_t,t,c)-(\epsilon-x)\|^2$$

Work through the mechanism

The pinned code defines xₜ=(1−t)x+tε, target ε−x, and single-step clean estimate xₜ−t v. Sampling moves toward lower t. The demo follows this complete convention. Some paper formulas use a different sign presentation; do not mix them with the code update.

Code count needs a unit. FlowMo-Hi has 256 internal slots with 56 bits per slot. The result table groups these into 1024 codes of 14 bits. Both descriptions carry 14336 nominal bits. At 256² pixels that is 0.21875 bits per pixel, before entropy coding or packaging.

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.

Noise time t × 100

50

Reverse Euler steps

5

Reset

The straight path has known velocity. A learned model must estimate velocity; this demo does not measure its error.

Computed values for the current example

Quantity	Value
t	0.500
x_t = (1 − t) × 2 + t × (−1)	0.500
v = ε − x	-3
Δt = −t / steps	-0.100
One reverse update	0.800

Training is not sampling

Stage 1A jointly trains the tokenizer with flow loss, a one-step LPIPS-VGG term, and LFQ terms. Stage 1B freezes the encoder and backpropagates through an eight-step reconstruction chain using a ResNet perceptual objective alongside flow loss. A separate MaskGIT later models the codes. Thus no GAN does not mean only diffusion L2, and eight post-training steps are not automatically the evaluation step count.

Reported evidence, not a local reproduction

Setting	Metric	Value	Interpretation
ImageNet-256, Lo, 256 × 18 bits	rFID ↓ / PSNR ↑	0.95 / 22.07	50k validation images; paper main table
ImageNet-256, Hi, 1024 × 14 bits	rFID ↓ / PSNR ↑	0.56 / 24.93	Code unit differs from internal slot count

What this does not establish

The reported architecture is large, about 945M encoder-plus-decoder parameters in Lo/Hi. Quality comparisons must count sampling calls, random seeds, and post-training. Reconstruction fidelity does not demonstrate text alignment. The code uses multi-axis RoPE, but the broadcast latent index is a slot identity, not a physical 3D point.

Check your understanding

Does the flow decoder see the original clean image during reconstruction sampling?

It receives its compressed code and a noisy state. The clean image is a training target, not a sampling input.

Sources and next steps

• FlowMo · 2503.11056v3
• Method / pinned implementation source

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.