Masked Generative Nested Transformers with Decode Time Scaling
Authors: Sahil Goyal
Debapriya Tula
Gagan Jain
Pradeep Shenoy
Prateek Jain
Sujoy Paul
Venue: ICML 2025 poster (2025)
Keywords: generation
efficiency
nested models
kv caching
model scaling
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: Efficient visual generation using a decode time model scaling approach and cached computations, achieving similar performance with significantly reduced compute.

Abstract: Recent advances in visual generation have made significant strides in producing content of exceptional quality. However, most methods suffer from a fundamental problem - a bottleneck of inference computational efficiency. Most of these algorithms involve multiple passes over a transformer model to generate tokens or denoise inputs. However, the model size is kept consistent throughout all iterations, which makes it computationally expensive. In this work, we aim to address this issue primarily through two key ideas - (a) not all parts of the generation process need equal compute, and we design a decode time model scaling schedule to utilize compute effectively, and (b) we can cache and reuse some of the intermediate computation. Combining these two ideas leads to using smaller models to process more tokens while large models process fewer tokens. These different-sized models do not increase the parameter size, as they share parameters. We rigorously experiment with ImageNet256$\times$256 , UCF101, and Kinetics600 to showcase the efficacy of the proposed method for image/video generation and frame prediction. Our experiments show that with almost $3\times$ less compute than baseline, our model obtains competitive performance.

Masked Generative Nested Transformers with Decode Time Scaling
Sahil Goyal
Debapriya Tula 1
Gagan Jain'
Pradeep Shenoy'
Prateek Jain' Sujoy Paul
Abstract
2022: Van Den Oord al., 2017), and so on. In recent years, diffusion models and modeling visual tokens as lan- guage have been the de-facto processes used to generate high-quality images. While initially proposed with a CNN or U-Net based architectures (Rombach et al., 2022; Saharia et al., 2022), transformer models have become the norm for these methods (Peebles & Xie, 2023; Yu et al., 2023a).
Recent advances in visual generation have made significant strides in producing content of excep- tional quality. However, most methods suffer from fundamental problem a bottleneck of inference computational efficiency. Most of these algorithms involve multiple passes over a trans- former model to generate tokens or denoise inputs. However, the model size is kept consistent for all iterations, making computationally expensive. In this work, we aim to address this issue primar- ily through two key ideas (a) not all parts of the generation process need equal compute, hence we design a decode time model scaling schedule to utilize compute effectively, and (b) we can cache and reuse some of the intermediate computation. Combining these two ideas leads using smaller models to process more tokens while large mod- els process fewer tokens. These different-sized models do not increase the parameter size, as they share parameters. We rigorously experiment with ImageNet256x255 256, ImageNet128x 128, UCF101, and Kinetics600 showcase the efficacy of the proposed method for image/video generation and frame prediction. Our experiments show that with almost 3x less compute than baseline, our model obtains competitive performance.
The recent advancements in visual generation can be cat- egorized along two axes (a) different types of denoising processes in the continuous latent space (Ho et al., 2020;
Nichol & Dhariwal, 2021b), discrete space (Gu et al., 2022;
Lou et al.) or masking in the discrete space (Yu et al., 2023a; Chang et al., 2022), continuous space (Li et al., 2024a) (b) modeling tokens either auto-regressively (Kondratyuk et al., 2024; Esser et al., 2021; Yu et al., 2021) with causal atten- tion or parallel decoding with bi-directional attention (Gu et al., 2022: Yu al., 2023a; Chang et al., 2022; Zheng et al., 2022). To achieve high synthesis fidelity, both, denoising in diffusion models, and raster scan based auto-regressive token modeling require several iterations.
Recently, parallel decoding of discrete tokens have shown promise in generating high quality images with few itera- tions MaskGIT (Chang et al., 2022), MAGVIT (Yu et al., 2023a), MUSE (Chang et al., 2023), MaskBIT (Weber et al., 2024), TiTok (Yu et al., 2024b). These models are trained with Masked Language Modeling (MLM) type losses, and the generation process involves unmasking a few confident tokens every decoding iteration, starting from all masked tokens. They can even surpass diffusion models, given a good visual tokenizer (Yu et al., 2023b; Weber et al., 2024).
1. Introduction
The last decade has witnessed tremendous progress in image and video generation, under diverse paradigms- generative adversarial networks (Brock, 2018; Sauer et al., 2022), de- noising processes such as diffusion models (Ho et al., 2020:
Although MaskGIT reduces decode complexity signifi- cantly, parallel decoding still includes several redundant computations. First, the need for same capacity model for all steps needs to be investigated. Second, unlike auto- regressive models, which eache its computation in all steps, parallel decoding performs re-computation for all tokens. We empirically find that smaller model can generate good- quality images quite fast, but its performance saturates after point with more decoding iterations. A bigger model can perform finer refinement and generate better-quality images.
2022b; Dhariwal & Nichol, 2021; Rombach et al., 2022: Gu et al., 2022), image generation via vector quantized tok- enization (Razavi al., 2019; Esser et al., 2021; Ge et al.,
"qqual contribution, work done while at Google DeepMind 'Google DeepMind University of California, Los Angeles. Cor- respondence to: Sahil Goyal <ggyalsahil@e caanhhil@ooogl..ooo google.com>, Sujoy Paul <sujoyp@ google.com>.
Motivated by these observations, we present Masked Generate Nested Transformers with Decode Time Scaling (MaGNeTS). We design a model size curriculum over the
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

MaGNeTS
MaskGIT++ (FID=2.3)
-3x FLOPs Reduction
MaGNeTS (FID=2.9)
Figure 1: Class-conditional image generation on ImageNe1256xx266. 256. Comparing MaskGIT++ and MaGNeTS (size: L).
decoding process, which efficiently utilize compute. MaG- NeTS gradually scales the model size up to the full model the decoding iterations instead of using single large model throughout. Operating discrete tokens, we cache key-value pairs of unmasked tokens and reuse them in later iterations. A combined effect of these two techniques leads to processing more tokens with smaller and fewer tokens with larger models. The heterogenous sized models share parameters as MatFormer (Kudugunta al., 2023). We build MaGNeTS on top of MaskGIT. We find that MaskGIT be drastically improved using classifier-fie free guidance, specifically when trained with it. We call this MaskGIT++ and use this as the improved baseline, presenting all results on top of it. We also conduct preliminary inference-time time ex- periments on diffusion models to demonstrate our method's ss generalization capabilities.
ct al., 2012) and Kinetics600 (Carreira et al., 2018). To summarize, the main contributions of this work are:
We introduce the concept of model size scheduling during the generation process to significantly reduce compute requirements.
We show that like auto-regressive models, KV-caching can also be used in parallel decoding, which can effec- tively reuse computation when refreshed appropriately.
We introduce nested modeling in image/video genera- tion to exploit the above ideas effectively.
Extensive experiments show that MaGNeTS offers 2.5 3.7x compute gains across tasks.
2. Related Work
On ImageNet, with 3x less compute, MaGNeTS generates images of similar quality as MaskGIT+ (see Figure 1). It is also comparable to state-of-the-art methods, which need orders of magnitude more compute. We also show MaG- NeTS's efficacy on video datasets like UCF101 (Soomro
Efficient Visual Generation. Image generation literature has seen significant improvements in the past years gener- ative adversarial networks (Brock, 2018; Sauer et al., 2022), discrete token based models (Chang et al., 2022; Yu et al., 2023a), diffusion-based models (Kingma & Gao, 2023:
2

---

MaGNeTS
Hoogeboom et al., 2023), and more recently hybrid models (Peebles & Xie, 2023; Yu et al., 2024c), but they often guz- zle computing power. Researchers tackle this bottleneck of computational costs with efficient model architectures and smarter sampling strategies.
Learning (Kusupati al., 2022) introduces the concept of nested structures into embedding dimensions, making them elastic. MatFormer (Kudugunta et al., 2023) applies the same concept to the MLP hidden layer in each transformer block, enabling extraction of multiple sized models. MoNE (Jain et al., 2024) and Flextron (Cai et al., 2024) learn to route tokens to variable sized nested models leading to com- pute efficient processing. Stochastic Bottleneck (Koike- Akino & Wang, 2020), MQT (Hu et al., 2024), One-D-Piece (Miwa et al., 2025), FlexTok (Bachmann et al., 2025) and Semanticist (Wen et al., 2025) explore nesting across the latent sequence dimension. In this work, we demonstrate how different stages of a multi-step task like image/video generation, can be efficiently handled by nested models in- stead of relying on the full model at every step, significantly reducing computation without sacrificing quality.
In diffusion model literature, there have been some work to reduce the number of sampling steps, by treating the sampling process like ordinary differential equations (Song et al., 2022; Lu et al., 2022; Liu et al., 2022), incorporating additional training process (Kong& & Ping. 2021; Nichol & Dhariwal, 2021a; Salimans & Ho, 2022; Song et al., 2023), sampling step distillation (Salimans& & Ho, 2022;
Song et al., 2023; Berthelot et al., 2023; Meng et al., 2023; Feng et al., 2024), sampling and training formulation mod- ifications (Esser al., 2024; Song al., 2023), and more. Recently, there has been growing interest in understanding how each step in the diffusion sampling process contributes (Choi et al., 2022; Park et al., 2023; Lee et al., 2024). These approaches analyze sampling steps leveraging distance met- rics such as LPIPS, fourier and spectral density analysis. Building on these explorations researchers have designed methods based on optimal sampling steps (Watson et al., 2022; Lee al., 2024), weighted training loss (Choi al., 2022), and step-specific models (Li et al., 2023; Yang et al., 2024; Lee et al., 2023). These step-specific models use computationally expensive evolutionary search algorithms, directly optimizing the quality metrie, FID (Heusel et al., 2017). Concurrently, researchers are actively addressing the inherent architectural costs of diffusion models, particularly those associated with transformer attention mechanisms (Yuan et al., 2024; Yan et al., 2024).
3. Preliminaries
Parallel Decoding for Image Generation. Masked Gener- ative Image Transformer (MaskGIT) (Chang et al., 2022) introduces a novel approach to image generation that sig- nificantly differs from traditional autoregressive models. In autoregressive decoding, images are generated sequentially, one pixel/token at time, following a raster scan order (Esser et al., 2021; Kondratyuk et al., 2024; Wang et al., 2024; Yu et al., 2024a; Li et al., 2024b). This sequential approach is computationally inefficient, as each token is conditioned only on the previously generated tokens, lead- ing to a bottleneck in processing time. MaskGIT generates all tokens of an image simultaneously, while iteratively re- fining them. This method enables significant acceleration in the decoding process. The tokens are discrete and ob- tained using Vector Quantized (VQ) autoencoders, Iearned with self-reconstruction and photo-realism losses (Yu et al., 2023a). The iterative parallel decoding process is repre- sented as:
On the other hand, certain works focus on building efficient and improved tokenizers. IDDM (Rombach et al., 2022) takes diffusion models from pixel to compressed latent space for efficient and scalable generation. FSQ (Mentzer et al., 2023), LFQ (Yu et al., 2023b) and BSQ (Zhao et al., 2024) explore certain vector quantization techniques in the discrete tok- enization process without explicitly learning the codebooks. VAR (Tian et al., 2024) explores multi-scale tokenizer to improve the generation quality. Recently TiTok (Yu et al., 2023b), FlowMo (Sargent et al., 2025) relax the topology of latent space and reduce the grid dimensionality of latents from 2D to ID. Further FleXTok (Bachmann et al., 2025) enables adaptive sequence length in ID tokenization for efficient generation. Instead of sampling tokenization process optimization, we tackle an orthogonal problem of efficient compute allocation over the multi-step generation process. This makes our approach usable with a variety of tokenizers, model architectures and sampling schemes.
X Mask Sample (M(X)=1,c),k)
(1)
where X E zzy, are the input tokens, N is the number of tokens, & E [1, K] denote the iteration number, with K be- ing the total number of iterations, X, is either completely masked for full generation, and partially masked for condi- tional generation tasks like frame prediction, is the cate- gory of image/video under generation. The Sample function utilizes logits predicted by the model (-). introduces cer- tain randomness, and sorts them by confidence, unmasking only top-k tokens while masking the rest. We follow this process as in (Chang et al., 2022; Yu et al., 2023a).
Nested Models. The core of our algorithm for inference- efficient decoding relies on variable-sized nested mod- els for efficient parameter-sharing. We use MatFormer's 's (Kudugunta et al., 2023) modeling approach to extract mul-
Nested Models. Rippel et al. (2014) introduced nested dropout to learn ordered representations of data that improve retrieval speed and adaptive data compression. Matryoshka
3

