Paper: Learning Stackable and Skippable LEGO Bricks for Efficient, Reconfigurable, and Variable-Resolution Diffusion Modeling (ICLR-2024)

Reviewer comment (Novelty):
It seems that the design of LEGO bricks borrows a lot from DiT, so it is a bit unclear to me how much additional contribution w.r.t. network design made in this work.

Author response:
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.

| Model                               | Unet        | DiT         | LEGO                      |
|-------------------------------------|-------------|-------------|---------------------------|
| Backbone                            | Convolution | Transformer | Convolution / Transformer |
| Attention                           | Global      | Global      | Local + Global            |
| Local feature                       | Multi-scale | Fixed-scale | Multi-scale               |
| Hierarchical structure | Downsample -> Upsample          | No hierachy          |   Patch decomposition (small -> large or large -> small)                   |
| Multi-variate resolution generation | No          | No          | Yes                       |


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.

Reviewer comment (Experiments):
It is also clear if the idea of LEGO (skippable and stackable backbone), is specific to DiT, or it is general enough to be applied to other types of backbone for diffusion models?

Author response:
> 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.

Reviewer comment (Experiments):
It is mentioned that the optimization of patch sizes for the LEGO bricks can be future work: can you give some intuition on how to choose the right patch sizes?

Author response:
> 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.

Reviewer comment (Reproducibility):
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?

Author response:
> 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: 

| Model             | FID  | sFID          | IS              | Precision | Recall | Iterated Imgs (M) |
|-------------------|------|---------------|-----------------|-----------|--------|---------------|
| DiT-XL/2          | 2.27 | 4.60          | 278.24          | 0.83      | 0.57   | 1792          |
| MDT-XL/2          | 1.79 | 4.57 | 283.01          | 0.81      | 0.6    | 1664          |
| MaskDiT           | 2.28 | 5.67          | 276.56          | 0.80      | 0.61   | 521           |
|-------------------|------|---------------|-----------------|-----------|--------|---------------|
| LEGO-XL-PG (ours) | 2.05 | 4.77          | 289.12          | 0.84      | 0.55   | 512           |
| LEGO-XL-PR (ours) | 2.35 | 5.21          | 284.73          | 0.83      | 0.60   | 512           |
| LEGO-XL-U (ours)  | 3.24 | 4.29          | 337.85 | 0.82      | 0.52   | 230           |

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.

Reviewer comment (Presentation):
The majority of the experiments and results are placed in the appendix, while a very large portion of the main paper is dedicated on an extensive introduction, related works, and more context setting at the start of the method section. The authors should strongly consider making the first 4 pages of the main paper significantly more concise, thereby allowing the main findings to move from the appendix to the experimental section of the main paper.

Author response:
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.

Reviewer comment (Presentation):
The graphic shown in right panel in Figure 1 is interesting, but there may be a better choice of representing the same data. Drawing circles with FLOPs as their radius makes differences between models appear much more significant than they are (since the difference in circle size == circle area grows with the square of the radius), making the plot somewhat misleading. A more common choice is a plot showing FID on the y-axis and FLOPs on the x-axis, where each model would be a single dot.

Author response:
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).

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

Author response:
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.

Reviewer comment (Theory):
Page 1: "This requirement arises from the model's need to learn how to predict the mean of clean images conditioned on noisy inputs [...]". Could the authors provide a citation for this claim?

Author response:
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).

Reviewer comment (Novelty):
It seems the authors in essence propose a block-based hierarchical architecture which is not very different from a UNet. While there is lots of talk on modularity and skipping of blocks, these aspects are only explored in the appendix on 64x64 CelebA images, i.e. an easy data set at a resolution where inference speed ups are not very interesting.

Author response:
> 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.

| Model                               | Unet        | DiT         | LEGO                      |
|-------------------------------------|-------------|-------------|---------------------------|
| Backbone                            | Convolution | Transformer | Convolution / Transformer |
| Attention                           | Global      | Global      | Local + Global            |
| Local feature                       | Multi-scale | Fixed-scale | Multi-scale               |
| Hierarchical structure | Downsample -> Upsample          | No hierachy          |   Patch decomposition (small -> large or large -> small)                   |
| Variable resolution generation | No          | No          | Yes                       |

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.

Reviewer comment (Novelty):
Generating images larger than the training resolution is demonstrated with a few examples, which might also be obtained by cleverly leveraging prior works.

Author response:
> 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: 


| Crop    | Grid - 256   | Random - 256 | Grid - 512   | Random - 512 |
|---------|--------------|--------------|--------------|--------------|
| DiT     | 0.15 &plusmn; 0.07 | 0.66 &plusmn; 0.13 | 0.55 &plusmn; 0.07 | 0.76 &plusmn; 0.06 |
| LEGO-PG | 0.14 &plusmn; 0.07 | 0.21 &plusmn; 0.17 | 0.36 &plusmn; 0.07 | 0.37 &plusmn; 0.06 |
| LEGO-PR | 0.15 &plusmn; 0.07 | 0.25 &plusmn; 0.15 | 0.55 &plusmn; 0.07 | 0.55 &plusmn; 0.05 |




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.

Reviewer comment (Experiments):
Given all these points it feels like the contribution is not as significant as suggested at the beginning of the paper. 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.

Author response:
> 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.

Reviewer comment (Evaluation):
How is the FID computed? FID numbers can differ substantially depending on the implementation, and whether the train or validation set is used as a reference. For example DiT and ADM both use the ADM Tensorflow evaluation suite.

Author response:
> 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.

Reviewer comment (Writing):
What is a “linear MLP” (page 7)? How does it differ from a linear layer?

Author response:
> 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."

Reviewer comment (Experiments):
How do the parameter counts of the proposed model compare with the baselines for the ImageNet 256 and 512 pixels experiments? I could only find parameter counts for the 64 pixel models.

Author response:
> 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:

| Model   | Total Param  | Param with skipping (PG) | Param with skipping (PR) |
|---------|--------------|--------------------------|--------------------------|
| LEGO-S  | 35           | 16                       | 26                       |
| LEGO-L  | 464          | 223                      | 331                      |
| LEGO-XL | 681          | 327                      | 555                      |

Reviewer comment (Reproducibility):
Are the bricks of a given stage shared across space, or are they specialized per patch?

Author response:
> 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.