AsCAN: Asymmetric Convolution-Attention Networks for Efficient Recognition and Generation

We thank the reviewer for the positive feedback on this work.

As already pointed out by the reviewer, ablations on large-scale text-to-image generation tasks are computationally very expensive. Hence, we performed small ablations on the ImageNet-1K dataset. We list these ablations below. We will include a detailed discussion of these characteristics in the final version.

Training algorithm ablations.

• We analyzed the performance of the model with and without pre-training on the ImageNet-1K task. Appendix Figure 8 shows the evolution of the FID score with the number of training iterations. It shows that training the model without pre-training on the ImageNet-1K task results in slower convergence, and the final FID achieved by the model without pre-training is worse than the one achieved by pre-training on this task. Thus, we demonstrate that pre-training on this task helps speed the convergence.

• We analyzed the importance of noise levels between $256\times256$ and $512\times512$ resolution stages in Appendix Figure 7. It shows that different resolutions prefer different input noise for the diffusion process. It helps us decide the use of noise=0.01 for $256$ resolution and noise=0.02 for $512$ resolution. We also ablated on other hyper-parameters (learning rate, weight decay) while training on $256\times256$ and $512\times512$ resolution stages.

Architecture ablations.

• We train on the class conditional $256\times256$ generation on ImageNet-1k. As seen in Table 3, our asymmetric architecture achieves a similar FID score at less than half FLOPs under the same training setup as other baselines. We provide additional results on this task in the rebuttal pdf (see Table 1), including inference latency and generation on the $512\times512$ resolution. It shows that our FLOP gains also translate to the inference latency gains.

We appreciate the reviewer for reading the paper thoroughly and providing invaluable feedback. Below, we have tried to answer their questions and concerns. While we answered some questions in the main rebuttal, we reiterate their highlights for completeness.

Originality: Limited Novelty.

• Our main contribution is the asymmetric distribution of convolution and transformer blocks in various stages in a hybrid architecture.
• We show that simple design choices (Sec. 3.1) yield architectures with existing blocks that achieve state-of-the-art performance and latency trade-offs. We demonstrate that this design is easily applicable to various applications.
• Many works design architectures based on parameter count and FLOPs, but this typically does not translate into inference throughput gains. Some of these issues come from using operators that do not contribute to parameter count and FLOPs but require non-trivial runtime, such as reshape, permute, etc. Others originate from the lack of efficient CUDA operators for these specialized attention and convolutional operators.
• In contrast, our proposal directly measures the inference throughput on different accelerators and incorporates building blocks that yield higher throughput.

Quality: Justification of C before T constraint. Our intuition behind this constraint comes from the following observations:

• Interleaving C and T in a stage performs many tensor reshape operations that do not add any FLOPs. However, these operations count towards runtime, and many such operations lower the model throughput.
• Given a feature map, C blocks can capture local and scale-aware features, while T blocks try to work out the dependency between all feature values. Thus, it would be more beneficial to perform a convolutional operator first to capture these local and scale features, followed by pairwise dependencies between all tokens.
• It helps to narrow the search process since interleaving these blocks would result in many possibilities and be hard to evaluate computationally.
• Further, to validate our assumptions, we have included configurations where T appears before C in Table 2 in the rebuttal pdf. In the models where the first stage consists of T, the throughput is significantly lower than in instances where C is the first stage. Similarly, models with T before C do not achieve similar accuracy vs performance trade-off as the configurations where C appears before T.

Quality: Missing throughput in the class conditional generation. Thank you for pointing this out. We benchmark the throughput (images generated per second) on an A100 GPU for all the baselines. In the attached rebuttal pdf, table 1 shows the throughput for one forward pass of batch size $64$ for all the models. To achieve an FID score of $2.23$, our $52$G FLOPs model achieves a $556$ samples/sec throughput while state-of-the-art DiT-XL/2-G with $118$G FLOPs achieves $293$ samples/sec. It shows that our asymmetric model still has nearly double the throughput compared to other models while achieving similar FID.

Clarity. We have included your comments in the manuscript. These will be reflected in the final version.

Significance.

• EfficientNetV2-M vs AsCAN-L. There seems to be a bit of confusion in reading EfficientNetV2-M numbers. This model has been trained with an input resolution of $384\times384$ and evaluated at an input resolution of $480\times480$. Thus, even with a smaller parameter count, it has much higher FLOPs than other models with nearly $50$M parameter range. Further, all the hybrid models have been trained and evaluated at $224\times224$ resolution. For instance, with train/test input resolution as $224$, to achieve close to $85.1\%$ top-1 accuracy, MaxViT-L requires 212M parameters, FasterViT-4 requires 425M parameters, MOAT-3 requires 190M parameters. When we follow the EfficientNetV2-m training strategy, AsCAN-L achieves $86.2\%$ top-1 accuracy, while the larger EfficientNetV2-L variant with $120$M parameters achieves $85.7\%$ top-1 accuracy.
• How much do the results depend on hardware? The impact of C and T blocks in a hybrid architecture results in non-linear behavior across accelerators and batch sizes. This is precisely the reason we do not rely too much on the number of floating point operations to estimate inference latency. We choose 16GB V100 and 80GB A100 GPUs as representatives of two popular RAM and accelerator designs. We expect the trend on H100 to be similar to A100 and other lower-end GPUs (such as A10G, L4, etc.) to be similar to V100. While currently we do not have access to these accelerators, we will try to include benchmarks on these accelerators in the final version.
• T2I Speed Up. PixArt-$\alpha$ is a purely transformer architecture while AsCAN is a hybrid architecture involving convolutional blocks in early stages. As we have observed in our ablative experiments over C and T distribution in Table 2 as well as configurations where T appears before C in rebuttal pdf, transformer layers in the early part of the network significantly reduce the throughput. We will add further investigations into PixArt-$\alpha$ in the final version.

Questions: User study review process. Thank you for noticing this. We did get approval for the user study. We did not disclose these details to preserve anonymity. We will disclose the review process in the final version.

Questions: NSFW evaluation. We used a set of internal prompts that are suggestive in nature, i.e., they do not explicitly ask for generating unsafe images, but rather hide these details in words. We generated images with SDXL and our model. We compared the amount of NSFW images in these two models. Almost all of the SDXL generations are NSFW while our generations are safe and do not include any nudity.

Limitations. We have included the inference memory consumption in Table 3. We will include the discussion on parameter aspects in the final version.

Dear Reviewer R6Bu,

Thank you very much for your valuable feedback and the positive evaluation of our work. We have included detailed explanations in response to your questions. As the deadline for the discussion period approaches, we would appreciate your review of these explanations to confirm that they resolve any remaining concerns. Let us know if you have any other questions.

Thank you once again for your insightful review.

Best regards, Authors

Thank you very much for your willingness to address my comments. I will read the changes in the main text as soon as they become available, and update my review accordingly.

Thank you very much for the provided discussion, I updated my rating assuming that this analysis will be part of the final manuscript.

We thank the reviewer for their feedback. We provide additional clarifications in this response to the best of our understanding of the reviewer's concerns.

(Q.1) Lack of theoretical analysis. We propose a new design for hybrid convolutional-transformer architectures and apply this new architecture to different applications (image recognition, class conditional generation, text-to-image generation, etc.). We show that our proposal achieves significantly better performance vs latency trade-offs. We analyze the throughput on various accelerators and show the theoretical FLOPs required by all the models. Our asymmetric architectures consist of transformer blocks with a quadratic complexity in the input sequence length. It would be good if the reviewer could elaborate on the theoretical analysis required in this neural network architecture design.

(Q.2) Marginal performance enhancements. It is unclear why the reviewer thinks performance enhancements are marginal in this work. Other reviewers have already acknowledged that performance improvements are non-trivial (see strengths mentioned by reviewers 3CqY, R6Bu, hpXB, and fxge). To be concrete,