---

MaGNeTS
Generated Imege
Empty Cache
17g
KV Cache
7P2
All Masied Tokens
UnMasked
Masked
Takons
Takens
Figure 2: MaGNeTS Decoding. We start from the smallest nested model with an empty cache and gradually move to bigger models over the decoding iterations. We iterate using particular sized model for few iterations, beforc moving the next model size. As we cache the key-value pairs the unmasked tokens, the KV cache size also increases over time. We also refresh cache when we switch models, hence its dimension also increases over decoding iterations.
Iter
0.01
0.01
0.02
0.03
0.04
0.05 0.05
0.07
0.06
0.09
0.08
0.11
0.120.10
0.12 0.11 0.120.10
0.130.08
0.15
0.15
Figure 3: Unmasked Token Density visualization in each decoding iteration averaged over 50k generated samples on ImageNet. Yellow represents higher density. Each pixel represents a token from 16 x 16 latent token space. (See Appendix A for category-wise token density).
tiple nested models, from a single model, without increasing the total parameter count. Given a full transformer model M, MatFormer defines nested models (m1..... mc), such that mi m....... mc M. Each m. has fewer parameters and reduced compute. The core idea of extracting nested models is that in transformer block, reduced computation using parameter subspace is performed via a sliced matrix multiplication. Assuming a parameter matrix W E R xd and feature vector x IRd, then the computation y Wx is partially obtained by computing V1:51 W o
same strategy is applied to the MLP layer. This process gives us models with close to linear reduction in parameter count and inference compute with the downscaling factor p.
4. Method
Given the preliminaries, here we introduce the core algo- rithm. We first discuss the idea of scheduling models of different sizes over decode iters of MaskGIT. Then, we dis- cuss the process of caching key-value parallel decoding, followed by how to refresh them to improve performance. We finally discuss the nested model training method. A pictorial overview of our method is presented in Figure 2.
if y is desired to be partial and y W T [ if input is par- tial. With such partial computations throughout the network, we can obtain nested models which share parameters.
While MatFormer (Kudugunta et al., 2023) obtained sub- models with partial computation only in the MLP layer, also do it in the Self-Attention layer, specifically in obtaining the Q K. V features. These features are of dimension x du where n, is the number of attention heads, d, is the p. feature dimension, and p is the model downscaling factor. We choose four downscaled models C 4, with pe E {1, 2, 8) in this work. After attention computation, this gives features that also p times downscaled. Then, it is projected back to the full model dimension d using partial computation, as the input features are partial. The
Decode Time Model Schedule. In iterative parallel decod- ing (Chang et al., 2022: Yu et al., 2023a), the same-sized model is used for all steps, starting with all tokens being masked. However, we hypothesize that certain stages the generation process might be easier than others. For example, in the initial steps, the model only needs capture coarse global structures, which can be achieved efficiently using smaller models. In the later steps, the model must refine finer details, which requires larger models. This hypothesis is bolstered with Figure 3, which shows that the generation
process starts unmasking tokens from the background and shifts the middle of the image the later iterations (more

---

MaGNeTS
categorical examples in Appendix Figure 8).
enhances efficiency. In auto-regressive models, caching progressively happens in one fixed direction. However, in parallel decoding, caching must depend on which tokens are unmasked over the iterations.
Deeadng Bessn
Concretely, starting from an empty cached set, we keep adding keys and values to the set for the tokens that are unmasked after the Masko steps (see Section 3). We do not update the predicted token indices for these unmasked tokens in the subsequent iterations. Hence, the cached key and values for the unmasked tokens are the only features the other masked tokens need. In every decoding iteration, we can categorize tokens into three main categories: unmasked tokens (for which we have cached KV), masked tokens that will be unmasked during the current iteration, and the rest of the masked tokens. Note that the KV cache for the second category tokens cannot be used in the next iteration but only in the iteration after that once we know their token indices in the current forward pass. We cache them in the next iteration for use in the immediately next iteration.
Figure 4: Nested Models at difffrent decoding iterations. Dif- ferent values of the downscaling factor correspond to the nested models. The diameter of the blobs indicates #iterations.
Our hypothesis is further motivated by Figure 4, which presents the generation quality (FID) over iterations of paral- lel decoding for different-sized models. The smallest model reaches a reasonably good FID score with very low FLOPs compared to the biggest model. However, it saturates after point, and the larger models surpasses the smaller ones in performance, demonstrating their ability to capture finer details and generate higher-quality images when provided with sufficient compute. This trend suggests that dynami- cally scaling the model size during decoding can exploit the varying task difficulty and achieve compute efficiency.
Caching is even more useful for decode time model sched- ules. For a schedule that progressively scales up the model size as decoding progresses, smaller models process more to- kens, while the larger models process fewer tokens, leading to an efficient yet good quality image generation process.
Algorithm MaGNeTS Decoding Algorithm
We use nested models to extract multiple models rather than using models with disjoint parameters. Nested models do not increase the parameter count and it also helps in better alignment of hypothesis when we shift model size over decode steps. The decode time model schedules can be generalized and represented as making the model choice in Equation (1) dependent on the iteration index as follows:
Input: Xo (Initial Tokens), K (8sepp)) N (Atokens), M (Nested Model Schodhle), e (class)) Instalize: U; cache (): Nore: No iica
token (Mask token id -1
while sep l e K 'dd
If k U and My + MM then Clear euche.
end
Get uncached tokens: X [dyg |s c Xxx & cuchee)
Compute prediction probubiiities and key-anlses p² (ke) - MMx cache)
Sampli tokens using current peodictions p², withor modifying presious predictions
X + Mask o Sample(MM((xx11..))
k MukkCTT-Smmplepp"
New indives to cache: C - iil cache, xf+1_- 1-1) Updase the ky cache: cache - cache (kv)) &
M ((mp))) (mpp)) (mpm) s.t. Σ k; K (2)
end return X
where PI, denote the downscaling factors of the corresponding nested models, and (m)* denotes that model m will be executed for & iterations. K represents the total number of iterations. We can think of different model sched- ules- (a) downscaling (starting with the full model and then gradually moving to the smallest model), (b) upscaling, (c) intermittently switching among a few models, and so on. We can also modify the integers k; to choose the number of times we stick to a model before switching. However, as intuitively discussed before, we empirically validate that gradually upscaling the model size gets the best trade-off between the compute and generation quality.
Intermittent Cache Refresh. Caching the key-value pairs for the unmasked tokens helps reduce computation, but it can slightly degrade performance. This happens because- (a) when we cache, the unmasked tokens are not updated in the subsequent iterations. (b) when we shift model size during generation, the attention layer, the query size dif- fers from the cached KV (see Section 3). While technically, we can zero-pad the KV to be compatible with the current model's 'ss query dimension, the model remains unfamiliar of such feature discrepancies between query and key-value.
Cached Parallel Decoding, Inspired by caching key-value pairs in ruto-regressive models, we explore caching in par- allel decoding, which retains relevant computations and
To remedy this, we strategically refresh the cache while changing the model size. Refreshing involves discarding the cached KV for that iteration and caching a newly com-
5

---

