Pixel-space video generation with U-JiT

Théophane Vallaeys; Jakob Verbeek; Matthieu Cord

cs.CVNov 3, 2026

Pixel-space video generation with U-JiT

Authors: Théophane Vallaeys, Jakob Verbeek, Matthieu Cord

Organizations: FAIR, Meta · Sorbonne University, Paris

Abstract

Latent diffusion and flow matching models are ubiquitous in state-of-the-art generative video models. While highly effective, relying on pretrained autoencoders to compress the data into a lower-dimensional latent space also comes with drawbacks. First, training proceeds in two stages: autoencoder pretraining, with compression and reconstruction objectives, followed by latent generative model training, precluding fully end-to-end training. Second, autoencoders also lead to an inherent loss of signal, and the trade-off between reconstruction quality and generative performance in the latent space is still an active area of research. In this paper we explore the viability of generative video models directly in pixel space. We analyze recent approaches in pixel-space generative modeling of images that rely on basic transformer designs, and identify several shortcomings that limit their effectiveness for the high-dimensional pixel-space video generation. To overcome these limitations, we propose a UNet-like transformer architecture along with improvements to the training recipe and sampling strategy. We conduct controlled experiments, training latent and pixel-space models from scratch on Kinetics-600. For the standard frames-to-video task at 128 resolution, our U-JiT model improves generative FVD by 27% over prior state-of-the-art results, and by 14% against the JiT pixel-space baseline. Scaling to higher 256 resolution, we find the improvement over pixel-space and latent-based models grows as we increase the spatiotemporal compression rate up to $64 \times 64 \times 8$ . Finally, we find that U-JiT improves per-frame FID by 27% over latent models and 61% over pixel-space models on class-to-video generation, with respectively 9% and 56% reduction in FVD.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison of JiT and our proposed U-JiT on video generation. While JiT processes the video with a single stack of large-patch blocks, U-JiT varies the patch size along model depth. In this manner the model benefits from small patches in the outer layers to capture high resolution detail, as well as high channel depth in the inner layers to model high-level semantic structure. On the right, frames generated by each model, and the PCA spectrum of generated vs . real videos.

Figure 2: PCA spectrum over

16\times 16\times 4

patches, computed over 1k real and sampled videos.

Panel of Figure 3 — Figure 3: U-JiT architecture overview. Encoding and decoding blocks are parametrized to match standard JiT blocks used in the main backbone in computational cost.

Panel of Figure 4 — Figure 4: U-JiT block, with attention tiling and patchification blocks.

Figure 5: Band power retention in decoded videos, depending on latent capacity.

Figure 6: Reconstruction quality over latent jittering.

Figure 7: Effect of two types of noise on metrics.

		CFG
	$\sigma_{t}$	1	1.5	2	2.5	3	3.5
FVD@5k $\downarrow$	1	3.43	3.21	3.15	3.15	3.16	3.17
	2	3.37	3.17	3.14	3.13	3.13	3.15
	4	3.33	3.14	3.12	3.11	3.12	3.15
	8	3.29	3.14	3.10	3.14	3.20	3.29
	16	3.32	3.15	3.14	3.26	3.47	3.78

Table 1: U-JiT-B sampling sweep on Kinetics-600, at 200 epochs. A CFG of 1 reverts to non-CFG sampling, and

\sigma_{t}\!=\!1

to linear time sampling.

Model (f8t4)	Pixel	#Prms.	FLOPs	$\downarrow$ FVD	$\downarrow$ FID
Video Diffusion Models ( Ho et al., 2022b )	✓	1.1B	4.1T	16.2	-
RIN ( Jabri et al., 2023 )	✓	411M	386G	10.8	-
MAGVIT-v2 ( Yu et al., 2024 )	$\times$	307M	-	4.3	-
DeRA ( Guo et al., 2025 )	$\times$	632M	-	4.1	-
W.A.L.T ( Gupta et al., 2025 )	$\times$	313M	-	3.3	-
Unified Latents ( Heek et al., 2026 )	$\times$	-	-	1.3	-

Table 2: Frames to video: 5 conditioning frames, 12 generated frames, at resolution

128\!\times\!128

