Learnings from Scaling Visual Tokenizers for Reconstruction and Generation

The Variational Auto-Encoder (VAE) (Kingma and Welling, 2013) is a framework that takes a visual input $X\in\mathbb{R}^{T\times H\times W\times 3}$ (where $T=1$ for images and $T>1$ for videos) is processed by an encoder $f_{\theta}$, parameterized by $\theta$. This encoder performs a spatial-temporal downsampling by a factor of $q\times p\times p$, producing a latent code. The encoder outputs parameters for a multivariate Gaussian distribution—mean $z_{m}$ and variance $z_{v}$ with $c$ channel size.:

The sampled latent vector $z$ is then fed into a decoder $g_{\psi}$, with parameters $\psi$, which reconstructs the input image $\hat{X}=g_{\psi}(z)$. The primary objective of the auto-encoder is to minimize the mean squared error between the reconstructed and original images, $\mathcal{L}_{\text{REC}}(\hat{X},X)$. To regularize the latent distribution to a unit Gaussian prior which is necessary to recover the variational lower bound, a KL divergence regularization term is added which we refer to as $\mathcal{L}_{\text{KL}}$. Recent advancements in VAEs used for downstream generation tasks (Esser et al., 2021; Rombach et al., 2022) incorporate additional objectives to improve the visual fidelity of the reconstructions. These include a perceptual loss based on VGG features (Johnson et al., 2016) $\mathcal{L}_{\text{LPIPS}}$ and an adversarial GAN objective, $\mathcal{L}_{\text{GAN}}$ (Goodfellow et al., 2014). The comprehensive loss function for the auto-encoder, $\mathcal{L}_{\text{AE}}(\hat{X},X,Z)$, is formulated as:

where $\beta$, $\eta$, and $\lambda$ are weights that balance the contribution of each term to the overall objective. We largely utilize the same overall loss, but ablate on the impact of each term in Section 3.4.

We now develop our visual tokenizer and pinpoint bottlenecks that we explore further in Section 3. The basic structure follows that of a variational auto-encoder (VAE) (Kingma and Welling, 2013) with an encoder-decoder architecture, but rather than relying on CNNs, we adopt a Vision Transformer (ViT) (Dosovitskiy et al., 2021) approach for better scalability. Our method builds on the ViViT framework (Arnab et al., 2021) to handle both images and videos. Specifically, a 3D convolution with kernel and stride size $q\times p\times p$ first tokenizes the input $X$ into a sequence $X_{\text{embed}}\in\mathbb{R}^{B\times L\times C_{f}}$, where $L=\frac{T}{q}\times\frac{H}{p}\times\frac{W}{p}$ and $C_{f}$ is the transformer’s feature dimension. A ViT encoder then processes $X_{\text{embed}}$, and a linear projection reduces the channel width to produce a compact representation $Z=f_{\theta}(X_{\text{embed}})\in\mathbb{R}^{B\times L\times 2c}$. Following the VAE formulation (Section 2), we recover $z\in\mathbb{R}^{B\times L\times c}$. We define

which effectively controls our compression ratio by specifying the total dimensionality of the latent space. As Section 3 highlights, $E$ is pivotal in predicting reconstruction performance. Both $c$ and $E$ are very important for generative performance as well. Though $E$ can be influence also by the number of tokens $L$, so we can potentially keep $c$ low while increasing $L$ for increased $E$.

For the decoder, a linear projection upsamples $z$ from $c$ to $C_{g}$ channels, after which a ViT decoder processes the tokens to predict $\hat{X}_{\text{embed}}$. Finally, a 3D transposed convolution recovers the original input resolution, producing $\hat{X}$. This covers the high level process of Vision Transformer Tokenizer or ViTok. Figure 1 illustrates this process. We denote ViTok configurations by specifying their encoder size, decoder size, and patch/stride parameters $(q,\,p)$. For instance, ViTok S-B/4x16 indicates a small encoder, a base decoder, and a patch stride of $q=4$, $p=16$. Table 1 provides details on the ViTok sizes.

