ENAT: Rethinking Spatial-temporal Interactions in Token-based Image Synthesis

W1: Relationship with MAE, CrossMAE

Thanks for your valuable suggestion. We are happy to include more discussions in our revision.

Firstly, we kindly note that our approach includes two key aspects: a disentangled architecture and a computation reuse mechanism. To the best of our knowledge, the latter is novel in NAT and highly effective, yet has not been explored in MAE/CrossMAE since they focus on representation learning, which does not involve a multi-step process that exists in generation tasks.

Secondly, while our disentangled architecture may appear similar to MAE and CrossMAE, it contains subtle yet important differences due to the distinction between representation learning and generation tasks. Specifically, MAE inputs all tokens into the decoder and employs full attention on them ([M] & [V] → [M] & [V]), and CrossMAE retains a single "[V] → [M]" information flow. In contrast, we adopt a "[M] & [V] → [M]" information flow in the decoder attention:

The results show that directly using MAE-type full attention merely adds computational overhead without improving performance, whereas implementing CrossMAE-type attention degrades performance notably.

This also aligns with our observation in Sec. 4.1 regarding generation tasks:

1. Visible tokens primarily provide information. Since the "[V] → [M]" component in our "[M] & [V] → [M]" attention already fulfills the information-providing role of visible tokens, additionally incorporating the visible tokens into the decoder input for full attention (as in MAE) provides no benefit.

2. The attention between masked tokens cannot be ignored (as opposed to CrossMAE). Intuitively, this allows masked tokens to be aware of each other during decoding, thus preventing incongruencies in the generated results. This further highlights the distinction between a generation-focused approach and a representation learning-focused approach.

Finally, we would like to clarify that our contributions are not limited to proposing a new architecture/mechanism. We have revealed several important scientific findings and insights in NATs, which are also novel contributions. We believe these findings can provide new perspectives for the research community.

W2: Novelty of SelfCrossAttention Module

Thanks for pointing out the related work [D], and we are happy to update our revision with this reference.

Additionally, we would like to kindly clarify that the SC-attention, as briefly introduced in Section 4.1, is not presented as our main contribution. Our primary contribution lies in uncovering new scientific findings in NAT generation, specifically regarding its spatial and temporal properties, that enhance our understanding and enable a better ENAT design.

W3 & Q2: Choices and Scores of Baseline Models

Thanks for your suggestion. We are happy to add GAN-type models like StyleGAN-XL in our revision, and we have already incorporated it in Fig. 3 of our global response.

For the choice of baselines in our work, since ENAT focuses on inference efficiency, we aim to compare ENAT with other models in a lightweight, low-FLOPs scenario. However, while the inference efficiency of generative models is important, it is generally under-explored in the original papers of state-of-the-art diffusion models (e.g., DiT, MDT, SiT), which mostly focus on enhancing generation performance. The official results of them are primarily obtained with hundreds of inference steps, making direct comparisons with ENAT challenging. For instance, as shown in Fig. 3, the official results of DiT, MDT, and SiT all concentrate at the high end of overall inference costs, requiring hundreds of times more computation than ENAT.

Fortunately, there are well-established fast sampling techniques (e.g. DPM-Solver) for accelerating diffusion models, and we have identified several diffusion models that are both highly competitive and officially supported* by these acceleration techniques. This allows us to reduce their sampling steps and compare them with ENAT in a fairer setting.

For the mentioned stronger contenders, we plotted MDT and SiT in Fig. 3. Simple diffusion is omitted because its computational cost is neither reported nor open-sourced for measurement. The strongest MDT model is also integrated with fast sampler for fewer sampling results. Few-step fast sampler results from the original U-ViT paper are also included for reference.

In our revision, we are willing to add more discussions on our baseline results and setups. Stronger baseline models like MDT and SiT will also be included.

---

*: U-ViT has DPM-Solver integrated in their Github repo. DiT and LDM are in the Hugging Face diffusers library, where DPM-Solver is also officially supported.

W4: More Qualitative Results

Thanks for your suggestion. In our global response, we have added qualitative comparisons with DiT-XL in Fig. 4 and included many more qualitative samples for all evaluated datasets in Fig. 5 & 6. More samples will also be added in our revision.

Q1: U-ViT as a Worse Starting Point?