with spatiotemporal compression f8t4. B- and XL-sized models. FLOPs are computed for a single forward pass of the model, not accounting for the decoder cost of latent models. Pixel column indicates pixel-space (✓) and latent-space (

\times

) generative models.

Figure 8: Scaling the spatiotemporal compression rate. FVD@5k (left) and FID@5k (right) as functions of the compression. Evaluated at 400 epochs, 1 to 15 frames, resolution

256\!\times\!256

Model (f32t8)	#Prms.	$\downarrow$ FVD	$\downarrow$ FID
$l$ -JiT-XL (Cosmos Tokenizer)	644M	3.94	12.41
$l$ -JiT-XL (Wan2.1 Tokenizer)	644M	2.73	12.86
$l$ -JiT-XL (Wan2.2 Tokenizer)	644M	2.66	9.97
$l$ -JiT-XL (LTX Tokenizer)	643M	5.09	21.06
JiT-XL (base sampling)	700M	8.98	52.53
JiT-XL	700M	5.86	30.27

Table 3: Image to video, strong compression . 1 to 16 frames, at resolution

256\times 256

with spatiotemporal compression f32t8.

Model (f32t8)	#Prms.	$\downarrow$ FVD	$\downarrow$ FID
$l$ -JiT-XL (Cosmos Tokenizer)	647M	4.78	17.44
$l$ -JiT-XL (Wan2.1 Tokenizer)	647M	3.31	21.52
$l$ -JiT-XL (Wan2.2 Tokenizer)	647M	3.55	18.08
$l$ -JiT-XL (LTX Tokenizer)	646M	3.96	22.24
JiT-XL (base sampling)	707M	9.45	65.49
JiT-XL	707M	6.88	32.60

Table 4: Class to video . 65 frames generated from a class condition, at resolution

256\times 256

with strong spatiotemporal compression f32t8.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

	JiT-B	JiT-L	JiT-XL
architecture
depth	12	24	28
hidden dim	768	1024	1152
heads	12	16	18
dropout	0
register tokens	32

Appendix

Table S1: Configurations of experiments.

Setting
Video resolution	$128\times 128$		$256\times 256$
Video length	17 frames		17 frames	65 frames
Tokenizer patch size	f8t4 ( $8\times 8\times 4$ )		f32t8 ( $32\times 32\times 8$ )
Model size	B	XL	XL	XL
Pixel-space noise scaling	2		2	4
First level

Appendix

Table S2: U-JiT blocks parametrization. Per-setting configuration of the U-JiT multi-level blocks across video resolutions, lengths, tokenizers, and model sizes, together with parameter counts and per forward-pass FLOPs compared to JiT .

		CFG
	$\sigma_{t}$	1	1.5	2	2.5	3	3.5
FVD@5k $\downarrow$	1	3.17	3.02	3.11	3.20	3.25	3.27
	2	3.13	3.02	3.12	3.21	3.26	3.28
	4	3.08	3.01	3.13	3.22	3.28	3.29
	8	3.05	3.01	3.15	3.22	3.29	3.34
	16	3.13	3.10	3.25	3.39	3.58	3.82

Appendix

Table S3:

l

-JiT-B (Wan2.2 Tokenizer) latent sampling sweep on Kinetics-600 (epoch 200), on 5-to-11 frames video generation. FVD@5k and FID@5k over

\sigma_{t}

(time shift) and CFG. Lower is better; best per metric in bold .

	Model (f32t8)	FLOPs		FVD@5k		FID@5k
Default configuration
(1)	U-JiT-L (ours)	137G		5.35		19.20
Loss ablations
(2)	w/o LPIPS	137G	–	6.33	$+$ 0.98	35.47	$+$ 16.27
Architecture ablations
(3)	w/o attention tiling	358G	$+$ 221G	5.32	$-$ 0.03	19.13	$-$ 0.07

Appendix

Table S4: Architectural and sampling ablations. Variants of U-JiT-L trained for 400 epochs on 1 conditioning frame to 15 evaluated frames, at resolution

256\!\times\!256

with spatiotemporal compression f32t8. The first block ablates the loss, the second the architectural choices (attention adaptations, patchification levels and network structure), and the third the sampling-time methods. FLOPs are computed for a single forward pass of the model.