We detail the training process for ViTok that will enable our exploration in Section 3.

In prior discrete cases performance depends on the number of tokens ($L$) and the size of the discrete codebook per token (Oord et al., 2017; Mentzer et al., 2023). For ViTok, the analogous factor is $E$ (Equation 2), which proves to be the critical determinant of reconstruction performance. The bottleneck $E$ is related to the number of pixels per floating point, $\frac{T\times H\times W\times 3}{E}$, representing the degree of compression applied.

To fully understand how $E$ functions as a bottleneck, we performed an extensive sweep through various configurations of ViTok investigating performance on 256p image reconstruction. For our first experiment, we look to explore all combinations of patch size $p=\{32,16,8\}$ and channel widths $c=\{4,8,16,32,64\}$ which gives various $E$ between $2^{8}$ to $2^{16}$. The patch size influences $L=\frac{H\times W}{p^{2}}$ and the amount of flops expended by the model due the quadratic nature of attention, while $c$ dictates the extent of the bottleneck between the encoder and the decoder. For these experiments, we fixed the encoder size to Small and the decoder to Base (Table 1). Our findings on scaling $E$ with 256p images are summarized in Figure 2. We provide more details and results in Appendix 8.

Figure 2 illustrates a strong correlation between $E$ and rFID/rIS/rSSIM/rPSNR. This indicates that $E$ is a significant predictor of the quality of the reconstruction, regardless of the shape of the code. Also, the behavior between different datasets reconstruction performance is similar with rFID changing slightly due to the size of the validation set difference (50k for ImageNet-1K vs 5k for COCO). Furthermore, for the same $E$, different patch sizes ($c=\frac{E\times p^{2}}{H\times W}$) yield similar performance. This suggests that increasing FLOPs for a fixed $E$ does not enhance performance, establishing $E$ as the most critical bottleneck in reconstruction performance for a given encoder. Figure 3 compares visualizations for different $E$ values on 256p images. As $E$ decreases, high-frequency details are lost, and when $E<4096$, significant texture and pattern information is degraded, although the overall image structure remains intact.

One potential source of concern is the precision of $E$ could effect reconstruction performance, therefore it should be shown via bits per pixel. We train ViTok S-B/16 at float32 precision and compare to bfloat16 precision in Table 2. There are almost no differences in performance, which shows that the precision of $E$ does not necessarily affect the reconstruction performance.

In this section, we investigate how $E$ influences performance in generative tasks by following the training protocol from Section 2.3 and using the same set of tokenizers evaluated in Figure 2. The results are in Figure 5.

The generative results exhibit a different trend compared to reconstruction, showing little to no linear correlation between $\log(E)$ and the generative metrics $\log(\text{gFID})$ or gIS. Figure 5 reveals that each patch size has an optimal $E$, leading to a second-order trend. The optimal configurations are $p=16$, $c=16$, $E=4096$; $p=8$, $c=4$, $E=4096$; and $p=32$, $c=32$, $E=2048$ for their respective patch sizes. Additionally, higher CFG settings tend to minimize the differences in gFID across various $E$ values. However, for gIS, higher channel sizes ($c>32$) and variants with poor reconstruction quality still result in poorer image quality, indicating that excessive channel sizes negatively impact performance despite CFG adjustments.

Closer analysis reveals that a low $E$ often bottlenecks the generative model, as the auto-encoder struggles with effective image reconstruction. Conversely, a high $E$, primarily driven by larger channel sizes ($c$), complicates model convergence and degrades both gFID and gIS metrics. These findings are corroborated by concurrent work that details a trade off between rFID and gFID in latent diffusion models (Yao and Wang, 2025). This highlights a critical trade-off in current latent diffusion models: $E$ and $c$ must be kept as low as possible to enhance generation performance while maintaining it high enough to ensure quality reconstructions. We provide generation visualizations for each tokenizer and trained DiT model in Appendix 9.

We aim to explore how scaling impacts auto-encoding in both reconstruction and generation tasks using ViTok. To test this, we fix the parameters to $p=16,c=16,L=256,E=4096$ for ViTok. We then conduct a sweep over different encoder and decoder sizesS-S, B-S, S-B, B-B, S-L, B-L, L-L defined in Table 1, following the same training protocol as described in Section 2.3. The results are reported in Figure 6 and 7.

As illustrated in Figure 6, the size of the encoder is not correlated with the reconstruction performance. In contrast, Figure 7 shows that the size of the decoder is positively correlated with the reconstruction performance. However, $E$ remains the dominant factor as doubling the decoder size does not provide the same effects as doubling $E$. For example, increasing the decoder size from Base to Large drops the rFID from 1.6 to 1.3 for $E=4096$, but doubling $E$ to $8192$ brings the rFID to 0.8 (Figure 2) for a decoder size Base. Overall, while scaling the decoder might be advantageous, scaling the encoder of a visual auto-encoder is unhelpful.

Figures 8 and 9 explore the effects of scaling the encoder and decoder on generation performance. In Figure 8, a slight negative correlation is observed between encoder size and generation results. This suggests that increasing the encoder size either has little to no impact on performance or may even detrimentally affect it, all while imposing additional computational burdens.

Similarly, Figure 9 shows that scaling the decoder exhibits minimal correlation with generation performance, indicating that enlarging the decoder offers limited benefits. Unlike reconstruction tasks, expanding the encoder or decoder does not significantly enhance generation quality; instead, it primarily increases training and inference costs. Notably, a 129M-parameter auto-encoder performs adequately (ViTok S-B/16), suggesting that future scaling efforts should focus on the generation model itself rather than the auto-encoder.

With the findings so far, we believe simply scaling the current auto-encoding (Esser et al., 2021) based tokenizers does not automatically lead to improved downstream generation performance. Therefore for generation, it is more cost-effective to concentrate scaling efforts on the generator itself, rather than the tokenizer.

As shown in Section 3.3, increasing the size of the decoder improves reconstruction, suggesting that the decoder behaves more like a generative model for the input $X$ and thus needs more computation than the encoder. To illustrate this, we compared how different losses balance traditional compression metrics (SSIM/PSNR) against generative metrics (FID/IS). SSIM/PSNR measure visual fidelity or how much of the original information is preserved, while FID/IS focus on visual quality and how closely outputs match the real dataset. This comparison shows how different choices of losses can shift the decoder’s role from strictly reconstructing to more actively generating content.

We conducted these experiments on ViTok by fixing $p=16$, $c=16$, and $E=4096$. We then trained with stage 1 and varied the LPIPS loss weight $\lambda\in\{0.0,0.5,1.0\}$ combined with the choice of L1 or L2 reconstruction loss (Equation 1). We also include our Stage 2 results following Section 2.3 to see the effect of the generative adversarial loss.

Figure 10 shows a clear trade-off among these losses. Without perceptual loss, we get worse rFID/rIS scores but better rSSIM/rPSNR, indicating that a strict MSE-based approach preserves the most original information. Increasing $\lambda$ gradually lowers SSIM/PSNR while improving FID/IS. Finally, fine-tuning the decoder with a GAN pushes these generative metrics further, achieving an rFID of 0.50 at the cost of lower SSIM/PSNR.

In addition, including the GAN also improves DiT’s downstream generation results. For instance, the Stage 1 model at $p=16$, $c=16$, and $E=4096$ reaches a gFID of 5.5 and a gIS of 160 at 500k steps with CFG=1.5 (Section 3.2). In comparison, the Stage 2 model achieves a gFID of 4.9 and a gIS of 210 at 500k steps, reflecting the same trade-offs seen in reconstruction. This underlines how strengthening the decoder’s generative capacity boosts overall performance in generation tasks, which makes the scaling benefits of decoders more complex than indicated in Section 3.3.

These results demonstrate that at a fixed $E$, aiming for higher visual quality requires sacrificing some traditional compression fidelity. This underscores that the decoder effectively acts as an extension of the generation model, creating visually pleasing results from the compressed representation. We provide more evidence of the decoder as a generative model as well as specific GAN ablations in Appendix 8.

We extend the application of ViTok to video tasks to examine the impact of $E$ on video reconstruction and to investigate redundancy in video data. To enable a direct comparison with our image results, we maintain a resolution of 256p and utilize 16-frame videos at 8 fps for both training and evaluation. Tokenizing videos can result in very large sequence lengths; for example, a tubelet size of 4×8 (with temporal stride $q=4$ and spatial stride $p=8$) for a video of dimensions 16×256×256 yields a sequence length of 4096 tokens. Therefore, based on our previous analysis of encoder and decoder sizes in Section 3.3, we use a small ViTok S-B variant to reduce computational burden, as $E$ is likely the more critical factor in this context.

To test how $E$ effects video we sweep over patch sizes $p\in\{8,16,32\}$ and temporal strides $q\in\{1,2,4,8\}$ following the protocol depicted in Section 2.3. As illustrated in Figure 11, the relationship between $E$ and the metrics rFVD/rFID mirrors the patterns observed in image tasks (Figure 2), where $\log(E)$ strongly correlates with reconstruction metrics. Figure 12 focuses on the UCF-101 dataset and demonstrates that, regardless of the selected spatial or temporal stride, $E$ remains the predominant factor influencing reconstruction performance. Consequently, adjusting spatial or temporal compression offers minimal advantages when $E$ is held constant for video reconstruction.

Comparing videos to images reveals that reaching similar rFID values requires $E\approx 16384$ to $E\approx 32768$ to achieve an rFID of 2.0, whereas for images $E=4096$ suffices. This difference, which is smaller than the naive 16$\times$ factor from frame-by-frame considerations, highlights that videos are more compressible than individual frames, and showing how ViTok can leverage this advantage.

We evaluate our models on image reconstruction and class-conditional image generation tasks using the ImageNet-1K (Deng et al., 2009) and COCO-2017 datasets at resolutions of 256p and 512p. For image reconstruction, we compare our continuous tokenizer-based models against several state-of-the-art methods, including SD-VAE 2.x (Rombach et al., 2022), SDXL-VAE (Podell et al., 2023), Consistency Decoder (OpenAI, 2023), and COSMOS (NVIDIA, 2024). It is important to note that discrete tokenizers present challenges for direct comparison with continuous tokenizers; therefore, our focus remains primarily on continuous tokenizers.

As shown in Table 3, our S-B/16 variant demonstrate highly competitive performance, achieving state-of-the-art (SOTA) rFID scores on both ImageNet-1K and COCO datasets. Furthermore, our models maintain competitive metrics in rSSIM and rPSNR. When scaling up the decoder size to L, the metrics improve further showing how scaling the decoder can be helpful in ViTok for reconstruction. Most importantly both variants of ViTok reduce the required FLOPs over prior CNN methods, which highlights the efficiency of ViTok. For the 512p image reconstruction results presented in Table 4, ViTok achieves SOTA reconstruction performance with a notable reduction in computational FLOPs over prior state of the art methods. In general, ViTok performs strongly in reconstruction benchmarks compared to prior methods in both FLOPs and performance.

Subsequently, we assess our auto-encoders on class-conditional image generation tasks at both 256p and 512p resolutions using the ImageNet-1K dataset. We follow the DiT training protocol outlined in Section 2.3, where we train a DiT-XL (675M parameter) model for 4 million steps paired with ViTok S-B/16 using 256 tokens for 256p generation and 1024 tokens for 512p generation. The results, summarized in Table 6, indicate that ViTok maintains competitive performance compared to the traditional SD-VAE trained with DiT and other continuous tokenizers in image generation. In 512p generation, ViTok performs on par with other methods, demonstrating ViTok’s efficacy at higher resolutions. Examples of generated images using our 256p and 512p tokenizers are illustrated in Figures 14 and 15, respectively.