We would like to clarify that ENAT only follows the training pipeline of [2], rather than its diffusion transformer model architecture. The pipeline of [2] is relatively standard, and our goals are twofold: 1) to validate the gains solely attributable to our design under fair conditions, and 2) to demonstrate that our method does not rely on specific training techniques. Exploring more advanced training techniques could potentially lead to further improvements, and we are happy to explore this in the future.

Q3: FID Evaluation Protocol

We follow the widely adopted OpenAI's evaluation protocol to compute FID against the ImageNet training set.

Thank you for the detailed answers and important clarifications.

I don't really have major concerns left, and will update my rating accordingly.

W1: How the Projection Module $f$ is Trained?

As discussed in line 221-224, we execute the "reuse mode forward" (Eq. (5)) with some probability during training, thus involving the parameters of $f$ in the computation graph. More specifically, the "reuse mode forward" during training is achieved through the following steps:

1. Masking Tokens: Given the current input token map $\mathbf{v}$, we mask a random subset (e.g., 30%) of visible tokens in $\mathbf{v}$ to create $\mathbf{v}^{\text{prev}}$.
2. Feature Extraction: Feed $\mathbf{v}^{\text{prev}}$ into the NAT model to obtain its features $\mathbf{z}^{\text{prev}}$.
3. Forward Pass and Loss Computation: Use Eq. (5) to forward the current input token map $\mathbf{v}$ along with $\mathbf{z}^{\text{prev}}$ obtained in Step 2, and compute the loss. In practice, a stop gradient operation is applied before feeding $\mathbf{z}^{\text{prev}}$ into Eq. (5).

By following Steps 1 and 2, we construct $**z**^{\text{prev}}$ and feed it into the projection module $f$ in Step 3, thus involving $f$ in the forward propagation. Since all operations are differentiable, gradients can then be effectively back-propagated to update the parameters of $f$.

We will elaborate further in our revision to make this process clear.

W2 & W3: Scalability of ENAT

In line with findings from other token-based models [A, B], ENAT is also scalable with increased model size. Due to limitations in time and computational resources, we conducted a preliminary experiment using a 1B parameter version of ENAT, referred to as ENAT-XL (28 layers, 1280 width):

As the results show, scaling up ENAT can further improve performance.

An interesting finding is that, the performance ENAT-XL is already approaching the reconstruction FID of our VQ-autoencoder, which measures the quality of images decoded from ground truth visual tokens (thus being an indicator of the “performance upper-bound” for generative models in this VQ space).

Therefore, for even further improvements over ENAT-XL, we may consider simultaneously scaling up the generative model and the VQ-autoencoder, which we are actively working on.

---

[A] Tian, Keyu, et al. "Visual autoregressive modeling: Scalable image generation via next-scale prediction." arXiv preprint arXiv:2404.02905 (2024).

[B] Chang, Huiwen, et al. "Muse: Text-To-Image Generation via Masked Generative Transformers." International Conference on Machine Learning. PMLR, 2023.

W4: More Visualization Results Revealing How the Method Works

Thanks for your suggestion. We have included additional visualization results in our global response.

Specifically, Fig. 1 visualizes the intermediate generation results throughout the process. The results suggest that ENAT generates superior images as it converges to decent quality results faster, and continues to refine for even better image quality till the end.

To further corroborate this point, we conducted a "one-step decoding" experiment. In this experiment, we applied various mask ratios (simulating inputs at different stages of generation) to ground truth image token maps and measured the quality of the generated images using FID for one-step decoded results (i.e., the model directly predicts all masked positions in one forward pass).

As the results indicate, ENAT consistently outperforms the standard NAT design. The results for the 0.2 mask ratio setting are also visualized in Fig. 2. This demonstrates that ENAT's predictions for masked positions are more reliable at each decoding step, allowing it to converge to high-quality results faster and produce superior images.

Finally, we have included more qualitative comparisons and samples in Fig. 4, 5, and 6 to further validate the effectiveness of our method.

W1: Relationships with MAE and MAGE

Firstly, we kindly clarify that our approach includes two key aspects: a disentangled architecture and a computation reuse mechanism. To the best of our knowledge, the computation reuse mechanism is a novel contribution that has not been explored in NAT, and has demonstrated great promise and effectiveness.

