Masked Generative Nested Transformers with Decode Time Scaling

We thank the reviewer for their valuable time and constructive reviews. We are glad to see that the reviewer appreciates the effectiveness of decode-time model scaling and caching, to significantly reduce computational costs in visual generation while maintaining competitive performance and demonstrating practical implications across image and video tasks.

Training complexities in nested models

As discussed in the related works section, several methods in the literature have successfully demonstrated the use of nested architectures for different tasks like language modeling (MatFormer, Kudugunta et al., 2023 and Flextron, Cai et al., 2024b), discriminative tasks like image/video classification (MoNE, Jain et al., 2024) as well as in representation learning (Matryoshka Learning, Kusupati et al., 2022) and some others discussed in Line 144 onwards (left column). These works show that nestedness does not introduce any training complexities. However, when we keep adding a lot of nested models, that might constrain the smaller models, making their performance inferior. To avoid that, in this work we propose a curriculum based distillation (Line 266 right column), which helps to improve the performance of the smaller models as shown in Table 7 of supplementary material. To further assess the training complexity, we ablate the effect of the number of nested models as suggested by Reviewer fe9H. Please refer to our response on “Impact of the number of nested models” to Reviewer fe9H for more details.

Quality of generated samples

Our work primarily focuses on enhancing efficiency of a baseline model, rather than directly improving its generation quality. Having said that, the compute-performance trade-off curve in Fig 6 shows that given a certain compute budget, our method can reach a better FID than the baseline.

In Table 1, we report results for a fixed model schedule of (3, 3, 3, 3), i.e. 3 iterations of each nested model. It takes a small hit in performance (0.6 FID) compared to the baseline, but being 2.65x inference efficient. To recover the small difference in performance we can change the model schedule to use bigger nested models like (0, 0, 6, 6), i.e., using 6 iterations of only the two largest nested models. With this, we achieve an at par quality with baseline (FID of 2.6) with 745 GFLOPs. This is 1.7× compute efficiency. Fig 6 also shows that the compute-performance tradeoff is even better for a smaller GFLOPs budget. | Method | Schedule | FID | # params | # GFLOPs | |---|---|---|---|---| | MaskGIT | NA | 6.2 | 303M | 647 | | MaskGIT++ | NA | 2.5 | 303M | 1.3k | | MagNeTS (ours) | (3, 3, 3, 3) | 3.1 | 303M | 490 | | MagNeTS (ours) | (0, 0, 6, 6) | 2.6 | 303M | 745 |

We thank the reviewer for their valuable time and constructive reviews. We are glad to see that the reviewer found this to be an efficient approach for image and video generation, experimentally demonstrating reduced computational costs while maintaining generation quality and offering an unexplored perspective compared to existing methods. We answer the reviewer's question below.

Inference Time

We would like to highlight that in Table 8 in our Supplementary Material we present practical efficiency like throughput (in img/sec) of our method as compared to the corresponding baseline. For the sake of completeness, we also present the latency below. As we can see the proposed method is 2.5x faster than the baseline in this setting. | Algorithm | Baseline (MaskGIT++) | MaGNeTS | |---|---|---| | Images/Sec | 22.5 | 56.3 | | Latency (ms) | 712 | 285 |

Generalization to other generative modeling frameworks

We would like to highlight that the idea of model size scheduling over decoding iterations is generic enough to be applied to other multistep processes like diffusion. The core idea is that some parts of the decoding/denoising process might be easier than others, hence allowing for a step-wise allocation of model capacity. While we explore fixed schedules in this work, the idea can be further extended to input-adaptive schedules, i.e., some images might be easier to generate than others and based on the input we can decide which model to use for a certain step.

To support this broader applicability of our algorithm, we conducted preliminary experiments on diffusion models. Following the above discussion, we did some experiments using model schedules for diffusion. Due to time constraints, we were not able to train a new diffusion model in nested fashion, rather use UViT’s [A] pretrained checkpoints [B] on ImageNet 64×64. We use two models - U-ViT-L/4 and U-ViT-M/4 - to demonstrate our idea on model schedule during inference. Some key details of experiment -

• We use the default number of sampling steps = 50 and batch size = 500 in all experiments. We use a single A100 GPU.
• We do not use classifier-free guidance. We do not use any caching for these experiments (due to the continuous nature of the input) and only demonstrate the generalizability of the model scheduling idea.
• Since the initial denoising steps play a crucial role in shaping the final output of the reverse diffusion process, we utilize the L model for these early stages and transition to the M model for the later denoising steps.
• Given that the L model has greater denoising capacity than the M model, we customize the noise schedule with larger denoising step sizes for L and smaller step sizes for M, balancing efficiency and performance. | Method | FID (50k) | # steps | time (sec/iter) | |---|---|---|---| | U-ViT-M/4 | 5.92 | 50 | 17.12 | | U-ViT-L/4 | 4.21 | 50 | 32.34 | | Ours (Model Sched) | 4.58 | 50 | 21.10 |

As we can see, with only model scheduling, we are able to achieve ~1.53x inference compute gains with almost similar performance as baseline. Exploring better schedules, training the models with nesting and distillation will offer better compute gains. This shows that the proposed method of model scheduling over multistep decode process in image/video generation is generic enough to be applied to different modeling approaches

Typos

Thank you for bringing these typos to our knowledge. We will fix them in our final revision.

References

[A] F. Bao, S. Nie, K. Xue, Y. Cao, C. Li, H. Su, and J. Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22669–22679, 2023.

[B] https://github.com/baofff/U-ViT