High-Resolution Video Generation

See all 16 papers

Sep 29, 2026

PixelUMM: Encoder-Free Unified Image and Video Understanding and Generation

Unified Multimodal Models (UMMs) often rely on separate visual representations for understanding and generation, increasing visual context length and complicating integration with established vision-language pretraining pipelines. Recent advances in pixel-space modeling offer an encoder-free alternative, but extending this paradigm from images to videos is non-trivial: video understanding and generation adopt different temporal representations, leaving the design of a unified visual interface an open question. We present PixelUMM, an encoder-free model for unified image and video understanding and generation directly in pixel space. PixelUMM represents images as spatial patches and videos as spatiotemporal tubelets, connecting raw pixels to a shared multimodal backbone through single-layer linear projections. Its Mixture-of-Transformers architecture combines shared attention with task-specific parameters and extends clean-pixel prediction to video generation, jointly supporting autoregressive text prediction and pixel-space flow matching. Experiments show that PixelUMM achieves competitive performance across image and video understanding and generation tasks. We further conduct empirical studies of key design choices, including decoder design and spatial-temporal patch size, providing insights for future pixel-space unified multimodal models.

Cong Wei, Xuanchi Ren, Bryan Chu +8

NVIDIA

Unified Multimodal Models Multimodal Generation Video Generation Multimodal Learning Multimodal Understanding

Ask PDF

Jul 20, 2026

Pixel-Space Diffusion Transformers

Latent diffusion models (LDMs) enable efficient high-resolution image synthesis by denoising in a VAE-compressed latent space. However, fixed visual tokenizers can discard fine textures and structural details, while separate representation and diffusion training creates a mismatch between reconstruction and generation objectives. These limitations have renewed interest in pixel-space diffusion, which models raw pixels directly, removes the VAE bottleneck, and supports end-to-end optimization. This formulation better matches the demands of high-fidelity generation but introduces challenges in high-dimensional modeling, including noise scheduling, loss weighting, token efficiency, and scalable architecture design. Pixel-space modeling also offers a promising basis for unified multimodal systems: raw pixels, text, and task conditions can be represented in a shared token space and jointly processed by a single Transformer, narrowing the gap between visual understanding and generation. This paper reviews Pixel-Space Diffusion Transformers (pDiTs) from the perspectives of model architecture, continuous generative mechanisms, and unified multimodal modeling. We summarize representative methods, identify key technical challenges, and discuss future directions toward high-fidelity, end-to-end vision foundation models that integrate generation and understanding.

Renye Yan, Jikang Cheng, You Wu +8

Peking University, China · Nanjing University, China · Stanford University, USA +3

Diffusion Transformer Pixel-Space Diffusion Models Diffusion-Based Image Generation Multimodal Diffusion Models

Ask PDF

Sep 16, 2026

vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes

Open-source video generative models ship almost exclusively as PyTorch/CUDA reference implementations. This leaves Cloud TPU pods without a production-ready inference path, despite offering large, cost-effective accelerator memory pools ideal for long-sequence spatiotemporal attention. We present vidax, an open-source JAX/Flax inference engine and zero-copy PyTorch-to-JAX weight translator for modern video generation architectures. vidax covers a diverse set of spatiotemporal models --- including Diffusion Transformers, omnimodal Mixture-of-Transformers, 3D VAEs, text encoders, and native samplers --- with zero PyTorch dependency in the execution path. The framework unifies 1D tensor parallelism with DeepSpeed-Ulysses sequence parallelism on a single JAX sharding mesh, integrates TPU flash-attention kernels, and implements per-layer weight offloading to support reference resolutions that exceed single-device memory. We benchmark compile times, latency, and peak memory utilization on TPU v4-8 hardware, and document real-world numerical bugs surfaced during checkpoint translation. vidax is released open-source as a baseline for JAX and TPU video generation research.

Congyue Deng

Massachusetts Institute of Technology

Video Diffusion Models Tensor Parallelism Video Generation Distributed Inference Diffusion Model Inference Acceleration

Ask PDF

Abstract

High-Resolution Video Generation

Explore similar work

PixelUMM: Encoder-Free Unified Image and Video Understanding and Generation

Pixel-Space Diffusion Transformers

vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes