Hollowed Net for On-Device Personalization of Text-to-Image Diffusion Models

Thank you for your thoughtful feedback. Please note our top-level comment with additional experimental results. Below we address specific questions.

1. Clarification on assumption and experimental details

The experiment depicted in Figure 2 of the manuscript, as detailed in Section 4.1, was conducted to identify less significant layers within the U-Net by analyzing the changes in LoRA weights before and after personalization. Specifically, we measured the extent of weight changes across different blocks of the U-Net, where smaller weight changes in certain layers after personalization indicate that these layers are less important for the personalization process. This assumption is inspired by previous work, notably Li et al. [22], Kumari et al. [7], and Shah et al. [6], which demonstrated similar analyses of weight changes.

Figure 2 of the manuscript illustrates the results obtained when personalizing the "monster toy" subject from the DreamBooth dataset, starting from the pretrained weights of the Stable Diffusion V2.1 model. We injected LoRA into various layers of the U-Net's attention and ResNet blocks, and then fine-tuned the model for 1000 steps with a learning rate of 1e-4, using a LoRA rank of 128. We observed that the overall trends in weight changes remained consistent across different learning rates, steps, LoRA ranks, LoRA injection locations, and personalizing subjects.

Additionally, to substantiate our findings, we conducted further experiments, the results of which are described in Figure 1 of the attached document. We analyzed more diverse subjects (30 subjects from the DreamBooth dataset and 101 subjects from the CustomConcept101 dataset) and across different models (Stable Diffusion V2.1 and SDXL). The x-axis of each plot represents the specific U-Net blocks, while the y-axis shows the changes in LoRA weights before and after personalization. For each dataset, we averaged the weight changes across subjects and provided error bars to indicate the statistical variance within each dataset.

Similar to the original experiment in Figure 2 of the manuscript, these extended experiments revealed that weight changes around the central blocks tend to be minimal, confirming that these layers are less involved in personalization. This suggests that our assumption holds across various personalization tasks and models. We plan to include these additional findings in the supplementary materials.

2. Discussion on the efficiency of Hollowed Net

The training memory used for training LoRA is 5.226GB, and more than 65% (3.404GB) comes from holding the entire model on GPU, and our aim is to remove the part of the models during training, in order to reduce the memory consumption during training without undermining the performance.

In particular, we assume an extremely resource-constrained environment where even running an inference (costing approximately 3.49GB in our experiments) can exhaust most of the available memory. This assumption is reasonable given that the total memory capacity of recently released mobile phones (e.g., iPhone) ranges from 6 to 8GB, which must be shared with the operating system and other background apps. In such cases, a 1.35GB memory usage reduction can be substantial, directly affecting the feasibility of performing fine-tuning.

3. Comparison with ControlNet

The main differences between ControlNet and Hollowed Net arise from their distinct goals. First, ControlNet focuses on adding spatial conditioning controls to diffusion models. It leverages a cloned copy to use the pre-trained knowledge of diffusion models as initialization. However, ControlNet does not address memory efficiency concerns, as maintaining both the original and cloned networks simultaneously for training and inference requires substantial GPU memory.

In contrast, Hollowed Net is designed to achieve personalized LoRA parameters in a memory-efficient manner. We are the first to introduce a two-step fine-tuning strategy that enables efficient personalization of diffusion models with memory demands as low as inference. Compared to ControlNet, Hollowed Net uses a cloned copy to ensure that the personalized LoRA parameters can be transferred to the original network. Unlike ControlNet, which employs a cloned copy to complement the original network after training, Hollowed Net utilizes the cloned copy as a proxy for training LoRA parameters, without intending for the cloned copy to be used concurrently with the original network.

4. Experiments with large datasets

We agree with you for the necessity of testing with larger datasets. Following your suggestions, we include the analysis with CustomConcept101 datasets, which you can find the analysis in the top-level comment.

Thank you for your thoughtful comments and for taking the time to review our paper. We have carefully considered your feedback and would like to provide some clarifications regarding our work.

1. Computational Advantage

While a reduction of 1.35 GB in GPU memory may seem modest under standard workstation settings (e.g. A100; 80GB), it represents a significant advantage in the context of on-device training (e.g. IPhone15; 6~8GB). For training large generative models with on-device chipsets, even a small reduction in GPU memory can have a substantial impact on feasibility. Even an increase in GPU memory of less than 1 GB can destabilize chipset functioning and make computation infeasible. In this context, our approach offers a crucial advantage in on-device fine-tuning of diffusion models by reducing training memory to a level as low as that required for inference, while standard LoRA consumes 35% more memory than Hollowed Net in proportional terms.

2. Novelty of the Work

We would like to emphasize that there is a clear distinction between our approach, Hollowed Net, and traditional layer pruning methods. As presented in recent work [1], accepted at ECCV’24, layer pruning involves completely removing selected layers, which necessitates extensive pre-training with large datasets to recover lost information and restore functionality. Diffusion models, however, often suffer from significant performance degradation post-pruning, as lost information may not be fully recoverable through pre-training.

In the table below, we present experiments with layer-pruned Stable Diffusion models by recent work (BK-SDM [1]), accepted at ECCV’24, using rank-128 LoRA. Compared to the results in Table 1 of our paper, these models achieve memory usage comparable to Hollowed Net but show significant performance degradation, generating malformed images for subjects and prompts outside the pre-training data. Despite the extensive pre-training, the performance of these models is compromised.

In contrast, Hollowed Net does not completely remove the middle layers and does not require any extensive pre-training. We temporarily exclude the selected layers during fine-tuning, while preserving essential information from those excluded layers by incorporating an additional buffering stage. Despite this buffering, the computational load for the entire training process of Hollowed Net (2052 TFLOPs) remains more efficient than LoRA fine-tuning (2148 TFLOPs).

Our approach of excluding layers for fine-tuning is a novel technique beyond traditional layer pruning, and merits further discussion within the community. In fact, Meta AI has recently released an approach for LLMs, shortly before our submission [2], which suggests that for fine-tuning LLM models, one can remove up to 40% of middle layers of the models and perform fine-tuning, and still achieve the comparable results. However, their approach differs from ours due to the characteristics of LLMs versus diffusion U-Nets. Meta AI's method involves complete removal of middle layers for both fine-tuning and inference, arguing that these layers do not store critical knowledge compared to shallow layers.

In contrast, our research shows that the deep layers of diffusion U-Nets contain crucial high-level image features, and their complete removal can lead to severe performance degradation, even with additional pre-training (as discussed for the case of [1] above). This underlines the significance of our two-stage fine-tuning strategy, which maintains performance while reducing memory usage without requiring additional pre-training.

[1] Kim, Bo-Kyeong, et al. "Bk-sdm: A lightweight, fast, and cheap version of stable diffusion." arXiv preprint arXiv:2305.15798 (2023).

[2] Gromov, Andrey, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

3. Discussion on Layer Selection

Regarding layer selection, we do not base our choices on computational expense of certain layers. Instead, we leverage the skip connections inherent in the U-Net architecture to determine which layers to remove. For our main results, we choose the third layer of the second down block, the entire third down block, the entire mid block, and the entire first up block to be hollowed during fine-tuning, which corresponds to 40% of the U-Net's parameters.

Our findings suggest that removing those layers during fine-tuning is feasible because these deep layers are more involved with high-level abstract features (e.g., the concept of a dog) rather than pixel-level personalization details of a specific instance (e.g., “a [V] dog”), which leads to relatively smaller change of weights in LoRA parameters for personalization in those layers (commonly observed for different models and datasets as presented in the pdf of the top-level comment). This characteristic allows us to use intermediate activations from the frozen original network, acquired with a class-level prompt (e.g., “a dog”) during the collection stage, to update the layers of Hollowed Net during the fine-tuning stage.

4. Scalability of the Work

To demonstrate the scalability of Hollowed Net, we present experimental results with CustomConcept101, the largest dataset available for T2I personalization, and results with SDXL, one of the largest publicly available diffusion models (included in the top-level comment and pdf), as well as user studies to validate the consistency of metrics with human evaluation (included in the comment for Reviewer SH8P). These extensive results consistently show that Hollowed Net does not require any trade-off between performance and computation efficiency compared to LoRA, for the removal of up to 40% of parameters. Our approach maintains the same or even slightly better performance compared to LoRA, while achieving significant computational advantages for on-device learning. As our novel two-stage fine-tuning strategy allows to maintain the knowledge of (temporarily) hollowed layers, it enables effective personalization for over 100 subjects with various prompts, including long range image editing like background editing and many others.

We hope this detailed explanation addresses your concerns regarding the novelty and effectiveness of our work. We are open to further discussion and clarifications if needed, and we sincerely hope these points will contribute to a reconsideration of the final rating.

Thank you for your time and consideration.

Thank you for your thoughtful feedback. Please note our top-level comment with additional experimental results. Below we address specific questions.

1. Experimental evaluation with user studies

