Architecture

The Cosmos Tokenizer is designed as an encoder-decoder architecture. Given an input video \(x_{0:T} \in \mathbb{R}^{(1+T) \times H \times W \times 3}\) with height \(H\), width \(W\), and frames \(T\), the encoder (\(\mathcal{E}\)) compresses it into tokens \(z_{0:T'} \in \mathbb{R}^{(1+T') \times H' \times W' \times C}\) with spatial and temporal compression factors \(s_{HW}\) and \(s_T\), respectively. The decoder (\(\mathcal{D}\)) reconstructs the video as \(x'_{0:T} \in \mathbb{R}^{(1+T) \times H \times W \times 3}\), mathematically as:

\[x'_{0:T} = \mathcal{D}\Big(\mathcal{E}\big(x_{0:T}\big)\Big).\]

Our architecture employs a temporally causal design, processing only current and past frames at each stage. The tokenizer operates in the wavelet transform space, where a 2-level wavelet transform downsamples inputs by a factor of four, grouping frames as \(\{g_0, g_1, \dots\} = \{x_0, x_{1:4}, x_{5:8}, \dots\}\). Subsequent stages encode these groups causally, ultimately outputting tokens \(z_{0:T'}\). This causal design is essential for physical AI applications, which often operate in a similar causal setting, and the wavelet transform enables effective compression by emphasizing semantic content.

(a) Temporal Causality: Illustration of the temporal causality mechanism, where inputs \(x_0, x_1, \dots, x_{12}\) are processed through grouped intermediate outputs via wavelet transform \(g_0, g_1, \dots\), and further refined by causal spatio-temporal convolution and attention operations.

(b) Network Architecture: The encoder-decoder network structure includes a 3D Haar wavelet, causal residual, downsampling, and spatio-temporal attention blocks. The decoder mirrors the encoder's structure.

Figure 4: Cosmos Tokenizer architecture.

Our encoder stages (after the wavelet transform) consist of residual blocks interleaved with down-sampling blocks. Each encoder block applies a spatio-temporal factorized 3D convolution: a 2D convolution of \(1 \times k \times k\) captures spatial information, followed by a temporal convolution of \(k \times 1 \times 1\) to capture dynamics, with left padding of \(k-1\) ensuring causality. Long-range dependencies are captured through a spatio-temporal factorized causal self-attention with a global support region, e.g., \(1+T'\) in the last encoder block. The decoder mirrors the encoder but replaces downsampling with upsampling. Figure 4 shows the Cosmos Tokenizer architecture.

We use a vanilla autoencoder (AE) to model continuous tokenizers' latent space and Finite-Scalar-Quantization (FSQ) (Mentzer et al., 2023) for discrete tokenizers. The latent dimension for continuous tokenizers is 16, while for discrete tokenizers it is 6, with FSQ levels of (8, 8, 8, 5, 5, 5), resulting in a vocabulary size of 64,000.

Training

We use a joint training strategy, alternating mini-batches of images and videos at a preset frequency. We supervise only the final output of the tokenizer's decoder, without auxiliary losses in the latent spaces (e.g., commitment or KL prior losses). For instance, a VAE for continuous tokenizers would require a KL prior loss, and a VQ-VAE for discrete quantization would require a commitment loss, unlike our vanilla AE and FSQ approach.

We employ a two-stage training scheme. In the first stage, we optimize with the L1 loss that minimizes the pixel-wise RGB difference between the input and reconstructed video (\(\hat{x}_{0:T}\)), given by \[\mathcal{L}_1 = \left\lVert \hat{x}_{0:T} - x_{0:T} \right\rVert_1,\] and the perceptual loss based on the VGG-19 features (Simonyan and Zisserman, 2014), given by, \[\mathcal{L}_{\text{Perceptual}} = \frac{1}{L} \sum_{l=1}^{L} \sum_{t} \alpha_l \left\lVert \texttt{VGG}_l (\hat{x}_t) - \texttt{VGG}_l (x_t) \right\rVert_1,\] where \(\texttt{VGG}_l (\cdot) \in \mathbb{R}^{H \times W \times C}\) is the features from the \(l\)-th layer of a pre-trained VGG-19 network, \(L\) is the number of layers considered, and \(\alpha_{l}\) is the weight of the \(l\)-th layer. When \(x\) is a video (i.e. \(T>0\)), we compute the perceptual loss as the average of the per-frame perceptual losses.

In the second stage, we use the optical flow (\(\texttt{OF}\)) loss (Teed and Deng, 2020) to handle the temporal smoothness of reconstructed videos, \[\mathcal{L}_{\text{Flow}} = \frac{1}{T} \sum_{t=2}^{T+1} \left\lVert \texttt{OF}(\hat{x}_{t}, \hat{x}_{t-1}) - \texttt{OF}(x_{t}, x_{t-1}) \right\rVert_1 + \frac{1}{T} \sum_{t=1}^{T} \left\lVert \texttt{OF}(\hat{x}_{t}, \hat{x}_{t+1}) - \texttt{OF}(x_{t}, x_{t+1}) \right\rVert_1,\] and the Gram-matrix (\(\texttt{GM}\)) loss (Gatys et al., 2016) to enhance the sharpness of reconstructed images, \[\mathcal{L}_{\text{Gram}} = \frac{1}{L} \sum_{l=1}^{L} \sum_t \alpha_l \left\lVert \texttt{GM}_l (\hat{x}_t) - \texttt{GM}_l (x_t) \right\rVert_1.\]

Additionally, we use adversarial loss in the fine-tuning stage to further enhance reconstruction details, particularly at large compression rates.

Evaluation

We extensively evaluated our Cosmos Tokenizer suite on diverse image and video benchmark datasets. For image tokenizers, we followed previous work and evaluated on MS-COCO 2017 (Lin et al., 2015), ImageNet-1K (Deng et al., 2009), FFHQ (Karras et al., 2019), and CelebA-HQ (Karras, 2017).

Figure 5: Example videos from TokenBench, including egocentric, driving, robotic manipulation, and web videos.

TokenBench To address the lack of standardized benchmarks for high-resolution, long-duration videos, we introduced TokenBench, covering various domains such as robotics, driving, egocentric, and web videos. TokenBench includes 100 randomly sampled, 10-second clips from each of BDD100K (Yu et al., 2020), EgoExo-4D (Grauman et al., 2024) BridgeData V2 (Walke et al., 2023), and Panda-70M (Chen et al., 2024). We manually filtered low-quality videos in Panda-70M and selected 100 scenes from EgoExo-4D, with one egocentric and one exocentric video per scene, totaling 500 videos. Examples of TokenBench are shown in Figure 5, and we have made it publicly available for the community at github.com/NVlabs/TokenBench. We also evaluated our video tokenizers on the DAVIS (Perazzi et al., 2016) dataset at 1080p.

Baselines and Evaluation Metrics We evaluated our tokenizers across different compression rates, comparing them with state-of-the-art (SOTA) image and video tokenizers (see Table 1 for details). The evaluation metrics include PSNR, SSIM, reconstruction Fréchet Inception Distance (rFID) for images, and Fréchet Video Distance (rFVD) for videos.

Qualitative Results

Figure 6: Comparison of video reconstruction with continuous video tokenizer with 4×8×8 compression (CV4×8×8).

Figure 6 and 7 show the reconstructed video frames using continuous video tokenizers, covering diverse scenarios such as natural scenes, autonomous driving, robotic manipulation, and egocentric videos. Figure 9 shows reconstructed images with different discrete image tokenizers, while Figure 8 presents error maps for continuous image tokenizers to highlight reconstruction differences. Compared to previous methods, Cosmos Tokenizer more effectively preserves structural and high-frequency details (e.g., grass, tree branches, text) with minimal visual distortion (e.g., faces, text) and ghosting.

Cosmos

Omni

Cosmos

Omni

Cosmos

Omni

Cosmos

Omni

Figure 7: Comparison of video reconstruction with discrete video tokenizers with 4×8×8 compression (DV4×8×8).

These qualitative results show that Cosmos Tokenizer is capable of encoding and decoding a wide range of visual contents and has the capability to preserve the highest visual quality of images and videos.

Figure 8: Comparison of image reconstruction error maps with continuous image tokenizers with 8×8 compression (CI8×8).

Figure 9: Comparison of image reconstruction with discrete image tokenizers with 16×16 compression (DI16×16).

Quantitative Results

Tables 2 and 3 present average quantitative metrics for continuous and discrete video tokenizers on various benchmarks. Cosmos Tokenizer achieves state-of-the-art performance on both the DAVIS and TokenBench datasets at a 4×8×8 compression ratio. Even at higher compression rates (8×8×8 and 8×16×16), Cosmos Tokenizer outperforms previous methods, demonstrating an excellent compression-quality trade-off.

Similarly, Tables 4 and 5 show that Cosmos Tokenizer consistently achieves superior results on various image benchmarks at 8×8 compression rate. Notably, at 4× larger compression rate of 16×16, the image quality of Cosmos Tokenizer often matches or surpasses prior art at 8×8. These results confirm Cosmos Tokenizer's capability to represent visual content effectively with high spatio-temporal compression.

Extended Visual Results

To further demonstrate the effectiveness of Cosmos Tokenizer, we provide comparison videos across driving, robotics, and sport videos. Pause the videos at desired temporal locations to closely inspect image quality differences.

Runtime Performance

Table 6 compares the parameters and average encoding and decoding times per image or video frame on a single A100 80GB GPU. Cosmos Tokenizer achieves 2× to 12× faster speeds than prior state-of-the-art tokenizers while maintaining the smallest model size, demonstrating high efficiency in encoding and decoding visual content.

Together with the qualitative results, quantitative metric evaluations, and the efficiency evaluation, these results demonstrate that Cosmos Tokenizer is well placed to serve as an effective and efficient building block in both diffusion-based and autoregressive models for image and video generation and the development of a world model. You can try out Cosmos Tokenizer on your projects using the source code and pre-trained models, which are now publicly available.

Contributors

Fitsum Reda, Jinwei Gu, Xian Liu, Songwei Ge, Ting-Chun Wang, Haoxiang Wang, Ming-Yu Liu