MaGNeTS
puted KV for the immediate next iteration. We empirically find that it bridges the performance gap that arises due to caching, The proposed decode time model scaling algo- rithm is presented in Algorithm 1, which uses MaskGIT's sampling strategy (Chang et al., 2022; Yu et al., 2023a) to sample tokens from logits predicted by the network.
Datasets. We evaluate our model on ImageNet256x2566 and
ImageNer128x 128 (Deng al., 2009) for image generation, UCFI01 (Soomro et al., 2012) for video generation and Kinetics600 (Carreira et al., 2018) for frame prediction (5- frame condition).
Implementation Details. We utilize the pretrained tokeniz- ers from MaskGIT (Chang et al., 2022) (for images) and MAGVIT (Yu et al., 2023a) (for vidcos) with the codebook size of 1024 tokens. We train different models for image sizes 256 256 and 128 128. Respective tokenizers com- press them to 16 16 discrete tokens. For videos, learn models for 16 128 x 128, where the tokenizer outputs 4 x 16 16 tokens. Following MaskGIT, we utilize the Bert model (Devlin et al., 2019) as a transformer backbone. We perform experiments at several model scales to understand the scaling behaviors of our algorithm. We utilize the same training
Training Nested Models. MatFormer (Kudugunta et al., 2023) opts for joint optimization of losses w.r.t. ground- truth from all models with equal weights. While this mode of training works for a small range of model downscaling, we found it to hurt performance with larger downscaling factors p. We introduce a combination of ground truth and distillation loss to address this issue. We perform online distillation progressively, where the teacher for model m, is model m4+1- The full model my((= M) is trained with only ground truth loss. This provides simpler optimiza- tion for the smaller nested models while maintaining the overall objective. Progressive distillation also reduces the teacher-student size gap, which can otherwise hurt distilla- tion performance (Stanton et al., 2021; Beyer et al., 2022: Mirzadeh et al., 2019). Given input X, ground truth label Y and loss function L, our training loss is expressed as:
to train our nested models as these baselines. We train our model for 270 epochs for all the experiments. Unless otherwise mentioned, throughout the paper, we employ same number of steps per model before switching to the next model, i.c., h = k2 = = kn, We fol- low a cosine schedule of unmasking tokens during inference. For image generation and frame prediction, we use classifier- free guidance for both MaGNeTS and respective baselines. We drop input class condition labels for 10% of the training batches in image generation to better facilitate classifier free guidance during image generation. We mention the details of sampling hyperparameters in Appendix B.
Lirain- L (C(mx(()) Y)+ +
N-1 () + (1-())))))))((( (x)
Evaluation Metrics. Following previous baselines, we use Fréchet Inception Distance (FID) (Heusel et al., 2017; Dhari- wal & Nichol, 2021) for image generation, Fréchet Video Distance (FVD) (Unterthiner et al., 2019) for the video gen- eration tasks, Inception Score (Salimans et al., 2016) for both tasks, as well as precision and recall for image genera- tion. We compare algorithms using inference-time GFLOPs. Refer Appendix D for GFLOPs computation details.
where a, controls the weight between the distillation and ground truth loss, which linearly decayed from 1 to O as training progresses. Note that a stop gradient is applied during distillation on m4.1 in the third term of the equation.
Classifier- Free Guidance. Following literature (Ho & Sal- imans, 2022: Yu et al., 2023b), we also utilize classificr-fres guidance during the generation process. Following the same motivation as decode time model scaling discussed above, which shows that the initial decoding iterations focus the background region, and gradually moves to the main object/region of interest in the final decoding iterations, we apply guidance to only a few final decoding iterations. We find that doing this offers similar quality images as applying guidance to all iters (refer Figure 9b). See Appendix for detailed analysis.
5.1. Image Generation
Comparison with Baselines. In this section, compare MaGNeTS with state-of-the-art methods in the literature for image generation. We list the results for 256 x 256 and 128 x 128 image generation on ImageNet- lk in Table I and Table 2 respectively. Table 1 shows that MaGNeTS can speed up the generation process by 2.65-33 3x (depending on total step count) compard MaskGIT++. with a negligible drop in FID. Refer Appendix D for real-time gains. Figure 5 illustrates that MaGNeTS significantly accelerates parallel decoding, which gets more pronounced as image resolu- tion grows. Figure I and Figure 10 show generated images from MaskGIT++ and MaGNeTS (ours). As shown in re- cent literature, using a superior tokenizer (Yu et al., 2023b; Weber et al., 2024) or optimized training/inference configu-
5. Experiments and Results
We conduct extensive experiments to demonstrate the ef- ficacy of our approach on three distinct tasks: class- conditional image generation, class-cona video gen- eration, and frame prediction.
6

---

MaGNeTS
Model
15
Bec
pacums
steps
Goops
digGANddopp (Brock,
87
StllaiAN-XI (Sneer 2022)
2.3
265.1
166M
Ingroved DDFM? (Nichol Dhatwal, 2021b)
12.3
70
62
280M
250
>130K
ADM Opeamlee (Dtariwal & Nichol, 2021)
3.9
215.8
83
53
55-M
250
371k
LDM-4-9 (Rombach et 2022)
3.6
247.7
400M
250
51.5k
D/ (Prahlees
2003)
278
81
MDF et 2023)
283.0
MaskDiT's> (Zheng et al.. 2023)
2.3
275.6
80
736M
250
>28k
CDM al. 2022))
4,9
158.7
8100
RIN (Jabri al. 2022)
3,4
182.
410M
1000
3341
Simple Diflusio (Bongchoom 2023)
2.4
VDM-+ (Kingmo
2023)
2.1
2B
Figure 5: Compute Comparison between uni- form model schedule (MaskGIT) and MaG- NeTS, for 12 decode iters.
EDDF (Hlag al., 2024)
2.1
450M
50
119k
ADM* (Wing et 2023)
2.7
50
7.8k
MAR eti 2024b)
1.8
296.0
81
479M
128
VOVAE-2 (Ruzavi
2019)
36
13,58
VQGAN" (Esee et 2021)
15.8
VQGAN (archilecture)++ MaskGfT (seup)
18.7
80.4
78
26
227M
256
34-
(Chag et al. 2022)
6.2
182.1
80
51
227M
647
32.
Scresuee
Mo- MO-VOGAN (7hing et al. 2022)
7.2
1301
72
55
389M
12
-Ik
Mask (Wisee at 2024)
10.3K
PAR-4x (Wing 20234)
PAR-10x" (Wing al., 2024)
2.9
262.5
82
3.18
MaskGETea
2,5
260.3
83
54
303M
12
1.3k
MaskGEI++44
2,5
290.6
84
51
303M
15
1.8k
MaGNeTT (ours)P
3.1
254.8
85
303M
490
(oumysh
Table I: Class-conditional Image Generation on ImageNet 256 256. "# steps" refers the number of neural network runs. denotes values taken from prior publications. indicates usage of extra training data. g denotes use of classifier- free guidance (Ho & Salimans, 2022) for all steps. J= represents use of guidance only for final x steps.
Figure 6: Compute Scaling Curve. Genera- tion performance vs compute for different model sizes. The blob size indicates parameter count.
rations (Ni et al., 2024a;b) can further boost MaGNeTS's performance.
Figure 6. shows the compute vs performance different models, with the blob size denoting the model size. For a certain parameter count, the baseline uses the full model for all 12 decoding steps, while the scheduled routines use sequence of nested models with downsampling factors 8, 4, 2, 1 for 3 steps each. Scaling up model size lead to much cheaper compute scaling of MaGNeTS than the baseline, with almost 3x compute reduction.
Note that, direct comparison with several recent diffu- sion methods isn't t feasible, as they typically report results on ImageNet64x64. We report some numbers in Table for diffusion models at 256 x 256 resolution for compari- son, though they may use different tokenizers or genera- tion architectures. Exploring these variations is beyond the scope of this work. In addition to our experiments with conduct preliminary
5.2. Video Generation
non-autoregressive transformers, experimentation on applying the core idea of decode-time model scaling to diffusion transformers (see Appendix C.I).
We use the MAGVIT (Yu et al., 2023a) framework to train parallel decoding based video generation and frame predic- tion models. Figure 11 shows generated videos of UCF101, We summarize the results for class-conditional video gen- eration UCF101 in Table 3 and for frame prediction on Kinetics600 in Table 4. Despite the challenging nature of video generation relative to image generation, results indi- cate that the decode time scaling of model size holds true even for video generation. MaGNeTS remains competi- tive to MAGVIT for frame prediction with ~ 3.7x lower compute.
Method
FID
# params
# steps
GFLOPs
DPM Stlver (I. at,, 2022)
4.1
422M
>3k
MaskGGT+44t
3.2
303M
12
1.3k
MaGNeTS (ours)"4
3.9
303M
490
Table 2: Class-conditional Image Generation on ImageNetl22x 12 188x122 128. "# steps" refers to the number of neu- ral network runs. denotes values taken from prior publications. g denotes of classificr-free guidance (Ho& & Salimans, 2022) for all steps. Iu represents use guidance only for final z steps.
5.3. Ablation Studies
Scaling Analysis. To understand the scaling properties of MaGNeTS we train models of different sizes S (22M), B (86M), L (303M) and XL (450M) for both the bascline as well as nested models needed for our algorithm. We use the same hyper-parameters for all, such as learning rate, epochs, weight decay, etc. We present the results in
Impact of Decode Time Model Schedule. We study the effect of different model scheduling choices. As discussed previously, we can think of different model schedules seal- ing up model size, scaling down, periodic scaling up and
7

---

MaGNeTS
5.0
Reated Hooub
Small Hode - Dig Hoad
5.0
45-
8
3.0
6680
4EA
35-
3A30
1A15
13.35
23-
Zio
110
206
752
(a) Scaling Up Schedules
(b) Scaling up vs down
Figure 7: Scheduling Options. (a) This shows the cempute-performanee trade-off for different schedule options while always scaling up model size over generation iters. The four numbers for cach point denote the number of iters cach model size operates in the order of downsampling factor p= (8, 4, 1). (b) This shows the benefit of scaling up model size compared to scaling it down during decoding.
Medool
Class FVD1
151
E auaaatt
Vilegs
E Gillpps
Method
FVD,
-
params
stops # GiPlopa
(ftegge
Cog (llong at. 2022)
109.2
9:40
DIGAN 2022)
1S 1007)
577.21
32.701 0.35
CCVS (Le Moing al. 2021)
S5 $5.0e1.0
DVD-GAN GAN (Cluk 2019)
32.97 1.70
Phenaki (Villegas al., 2022)
36.4
Video DriltazeeP uloeeral, alall 202200
57.00s 0152
1.18
256
THVD GADD-GAN-PP RPP (Lac et al., 2020)
25.7 0.7
12.54 a 0.06
TAIS" (Ge 2002)
420w
52.63 01234
121M
1024
Tranframee' (Nash al.,, 2022)
25,4
662M
OCVS<SgNGANN (Ia Molag et al, 2021)
3M6 15
24.67
0:13
RaMViD (Hippe et
DATS (Ge dta. 2002)
Mate-A-Vidarr atal 2000
3321 18
551
79.264 0.38
13.00
22154
1024
Vidao Dillusiu et at., 20026)
16.2 0..
1564
1.18
OugVide (Blos 20120
622
50.45
MAGVIT B
24.5s 0.9
87M
12
-1.3
(Singee
MAGVIT-L
7.2 a 6.1
16.48 0.01
306M
12
> 4.33
pllogger 2024)
MAGVIT-L
6.6 0.1
16,29
0.01
306M
PAR 16
100.4
M&GNeTS (OURS)
1088
0.00
MAOVEzB" cl. 2022g
1994
83.55 0.14
87M
12
-1.k
MaGNeTN (ousspi
9.6 0.1
16.25 0.01
306M
12
-1.4k
MAGVITI et. 2013a)
54+ 0:21
T00M
17
MsGNNTS (ni))
96 4+2
88 95-50-330 20
10MM
12
-17
Table 4: Frame prediction on K600. D denotes values taken
Table 3: Class-conditional Video Generation on UCF-101. Methods in gray are pretrained on additional large video data. Methods with / in the Class column are class-conditional, while the others are unconditional. Methods marked with use custom resolutions, while the others are at 128x 128. denotes values taken from prior publications. No guidance used for UCF101.
from papers. ga denotes use of guidance only for final x steps.
formance than scaling down model size. This shows that bigger models are better utilized in the later iters.
Impact of Caching and Refresh. We now discuss the im- pact of caching and its refresh. For this analysis, we use a uniform model schedule: h1 = l2 = k3 = K4 = 3. We also perform eaching and refresh on the baseline model, which has not been trained with any nesting with the same model applied for all iterations. For the baseline, we also refresh the cache at exactly the same steps the scheduled model. We present the results in Table 5. The columns "Baseline" and "Scheduled" do not involve any cache. While caching degrades the performance a bit, refreshing it intermittently can avoid the degradation. While refresh does have some compute overhead. it does help significantly to bridge the quality gap. Scheduling of models with caching and refresh has the best compute-performant trade-off
down, and so on. For this analysis, consider the L-sized model, with three nested models within it with parameter reduction by roughly 2: T 5. We can denote the number of times these four models are called during decoding as (ki, k2, A3, k4), s.t., Σi=1 K;- = 12. We drop the model no- tation of my in Equation (2) for simplicity and explicitly mention the model names in the text, as discussed next.
I
First, we evaluate all combinations of k; for which we al- ways scale up in Figure 7a in red and scale down in Figure 7b in blue. The green curve shows the performance of the indi- vidual nested models. We have the following observations (I) for a certain compute budget, the scheduling of models over generation iterations (red dots) can offer better perfor- mance than using a single nested model (green curve) for all steps. (2) Models that have smoother transitions in nested models, such as (3. 3,3, 3.3) or (0, O, 8, 4), offer much better performance than the ones which has abrupt model transi- tion such as (6, 0, 0, 6) or (3.0,0, 0, 0,9), i.e., directly jumping from the smallest to the biggest model. (3) Figure 7b shows that scaling up nested model size offers much better per-
Algorithm
Baseline
Cache
Refresh
Scheluled
+ Cache 4 Refresh
FID FLOP Gains (times)
2.5
3.4
266
311
48
3.1
21
Caching Ablation. Only caching performs inferior, which is bridged by refreshing it. A scaling up model schedule with caching and its refresh offers the best compute- porfuut--erffrmmacc tra Results are on ImageNet256x256 with L-sized model.
Table 5:
8