Based on you suggestion, we conducted user studies with 40 participants, each completing a set of 25 comparative tasks. For each task, participants were presented with one reference image, one prompt, and two generated images (A and B). They answered two questions: subject fidelity and text fidelity. Each pair of generated images, A and B, was created using Hollowed Net and LoRA FT, and the labels (A or B) were randomly assigned for each task. The table below displays the results of the user studies. These results confirm that users generally find the images generated by Hollowed Net and LoRA FT to be similar in terms of both subject fidelity and text fidelity, consistent with the main results presented in Table 1 of the paper.

2. More analysis on model size and latency

We evaluated our HollowedNet based on the Stable Diffusion v2.1 with 866M. We think that our fine-tuning approach of 1000 steps and 3.88GB given this backbone is an on-device-feasible solution with the following reasons:

2.1) [Memory] Recent mobile phones can support 8GB or 12GB memory (iPhone15 or GalaxyS23 Ultra), and thus the training memory 3.88GB can fit into this space.

2.2) [Latency] Though we couldn't measure the actual latency on a mobile device, we expect the latency based on the previous report. For example, the Table 1 in [1] shows around 7s-15s for an inference on a GalaxyS23 device. Compared to an inference, we observed that our fine-tuning needs 2.2x more computations than an inference (6.013 TFlops vs 2.717 TFlops). Based on this, for a single training step, roughly we can say it takes around 22 sec (2.2times x 10sec/inference). So, for 1000 steps, it would take 6 hours, and this can be done during charging time at night. Furthermore, combining with a HyperNetwork which provides initial LoRA parameters, we may also expect 40 times reduction in training steps (as reported in the Hyper-dreambooth paper), and this leads to less than 10min solution.

[1] Choi, Jiwoong, et al. "Squeezing large-scale diffusion models for mobile." arXiv preprint arXiv:2307.01193 (2023).

3. Memory for data buffer

For collection, 200 buffered samples are sufficient for achieving high-fidelity results Each set of pre-computed outputs requires approximately 1.35 MB. Therefore, $1.35 \times 200 = 270$ MB of additional storage is needed for keeping all data required later during fine-tuning. Whether all this data is held in RAM simultaneously rather than SSD will depend on the specific implementation.

4. Working with quantized/flow models

We believe that there is no clear reason why Hollowed Net cannot be applied to quantized/flow models, but some modifications may be necessary depending on the technical differences in the model architecture.

5. LoRA model that does not finetune all the attention layers

The table below shows the results of selectively applying LoRAs to the layers, excluding those removed in the Hollowed Net. We find that selectively applying LoRAs leads to a slight decrease in performance. This may be due to the cross-attention layers in the middle receiving a personalization prompt (e.g., "a [V] dog") but not learning how to process the unique token "[V]". More importantly, in terms of memory efficiency, using fewer LoRAs results in only a minimal decrease in training memory. As shown in the table, the main advantage of the Hollowed Net is its ability to avoid holding the hollowed layers in GPU memory during fine-tuning, resulting in a 35% reduction in GPU memory usage compared to LoRA.

6. Discussion on conditional input/skip connection

How to handle inconsistency and leverage it to produce better personalization results is what the model learns during fine-tuning. The deep layers in the middle are more involved with high-level abstract features (e.g., the concept of a dog) rather than the pixel-level personalization details of a specific instance (e.g., a [V] dog). Thus, at this level, as long as the categories are well-defined, the deep features of a specific instance can be replaced with the deep features of its corresponding class (e.g., a dog) through proper fine-tuning. In fact, both the results in Table 1 of the paper and the user study presented above show that Hollowed Net can perform slightly better than LoRA fine-tuning in some cases. We assume this is because using class-level inputs for the middle layers can result in positive regularization for effective personalization to a certain extent. This suggests that the resolution of defining class tokens can lead to better or worse results compared to LoRA.

7. Expanded explanation of analysis on LoRA weight changes

Please refer to the answer for the first question of the last reviewer: Jp6t.

8. In figure 3, how is the red arrow from block2 to block3 used?

Thank you for your keen observation and a thorough review. The red arrow from block2 to block3 should be removed.

Thank you for your valuable feedback. Please note our top-level comment with additional experimental results. Below we address specific questions.

1. Further discussion on computational overload

In the table above, we provide a detailed analysis of computational load (FLOPs) and memory usage (peak GPU memory) for each stage of Hollowed Net and LoRA. Each number corresponds to one step of each stage: one forward pass for Collection and Inference, and one forward+backward pass for Fine-tuning.