Secondly, while our "disentangled architecture" may appear similar to MAE & MAGE, they are in fact different. Specifically, MAE & MAGE view the idea of encoding only visible tokens as a way to accelerate the training process for representation learning. In contrast, our disentangled design focuses on building an efficient architecture tailored for image generation. The distinct objective results in small but important methodological differences, primarily manifested in the decoder's attention. Specifically, both MAE and MAGE adopt full attention on all tokens ([M] & [V] → [M] & [V]) in their decoder. On the contrary, the decoder in ENAT only inputs the mask token and integrates the visible ones via SC-Attention, creating a tailored "[M] & [V] → [M]" information flow in the decoder. Experimental results show that naively adopting MAE/MAGE-type full attention only increases computational overhead without enhancing performance:

This aligns with our observation in Section 4.1: In NAT generation, visible tokens primarily provide information for mask tokens. Thus, the [V] → [M] part in [M] & [V] → [M] has adequately fulfilled the role of the visible token, making it unnecessary to incorporate the visible token into the decoder input and perform full attention.

We will add more discussions to clarify the relationships with these works in our revision.

W2: Intuition for Enabling Smaller Sampling Steps

The key intuition behind our better performance with smaller steps is that: each decoding step in ENAT is more effective. To illustrate this, we conducted an experiment by applying different mask ratios (mimicking inputs in different steps of generation) to ground truth images and measuring the quality of one-step decoded results (i.e., the model directly predicts all masked positions in one forward pass) using FID.

As the results show, ENAT consistently outperforms the baseline NAT design. This suggests that ENAT's predictions for masked positions are more reliable (i.e., of higher quality) at every decoding step, allowing them to take larger strides (decode more tokens per step) and producing high-quality generation with fewer refinement steps.

The qualitative results for the 0.2 mask ratio setting and the intermediate generation results of ENAT are also visualized in Fig. 2 and Fig. 1 of our global response to further illustrate this point.

W1: Questions on MSCOCO Results

Firstly, we kindly clarify that the two models, U-ViT and ENAT-B, may not be viewed as being "on the same level." Our primary objective is to improve inference efficiency. Therefore, FLOPs, as an important indicator of practical efficiency (see Fig. 6 of our original paper), is much more critical than model parameter count. Moreover, the increase in model parameters is in fact an intentional result of our disentangled architecture, which significantly improves performance without affecting efficiency (see Tab. 6a of our original paper).

Secondly, we have only trained our ENAT model on COCO text-to-image generation for 150K steps (see A.2 of our original paper) as it already outperforms U-ViT trained with nearly 7$\times$ more training budget (i.e. 1000K steps). In a fairer comparison where U-ViT is also trained for 150K steps, ENAT is still significantly better:

We will add more discussions on this in our revised experiment section.

Q1: Elaboration on ENAT's Training Process

Sure, the detailed process of training in reuse mode (Eq. (5)) is elaborated below:

1. Masking Tokens: Given the current input token map $\mathbf{v}$, we mask a random subset (e.g., 30%) of visible tokens in $\mathbf{v}$ to create $\mathbf{v}^{\text{prev}}$.
2. Feature Extraction: Feed $\mathbf{v}^{\text{prev}}$ into the NAT model to obtain its features $\mathbf{z}^{\text{prev}}$.
3. Forward Pass and Loss Computation: Use Eq. (5) to forward the current input token map $\mathbf{v}$ along with $\mathbf{z}^{\text{prev}}$ obtained in Step 2, and compute the loss. In practice, a stop gradient operation is applied before feeding $\mathbf{z}^{\text{prev}}$ into Eq. (5).

Since $\mathbf{v}^{\text{prev}}$ created in Step 1 contains less visible tokens than $\mathbf{v}$, it effectively simulates the "previous step input" that ENAT will encounter during generation, thereby equipping ENAT with the ability to correctly handle previous step's feature for reuse.

We will elaborate this process in more detail in our revision.

Q2: Setups for Spatial Disentanglement Experiments

Yes, we intentionally kept the FLOPs unchanged for a clear comparison. However, we achieved this by enlarging the network width, which we will further clarify in our revised manuscript. Enhanced performance at the same FLOPs generally equates to lower computational cost for the same performance, thus still demonstrating our method's efficiency.