• For image recognition tasks (see Figure 3, Table 7), AsCAN architectures achieve similar accuracy with up to $2\times$ throughput increase compared to other architectures such as MaxViT, FasterViT, ConvNeXt, CoAtNet, etc. We have incorporated other related works in the rebuttal pdf (see Table 3 and Figure 1) that show similar gains in performance.

• On class conditional generation (see Table 3), our architecture achieves similar FID scores with less than half the FLOPs. We have included latency and $512\times512$ generation tasks in the rebuttal pdf (see Table 1).

• Similarly, on the Text-to-Image generation task (see Table 4, Table 5, Table 11), our models achieve much better resource efficiency and image generation quality when compared to existing baselines.

(Q.3) Analysis on training characteristics. We have performed experiments to ablate and improve the training quality.

• We analyzed the model performance with and without pre-training on the ImageNet-1K task. Appendix Figure 8 shows the evolution of the FID score with the number of training iterations. It shows that training the model without pre-training on the ImageNet-1K task results in slower convergence, and the final FID achieved by the model without pre-training is worse than the one achieved by pre-training on this task.

• We analyzed the importance of noise levels between $256\times256$ and $512\times512$ resolution stages in Appendix Figure 7. It shows that different resolutions prefer different input noise for the diffusion process. It helps us decide the use of noise=0.01 for $256$ resolution and noise=0.02 for $512$ resolution.

Thanks for your rebuttal. In Tables 1 and 2, your architectures had significantly reduced the inference time on GPUs, having marginal performance enhancements. I understand the contributions in terms of reduced computations in GPUs. The ratio of enhanced throughput with batch=1 seems to be small compared to the case with batch= 16, and 64 in Table 7. I think that the reason of the throughput enhancements on GPUs could be explained in a structural point of view.

We thank the reviewer for reading the rebuttal and responding promptly. We would like to clarify some remarks below.

• We appreciate that the reviewer acknowledges we achieved significantly reduced inference time. In all our experiments, we analyze the trade-off between performance (top-1 accuracy / FID score) and latency (throughput measure in images processed per second) achieved by the architectures. We would like to humbly clarify that we do not claim that we significantly improve performance. Instead, we claim to achieve significantly better performance vs. latency trade-offs. This can be observed by our empirical evaluations (see Table 7 in the main text, Figure 1, Table 1, and 3 in the rebuttal pdf), where we obtain significantly better throughput to achieve similar performance as other models. Vice-versa, we achieve better performance at the same throughput (see Figure 1 in rebuttal pdf). Thanks for the comments from the reviewer. We will further improve the writing of this paper to make it more clear about the claims of this paper.

• Since we focus on designing architectures suitable for GPUs, we provide throughput numbers for different batch sizes (B=1, 16, 64). At this large scale, batched inference (with B>1) makes a lot more sense since lower batch sizes (e.g., B=1) do not utilize the GPU memory fully and end up returning a very non-linear behavior. Even at B=1, we still achieve much better performance vs. throughput trade-off than many baselines.

We hope the above response could help address the concern of the reviewer. Please let us know if the reviewer has other questions and we would be very happy to help answer.

We appreciate the reviewer for reading the paper thoroughly and providing invaluable feedback. Below, we have tried to answer your questions and concerns. While we answered some questions in the main rebuttal, we reiterate their highlights for completeness.

Explanation of higher throughput while maintaining or improving accuracy.

• Many existing works design architectures based on parameter count and FLOPs, but this typically does not translate into inference throughput gains. Some of these issues come from using operators that do not contribute to parameter count and FLOPs but require non-trivial runtime, such as reshape, permute, etc. Others originate from the lack of efficient CUDA operators for these specialized attention and convolutional blocks.

• For instance, MaxViT consists of MBConv and Axial-attention transformer blocks. MBConv block is more friendly for mobile devices due to separable convolutions but hurts throughput on high-end GPU accelerators. Similarly, Axial-attention heavily invokes the permute operations. It results in additional overhead and hurts throughput.

• In contrast, we utilize building blocks that are efficient for high-end GPU accelerators, and our experimental design directly measures the inference throughput on different accelerators to search for optimal performance vs throughput trade-offs. This search favors asymmetric design with more convolution blocks in the early stages and more transformer blocks in the later stages. This helps reduce the occurrence of reshape and permute operations appearing in instances where C and T blocks are interleaved repeatedly.

Scaling to larger architectures. There are many strategies to scale the asymmetric design. In one direction, we can train our existing models with larger input resolution to achieve better top-1 accuracy. For instance, training AsCAN-L trained with $384$ resolution results in $86.2$% top-1 accuracy, compared to training with $224$ resolution which yields $85.2$%. On the other direction, we can scale these base architectures similar to previous works (CoAtNet, MaxViT, EfficientNetV2, etc.) by proportionally scaling the width and block repetitions in the stages (S0, S1, S2, S3). For instance, we can scale to AsCAN-XL variant by increasing the number of C and T blocks as (CC, C$^{4}$T$^{2}$, C$^{8}$T$^{8}$, C$^{4}$T$^{8}$), resulting in a $340$M parameter model that achieves $86.7$% top-1 accuracy. We can similarly scale these configurations. We can obtain even larger models by scaling the stages S1, S2, and S3.

Scaling Text-to-Image Generation w.r.t. data size. While it would be good to understand the model performance scaling w.r.t. data size, unfortunately, we do not have compute resources to train multiple such text-to-image generative models at any higher data scale than $450$M image-text pairs.

Class conditional generation with $512\times512$ resolution. We trained our small UNet variant on the $512\times512$ generation task. We report its FLOPs, throughput, and FID scores in the rebuttal pdf (see Table 1). On this task, DiT-XL/2 with $525$G FLOPs and $51$ throughput achieves an FID score of $3.04$. Similarly, U-ViT-H/2 with $546$G FLOPs and $45$ throughput achieves an FID score of $4.05$. In contrast, our smaller model with $224$G and $130$ throughput achieves an FID score of $3.15$. Thus, it has more than twice the throughput with a similar FID score.

Class conditional $256\times256$ generation with 100G FLOPs. We created an additional class conditional generation model on $256\times256$ resolution with an asymmetric model using 103G FLOPs. It achieves a throughput of $360$ samples/sec with an FID score of $2.08$ while DiT-XL/2-G model achieves a throughput of $293$ samples/sec with an FID score of $2.27$. We will include higher FLOPs models in the final version.

Design choice ablations. We have included ablative results for configurations with T before C as well as where the first stage consists of T blocks (see Table 2 in rebuttal pdf). In the instances where the first stage consists of T, the throughput is significantly lower than in the instances where C is the first stage. Similarly, the configurations with T before C do not achieve similar accuracy vs performance trade-off as the configurations where C appears before T.

Generalizability of architecture search on ImageNet-1k. We use the ImageNet-1K task for the architecture search problem since it is computationally cheaper for ablations when compared to other tasks and is still representative enough for the vision domain. To show its benefits, we use this simple design in other tasks without any task-specific optimizations. In the class-conditional generation, our asymmetric architecture achieves similar FID as various state-of-the-art models with less than half the FLOPs (translating into half latency). We see similar improvements in other tasks like Text-to-Image generation. Thus, we can conclude that the search on the ImageNet-1k task generalizes to other tasks. Indeed, we would expect to achieve even better performance once we optimize these modules per task, but this process is computationally expensive.

Clarity in terminology. Thanks for pointing out the formatting and terminology issues. We will address these in the final version.

Dear Reviewer 3CqY,

Thank you very much for your valuable and constructive feedback. We have included detailed explanations and additional experiments in response to your questions. As the deadline for the discussion period approaches, we would appreciate your review of these explanations to confirm that they fully meet your expectations and resolve any remaining concerns.

Thank you once again for your insightful review.

Best regards, Authors

We appreciate the reviewer for reading the paper thoroughly and providing invaluable feedback. Below, we have tried to answer your questions and concerns. While we answered some questions in the main rebuttal, we reiterate their highlights for completeness.

Stem Design Choice. This is a popular stem design and has been utilized in related works such as MaxViT, CoAtNet, etc. We have also included ablations with some configurations that replace convolution blocks with transformers.

Rationale behind the asymmetric design. Existing hybrid architectures such as MaxViT, FasterViT, CoAtNet, etc., follow a symmetric design in which C and T blocks are uniformly distributed across stages or within a stage. In contrast, AsCAN architectures recommend an asymmetric distribution with a preference for more convolutional blocks in the early stages and more transformer blocks in the later stages. Further, since even this search space is large for a reasonable computational budget search, we restrict ourselves to some design choices to find reasonably performing architectures in this search space. In Table 2, we perform ablations on various combinations of these C and T blocks and show that asymmetric configurations yield better performance-latency trade-offs. It is further evaluated in Figure 3 and Table 7, which compare AsCAN with existing hybrid architectures on various accelerators.

Why modules optimized for recognition tasks generalize on generative tasks. We use the ImageNet-1K task for the search problem since it is computationally cheaper for ablations when compared to other tasks and is still representative enough for the vision domain. To show its benefits, we use this simple design in other tasks without any task-specific optimizations. In the class-conditional generation, our asymmetric architecture achieves similar FID as various state-of-the-art models with less than half the FLOPs (translating into half latency). We see similar improvements in other tasks like Text-to-Image generation. Thus, we can conclude that the search on the ImageNet-1k task generalizes to other tasks. Indeed, we would expect to achieve even better performance once we optimize these modules per task, but this process is computationally expensive.

Architectural results presented in Table 2 can be considered expected outcomes based on the known knowledge.

• We already compare against SMT models ([2] Scale-Aware Modulation Meet Transformer, SMT-B, and SMT-S variants), see Figure 3 and Table 7 for comparison. For instance, AsCAN-B architecture achieves $84.73\%$ top-1 accuracy with a throughput of $590$ samples per second on a V100 GPU, while SMT-B architecture achieves $84.3\%$ top-1 accuracy with a throughput of $243$ samples per second. Notice that many operations in SMT architectures are not high-end accelerator friendly and thus their lower FLOPs do not translate to higher throughput. Further, SMT architecture has a symmetric distribution of convolution and transformer blocks (initial stages have only C and later stages only have T).

• Similarly, MobileNetV4 architectures are targeted toward mobile device applications and hence focus on layers that are mobile-friendly. For instance, having depth-wise separable convolutions does not utilize the full accelerator capacity on high-end GPUs. Further, this architecture only leverages transformer blocks sparingly in a few configurations.

• Besides, we already compared against many hybrid architectures that have such a design (either C followed by T in different stages) or uniform mixing of C and T blocks within a stage. Instead, we are proposing an asymmetric distribution of these blocks in different stages.

Provide memory consumption for the proposed architectures. We have included the inference memory consumption along with the numbers for the remaining new architectures requested by the reviewer. As illustrated in Table 3 (see attached rebuttal pdf), AsCAN architectures have lower memory consumption compared to other architectures. For instance, for batch-size $64$, the MogaNet-XL model consumes $74.9$GB memory while AsCAN-L consumes $21.2$GB memory. This is still lower than the memory consumed by RDNet-L model which is a purely convolutional architecture and requires $26.2$GB memory.

Compare with recently proposed works. Thank you for listing the missing related works. We have included a performance vs throughput comparison with these works in Table~3 in the attached rebuttal pdf. We have also included these architectures in Figure 1 for the performance-latency trade-off on V100 and A100 GPUs. In this list, there are architectures that have significantly lower FLOPs but this does not result in significant gains in the throughput when measured on high-end accelerators. This is due to a large presence of operators such as permute which do not add floating point operations but require additional runtime. In contrast, AsCAN architectures still outperform these works on inference latency vs top-1 accuracy trade-offs.