For fine-tuning, $\sim$1000 steps are required, totaling $2.004 \times 1000 = 2004$ TFLOPs of computation. For Collection, 200 buffered samples are sufficient for achieving high-fidelity results, which requires $0.238 \times 200 = 47.6$ TFLOPs of additional computation. Thus, the total computation required for training with Hollowed Net is $2004 + 47.6 = 2051.6$ TFLOPs, which is lower than the $2.148 \times 1000 = 2148$ TFLOPs needed for fine-tuning LoRA.

For inference, Hollowet Net requires around 28% more computation than LoRA, as it needs to repeat a part of early downblocks to reproduce the path used in training. Considering that recent works show that on-device inference of Stable Diffusion models can be done within 7-15 seconds, those additional computational cost for inference may not be a significant issue, considering the advantage of significantly lowering the memory cost of fine-tuning to a feasible level.

However, if inference latency is a concern, a variant of the inference method for Hollowed Net can be considered. This variant involves a single inference path with LoRA, where the cross-attention layers hollowed during training receive a non-personalized prompt (e.g., "A dog") as input. As shown in the table below, this variant (Single Path) requires the same amount of computational FLOPs as LoRA but achieves a level of personalization equivalent to that of the default inference method (Dual Path). One potential explanation for the effectiveness of this efficient method is that the high-level features in the middle layers are focused on processing abstract class-level characteristics of an image rather than capturing personalization details at the pixel level. Therefore, even if the middle layers are fed features of a personalized dog (in Single Path) instead of a non-personalized dog (in Dual Path) from the early layers, the extraction of high-level features for "A dog" is not affected.

2. Further discussion on space consumption

The Peak Usage reported in Table 1 of our paper represents the maximum GPU memory (VRAM) consumption during the fine-tuning stage. As shown in the table above, the VRAM usage during the Collection stage is significantly lower because only part of the upper blocks needs to be retained. Each set of pre-computed outputs requires approximately 1.35 MB. Therefore, $1.35 \times 200 = 270$ MB of additional storage is needed for keeping all data required later during fine-tuning. Whether all this data is held in RAM simultaneously rather than SSD will depend on the specific implementation.

3. Results with LoRA rank between 1 and 128

In the table below, we present the results using LoRA with different ranks (4 and 16). While the default rank of 4 in the diffusers library is often used, we have found that it may oversimplify personalization details or fail to effectively handle a range of challenging subjects and prompts. Increasing the rank from 4 to 16 improves subject fidelity. However, to achieve personalization quality comparable to full fine-tuning across all subjects and prompts in the DreamBooth dataset, we find that a rank of 128 is necessary.

4. Discussion on the design of Hollowed Net

This is indeed a very interesting suggestion. Modifying the gradient flow by stopping gradients from passing through the hollowed part could replicate backpropagation in Hollowed Net while maintaining the hollowed layers. However, the primary advantage of Hollowed Net is that it avoids holding these hollowed layers in GPU memory during fine-tuning, which saves approximately 1.29GB of additional memory. This reduction in memory usage is a significant benefit, making it challenging to eliminate the extra buffering stages.

Thank you for your kind response. We are sincerely grateful for the time and effort you dedicated to thoroughly reviewing our paper and participating in the discussion. We sincerely hope this additional information addresses your concerns and contribute to a more favorable assessment of our work. We will carefully consider all the feedback you provided during the rebuttal and discussion periods as we update the paper to enhance its completeness. Once again, thank you very much for your thoughtful and thorough review.

Thank you for your thoughtful feedback. Please note our top-level comment with additional experimental results. Below we address specific questions.

1. 1.$\sim$GB memory usage save, while beneficial, is still marginal

In this paper, we are assuming an extremely resource-constrained environment where even running an inference (costing approximately 3.49GB in our experiments) can exhaust most of the available memory. This assumption is reasonable given that the total memory capacity of recently released mobile phones (e.g., iPhone) ranges from 6 to 8GB, which must be shared with the operating system and other background apps. In such cases, a 1.35GB memory usage reduction can be substantial, directly affecting the feasibility of performing fine-tuning.

The strength of Hollowed Net lies in its ability to enable model fine-tuning with memory usage similar to or even less than that required for inference, depending on the proportion of hollowed layers. Furthermore, in proportional terms, applying Hollowed Net results in a 35% memory reduction compared to LoRA, which is a significant improvement. In addition, more memory reduction can be achieved with more hollowness depending on the usecases and tasks. We believe this substantial reduction is crucial in environments with limited resources and makes our approach highly beneficial.

2. Comparison with other memory-efficient approaches

