OmniTokenizer: A Joint Image-Video Tokenizer for Visual Generation

Summary

This paper proposes OmniTokenizer, a tokenizer that can be used to tokenize both image and video data. To train the OmniTokenizer, they devise an architecture that decouples spatial and temporal axis, offering efficiency. Next, they perform a progressive training scheme that initializes from image-only training at fixed resolution that is followed by a joint image & video training at multiple resolutions. The resulting VQ tokenizer can be used with autoregressive modeling objective to perform generation. Optionally, it can also be fine-tuned with KL loss to be used as a diffusion tokenizer for latent diffusion models. The evaluations are performed on several image and video datasets to demonstrate the effectiveness of the tokenizer.

Strengths

Having a joint tokenizer for image and video modalities is an important step. From the results, the paper seems to do a decent job at it and the resulting tokenizers could be useful to the community. The ideas of decoupling spatial & temporal axis and progressive training are intuitive and seem effective from the provided ablation results. The paper is clearly written.

Weaknesses

My main weakness point with the paper is the lack of comparisons to the relevant baselines. I think MagVITv2 is the closest baseline to your method, but there is no in-depth comparison to it. I can see in Table 4 that MagVITv2 was reported in NAR setting while yours is in AR setting. I think it's important to have an apples-to-apples comparison to understand the effectiveness of your method. Same applies to image generation results in Table 3. For the comparison, I wouldn't limit the comparison to only generation results but also other relevant & important aspects such as video compression and the quality of learned representations. Other important and missing baselines are MAGE and MaskGIT. There are also recent advancements in the literature such as VAR but I wouldn't penalize for them as they are quite recent. Still, having a discussion & comparison to them would be useful. Other questions and comments: - How the trends would look like for 512x512 resolution? - For the reconstruction results, having other metrics such as PSNR, SSIM, and IS would be helpful. - Seeing the qualitative video results on actual videos would be really helpful. From the provided visuals it's unclear. - Are you planning to release the checkpoints & code for OmniTokenizer? - How do you explain the small performance gap for generation results with diffusion models? Can we say that OmniTokenizer is more effective with an autoregressive Transformer? - It was mentioned that the decoder is symmetric with encoder. Could you clarify that and provide a full architecture visualization? Could you also clarify the 2D and 3D patch embeddings? - What would be the reconstruction results for diffusion models on Table 1 & 2? You reported SD1.4 in the appendix. There are some numerical differences though (e.g. yours rFID is 0.69 in the main paper and 0.68 in the appendix), what is the reason for this? What are the reconstruction results for other latent diffusion models like SD3? - In the ablation results (Table 7) what would be the performance if you skip the fixed resolution image training stage and perform only the second stage on image & video jointly?

Questions

Please see the weaknesses for details. In particular, I am wondering an in-depth comparison against relevant baselines e.g. MagVITv2, MAGE, and MaskGIT. I am also wondering qualitative results on videos and some training & experimental clarifications I discussed in the weaknesses. I will update my final score accordingly.

Rating

6

Confidence

5

Soundness

3

Presentation

2

Contribution

2

Limitations

Yes.

Summary

This paper proposes OmniTokenizer, a transformer-based visual tokenizer model that processes both image and video input and achieves state-of-the-art reconstruction quality. OmniTokenizer's core designs are a decoupled spatial-temporal attention mechanism and a progressive training schedule. Two OmniTokenizers are trained for autoregressive generative models and diffusion models, respectively. Extensive experiments further show the new tokenizers effectively improve image and video generation performance, consolidating OmniTokenizer's effectiveness.

Strengths

1. Visual tokenizers are essential components for representation learning and generation, laying the foundation for many cutting-edge techniques, including image/video diffusion models and multimodal large language models. This paper's efforts in improving the tokenization quality and reconstruction performance of visual tokenizers are appreciated. 2. The proposed OmniTokenizer employs image and video data jointly with a properly designed two-stage training paradigm. Ablation studies show that the joint two-stage training strategy effectively improves the reconstruction of both image and video data. 3. Two variants of OmniTokenizers are trained, one with VQ and one with the vanilla VAE's KL regularization. The paper further trains autoregressive and diffusion-based image/video generative models with the OmniTokenizers, demonstrating superior results across several image and video datasets.

Weaknesses

1. Although there are extensive quantitative results on image/video reconstruction and generation, the qualitative comparisons are insufficient. Much more video reconstruction/generation comparisons can be provided in the supplementary materials to better demonstrate the effectiveness of OmniTokenizers. 2. The paper does not elaborate on the intuitions for designing decoupled window attention and causal attention. Table 8 investigates different architectural design choices, but there are no analyses on why the proposed design leads to better reconstruction quality than the other two variants. 3. Section 3.1.1 of the paper is too sketchy. To make the architecture reproducible and the paper self-contained, at least some concrete description or illustration of the window/causal attention mechanisms should be provided. Similarly, there are not many details about the image/video generative models trained with OmniTokenizers— the concrete model architecture/specifications and training setups could be provided in the appendix. 4. How is the order of VQ and KL training determined? There is no sufficient explanation on why the KL fine-tuning is conducted after the VQ training, and it's also unclear if the two variants of OmniTokenizers use two separate decoders or share the same decoder.

