Learning Stackable and Skippable LEGO Bricks for Efficient, Reconfigurable, and Variable-Resolution Diffusion Modeling

Thank you for your positive review and constructive feedback. We're delighted that the concept of LEGO bricks for a reconfigurable diffusion backbone has been well-received. Your appreciation for the paper's clarity and the performance of our design is encouraging. Now, to address your concerns and questions:

Q: What additional network design contributions does LEGO offer compared to DiT?

To better address this question, we summarize the characteristics of Unet, DiT and LEGO in the table below.

The primary contribution of our work is the development of the LEGO framework itself. Unlike DiT, which follows a traditional transformer architecture, LEGO introduces a novel, modular approach and is open to deploy both convolutional and transformer basic blocks. This modularity allows for individual components or 'bricks' to be easily skipped or stacked, enabling dynamic architecture reconfiguration based on specific requirements, a feature not inherent in DiT.

In the design of a single LEGO brick, we mainly consider two perspectives: 1) the modular brick needs to be able to handle local-level feature; 2) the local-level processed information needs to be aggregated seamlessly with attention mechanism. Considering these perspectives, DiT blocks well meets these requirements, and we borrow the successful design in DiT, provides a solid base for image processing, to form each LEGO brick.

The entire LEGO model design is however different from DiT, where we can see the multi-scale local feature and local+global attention in the entire model. LEGO optimizes processing power and reduces computational load, particularly in the scenarios where full engagement of all network layers is unnecessary. Moreover, another key aspect of our work is the adaptability of LEGO bricks. LEGO extends this by offering a versatile framework that can be adapted to various types of backbones beyond transformers, such as convolutional networks.

Q: Is the idea of LEGO (skippable and stackable backbone) specific to DiT, or it is general enough to be applied to other types of backbone for diffusion models?

We designed the LEGO framework with versatility in mind, aiming to make it applicable beyond just the DiT blocks used in our main experiments. To demonstrate the general applicability of the LEGO framework, we conducted experiments with the integration of the LEGO concept with a Unet backbone, as shown in Table 4 located in Section 4.3 of the revision.

Q: Some intuition on how to choose the right patch sizes?

We have conducted comprehensive studies in how to choose patch size and the attention length, as detailed in our experimental studies in Appendix C.2. Our core objective was to identify a patch size that maximizes model performance while minimizing computational costs. Through extensive experimentation, we developed a general rule: incrementally increasing the patch size by a factor of four, resulting in a summarized recipe in Tables 4 and 5 of our manuscript.

Q: There are two spatial refinement settings mentioned in the paper: PG and PR. So does that mean a model can only take one of these two configurations? or they can be used interchangeably in one model?

Thank you for your insightful question regarding the spatial refinement settings. Indeed PG and PR configurations, while distinct in their approaches, are not mutually exclusive and can indeed be combined within a single model to leverage their respective strengths.

We have been actively developing LEGO-U, a variant that combines the features of LEGO-PR and LEGO-PG. This model processes image patches starting with larger resolutions, transitioning to smaller ones, and then reverting back to larger resolutions. This approach draws inspiration from the architecture of the U-Net, which similarly operates with multiple downsampling and upsampling stages to have multi-scale representations.

We provide some preliminary results of LEGO-U into table below:

Despite the experiments still being underway and not fully completed, these initial findings already demonstrate competitive performance compared to the baselines, LEGO-PR and LEGO-PG. Notably, LEGO-XL-U achieves the best inception score and sFID, indicating the capacity to generate high-quality images.

We regard LEGO-U as a distinct project meriting further exploration and study. We welcome your guidance on whether you think it should be included in the current paper.

Dear Reviewer ky4P,

Thank you for your insightful review and valuable suggestions. We appreciate it if you could review our response by November 22nd, as we will be unable to participate in further discussions after this date. If you have any follow-up questions or require additional information, please feel free to let us know. We greatly appreciate your expertise and the time you've dedicated to our manuscript.

Warm regards,

Authors

Thank you for your detailed and constructive review. We are grateful for your acknowledgment of LEGO's originality and your suggestions in polishing the presentation. We have made revision accordingly, and provide our response to your concerns point-to-point.

Q: The method presentation could be improved. The majority of the experiments and results are placed in the appendix.

We have re-organized the presentation of our paper, making more space to the experiments. Specifically, we move some high-level review and discussions to Appendix A. We move the study of the sampling with skippable LEGO bricks and the study in terms of the versatility of LEGO in section 4.2 and 4.3.

Q: The graphic shown in right panel in Figure 1 can be improved.

Following your suggestion, we update the right figure of Figure 1 (Figure 2 in the revised paper). The updated figure now features a plot with FLOPs on the x-axis (presented with log scale).

Q: How many patches are selected (for a single LEGO brick)? During training time & inference.

In our main experiments, 50% and 75% of all non-overlapping patches are used in the training of patch-bricks on CelebA and ImageNet datasets. In the inference time, as we have no patches from real images that can be leveraged, we need to use all patches for the generation, but we can skip bricks to save computation cost.

Q: Reference for: Page 1 - "This requirement arises from the model's need to learn how to predict the mean of clean images conditioned on noisy inputs [...]".