Unfortunately, SVDiff does not have an official code release. We attempted to use a third-party reproduced code from GitHub, but it was not implemented in a memory-efficient manner, requiring over 14GB of peak GPU memory for processing a single image.

Theoretically speaking, the mechanism of SVDiff is not significantly different from LoRA, in terms of memory usage. As described in the paper, SVDiff uses 2200 times fewer parameters compared to vanilla DreamBooth, which suggests that SVDiff updates around 391K parameters. This is nearly double the number of parameters updated by rank-1 LoRA, presented in Table 1 of our paper. Consequently, we estimate that the minimal memory usage achievable by SVDiff would be at least 1GB more than Hollowed Net. In fact, all fine-tuning methods based on reducing updating parameters like SVDiff have a fundamental limitation: they can reduce the number of updated parameters but cannot reduce the actual model size, which takes a significant portion of memory usage when performing backpropagation over entire backbone network. In this context, rank-1 LoRA demonstrates the minimal baseline of memory usage that these methods can achieve.

Another memory-efficient approach to consider is model compression, which reduces the backbone model's size through pre-training with a large dataset. We present results using BK-SDM, a recent method accepted at ECCV '24, in the table below. We fine-tuned the compressed Stable Diffusion models of BK-SDM using rank-128 LoRA on the DreamBooth dataset. Our experimental results show that while these models achieve memory usage comparable to Hollowed Net, they suffer from significant performance degradation. These models' capacity is limited to the data used during model compression, and for subjects and prompts outside this data, they can even generate seriously malformed images.

In contrast, Hollowed Net achieves both low memory usage, comparable to compressed models, and high-quality results, similar to full fine-tuning of diffusion models. This is achieved through our novel approach, which effectively removes layers during training while preserving the knowledge from those hollowed layers throughout both training and inference stages.

3. Discussion on scalability and generalizability

To demonstrate the scalability and generalizability of our method, we include additional results and analysis in top-level comments. Please also refer to the answer for the first question of the last reviewer: Jp6t.

Dear Reviewer cT56,

As the discussion period draws to a close, we are mindful of the limited time available to address any further questions or requests for clarification. To ensure we provide all necessary information for your final recommendation, we would like to further clarify our previous rebuttal, particularly regarding scalability and generalizability, which we referenced other comments we made but did not directly address.

To address your concerns on the scalability and generalizability of Hollowed Net, we have presented experimental results with SDXL, one of the largest publicly available diffusion models, and results with CustomConcept101, the largest dataset available for T2I personalization (included in the top-level comment and pdf), as well as user studies to validate the consistency of metrics with human evaluation (included in the comment for Reviewer SH8P). Across all these experiments, Hollowed Net consistently demonstrates its ability to enable effective personalization for over 100 subjects with various challenging prompts of different applications including background editing, artistic rendition, accessorization, and property modification. This performance is comparable to Full / LoRA fine-tuning of diffusion models, while significantly reducing the required training memory.

Our findings suggest that Hollowed Net’s approach of excluding middle layers during fine-tuning is feasible because these deep layers are more involved with high-level abstract features (e.g., the concept of a dog) rather than pixel-level personalization details of a specific instance (e.g., “a [V] dog”), which leads to relatively smaller change in LoRA parameters during personalization in those layers, a trend commonly observed across different models and datasets (as presented in the pdf of the top-level comment). This characteristic allows us to use intermediate activations from the frozen original network, obtained with a class-level prompt (e.g., “a dog”) during the collection stage, to update the layers of Hollowed Net during the fine-tuning stage.

As elaborated in our rebuttal above, Hollowed Net is a novel method that overcomes the fundamental limitation of fine-tuning methods based on reducing updating parameters, such as SVDiff, which cannot reduce the actual model size that requires a significant amount of GPU memory during fine-tuning. While a reduction of 1.35 GB in GPU memory may seem modest under standard workstation settings (e.g. A100; 80GB), it represents a significant advantage in the context of on-device training (e.g. IPhone15; 6~8GB). For training large generative models with on-device chipsets, even a small reduction in GPU memory can make a substantial impact on feasibility. Even an increase in GPU memory of less than 1 GB can destabilize chipset functioning and make computation infeasible. In this context, our approach offers a crucial advantage in on-device fine-tuning of diffusion models by reducing training memory to a level as low as that required for inference, while standard LoRA consumes 35% more memory than Hollowed Net in proportional terms.

We sincerely appreciate the time and effort you have invested in reviewing our work. We hope these clarifications and additional details address your concerns and contribute to a more favorable assessment of our work.