---

MaGNeTS
The efficiency of using nested models. In MaGNeTS we nested models instead of separately trained smaller sized models. This has two advantages (a) parameter sharing, which limits the number of parameters to just that of the full model, compared to 1.875x (= 1+1/2+1/45 1/4+1/8) for disjoint models. Increasing the parameter count will increase mem- ory requirements. (b) Nested models are trained efficiently in just a single training run. When trained with distillation, they generate better models than training standalone mod- els (refer Table 8) of the same size the nested models. For performance comparison, we trained standalone (I.- sized) models of the same size as the nested models for both UCF101 and ImageNet. The results are presented in Table 6. Nested models can efficiently share parameters without loss in performance (ImageNet) and offer constraints that help in better performance (UCF101) than using standalone models.
Inference Model Schedules
# TTained Nested Models +
3
4
5
6
(0,0,0,0,00 12)
2.3
2.4
2.4
2.4
2.4
(0,0,0,0, 12,0)
2.6
2.7
2.7
2.8
2.8
(0,0,0, 12,00
3,4
3.5
3.7
3.8
12,0,12,0,0.00
5.2
5.7
6.3
(0,12,0.0,0.0)
8.9
10.8
(0,0,0,0,6,00
2.6
2.6
2.6
2.7
2.7
(0,0,0,4,4,)) 4,4,4
2.7
2.8
2.8
2.8
3,,0,3,3,,,3
3.1
3.2
3.2
3,
6.3
6.6
Table 7: Ablation of number of nested models. Experiments ImageNet256 256 on L-sized model.
Datasct
Training Algo.
p=1
p=
2
p=
4
p= 8
Schoduled
Only GT Only Distill
2.9
3.9
Dataset
Nested Models
Standalone Models
ImageNot
Traning Diverge
GT Distill
2.7
3.5
5.2
ImageNet (FID)
3.1
3.1 115.0
Only GT
80.0
101.3
143.8
221.8
112.6
UCFI01 (FVD)
96.4
UCFI01
Only Distill
Tramnnng Diverpl
GT - Distill
78.3
91.2
115.4
164.4
96.4
Table 6: Nested vs Standalone Models. This table presents the performance comparison between using nested models vs. standalone, independent models without parameter sharing in the decode-time scheduling algorithm of MaGNeTTS.
Table 8: Distillation Ablation. This shows the impact of different training losses used for the nested models on ImageNet256x255 (size: L)) and UCF101 (size: L.. Using only distillation diverges while using only ground-truth losses performs worse than our ap- proach (third row), where we combine ground-truth and distillation losses with a linear decay from the former to the latter.
Impact of number of nested models. We train different settings p- = (1,2) (Iwo models), p (1,2,4) (three mod- p p= (1, 4, 8 (four models), p 21,2, 8, 16) (five models), and p 2, 16, 32) (six models). We ob- serve that for all of these cases, the biggest model perfor- mance remains almost same. However, the performance of the smaller models degrades as shown in Table 7.
the number of attention heads (n₂), applying the same par- tial computation strategy as discussed in Section 3. However, this performed worse than nesting along the head feature dimension (which we use for attention parameter nesting).
We hypothesize that the drop in performance of smaller models is due to their lower representational power. As we add more nested models, the task complexity of the shared representation increases, and burdens the mid-sized model. However, this drop in performance does not significantly impact the performance of model scheduling, the larger models dominate the final results. Note that all of these results are on top of models trained with progressive distilla- tion (see Section 4), which helps to retain the performance to some extent as shown in Table 8.
6. Conclusion
In this paper, we propose MaGNeTS, allocate variable compute along image/video generations steps. We show that instead of always using the same sized model for all decoding steps, we can start from a model which is nested and fraction of its full size, and then gradually increase model size. This along with key-value caching in the par- allel decoding paradigm obtains significant compute gains. We believe that our exploration of dynamic compute opens exciting new research directions in efficient generative mod- els. In future works, we plan to explore token-dependent model schedules for further compute gains.
Impact of Distillation. We use two types of losses to train the nested models loss w.r.t the ground-truth tokens and distillation loss using the next bigger model as the teacher. The weight between the two losses is also linearly interpo- lated from the former to the latter. We compare this training strategy with the two extremes only ground truth loss and only distillation loss and present the results in Table 8. As we can see, using only distillation loss results in divergence. Using ground-truth loss is also inferior to linearly annealing.
Impact Statement
Our work mainly focuses on improving the inference time compute efficiency of state-of oh-the-rrt visual generative mod- els in literature. As the work have been conducted on pub- licly available datasets, we do not see any potential ethical or societal concerns.
Nested Attention Heads We also investigate nesting along

---

MaGNeTS
References
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding, 2019. URL pettps://rrxie org/abs/ 11100.111... 04805.
Bachmann, R., Allardice, J., Mizrahi, D., Fini, E., Kar, O. F. Amirloo, E., El-Nouby, A., Zamir, A., and Dehghan, A. Resampling images into 1d token sequences of flexible length, 2025. URL https:: wwxiv.ovoo........oo org/ abs/2502.132 13967.
Flextok:
Dhariwal, P. and Nichol, Diffusion models beat gans image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
Bao, F., Nie, S., Xue, K., Cao, Y., Li, C., Su, H., and Zhu, All are worth words: A vit backbone for diffusion mod- cls, 2023. URL https:://// /atpii//aaivv.v/ www.aaiwww....a.eeeeoooo absgg//b///a ar2/2 12152.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for nigh-resolution image synthesis. CVPR, pp. 12873- 12883, 2021.
Berthelot, D., Autef, A., Lin, J., Yap, D. A., Zhai, S., Hu, S., Zheng, D., Talbott, W., and Gu, E. Tract: Denoising diffu- sion models with transitive closure iime-iitillattion, 2023. URL https:tp///.. // /arxiv..... org/ abs/ 2333 04248.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, E., Podell, D., Dockhorn, T,, English, Z., Lacey, K., Goodwin, A., Marek, Y., and Rombach, R. Scaling recti- fied flow transformers for high resolution image synthe- sis, 2024. https:: /arxiv.. org aaa//b//24000 03206.
Beyer, L., Zhai, X., Royer, A., Markeeva, L., Anil, R., and Kolesnikov, A. Knowledge distillation: A good teacher is patient and consistent, 2022. URL https: //arxiv. org/ 2216. 05237.
Feng, W., Yang, C., An, Z., Huang, L., Diao, B., Wang, F., and Xu, Y. Relational diffusion distillation for efficient image generation, 2024.
Brock, A. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809 1096, 2018.
abs / 241 0 07679.
Cai, R., Muralidharan, S., Heinrich, G., Yin, H., Wang, Z., Kautz, J., and Molchanov, P. Flextron: Many- in-one flexible large language model. arXiv preprint arXiv:2406. 10260, 2024.
Gao, S., Zhou, P., Cheng, M.-M., and Yan, S. Masked diffusion transformer is a strong image synthesizer. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 24173333,, 2023.
Carreira, J., Noland, E., Banki-Horvath, A., Hillier, C., and Zisserman, A. A short note about kinetics-600, 2018. URL https:: / arxii. www.i..org///....vo.oo....oooov. abs 808. 01340.
Ge, S., Hayes, T., Yang, H., Yin, X., Pang, G., Jacobs, D., Huang, J.-B., and Parikh, D. Long video generation with time-agnostic vegan and time-sensitive transformer, 2022, URL https: / aax wiv. org/ abs// 22002 1.2204..20 03638.
Chang, H., Zhang, H., Jiang, L., Liu. C., and Freeman, W. T. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 11315-11325, 2022.
Gu, S.. Chen, D.. Bao, J., Wen, Zhang, B., Chen, D., Yuan, L., and Guo, B. Vector quantized diffu- sion model for text-to-image synthesis, 2022. URL https:/ /ttp/aaxviv//v////////eea/.a.ava http://aawwvv....aaeeeeeee wtgg/os////////.////
Chang. H., Zhang, H., Barber, J., Maschinot, A., Lezama, J., Jiang, L., Yang, M.-H., Murphy, K., Freeman, W. T., Rubinstein, M., et al. Muse: Text-to-image genera- tion via masked generative transformers. arXiv preprint vrXiv:2311.00700 2023.
Hang. T., Gu, S., Li, C., Bao, J., Chen, D., Hu, H., Geng, X., and Guo, Efficient diffusion training via min- weighting strategy, 2024. URL htttt:///aaatt org /a//abbb//3333 09556.
Choi, J., Lee, J., Shin, C., Kim, S., Kim, H., and Yoon, S. Perception prioritized training of diffuusin models, 2022. URL https: arxiv..iiivv//iii//aaaaaaaaaavva org/ abs/ 2204. 00227.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Clark, A., Donahue, J., and Simonyan, K. Adversarial video generation on complex datasets. arXiv preprint arXiv: 100.006177.0 2019.
Ho, J. and Salimans, Classifier-free diffusion guid- ance, 2022. URL https://arxiv.orrt 12598..
Deng. J., Dong, W. Socher, R., Li, and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255 leee, 2009.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840066855 2020.
10

---

MaGNeTS
Ho, J., Saharia, C., Chan, W., Fleet, D. J., Norouzi, M., and Salimans, T. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47):1-33 2022a.
Farhadi, Jain, P., al. Matformer: Nested transformer for elastic inference. arXiv preprint 10.007777, 2023.
Kusupati, A., Bhatt, G., Rege, A., Wallingford, M., Sinha, A., Ramanujan, V., Howard- Snyder. W., Chen, K., Kakade, S., Jain, P., al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233-30244 2022.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, J. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-866 2022b.
Hong, W., Ding, M., Zheng, W. Liu, X., and Tang, Cogvideo: Large-scale pretraining for text-to-video gener- ation via transformers. arXiv preprint arXiv:2205.15868, 2022.
Le Moing, G., Ponce, J., and Schmid, C. Cevs: Context- aware controllable video synthesis. Advances Neural Information Processing Systems, 34: 14:14404--4055, 2021.
Lee, H., Lee, H., Gye, S., and Kim. J. Beta sampling is all you need: Efficient image generation strategy for diffusion models using stepwise spectral analysis, 2024. URL https::tp//. // arxiv. .ww.v.orrrggooo..... aabs///////// 2207
Hoogeboom, E., Heek, J., and Salimans, T. simple diffu- sion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pp- 13213-13232. PMLR, 2023.
Lee, Y., Kim, J.-Y., Go, H., Jeong, M., Oh, S., and Choi, S. Multi-architecture multi-expert diffusion models, 2023, URL https: :ttttp::///aaiee http::::/aaaxvvva/aa/e
Höppe. T., Mehrjou, A., Bauer, S., Nielsen, D., and Dittadi, A. Diffusion models for video prediction and infilling. arXiv preprint ar arXiv:2206.07696 2022.
abs http.//rrrraa........../ 0499.
Li, L., Li, H., Zheng, X., Wu, J., Xiao, X., Wang, R., Zheng, M., Pan, X., Chao, F., and Ji, R. Autodiffusion: Training-free optimization of time steps and architectures for automated diffusion model acceleration, 2023. URL https: /ttp/vixiiv..///aa///e/////aaava http:////axxxi.../aabbaoeee org. 2209.
Hu, W., Dou, Z.-Y. Li, L. H., Kamath, A., Peng, N., and Chang, K.-W. Matryoshka query transformer for large vision-language models, 2024. URL https://aaax
Jabri, A., Fleet, D., and Chen, T. Scalable adaptive computation for iterative generation. arXiv preprint arXiv:2212.11972 2022.
Li, T., Tian, Y., Li, H., Deng, M., and He, Autoregres- sive image generation without vector quantization. arXiv preprint 2XXv:2446.111888 2024a.
Jain, G., Hegde, N., Kusupati, A., Nagrani, A., Buch, S., Jain, P., Arnab, A., and Paul, S. Mixture of nested ex- perts:
Li, T., Tian, Y., Li, H., Deng, M., and Hc, K. Autoregres- sive image generation without vector quantization, 2024b. URL https: //aaaaivvv./..////// /ttp//aawv.rrg/..a...////.. abb////////// ab//2/////444// 2406 11838.
Adaptive processing of visual tokens, 2024. URL https:/ / rxiv. org/ abs///b////// abs/ 19955.
Liu, L., Ren, Y., Lin, Z., and Zhao, Z. Pseudo numerical methods for diffusion models on manifolds, 2022. URL. https: /at::://xxivvvvii wttp://.wwgg/bbb//.....o 09778.
Kingma, D. P. and Gao, R. Understanding the diffusion objective as weighted integral of elbos. arXiv preprint :XXi::3003.0084,, 2, 2023.
Lou, A., Meng, C., and Ermon, S. Discrete diffusion mod- eling by estimating the ratios of the data distribution. In Forty-frst International Conference on Machine Learn- ing.
Koike-Akino, T. and Wang, Y. Stochastic bottleneck: Rateless auto-encoder for flexible dimensionality reduc- tion, 2020. URL https: /ttp:///aavvvvvvav org/abs 02870.
http::/:////xxiiooo/////sss ab000 2005.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps, 2022. URL https: / //arxiv. org abs 009.2...097.
Kondratyuk, D., Yu, L.,, Gu, X., Lezama, J., Huang, J., Hornung, R., Adam, H., Akbari, H., Alon, Y., Birodkar, V., et al. Videopoet: A large language model for zero-shot video generation. ICML, 2024.
Luc, P., Clark, A., Dieleman, S., Casas, d. I.., Doron, Y., Cassirer, A., and Simonyan, K. Transformation-bait adversarial video prediction farge-scale data. arXiv preprint arXiv:2003.04003 2020.
Kong, and Ping, W. On fast sampling of diffuusion prob- abilistic models, 2021. URL hittps://arxiv.. abs/2106. 00132.
Kudugunta, S., Kusupati, A., Dettmers, T., Chen, K., Dhillon, I. Tsvetkov, Y., Hajishirzi, H., Kakade, S.,
Meng, C., Rombach, R., Gao, R., Kingma, D. P., Ermon, S., Ho, J., and Salimans, T. On distillation of guided