We thank the reviewer for their valuable time and constructive reviews. We are glad to see that the reviewer acknowledged the work to be a promising and efficient approach for image/video generation by introducing model size scheduling and demonstrating the effectiveness of KV caching and nested models, supported by strong experimental results and a well-designed ablation study. We answer the reviewer's question below.

Impact of the number of nested models

Thanks for the interesting question. We analyse the effect of the number of nested models on performance. We train four more setting p=[1, 2] (two models), p=[1, 2, 4] (three models), p=[1, 2, 4, 8, 16] (five models) and p=[1, 2, 4, 8, 16, 32] (six models) apart from p=[1, 2, 4, 8] (four models) which we have in the paper. We observe that for all of these models, the biggest model performance remains the same for all cases. However, the performance of the smaller models, let's say $\frac{d}{8}$, degrades by 0.5 FID when we add $\frac{d}{16}$ and then further by 0.6 FID when we add the $\frac{d}{32}$ nested model. We hypothesize that the drop in performance of smaller models is due to its lower representational power. As we add more nested models, the complexity of the shared representation increases, and burdens the smaller model. This drop in performance does not impact the performance of model scheduling (MaGNeTS, FID=3.1), as the larger models dominate the final results. Note that all of these results are on top of models trained with distillation (Line 261 onwards, right column), which itself helps to retain the performance up to 4 distilled models. This can be seen from Table 7 of supplementary material, which shows that distillation helps to boost the performance of smaller nested models.

Qualitative Analysis

We would like to clarify that we have included additional qualitative results in the Supplementary Material (Fig 10 for image generation, Fig 11 for video generation) along with the main qualitative results in Fig 1. We also present some failure cases in Fig 12. We will add more results to the final supplementary material.

We thank the reviewer for their valuable time and constructive reviews. We are happy to hear that the reviewer acknowledged the novelty of the dynamic compute allocation approach in our work, which helps to reduce redundant computation, and experimental evidence supporting the claims. We answer the reviewer's question below.

Claim about KV Caching and Refresh for Inference Efficiency

We would like to clarify that the claim in our paper is not about KV caching alone improving efficiency without performance loss. Instead, it is the combination of KV caching with intermittent refresh that makes our approach inference-efficient without degrading performance, as discussed in the paper. All our efficiency claims explicitly include the compute required for cache refresh. For e.g.,

• In Line 109 (left column), we state “KV caching can also be used in parallel decoding, which can effectively reuse computation when refreshed appropriately”.
• In Line 240 (right column) we mention “Caching the key-value pairs for the unmasked tokens helps reduce computation, but it can slightly degrade performance”.
• Then in Line 251 (right column) we mention “To remedy this, we strategically refresh the cache while changing the model size.”

Trade-off analysis

It is interesting to note that compute efficiency is better represented as a compute-performance tradeoff curve as shown in Fig 6. This curve illustrates the relationship between compute/latency and model quality:

• Given a fixed compute or latency budget, it shows that the proposed method obtains the best quality.
• Conversely, given a desired quality requirement, it shows that the proposed method is the most inference-efficient one. This tradeoff analysis effectively captures the true essence of compute efficiency.

Comparison to Efficient Diffusion Models

As discussed in the related works section of the paper, efficient diffusion methods can be categorized around two main categories – (1) reducing the number of network calls, (2) designing better network architectures to reduce the computation of each call. Our method is complementary to both of these approaches and can be combined to further enhance efficiency. Having said that, we do compare with a bunch of efficient diffusion methods in Table 1 of the paper. Moreover, as mentioned in Line 356, several recent diffusion works (example - Feng et al 2024, Lee et al 2024, Meng et al 2023, Berthelot et al 2023, Song et al 2023, Zheng et al 2023) only report results on the low-resolution of ImageNet (64×64), and therefore a direct comparison is not possible, as all our experiments are on 256x256. In addition to the ones presented in Table 1 of the paper, below we add some more comparisons with efficient diffusion methods which report for image size 128 and 256 on ImageNet. As we can observe, while some diffusion models do perform well, it needs considerably more steps and hence FLOPs compared to MaGNeTS. We will add these comparisons to the main paper.

| Method | Image Size | FID | Params | Steps | GFLOPs | |---|---|---|---|---|---| | DPM-Solver (Lu et al 2022) | 128 | 4.1 | 422M | 12 | >3000 | | MagNeTS (Ours) | 128 | 3.9 | 303M | 12 | 236 | |---|---|---|---|---|---| | EDiff [A] | 256 | 2.1 | 450M | 50 | 119k | | LPDM-ADM [B] | 256 | 2.7 | - | 50 | - | | MagNeTS (Ours) | 256 | 3.1 | 303M | 12 | 490 |

Artifacts

We would like to clarify that our approach does not introduce new artifacts, but inherits these properties from the baseline (MaskGIT) on which we apply our model scheduling approach. We discuss these artifacts in Section E of Supplementary Material: Lines 757-761 (right column). These artifacts are also visualized in Figure 12, which show that failure cases (like faces of humans) in the baseline method (MaskGIT++) directly translate to failure cases in our method. However, we reemphasize that improving this aspect of generative modeling is orthogonal and beyond the scope of the current work.

References

[A] Hang, T., Gu, S., Li, C., Bao, J., Chen, D., Hu, H., ... & Guo, B. (2023). Efficient diffusion training via min-snr weighting strategy. CVPR, https://arxiv.org/pdf/2303.09556

[B] Wang, Z., Jiang, Y., Zheng, H., Wang, P., He, P., Wang, Z., ... & Zhou, M. (2023). Patch diffusion: Faster and more data-efficient training of diffusion models. NeurIPS, https://arxiv.org/abs/2304.12526