We have clarified that the concept of “diffusion models learn how to predict the mean of clean images conditioned on noisy inputs” can be derived from two theoretical foundations: either Tweedie’s formula (Robbins, 1992; Efron, 2011), as illustrated in Luo, 2022 and Chung et al., 2022, or the Bregman divergence (Banerjee et al., 2005), as revealed in Zhou et al. (2023).

Q: “Generating images larger than the training resolution is demonstrated with a few examples, which might also be obtained by cleverly leveraging prior works.”

We appreciate your suggestion to consider prior works in the context of generating images larger than the training resolution. Following your advice, we have now included a reference to the MultiDiffusion work in our manuscript. This work deploys techniques similar to those we described in Appendix C.4, particularly in generating large-content images using text-to-image diffusion models.

Since our case is to generate large-content with class-conditional model pre-trained with smaller resolution, we compare the large-content generation using pre-trained DiT model on ImageNet 256 x 256 and 512 x 512. To provide a thorough comparison, we generated images using our method and then compared the LPIPS distance (a perceptual score) with those of local crops:

The results of this comparison, which we have included in the revised manuscript, demonstrate that our method is effective in generating high-quality, large-resolution images by orchestrating the generated small patches.

Q: How is the FID computed?

We use the same evaluation suite, i.e., the ADM Tensorflow evaluation suite and their reference batch, ensuring that our evaluations are rigorous and reliable. We have added these details in Appendix D.

Q: What is a “linear MLP”?

Thank you for highlighting the ambiguity in our use of the term "linear MLP." To avoid confusion and ensure clarity in our manuscript, we have revised the text to specifically mention "a linear layer" instead of "linear MLP."

Q: How do the parameter counts of the proposed model compare with the baselines for the ImageNet 256 and 512 pixels experiments?

The parameter counts are listed below, we have included them in the Tables 5-6 of the revised manuscript:

Q: Are the bricks of a given stage shared across space, or are they specialized per patch?

In the current design of our model, each brick at a specific stage is indeed shared across the entire spatial domain. This means that the same brick processes all possible patches within that stage.

Thank you for your valuable feedback. In what follows, we provide detailed responses to each of your concerns.

Q: Similarity to UNet? The proposed block-based hierarchical architecture was viewed as not being significantly different from a UNet structure.

We summarize the characteristics of Unet, DiT and LEGO in the table below.

LEGO distinguishes itself from Unet in several key aspects:

• The LEGO bricks that constitute the LEGO network are distinct from the convolutional filters employed in UNet. Unlike UNet, which relies heavily on convolutional layers, the LEGO model is flexible to adopt either convolutional or transformer blocks.

• The self-attention deployed in the Unet is global-level, while LEGO adopts a different approach, which makes its self-attention length vary depending on the 'LEGO brick'. In some cases, attention is confined to specific image patches defined by the brick size, while in some other bricks, it extends over the entire image. This selective mechanism allows for more focused and efficient feature aggregation.

• Both models extract multiscale features, but they achieve this goal differently. The Unet uses a series of downsampling and upsampling stages, whereas LEGO leverages varied patch decomposition methods. This distinction in processing strategy results in diverse feature representations.

• A notable difference is that each 'brick' in LEGO can be trained to independently generate patch-wise outputs, resulting in variable resolution generation, a capability not inherently present in Unet's intermediate features. This property of LEGO bricks enhances their utility in diverse generative tasks.

• The variants LEGO-PG and LEGO-PR, as discussed in this paper, do not arrange the LEGO bricks in a manner akin to Unet. The LEGO-U variant, mentioned in the 'Overview of Revisions,' draws inspiration from Unet in its organization of local units. However, even in LEGO-U, the construction of these local units remains distinct from those in Unet.

Q: Limited Exploration of Modularity and Skipping Blocks? “The modularity and skipping of blocks are only explored in the appendix on 64x64 CelebA images, which is an easy data where inference speed ups are not very interesting. Experiments where skipping blocks at inference time yield substantial wallclock-time speedups for class conditional ImageNet generation at 256 or 512 pixels resolution would be more convincing.”

We move the figure presenting the on the ImageNet 256x256 data to Figure 4 (Figure 7 of the original submission). These experiments demonstrate similar effects shown in the experiments on CelebA, but also validate the effectiveness of our approach in a more complex and high-resolution setting.

To further substantiate our findings, we have now included additional analyses on ImageNet 64x64 and ImageNet 512x512 in Figure 8-9. These are similar in spirit to Figures 4. These results provide more comprehensive evidence of the efficiency gains achievable through our model, especially in terms of wallclock-time speedups during class-conditional ImageNet generation at resolutions of 64, 256 and 512 pixels.

Thanks to the authors for their detailed rebuttal.

I now have fewer concerns about the benefits of skipping blocks given the quantitative results on ImageNet highlighted in the rebuttal. Overall, I'm leaning towards acceptance.

I appreciate that the authors report the training wall clock time as a function of t_break. The only remaining evaluation that I think is really important for the final version of the paper is inference time (actual runtime on accelerator, not FLOPs) vs sample quality/FID, for different t_break.

Dear Reviewer x9WT,

As we will not be able to provide additional results or modifications to our manuscript after November 22nd, we are reaching out to see if our response adequately addresses your concerns.

We greatly value your comments and are committed to refining our work. Any further feedback would be helpful in polishing our final revisions.

Thank you for your time and expertise.

Sincerely,

Authors

Reviewers, PTAL at the rebuttal.