---

MaGNeTS
diffusion models, 2023. httpsttss///arxi abs/2210. 2210..3142..
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models, 2022. URL https: /ttp:s://axxvvoo/o/// abs 22t/22122..2 10752.
Mentzer, F., Minnen, D., Agustsson, E., and Tschannen, M.
Finite scalar quantization: Vq-vae made simple, 2023.
URL htllhttss://axiiiee http::://axxv..orrrr abs 2309. 15505.
Saharia, C., Chan, W., Saxena. S., Li, L., Whang, Denton, E. I.., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image dif- fusion models with deep language understanding. Ad- vances in neural information processing systems, 35: 36479-36494, 2022.
Mirzadeh, S.-I., Farajtabar, M., Li, A., Levine, N., Mat- sukawa, A., and Ghasemzadch, H. Improved knowledge distillation via teacher assistant, 2019. URL https: wrw.vorg//aaee abs/ 1902. 03393.
Miwa, K., Sasaki, K., Arai, H., Takahashi, T., and Ya- maguchi, Y. One-d-piece: Image tokenizer meets coality-controllabbee compression, 2025. URL https: //arxiv org. abs/ / 2501 10064. .00644
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models, 2022. URL https: / arxiv. rexiv.org/a prg/abs////e...o.. 2202. 222.222222222 00512.
Salimans, T., Goodfellow, Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for train- ing gans, 2016. URL 1 1606. 03498.
Nash, C., Carreira, J., Walker, J., Barr, L., Jaegle, A., Ma- linowski, M., and Battaglia, P. Transframer: Arbitrary frame prediction with generative models. arXiv preprint arXiv:2203..m 09494, 2022.
orgg org//bs/ /
Sargent, K., Hsu, K., Johnson, J., Fei-Fei, L., and Wu, J. Flow to the mode: Mode-seeking diffusion autoen- coders for state-of-the-art image tokenization, 2025. URL https: /atps://arxivv wwwpxiivvorr///ae wwt./b////...o..o..ee /250gg/bbb//33333 11056.
Ni, Z., Wang. Y., Zhou, R., Guo, J., Hu, J., Liu, Z., Song, S., Yao, Y., and Huang, G. Revisiting non-aetoregressive transformers for efficient image synthesis, 2024a. URL https: arxiv. org/ abs 2406. 05478.
Sauer, A., Schwarz, K., and Geiger, A. Stylegan-xl: Scaling stylegan to large diverse datasets, 2022. URL https: /
Ni, Z., Wang, Y., Zhou, R., Han, Y., Guo, J., Liu, Z., Yao, Y., and Huang, Enat: Rethinking interactions token-based image synthesis, 2024b. URL https: /ttas://araae //apxva...vv.///aa//aa/.aaaaiiia org/ abs/2411 06959.
220.. 00273.
Singer, U., Polyak, A., Hayes, T., Yin, X., An, J., Zhang, S., Hu, Q.. Yang. H., Ashual, O., Gafni, O., al. Make-a- video: Text-to-video generation without text-video data. arXiv preprint arXiv:2209. 14792, 2022.
Nichol, A. and Dhariwal, P. Improved denoising diffusion probabilistic models, 2021a. URL
orgg/bs//1000 /2 09672.
Skorokhodov, I., Tulyakov, S., and Elhoseiny, M. Stylegan- V: A continuous video generator with the price, image quality and perks of stylegan2. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 3626-3636, 2022.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffu- sion probabilistic models. In International conference on machine learning, pp. 8162-8171, PMLR, 2021b.
Park, Y.-H., Kwon, M., Choi, J., Jo, J., and Uh, Y. Un- derstanding the latent space of diffusion models through the lens of riemannian geometry, 2023. URL https: //arxiv. orgo abs/ 2307 12868.
Song, J., Meng, C., and Ermon, Denoising diffusion implicit models, 2022. URL https://azxiv.orrg abs / 221 0 02502.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Con- sistency models, 2023. URL abs abs/2303.0m 214.6 01469.
Peebles, W. and Xie, S. Scalable diffusson models with trans- formers, 2023. URL https://arxix: www.vvorrrreeooe abs/ 0212.09748.
Soomro, K., Zamir, A. R., and Shah, M. Ucf101: A dataset of 101 human actions classes from videos in the wild, 2012. URL httpps://axxi...rg//////// 1212. 0402.
Razavi, A., Van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with 2q-vae-2. Advances in neural information processing systems, 32, 2019.
Rippel, O., Gelbart, M., and Adams. R. Learning ordered representations with nested dropout. In International Con- ference on Machine Learning, pp. 1746-1754 PMLR, 2014.
Stanton, S., Izmailov, P., Kirichenko, P., Alemi, A. A., and Wilson, A. G. Does knowledge distillation re- ally work?, 2021. URL https: arrxi.iv wrt. /bs// 2106 05945.
12