Performance tradeoff like Fig. 3 on an A100 GPU. Thank you for raising this point. We already report throughput and performance in the Appendix (Table 7). We have included the requested performance trade-off figure for an A100 GPU in the rebuttal pdf. This follows a similar trend as the one for V100 GPU.

Confusion in Transformer block naming. Thank you for pointing this out, we will update the terminology in the paper.

Lack of held-out set in ImageNet experiments. We follow previous works in the experimental setup.

Semantic Segmentation results missing recent architectures. Thank you for the feedback. We will include other results for semantic segmentation in the final version.

Dear Reviewer hpXB,

Thank you very much for your valuable and constructive feedback. We have included detailed explanations and additional experiments in response to your questions. As the deadline for the discussion period approaches, we would appreciate your review of these explanations to confirm that they fully meet your expectations and resolve any remaining concerns.

Thank you once again for your insightful review.

Best regards, Authors

Dear Reviewer hpXB,

We would like to thank you again for your valuable feedback on our paper.

As the period for the Author-Reviewer discussion is closing very soon, we would like to use this opportunity to kindly ask if our responses sufficiently clarify your concerns. We sincerely appreciate your time and consideration.

Best Regards, Authors

Sorry for the very late reply to the responses. I greatly appreciate the detailed responses to my concerns, but my concerns still remain:

Existing hybrid architectures such as MaxViT, FasterViT, CoAtNet, etc., follow a symmetric design in which C and T blocks are uniformly distributed across stages or within a stage. In contrast, AsCAN architectures recommend an asymmetric distribution with a preference for more convolutional blocks in the early stages and more transformer blocks in the later stages.

Asymmetric architectures outperform many existing models that utilize specialized attention and convolutional operators. These works design models based on parameter count and FLOPs, but it typically does not translate into inference throughput gains. Some of these issues come from operators that do not contribute to parameter count and FLOPs but require non-trivial runtime, such as reshape, permute, etc. Others originate from the lack of efficient CUDA operators for these specialized attention and convolutional operators.

The above authors' responses still do not give me a rationale for the architectural choice: I am still curious about why the asymmetric distribution of convolutions and transformers works well. Even though performance improvements have been observed, I feel that the design concept might not be a universal option beyond ImageNet. Additionally, the CTTTTC block in the right middle of the generative model design (in Fig. 2) suggests that the authors' claim may vary depending on the specific scenario. I believe a NeurIPS paper should provide rationale/insights into the suggested method, enabling readers to apply it to other tasks or domains. More evidence of the transferability or universality of the discriminative architecture is certainly needed. I understand that time constraints may hinder additional experiments in this discussion round, but these could be considered in future revisions.

Another concern is that the authors claim higher acc vs. throughput trade-offs of a proposed architecture as a key contribution. However, from my perspective, the higher throughput largely depends on the use of FusedMBConv and the parallel Transformer architecture, as demonstrated in Table 1. Therefore, this may not be a unique contribution of this work, as it relies on existing GPU-friendly operations. Furthermore, Table 2 does not appear to be a controlled experiment in terms of the computational budget, which is the number of parameters. I believe that constraining the number of parameters would better highlight the trade-offs for "some" convolution-dominated counterparts, as these networks in the table seem to "typically" have more parameters to compute, which could put them at a disadvantage. I believe that more controlled experiments might lead to different conclusions, so the contribution of finding an asymmetric architecture in Table 2 is somewhat diluted.

I also noticed that Stage 1 (S1) could be merged into the stem (S0), which leads to the final number of stages to three and deviates from the standard four-stage design. Are there any particular intuitions behind this architectural choice? (I recall seeing something similar designs elsewhere, but I can't remember which work it was. and there were no clear explanations for it there either)

Finally, the authors made great efforts to address my concerns, and the contribution to generative modeling is noteworthy, so I am increasing my rating to 5. The authors are encouraged to reflect all the reviewer's concerns and the responses in to the final revision.