Paper: Language Model Beats Diffusion - Tokenizer is key to visual generation (ICLR-2024)

Reviewer comment (Experiments):
While the presented method is tailored for masked LM, many of the prevailing and powerful LLMs, such as LLaMA [A], employ an autoregressive approach. Incorporating results from AR-LM would greatly enhance the paper's relevance to the community.

Author response:
In the Appendix, we have added experiments on the autoregressive language model (AR-LM) where the proposed tokenizer outperforms the prior work MAGVIT [R3-A]. The results, also provided in Table R3.1 below, show our tokenizer can also work with the AR-LM.

It is worth noting that we have observed that achieving optimal performance requires a much larger AR-LM model size. But due to time constraints, we have not been able to train a larger AR-LM model with an adequate number of training steps. That also partly explains our current choice of not building our work on the autoregressive language model, and we leave it as future work.

*Table R3.1: Video generation results: class-conditional generation on UCF-101 with AR-LM models. We use the same transformer configuration as MLM experiments but without vocabulary factorization and weight tying. As a result, the AR-LM with our tokenizer uses more parameters in the embedding table and the softmax layer.*

| Tokenizer | FVD↓ | # Params | # Steps |
|---|:---:|:---:|:---:|
| MAGVIT | 265 | 306M | 1024 |
| *This paper* | **109** | 840M | 1280 |

Reviewer comment (Experiments):
In Table 4, despite the good action recognition performance showcased by the proposed method, it doesn't conclusively establish its efficacy as a viable self-supervised pre-training target. Notably, some pivotal baselines, like pixel colors and the image descriptor from MaskFeat [B], are missing.

Author response:
Following the reviewer’s suggestion, we have added pixel colors and HoG descriptors from MaskFeat [B] as baseline self-supervised pre-training targets into revised Table 4, with a distilled version in Table R3.2 below. As shown, pre-training with tokens from our model as the target achieves the best performance compared to the previous leading tokenizer MAGVIT as well as other baselines, including raw pixels, 3D VQ-VAE, and HoG descriptors. 

*Table R3.2 Video action recognition performance (classification accuracy↑ $\times$100) with different self-supervised pre-training targets as the transformer output.*

| Target | SSv2 |
|---|:---:|
| 3D VQ-VAE | 64.13 |
| MAGVIT | 67.22 |
| *This paper* | **67.38** |
| Raw pixel | 64.83 |
| HoG descriptor | 65.86 |

Reviewer comment (Experiments):
In Figure 1, it's highlighted that the VQ generation FID sees a pivotal change at a vocabulary size of 2^14, while the LFG generation FID consistently improves. I'm curious to understand how the LFG generation FID would respond to even larger vocabulary sizes.

Author response:
In Table R3.3, we showcase an instantiation of LFQ with a much larger vocabulary size at $2^{40}$. As shown, it yields consistent improvement over $2^{18}$ in both reconstruction and generation FID metrics.

*Table R3.3 Class-conditional image generation results on ImageNet 512$\times$512 with larger LFQ vocabularies. We use an MLM model and 12 decoding steps.*

| Vocabulary | Groups | # Params | rFID↓ | FID↓ |
|:---:|:---:|:---:|:---:|:---:|
| $2^{18}$ | 2 | 307M | 1.22 | 5.93 |
| $2^{40}$ | 4 | 312M | **0.80** | **5.15** |

Reviewer comment (Experiments):
The question “Why masked LM and not AR LM for image / video generation?” for evaluating the tokenizer was not clearly answered.

Author response:
We selected the masked language model (MLM) due to its competitive performance on benchmark datasets [R2-A, R2-B] compared to the autoregressive language model (AR-LM), where MLM also uses much fewer decoding steps than AR-LM. In Table R2.1, we also show experiments on the AR-LM where the proposed tokenizer outperforms the prior work MAGVIT [R2-A]. The results show our tokenizer can also work with the AR-LM, but the AR-LM lags behind the state-of-the-art. We have included this rationale in the revised Section 4.2 and added the new result to the Appendix.

It is worth noting that we have observed that achieving optimal performance requires a much larger AR-LM model size. But due to time constraints, we have not been able to train a larger AR-LM model with an adequate number of training steps. That also partly explains our current choice not to build our work on the autoregressive language model. We are investigating the integration of the tokenizer with the large AR-LM and will study this in our future research.

*Table R2.1: Video generation results: class-conditional generation on UCF-101 with AR-LM models. We use the same transformer configuration as MLM experiments but without vocabulary factorization and weight tying. As a result, the AR-LM with our tokenizer uses more parameters in the embedding table and the softmax layer.*

| Tokenizer | FVD↓ | # Params | # Steps |
|---|:---:|:---:|:---:|
| MAGVIT | 265 | 306M | 1024 |
| *This paper* | **109** | 840M | 1280 |

Reviewer comment (Reproducibility):
There is no mention of the decoder in VQ-VAE and the VQ-VAE loss used when no lookup is applied.

Author response:
Initially, the decoder was only briefly mentioned in Section 2 and in the caption of Figure 2. We have improved the clarity by elaborating on the additional details on model design and hyperparameters in Appendix A.2 with an architecture diagram in Figure 7, as well as by providing a discussion of the loss functions in Section 3.1.

Reviewer comment (Theory):
More motivation is needed on why the authors choose to use a causal encoder for the tokenizer when doing masked LM for image / video generation.