---

MaGNeTS
Tian, K., Jiang, Y., Yuan, Z., Peng, B., and Wang, L. Visual autoregressive modeling: Scalable image generation via next-scale prediction, 2024. URL https: /arxiv. org/abs/ 2404. 02905.
Yu, L., Cheng, Y., Sohn, K., Lezama, J., Zhang, H., Chang, H., Hauptmann, A. G., Yang, M.-H., Hao, Y., Essa, I., et al. Magvit: Masked generative video transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10459-10469, 2023a.
Unterthiner, T., van Steenkiste, S., Kurach, K., Marinier, R., Michalski, M., and Gelly, S. Fvd: A new metric for video generation. 2019.
Yu, I.., Lezama, J., Gundavarapu, N. B., Versari, L.., Sohn, K., Minnen, D., Cheng, Y., Gupta, A., Gu, X., Haupt- mann, A. G., et al. Language model beats diffusion- tokenizer is key to visual generation. arXiv preprint arXiv:2310.05737,222 2023b.
Van Den Oord, A., Vinyals, O., et al. Neural discrete rep- resentation learning. Advances in neural information processing systems, 30, 2017.
Yu, Q.. He, J., Deng, X., Shen, X., and Chen, L.-C. Ran- domized autoregressive visual generation, 2024a. URL
Villegas, R., Babaeizadeh, M., Kindermans, P.-J., Moraldo, H., Zhang, H., Saffar, M. T., Castro, S., Kunze, J., and Erhan, D. Phenaki: Variable length video generation from open domain textual descriptions. In International Conference Learning Representations, 2022.
/ttpss//arxiv..oo
Yu, Q. Weber, M., Deng, X., Shen, X., Cremers, D., and Chen, L.-C. An image is worth 32 tokens for reconstruction and generation, 2024b. URL https: //axxiv.rrg/ab//20000 07550
Wang, Y., Ren, S., Lin, Z., Han, Y., Guo, H., Yang, Z., Zou, D., Feng, J., and Liu, X. Parallelized autoregressive visual generation, 2024. URL htRp://t::/aatt //arxvv.///i////aaaaaaeeeeeaae org/ 15119.
Yu, S., Tack, J., Mo. S., Kim, H., Kim, J., Ha, J.-W., and Shin, J. Generating videos with dynamics-aware im- plicit generative adversarial networks. arXiv preprint arXiv:2202. 10571, 2022.
Wang, Z., Jiang, Y., Zheng, H., Wang, P., He, P., Wang, Z., Chen, W., and Zhou, M. Patch diffusion: Faster and more data-efficient training of diffusion models, 2023, URL https: arxiv. org/ abs/ 304 12526.
Yu, S., Kwak. S., Jang, H., Jeong. J., Huang, J., Shin, J., and Xie, S. Representation alignment for generation: Training diffusion transformers is easier than you think, 2024c. URL https: arxiv.. wwwii.oorr.. org/ abs /2400. 06940..0009..
06940.
Watson, D., Chan, W., Ho, J., and Norouzi, M. Learning fast samplers for diffusion models by differentiating through sample quality, 2022. URL https: /ttpx/x.xxv..///////..././.. abs/2202, 05830.
Yuan, Z., Zhang, H., Lu, P., Ning, X., Zhang, L.. Zhao, T., Yan, S., Dai, G., and Wang, Y. Ditfastattn: Attention compression for diffusion transformer models. arXiv 2024.
preprint arXiv:2406.0885
Weber, M., Yu, L., Yu, Q. Deng, X., Shen, X., Cre- mers, D., and Chen, L.-C. Maskbit: Embedding-free image generation via bit tokens, 2024. URL https: //arxiv. org/ abs /209 16211
Zhao, Y., Xiong, Y., and Krähenbühl, P. Image and video tokenization with binary spherical quantization, 2024. URL https: aa//aaxxvvvv../..//.aa /aa/ab///44 07548.
Zheng, C., Vuong, L. T., Cai,, J., and Phung, D. Movq: Mod- ulating quantized vectors for high-fidelity image genera- tion, 2022. URL https: pttttps///axivv /awxiv.iiii..///ii///eaa///iaiia org abs/2209. 09002.
Wen, X., Zhao, B., Elezi, Deng, J., and Qi. "principal components" enable new language of images, 2025.
URL htttstss://axii org/ abs/2503 685.
Yan, N., Gu, J., and Rush, A. M. Diffusion models without attention. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 8239- 8249, 2024.
Zheng, H., Nie, W., Vahdat, A., and Anandkumar, Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306.0985 2023.
Yang, S., Chen, Y., Wang, L., Liu, S., and Chen, Y. Denoising diffusion step-aware models, 2024. URL https: http://axxvvvv. org/ abss abss2310 0.310.. 03337.
Yu, J., Li. X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y. Vector-quantized image modeling with improved vggan. arXiv preprint arXiv:2/ 10. 10.04627, 2021.
13

---

MaGNeTS
A. Motivation for Decode Time Model Scaling
is the current decoding iteration and K is the total number of decoding iterations. We report the hyperparameters use in in Table 9. We use bias=0.5 and scale=0.8 for all experiments.
Our visualization token density averaged across 50k Ima- geNet samples reveals a dynamic pattern initial decoding iterations prioritize background regions. In contrast, later iterations focus on the center where foreground objects region of interest typically reside. This highlights the need to allocate resources efficiently during generation. To fur- ther investigate this behavior, we examine token density across various ImageNet categories (refer Figure 8). This category-wise analysis further motivates our focus on de- code time scaling. Figure 10 shows more qualitative results on ImageNet256 256 and Figure 11 shows samples UCF101.
Dataset
Method
gs
MTemp
MaskGIT++ MaGNeTS
65
ImageNet
65
UCF101
MAGVIT/ MaGNeTS
0
5
MAGVIT MaGNeTS
10
12.5
Kinetics600
5
10
Table 9; Best Sampling Hyperparameters.
B. prrrameerr Details
The MaskGIT algorithm has the following which we discuss next.
C. Additional Experiments
C.1. Preliminary Diffusion experiments
Guidance Scale (gs). It is used classifier-free guidance (Ho & Salimans, 2022) and governs the calculation of final logits during inference as shown in Equation (4).
We conduct initial experiments using model scheduling on diffusion models. Instead of training a new diffusion model with nesting and distillation, time experiments. We use publicly available pretrained checkpoints of UViT (Bao et al., 2023) on ImageNet64 x64. Specifically, we employ two models U-ViT-L/4 (large) and U-ViT-M/4 (mid) to investigate the impact of model scheduling during inference.
focus solely inference-
logitSpnal logitscoud + A gs (logitScondd uncona) (4)
from class-conditional input, lngot are from unconditional input, and A depends on the mask-ratio of the current decoding iteration.
where logitscond
Figure 8 shows that the initial decoding iterations of par- allel decoding focus on the background region, and focus gradually shifts the main object/region in the final de- coding iterations. Motivated by this, experimented with applying guidance to only few final decoding iterations and present our findings in Figure 9b. As we can see, most of the decoding iterations do not require guidance. We use guidance only for final few decoding iterations for class- conditional generation in ImageNet256x256 and frame pre- diction in Kinetics600. Following MAGVIT (Yu et al., 2023a), for class generation in UCF101 we do not use classifier-free guidance.
Implementation Details We use the default number of sampling steps of 50 and batch size of 500 in all experiments. We do not use classifier-free guidance. We do not use any caching for these experiments due to the continuous nature of the input. All experiments are run on single A A100 GPU.
Optimal Model Schedule Since the initial denoising steps play a crucial role shaping the final output of the reverse diffusion process, we utilize the model for these early stages and transition to the M model for the later de- noising 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.
Mask Temperature (MTemp). It controls the random- ness introduced on top of the token predictions to mask tokens.
Quantitative Results Refer Table 10 for results. With only model scheduling, we are able achieve ~1.53x in- ference compute gains with almost similar performance as baseline. Exploring more refined schedules, training the models with nesting and distillation will offer better com- pute gains. Additionally, nesting would enable parameter sharing, unlike the current setup, which relies on separate models. This shows that the proposed method of model scheduling over multi-step decode process in image/video
Sampling Temperature (STemp). It controls the ran- domness of the sampling from the categorical distribution of logits. Tokens are sampled from logits/STemp. STemp is calculated by Equation (5).
STemp bias scale- (1 - 1k)kk)
(5)
where bias and scale are hyperparameters (see Table 9), &
14