For our video comparison, our reconstruction metrics are computed on the UCF-101 training set and compared against state-of-the-art methods including TATS (Ge et al., 2022), LARP (Wang et al., 2024), and MAGViTv1/v2 (Yu et al., 2023b, a). The results are presented in Table 5. Our tokenizers demonstrate very competitive performance relative to prior work. Specifically, S-B/4x8 ($1024$ tokens) achieves state-of-the-art (SOTA) rFVD results compared to other CNN-based continuous tokenizers with the same total compression ratio. When applying further compression, the rFVD metrics show a slight degradation; however, they remain highly competitive with existing methods. Notably, our S-B/8x8 ($512$ tokens) variant matches the performance of LARP (Wang et al., 2024), which operates with $1024$ tokens. Additionally, our approach significantly reduces FLOPs compared to Transformer-based prior method LARP, underscoring the efficiency and versatility of ViTok.

We further evaluate our models on class-conditional video generation using the UCF-101 dataset. We train a DiT-L model across all compression variants for 500K steps on the UCF-101 training set, computing gFID and gFVD metrics with a batch size of 256 and a learning rate of $1\times 10^{-4}$. The results are summarized in Table 7. ViTok achieves SOTA gFVD scores at $1024$ tokens and maintains highly competitive gFVD scores at $512$ tokens ($\times 8$ by $\times 8$ compression), representing the most efficient level of token compression for any tokenizer so far. At $256$ tokens, ViTok’s performance experiences a further decline but remains competitive within the field. Example video generations using our $1024$-token configuration are illustrated in Figure 16.

We provide further detail of the ImageNet-1K validation reconstruction results from Figure 2 in Figure 17. Here we show different patch sizes and channels over $E$. This shows that regardless of patch size and FLOPs usage, $E$ is highly correlated with the reconstruction perforance

In Figure 18, we study how various loss settings affect finetuning of the GAN decoder. Our goal is to highlight the trade-off and the decoder’s transition toward more generative behavior. We use ViTok S-B/16 on 256p images, following the protocol in Section 2.3 for stage 2 fine-tuning from a model trained on stage 1.

We compare:

• •

  Finetuning the decoder with the same Stage 1 loss (no GAN).

• •

  Finetuning with discriminator learning rates ($\{1\times 10^{-5},2\times 10^{-5},4\times 10^{-5}\}$) and a GAN weight of 0.05.

• •

  Finetuning the full encoder/decoder with the GAN.

• •

  Using a higher GAN weight of 0.1 with a discriminator learning rate of $1\times 10^{-5}$.

From Figure 18, the best setting is a GAN weight of 0.05 and a discriminator learning rate of $2\times 10^{-5}$. A higher discriminator learning rate causes training instabilities, while a lower rate degrades performance. Full finetuning works but does slightly worse than just finetuning the decoder. Finetuning without a GAN shows no improvement, confirming that GAN training is the primary driver of better results.

Finally, we see an inherent trade-off: improving FID tends to worsen SSIM/PSNR, indicating that as the decoder focuses on visual fidelity, it shifts more toward generative outputs. This demonstrates the decoder’s evolving role as a generative model to enhance visual performance.

In this section, we describe two variants of ViTok that provide different potential directions for tokenization. First we describe and evaluate our latent variation that does 1D tokenization and can form more arbitrary code shapes, then we discuss and evaluate our masking variant that allows for variable, adaptive tokenization.

We include more video generation results in this section from Table 7 and show example generations at 512 and 256 tokens respectively.

Here provide generation examples from our sweep conducted in Figure 22. $p=16$ visuals are in Figure 23, $p=32$ visuals are in Figure 24, and $p=8$ visuals are in Figure 25.