Questions

The paper studies visual tokenizers and proposes an effective architecture with a joint image-video training framework. The thorough quantitative results convincingly demonstrate the advantages of the proposed OmniTokenizers, and further scaling up this framework could lead to more powerful tokenizers that benefit future research. However, the paper lacks several important information: i) qualitative comparisons (especially for videos), ii) details of the tokenizer architecture and insights on why it works better than other variants, and iii) details of the image/video generative models. The lack of these details/analyses hurts the clarity and reproducibility. I will increase the score if these concerns are properly addressed in the rebuttal.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The limitations and broader impacts are discussed in Sections 4 and 5 of the main paper.

Summary

The paper introduces OmniTokenizer, a transformer-based tokenizer designed for both image and video tokenization within a unified framework. This tokenizer employs a spatial-temporal decoupled architecture, using window attention for spatial and causal attention for temporal modeling. The approach leverages a progressive training strategy, starting with image data to develop spatial encoding skills and then extending to video data to handle temporal dynamics. Extensive testing on several datasets demonstrates that OmniTokenizer achieves good performance in reconstruction tasks and enhances the effectiveness of both AR-based and diffusion based generative models.

Strengths

- The progressive training strategy is intuitive, allowing the model to learn incrementally from simpler to more complex tasks. This strategy contributes to the model's performance and makes sense to me. - The KL fine-tuning for converting a VQ tokenizer to a continuous tokenizer seems interesting and few works explored that.

Weaknesses

- The model does not compare with magvitv2 thoroughly, which should be a strong baseline in terms of reconstruction and generation. - The architecture for single frame mostly follows ViT-VQGAN and the improvement comes from joint ImageNet/UCF training, when compared with single frame baselines in Table 1/3/4 the comparison seems unfair as the baseline tokenizers are trained on ImageNet only, but the UCF data can also be used as a single frame data source. - The factorized spatial-temporal attention is not novel as many previous works also employed that technique.

Questions

Overall I believe this paper shows some interesting points like progressive training and KL fine-tuning. However, the unfair comparison with baselines (magvitv2 and others) prevents me from giving higher ratings, and it is difficult to evaluate the significance of the performance given the current results

Rating

4

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

Summary

The paper introduces OmniTokenizer, a transformer-based image-video tokenizer designed for visual generation tasks. It adopts a spatial-temporal decoupled architecture, integrating window attention for spatial modeling and causal attention for temporal dynamics, allowing it to process both image and video data within a unified framework. A progressive training strategy is proposed, starting with image data to develop spatial encoding and then incorporating video data for learning temporal features across multiple resolutions. Extensive experiments demonstrate OmniTokenizer's state-of-the-art performance in reconstruction tasks on various datasets, outperforming previous methods.

Strengths

1. The proposed tokenizer can enhance the performance of both language model-based and diffusion model-based visual generation approaches. 2. The key of this work lies in its potential to provide a versatile (image and video) and efficient tool for translating complex visual data into compact latent representations. 3. This paper is generally well-written and clearly stated.

Weaknesses

1. The novelty is limited. The proposed tokenizer architecture is not new and is basically the same as previous methods (e.g., MagVit). The only difference is further regarding images as a 1-frame video to pretrain the network. I don’t think that counts as enough contributions. 2. The spatial-temporal decoupled architecture is a key aspect of OmniTokenizer. A deeper dive into the role and impact of different attention mechanisms on the model's performance could offer more clarity on design choices and potential improvements. 3. While the paper mentions the potential for scaling up the model capacity, it does not provide a detailed analysis of how the model scales with increased data volume or model size. Providing insights into scalability, such as computational complexity, training time, and resource requirements, would be valuable.

Questions

See weakness

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Additional Notes on MAGVITv2 Comparison

MAGVITv2 utilizes an exceptionally large codebook size (2^18), whereas we have used only 16k, which precludes a fair comparison. Additionally, since MAGVITv2 is not open-source, reproducing its work to achieve a fair comparison is challenging. In future work, we will compare performance using the same codebook size.

Thanks, rebuttal was helpful

Thank you for the detailed response and the new results you provided, they were helpful. Please incorporate all of them to your paper. I still think having an apples-to-apples comparison with MagVITv2 on visual generation will be useful to the practitioners to understand the capabilities of two approaches. Consider adding this to your paper as well. I'm raising my score to weak accept.