---

MaGNeTS
0.01
0.02
0.03 0.03
0.07 0.05
0.10..05
0:12 0.07
0.15
0.15 0.08
0.15 0.09
0.14
0.440:4
0.17 0.02
0.21
0.07
0.00 0.07
0.11
0.14::09
0.14
013010
014007
016
0.01
0.01 0.02
0.03 0.04
0.05 0.05
0.07 o 66
0. 0.07
0.10008
0.12
0.13.10
014010
014009
0.17
0.18
Figure 8: Visualization of token density unmasked in each iteration averaged over 10k generated samples different categories of ImageNet. The top example shows category volcano. Middle and bottom examples show cdishrag.dishclog and "goldfish, Carassius auratus", respectively. Yellow color represents higher density, and each pixel represents token from the 16x x 16 token space.
generation generic enough to applied to different mod- eling approaches.
Calculation of GFLOPs. We illustrate calculation of inference GFLOPs via Python pseudo-code in Table 12. We double the GFLOPs decoding iterations where classifier- free guidance (Ho & Salimans, 2022) is used. Note that always use cosine schedule to determine the number of tokens to be unmasked in every step.
Method
FID (50K)
puramns
# steps
Time (seefiterr
U-VIT-ME
L/V
4.1
287M
50
32.34
sched)
287)
Table
10:
Class-conditional Image Generation on "# steps" refers to the number of neural
Real-Time Inference Benefits. addition to the theoret- ical FLOP gains offered by MaGNeTS. here we want to analyze the real-time gains that it offers. We implement MaGNeTS single TPUv5 chip and present the results in Table 11.
ImageNet64x64. network runs.
D. Compute Gains
Algorithm >
Baseline (MaskGIT++)
MaGNeTS
Per-step FLOPs. Figure 9a illustrates the inference-time computational cost, measured in GFLOPs, per iteration for the baseline model and MaGNeTS. As amount of FLOPPs are drastically reduced using MaGNeTS. This is for a schedule with k k2 = 33 K4 3. The spikes after every 3 iterations are due to the cache refresh step. Mechanisms to get rid of the cache refresh can further reduce the total compute needed.
Images/Sec Latency (ms)
22.5
56.3
712
285
see the
Real-Time Inference Efficiency. These show the number of generated images per and latency. These results are ImagcNer256x256 with model size XL.
Table 11:
15

---

MaGNeTS
Model Scale:
Model Scale: XL
Baseline
MaGNiTS MaskGiT+++
MaGNeTS
MascCi++
8-
Iterations
Iterationn w/ Guidance
(a) GFLOPs per sampling step.
(b) Guidance Analysis.
Figure 9: (a) Inference GFLOPs per step baseline and MaGNeTS, (b) generation performance (FID) on ImageNet Number decoding iterations w/ guidance for different model scales. Note that we start from last decoding iteration. For example. "No, of iterations w/ Guidance 6" means we use guidance only for final six itcrations (out of total 16 iterations). This shows that using guidance only for few final iterations is enough in the parallel decoding setup.
E. Limitations.
While our approach demonstrates strong performance in image and video generation, we acknowledge certain lim- itations. Some artifacts inherent to MaskGIT++ may also appear in our generated outputs (see Figure for examples on ImageNet256 256). Such artifacts are common in mod- els trained controlled datasets like ImageNet. Moreover, the quality of the pretrained tokenizers (Yu et al., 2023b; We- ber et al., 2024) directly impacts our method's effectiveness; however, improving these tokenizers is beyond the scope of this work. Although, use of nesting and decode time scaling does not have any specific requirement for model architec- ture and sampling scheme, to unlock the further benefits of KV caching, the process needs to generate discrete tokens.
16

---

MaGNeTS
MaskGIT++ (FID=2.3)
~3x FLOPs Reduction
MaGNeTS (FID=2.9)
Figure 10: Class-conditional Image Generation. More qualitative results ImageNet. Comparing MaskGIT++ and MaGNeTS (size: L, epochs: 270).
17

---

Figure 11: Class-condittonall Video Generation on UCF101 16-frame videos are generated at 128x128 resolution 25 fps. Every third frame is shown for each video. The classes from top to bottom are Lunges, Bench Press, Handstand Pushups, Cutting In Kitchen.
MaskGIT++
MaGNeTS
Figure 12: Failure cases. Similar to existing methods, our system can produce results with noticeable artifacts.
18

---