Author response:
We have found that causal encoders facilitate frame prediction and result in improved generation FVD, as shown in Table R2.2. We have added this comparison into the revised Table 1 and Section 4.2. Additionally, this approach enables the representation of image and video within a unified token space, which could assist the community to explore joint image-video training.

Nevertheless, we agree with the reviewer that the causal design is also suitable for autoregressive language modeling. We have included a preliminary study of this in the Appendix and also in Table R 2.1. The result shows that our tokenizer also provides benefits for the autoregressive language model. Due to the reasons outlined in Q1, we chose to focus on using the masked LM as the backbone for generation.

*Table R2.2: Comparison of non-causal and causal tokenizer on the Kinetics-600 frame prediction task. A masked LM with 307M parameters is employed using 12 inference steps.*
 
| Causal | K600 FVD↓ |
|:---:|:---:|
| No | 11.6±0.6 |
| *Yes* | **5.2±0.2** |

Reviewer comment (Writing):
It is not clear why the authors tackle video generation if the aim was to understand the fundamentals of tokenization. It may be desirable for them to clearly motivate why they study videos and not images alone.

Author response:
Existing generation models, such as the stable diffusion model, have achieved remarkable results in image generation. Compared to image generation, video generation still faces substantial challenges in generating consistent and realistic motion. We are interested in exploring the capabilities of language models in tackling this unsolved challenge. Therefore, this paper focuses on the video tokenizer that can effectively model spatial-temporal dynamics for generative modeling. We have clarified this point in the beginning of Section 3 of the revised manuscript.

Reviewer comment (Evaluation):
It may be interesting for the reader to understand the computational complexity of both the tokenizer (encoder) and the detokenizer (decoder) and how they compare with video or audio codecs.

Author response:
Following the reviewer’s recommendation, we evaluate the computational complexity of our models and MAGVIT using MFLOPs (megaFLOPs) per pixel and frames per second, as shown in Table R2.3. At a similar complexity, our model achieves a higher compression ratio (lower bpp) while maintaining a similar quality (LPIPS) compared to MAGVIT [R2-A]. With a higher complexity, our model demonstrates significantly higher quality as measured by LPIPS and also human rater study shown in Figure 6.

We find it difficult to compare our method with the standard codecs because our model currently runs on accelerators with floating point operations while standard codecs are often implemented on CPU using integer operators or supported by hardware encoder/decoders. We leave this comparison to our future work.

We acknowledge that, judging from its network size and FLOPs, our current model may incur higher computational costs using TPUs compared to the common setups of standard codecs, e.g. HEVC Main Profile. We have added a short discussion in the revised Section 4.3. However, inspired by the impressive advancements in the acceleration of neural networks, we are cautiously optimistic about future research in this direction that could significantly improve the efficiency of the tokenizer.

*Table R2.3: Computational complexity of tokenizers at 360p resolution (480x360). Frames per second is measured with a single TPUv4 chip.*

| Model | bpp | LPIPS↓ | MFLOPs per pixel↓ || Frames per second↑ ||
|---|:---:|:---:|:---:|:---:|:---:|:---:|
|  |  |  | Tokenizer | Detokenizer | Tokenizer | Detokenizer |
| MAGVIT | 0.0384 | 0.1440 | 5.6 | 10.8 | 133 | 50 |
| *This paper* | **0.0147** | 0.1442 | 6.1 | 8.0 | 131 | 67 |
| *This paper* | 0.0532 | **0.0823** | 15.9 | 21.3 | 54 | 29 |

Reviewer comment (Experiments):
In the video recognition setup, it seems unnecessary to detokenize the visual tokens back to pixels since BEVT and BEIT can work with tokenized input. I understand one of the main reasons is that the underlying recognition model is the ViViT which takes raw pixels as input (as stated in the draft). However, you may also have a comparison with BEVT.

Author response:
We appreciate your understanding that ViViT `takes raw pixels as input`. The detokenization experiments used the **frozen** ViViT model that was trained on RGB pixels. For proper inference, the detokenizer is needed to convert the tokens back into pixel inputs, aligning with the input requirements of the frozen ViViT model. As shown in Table 4, our model shows superior performance compared to MAGVIT [R1-B] and is very close to raw pixels while only using ~1/500 data bandwidth. This setting was mentioned in Appendix A.3 of the original submission. And we have revised Section 4.4 to clarify this.

Following the reviewer’s suggestion, we also train a BEVT model with tokens as both the input and the output for action recognition. The result is shown in Table R1.1. We find that directly using token inputs yields suboptimal performance in part because the current architecture and pre-training setups are tailored for pixel inputs. Further research is needed to investigate this token-in-token-out pretraining paradigm for action recognition, where both the input and output are based on discrete tokens.

*Table R1.1: Video action recognition performance (classification accuracy↑ $\times$100) on SSv2 dataset.*

| Tokenizer | Input+Output | Output only |
|---|:---:|:---:|
| MAGVIT | 60.86 | 67.22 |
| *This paper* | **60.99** | **67.38** |

> > [R1-B] [MAGVIT: Masked Generative Video Transformer. In CVPR 2023.](http://openaccess.thecvf.com/content/CVPR2023/html/Yu_MAGVIT_Masked_Generative_Video_Transformer_CVPR_2023_paper.html)