MaGNeTS
Funstion get the GFlops for current decoding iteration
def ret_flops (num tokens cached,
Eokens_cached, num_tokens pam__ttens_pooceseedde model_id, parans, version)
num num_layers, hidden_size, mip_dim, nur_heads parama [version]
ękv 4 qun_tooknssproccesed hidden_size (hidden_size // model_id)
attn = nun_tokens_prroeesed (nun_tokens_procesee)
nun_tokens_cachell
sidden_size
mip nun_tokens_prrcest return (qkv + attn + m1p) num_layers // 1e9
(mlp_dim // model_idy hidden_sizz
Function get the total inference GFlops
10
def get_total. totalotllaa..o #lops (version, num sum_iters, use cache, refresh at, total_toxens, model_id_scheduie params, num_cond_tokens=0)] assert nun_cond_tokens totalltokens
refresh_cache_at [int (x) for in att if assert len (model sc_sccud iters num cached total_flops
14
15
16
MaGNeTS (curs) doesn't need to process the conditioned tokens in the frame prediction task total_tckens num num_cond_token
18
19
for in range (num (num_iterr)) itters) ratic i num iters
22
23
Cosine masking schedule num_processed np.cos inp.pi/2, * ratio) total_tokens
24
26
Even are performing caching, all tokens are processed in first iteration. iterations where cache refreshed
27
if == 0 or i in refresh_cache._eee and use iis_ccaceer
sarel__tokesstnu_momm tetal_flops += get_fleps tote model idee_idd_scheule params, version)
29
30
we always cache the conditioned tokens else:
11
tetal_flopp +- glo_f (nur_ccehad-mumnooooccoeee total _tooenes-uu_ccchheeee model Lodel_id_schedulle params, version)
33
34
if use cache:
35
num cached total tokens num processed
return total_flops
Sample function call for class-conditional image generation params is a dictionary of the form (version: (num_layers, hidden_size, mip_dim, nun_heada)
common (' verssinn'; 'I', 'num_iters': 12, 'total_tokens' 257, 'parans's params)
baseline d'use_cache" False, 'refresh_at's ", 'model "model_id_schedeeee scheduee''''' (1,) *ccommon)
ours = ( _ase_ccchee: True, 'refresh_at'; 13,6, 9', sccheuleds (8,) (4, (2,) *3 (1, x3, **common)
print (get_total, flops basellnee) get_tota.1 flops (k*ours))
total_tokens 1025 for class-condit onal video generation and frame prediction nun_cond_tokens 512 for frame prediction
Table 12: Python pseudo-code illustrating the calculation of inference GFLOPs.
19


Reviews:
[{"review_id": "T6UEZfv8A3", "reviewer": "Reviewer_NcBN", "summary": "The paper introduces MaGNeTS, a approach to improving the efficiency of visual generative models by dynamically scaling model size during decoding.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Claim about KV Caching and Refresh for Inference Efficiency**\n\n> 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.\nFor e.g.,\n> * In Line 109 (left column), we state “KV caching can also be used in parallel decoding, which can effectively reuse computation when refreshed appropriately”. \n> * 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”. \n> * Then in Line 251 (right column) we mention “To remedy this, we strategically refresh the cache while changing the model size.”\n\n**Trade-off analysis**\n\n> 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:\n> * Given a fixed compute or latency budget, it shows that the proposed method obtains the best quality.\n> * Conversely, given a desired quality requirement, it shows that the proposed method is the most inference-efficient one.\nThis tradeoff analysis effectively captures the true essence of compute efficiency. \n\n**Comparison to Efficient Diffusion Models**\n\n> 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](https://arxiv.org/pdf/2410.07679), [Lee et al 2024](https://arxiv.org/pdf/2407.12173), [Meng et al 2023](https://arxiv.org/pdf/2210.03142), [Berthelot et al 2023](https://arxiv.org/pdf/2303.04248), [Song et al 2023](https://arxiv.org/pdf/2303.01469), [Zheng et al 2023](https://arxiv.org/pdf/2211.13449)) 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. \n\n> | Method | Image Size | FID | Params | Steps | GFLOPs |\n|---|---|---|---|---|---|\n| DPM-Solver (Lu et al 2022) | 128 | 4.1 | 422M | 12 | >3000 |\n| MagNeTS (Ours) | 128 |  3.9 | 303M | 12 | 236 |\n|---|---|---|---|---|---|\n| EDiff [A] | 256 | 2.1 | 450M | 50 | 119k |\n| LPDM-ADM [B] | 256 | 2.7 | - | 50 | - |\n| MagNeTS (Ours) | 256 | 3.1 | 303M | 12 | 490 |\n\n**Artifacts**\n\n> 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.\n\n**References**\n\n[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 \n\n[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", "claims_and_evidence": "The claims made in the submission are largely supported by clear and convincing evidence, particularly through extensive experiments on ImageNet256×256, UCF101, and Kinetics600.\n\nHowever, the authors claim that KV caching improves efficiency without performance loss, but Table 4 shows that caching degrades FID, and only with cache refresh does performance recover.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria are appropriate for the problem. MaGNeTS is tested on ImageNet256×256, UCF101, and Kinetics600, which are standard benchmarks for image/video generation. Metrics like FID, IS, and FVD effectively measure generation quality. However, The model is compared to old baselines, and while compute savings (2.5–3.7×) are significant, FID scores slightly degrade, requiring further trade-off analysis.", "theoretical_claims": "The paper primarily focuses on algorithmic innovations and empirical results rather than extensive theoretical proofs.", "experimental_designs_or_analyses": "Yes, the experimental design and analyses were reviewed for soundness and validity.", "relation_to_broader_scientific_literature": "The key contributions of this paper build upon prior research in efficient visual generation, parallel decoding, and nested transformer models, while introducing novel improvements in compute efficiency through decode time scaling and KV caching.", "essential_references_not_discussed": "I think the paper has good references .", "other_strengths_and_weaknesses": "Strengths :\n\nThe paper introduces decode time model scaling, a novel dynamic compute allocation approach that progressively scales model size, reducing redundant computation.\n\nWeaknesses:\n\nComparison to Efficient Diffusion Models is Missing", "other_comments_or_suggestions": "Refer to Other Strengths And Weaknesses", "questions_for_authors": "Regarding artefacts in the generated results: The paper acknowledges that there may be artefacts in the results generated by MaGNeTS. Can the authors provide more specific details about these artefacts?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fv44NDdJmN", "reviewer": "Reviewer_fe9H", "summary": "This paper proposed a promising efficient approach for image/video generation. Specifically, this work introduces the concept of model size scheduling during the generation process to significantly reduce compute requirements. It demonstrates KV cache also works for parallel decoding. They used nested modeling to achieve these ideas. The experimental results show strong performance of the proposed efficient method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Impact of the number of nested models**\n\n> 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. \n\n**Qualitative Analysis**\n\n> 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.", "claims_and_evidence": "There are three key parts of the proposed method: 1) Decode Time Model Schedule; 2) Caching and Refresh; and 3) Nested models.\n\nThe experiments and ablation study demonstrate the effectiveness of these modules.", "methods_and_evaluation_criteria": "**Method:**\nAfter reading the first section of the supplementary material, I think the author's motivation is very natural. However, I have a concern about the number of the nest models. Is there any experiment to analyze how the number of the nest models effect the performance?\n\n**Evaluation:** It will be better to add some cases about the generated videos in the suppl.", "theoretical_claims": "There is no proof for theoretical claims.", "experimental_designs_or_analyses": "The quantitive experiments are well-designed, and the ablation study strongly demonstrates the efficiency of the proposed method", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Additional Weaknesses: There was relatively little qualitative analysis, and I would have liked to see more visual contrasts.", "other_comments_or_suggestions": "It is a promising work about the efficiency of the vision generative models. I think this paper deserves to be accepted.", "questions_for_authors": "Please refer to the \"Methods And Evaluation Criteria\" and \"Other Strengths And Weaknesses\".", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LFLxHlzWgl", "reviewer": "Reviewer_MdBq", "summary": "This paper introduces Nested Transformers for efficient image and video generation. The method progressively increases the model size during decoding to reduce computational costs in the early steps. Additionally, KV-caching is employed across decoding steps to further enhance efficiency. Experiments are conducted on both image and video generation tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Inference Time**\n\n> 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.\n| Algorithm | Baseline (MaskGIT++) | MaGNeTS |\n|---|---|---|\n| Images/Sec | 22.5 | 56.3 |\n| Latency (ms) | 712 | 285 |\n\n**Generalization to other generative modeling frameworks**\n\n> 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. \n\n> 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 - \n> - We use the default number of sampling steps = 50 and batch size = 500 in all experiments. We use a single A100 GPU. \n> - 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.\n> - 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.\n> - 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.\n| Method | FID (50k) | # steps | time (sec/iter) |\n|---|---|---|---|\n| **U-ViT-M/4** | 5.92 | 50 | 17.12 |\n| **U-ViT-L/4** | 4.21 | 50 | 32.34 |\n| **Ours (Model Sched)** | 4.58 | 50 | 21.10 |\n\n> 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\n\n**Typos**\n\n> Thank you for bringing these typos to our knowledge. We will fix them in our final revision.\n\n**References**\n\n> [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.\n\n> [B] https://github.com/baofff/U-ViT", "claims_and_evidence": "Experimental results support the claim that the proposed method reduces computational cost in terms of FLOPs while maintaining generation quality.", "methods_and_evaluation_criteria": "Using nested modeling and gradually increasing the model size is a reasonable approach to reducing the computational cost of generative models.\n\nFrom the evaluation perspective, the paper includes comparisons on both image and video generation tasks, which appropriately validate the proposed method.", "theoretical_claims": "This paper does not include any theoretical proofs.", "experimental_designs_or_analyses": "This paper evaluates model efficiency using parameter size, FLOPs. However, it would be beneficial to include inference time in the ablation study as a direct indicator of practical efficiency.", "relation_to_broader_scientific_literature": "This paper presents an efficient approach for generative models, building upon MaskGiT. It improves efficiency by approximately three times while keeping the generation quality, which is a significant advancement. \n\nThe contributions of this paper are related to distillation-based methods but from a nested modeling perspective, which has not been explored before.", "essential_references_not_discussed": "I do not notice any other essential references that require further discussion.", "other_strengths_and_weaknesses": "The design of the proposed method appears to be specifically tailored to the MaskGiT framework and may not generalize well to other generative modeling frameworks, which is a potential limitation.", "other_comments_or_suggestions": "typos:  \n1. missing space in Line 42 between \"and\" and \"video\".\n2. unexpected space in Line 238 and Line 309.", "questions_for_authors": "1. Can the proposed method be applied to other frameworks, or is it limited to the MaskGiT framework?\n2. Why was inference time not reported as part of the evaluation to validate the efficiency of the method?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cdEKzR6Hk8", "reviewer": "Reviewer_XpTJ", "summary": "Recent advances in visual generation have improved content quality but face challenges in computational efficiency during inference. Many algorithms require multiple passes over a transformer model, keeping a consistent model size that leads to high computational costs. This work proposes two strategies to address this: (a) implementing a decode-time model scaling schedule to allocate computational resources more effectively, and (b) caching and reusing computations. These approaches allow smaller models to handle more tokens while larger models process fewer, without increasing parameter size due to shared parameters. This results in competitive performance with significantly reduced computational costs.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n**Training complexities in nested models**\n\n> 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.\n\n**Quality of generated samples**\n\n> 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.\n\n> 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.\n> | Method | Schedule | FID | # params | # GFLOPs |\n|---|---|---|---|---|\n| MaskGIT | NA | 6.2  | 303M | 647 |\n| MaskGIT++ | NA | 2.5  | 303M | 1.3k |\n| MagNeTS (ours) | (3, 3, 3, 3) | 3.1 | 303M | 490 |\n| MagNeTS (ours) | (0, 0, 6, 6) | 2.6 | 303M | 745 |", "claims_and_evidence": "Supported by references to various paradigms, but could benefit from specific performance metrics to illustrate improvements. This claim is generally accepted.", "methods_and_evaluation_criteria": "The submission introduces MaGNeTS, which employs model size scheduling and KV-caching during the decoding process. This approach logically addresses the identified inefficiencies in parallel decoding and redundancy in computations. The gradual scaling of model size is a sensible strategy to optimize computational resources, making it relevant for high-quality image and video generation. \n\nThe use of benchmark datasets like ImageNet, UCF101, and Kinetics600 is appropriate for evaluating the performance of generative models. These datasets are widely recognized in the field and provide a robust basis for comparing the quality of generated outputs.", "theoretical_claims": "This paper does not cover theoretical claims.", "experimental_designs_or_analyses": "The submission outlines the use of benchmark datasets (ImageNet, UCF101, and Kinetics600) to evaluate MaGNeTS. This choice is sound as these datasets are well-established and relevant for image and video generation tasks.", "relation_to_broader_scientific_literature": "The caching mechanism proposed for reusing computations is reminiscent of techniques in self-attention models, where key-value pairs are cached to improve efficiency (Gu et al., 2022). This notion of leveraging previously computed results to enhance performance aligns with broader trends in machine learning focused on efficiency.", "essential_references_not_discussed": "I don't think there is any important related work that has not been discussed.", "other_strengths_and_weaknesses": "Strengths\n\nBy reducing computational costs by 2.5–3.7× while maintaining quality, the method has practical implications for real-time applications and resource-constrained environments. The validation across both image (ImageNet) and video (UCF101, Kinetics600) datasets underscores its versatility.\n\nWeaknesses\n\nThe nested architecture might introduce training complexities, such as balancing shared parameters across sub-models.\n\nGenerating using the method in this paper will result in a certain degree of degradation in generation quality. Is there a model with lower FID?", "other_comments_or_suggestions": "No other suggestions", "questions_for_authors": "No other questions", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes Masked Generate Nested Transformers with Decode Time Scaling (MaGNeTS), a novel approach for efficient visual generation that dynamically scales model size during decoding and utilizes KV caching with intermittent cache refresh. Through nested transformer architectures and compute-aware scheduling, MaGNeTS significantly reduces inference computational costs without compromising generation quality much, as shown across image and video benchmarks (ImageNet256, UCF101, Kinetics600). Reviewers particularly appreciate the strong empirical validation, practical impact, and generalizability of the proposed method. Initial concerns regarding KV caching performance trade-offs, qualitative analyses, and generalization beyond MaskGIT were effectively addressed in the authors' thorough rebuttal. Therefore, I recommend acceptance of this paper.
Decision: Accept (poster)