ScaleLong: Towards More Stable Training of Diffusion Model via Scaling Network Long Skip Connection

In diffusion models, UNet is the most popular network backbone, since its long skip connects (LSCs) to connect distant network blocks can aggregate long-distant information and alleviate vanishing gradient. Unfortunately, UNet often suffers from unstable training in diffusion models which can be alleviated by scaling its LSC coefficients smaller. However, theoretical understandings of the instability of UNet in diffusion models and also the performance improvement of LSC scaling remain absent yet. To solve this issue, we theoretically show that the coefficients of LSCs in UNet have big effects on the stableness of the forward and backward propagation and robustness of UNet. Specifically, the hidden feature and gradient of UNet at any layer can oscillate and their oscillation ranges are actually large which explains the instability of UNet training. Moreover, UNet is also provably sensitive to perturbed input, and predicts an output distant from the desired output, yielding oscillatory loss and thus oscillatory gradient. Besides, we also observe the theoretical benefits of the LSC coefficient scaling of UNet in the stableness of hidden features and gradient and also robustness. Finally, inspired by our theory, we propose an effective coefficient scaling framework ScaleLong that scales the coefficients of LSC in UNet and better improves the training stability of UNet. Experimental results on four famous datasets show that our methods are superior to stabilize training and yield about 1.5x training acceleration on different diffusion models with UNet or UViT backbones. Code: https://github.com/sail-sg/ScaleLong

Paper

Similar papers

Peer review

Reviewer Ehww6/10 · confidence 2/52023-07-07

Summary

They state that diffusion models using Unet suffer from unstable training and oscillations of features and gradients. They also state that while this is sensitive to coefficients related to scaling the skip connections of Unet. They set out to provide an explanation and more robust scaling methods for the skip connections to address this. Their methods are constant scaling based exponentially on the depth of the skip connection, and learnable scaling. The methods seem to consistently reduce FID versus the same model with regular scaling.

Strengths

I believe this paper’s contribution is simple and meaningful, and written clearly. The experiments seem thorough to changes made and the questions you would have reading it. I cannot fully speak either way to novelty in this area, but it seems like a sound paper.

Weaknesses

I did not see meaningful weaknesses in the paper.

Questions

N/A

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer VLJa6/10 · confidence 3/52023-07-07

Summary

U-Net is the most popular neural network backbone for diffusion models. In U-Net, the long skip connection (LSC) links the long-distant information near to the input and the intermediate network outputs. However, they suffers from unstable training, which is resolved by scaling down the LSC coefficients. This paper addresses the theoretical aspects of stabilizing LSC, robustifies the training process, and accelerates this. Both hand-crafted and learning-based parameterization are proposed.

Strengths

* The organization of the paper is clear, and the figures of motivating experiments are well demonstrated. * The method and simple yet efficient, improving both the training speed and sampling performance while requiring only some hand-crafted hyperparameters (CS) or lightweight learning module (LS).

Weaknesses

* This work does not deal with more recent diffusion model baselines such as EDMs, which is much more powerful than the methods that are compared in this work. * Prior, related works are not introduced, especially for those who are not familiar with the scaling of skip connections in other network architectures.

Questions

* This work focuses on scaling the skip connection coefficients U-Net-based diffusion models. Does this technique also hold for other U-Net based methods rather than diffusion models? * Can you add some figures or tables that addresses the optimal value of $\kappa$ in CS cases, like the $\kappa$ from LS in Figure 7? * Please write the full name of IE and SE in line 330. * Please address the number of sampling steps in each methods. * Does parameterizing LC to some additional module (Figure 4) and training the network end-to-end adequately learn $\kappa_i$? =========== Corrections * Line 85: serir --> series * Line 86: UNettraining --> UNet training

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors adequately addressed the limitations and introduced further directions.

Reviewer xZoY7/10 · confidence 3/52023-07-07

Summary

This paper proposes to scale the skip connection in diffusion model Unets by an exponential factor. The authors should that the feature norms of vanilla Unets oscillate across batches, and that their proposed method results in much smaller feature output oscillations. They conjecture that this method stabilizes training, and prove bounds on the oscillation which their method decreases. The proposed method shows faster convergence and shows improvements across multiple models, datasets, model sizes.

Strengths

1. The authors method follows prior work on scaling block output and/or the skip connection, where such scaling methods have been shown to increase stability/robustness and result in Lipschitz continuous models. 1. The learnt scales in LS method seems to somewhat mirror the author's theoretical exponential CS method, providing further validation. 1. The authors method results in faster convergence and shows improvements across multiple models, datasets, model sizes, and compares favourably to some other scaling methods. 1. The author's proposed method can be very easily adapted to existing models.

Weaknesses

1. The authors show that feature norms oscillate across samples, and they say this implies parameters must also be oscillating (line 47, line 172). While the authors demonstrate that feature norms do oscillate (and I am willing to agree the same may be true of gradient norms), this may not directly cause the parameters/updates to oscillate. 1. The authors works (and proposed exponential CS solution) is remarkably similar to [1], where a scaling of $b^l$ is proposed for the block output (compared to this paper's scaling of the skip connection by a similar value). An experimental comparison with this method should be added to Table 3, or perhaps the similarities/differences discussed. 1. Discussion of prior related works should be expanded. For example, [2] proposes scaling of the skip connection and/or the block output, and has similar exponential bounds. 1. The constants in theorem 3.1 are ignored (assumed to be 1) in line 270 to derive the values for the scaling parameter. If this constant is for example 10 or 0.1, that will dramatically change the proposed values for CS. Further investigation into this constant is required. [1] Hanin, B., & Rolnick, D. (2018). How to Start Training: The Effect of Initialization and Architecture. Advances in Neural Information Processing Systems, 31, 571–581. [2] Balduzzi, D., Frean, M., Leary, L., Lewis, J.P., Ma, K.W. &amp; McWilliams, B.. (2017). The Shattered Gradients Problem: If resnets are the answer, then what is the question?. <i>Proceedings of the 34th International Conference on Machine Learning</i>, in <i>Proceedings of Machine Learning Research</i> 70:342-350 Available from https://proceedings.mlr.press/v70/balduzzi17b.html.

Questions

1. The authors derive bounds for oscillations (eq 4), and they show that their methods decreases this bound, but how tight is this bound? Comparing the minimum and maximum norms predicted by this formula, with that in Figure 1., is important to show if this bound is indeed useful, or perhaps their proposed method has some other side effect that results in the performance improvements. 1. Direct evidence of parameter oscillation (mentioned in line 47, 172 can be shown), perhaps by comparing a per-parameter Adam style "m" and "v" across the baseline method and the author's method. If oscillation of parameters is indeed smaller, one would perhaps expect the norm of "m/v" to be larger in the author's case. Is this indeed the case? 1. Given that the exponential change in scaling with increasing total number of layers, does the performance of the method suffer compared to $\frac{1}{\sqrt{2}}-CS$ as number of layers increases to more depths (as large depth will cause the skip connection to essentially be scaled to zero)? Minor presentation issues (authors are not expected to respond to these, as they are minor issues that can be fixed later): 1. Line 85 "serir" should perhaps be "series"? 1. In table 1, comparisons to author's method "x+LS/CS" should perhaps be placed directly below "x", to make it easier to compare scores. 1. The values for CS for table 1 are mentioned in the supplementary, but perhaps they should be mentioned at line 272 or in 315. 1. Quotes around 'CS.py' in supplementary above Figure 2 are inverted. 1. $\frac{1}{\sqrt{2}}-CS$ reads as "1 by square root 2 minus CS", and for a moment I was trying to think what $CS$ are the authors subtracting from $\frac{1}{\sqrt{2}}$. Perhaps a better notation for this may be something along the lines of $(CS)\frac{1}{\sqrt{2}}$?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

-

Reviewer Ttzr5/10 · confidence 3/52023-07-13

Summary

In this paper, the authors focus on the challenge of the instability arising from the commonly adapted U-Net architecture for diffusion models. In particular, the authors start by theoretically analyzing the influence of the coefficients of long skip connects in U-Net-based diffusion models, specifically on the stability of the forward and backward propagation, along with the robustness of the network. Motivated by the theoretical discussions, the authors propose two corresponding scaling methods that adapt the coefficients of the long skip connections for more stable training, including a constant scaling method and a learnable scaling method, which is straightforward and ideally similar to that in dynamic neural networks and meta learning, but backed by solid theoretical analysis. Experiments on four datasets validate the effectiveness of the proposed two forms of scaling manners, with key codes provided in the supplement.

Strengths

- The proposed constant and learnable scaling methods are very straightforward, but with solid theoretical analysis and guarantee. The results shown in Fig. 1 are also promising that successfully alleviate the oscillations. - Extensive ablation studies including the robustness of LS to network architecture are provided. - Key codes are provided in the supplementary material for handy reproducibility.

Weaknesses

- Method: The second part of learnable scaling method in this paper is algorithmically related to meta learning and dynamic neural networks. However, there is a lack of discussions on related works. Also, it would be great if the authors could discuss some alternative ways for learnable scaling by borrowing the ideas from meta learning and dynamic neural networks. - Experiments: The authors are encouraged to provide some qualitative comparative results with existing methods, as a complement to the quantitative results reported at, for example, Tab. 1. The instability issue in diffusion models with U-Net is not my primary research area. As such, I may not be capable of correctly evaluating the novelty of this paper. I will refer to the comments of other reviewers for my final justification.

Questions

Please kindly refer to the weakness section for more details.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations are discussed in Sect. 6 of the paper.

Reviewer 7ePs6/10 · confidence 2/52023-07-25

Summary

The paper discusses the stability issues observed while training UNet in diffusion models, and theorizes on the role of Long Skip Connections (LSCs) in causing this instability. Diffusion models (DMs), lauded for their ability to model realistic data distributions, involve a forward and a reverse diffusion process. Most DMs use UNet as their backbone due to its use of LSCs which facilitate long-distance information aggregation and prevent vanishing gradient issues. However, despite the use of shared UNet for predicting injected noise at each step, instability is noticed during training. The paper's main contributions revolve around investigating this instability and deriving effective (and efficient) methods to address it. Specifically, it is theoretically proven that the coefficients of LSCs in UNet significantly affect the stability of forward and backward propagation as well as the robustness of UNet. The paper also proposes two coefficient scaling methods, Constant Scaling (CS) and Learnable Scaling (LS), designed to adjust the coefficients of LSCs in UNet for training stability. CS involves setting the coefficients as a series of exponentially-decaying constants while LS uses a small shared network to predict the scaling coefficients for each LSC.

Strengths

- One of the key strengths of the paper is its rigorous theoretical analysis. It provides a comprehensive understanding of the instability of UNet in diffusion models by focusing on the significant impact of the coefficients of Long Skip Connections (LSCs). It's an important insight that broadens the understanding of UNet's performance in diffusion models. - The paper proposes two novel scaling methods, Constant Scaling (CS) and Learnable Scaling (LS), designed to enhance the stability of the UNet training process. These simple methods address the identified problem and provide tangible ways to improve training stability, indicating a proactive approach to problem-solving.

Weaknesses

- The paper primarily focuses on the role of Long Skip Connections (LSCs) in causing training instability in UNet. While LSCs are a significant part of the model, the authors might have explored other potential contributing factors to this instability as well. Broadening the scope of their investigation could potentially have led to a more comprehensive understanding of the issue. - The placement of Figure 1 doesn't align well with its corresponding description, causing some disconnection. Ideally, figures on the first page should offer readers a quick, intuitive understanding of the paper's content. In its present state, Figure 1 seems to lack sufficient information to accomplish this. If this positioning is due to space constraints, please consider relocating it to a more suitable page.

Questions

-

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

-

Reviewer KfUJ7/10 · confidence 2/52023-07-26

Summary

The paper presents a study and algorithm for scaling UNet's long-range connections such that convergence and stability can be improved. The results are strong in the setup of training diffusion models.

Strengths

- The paper is very well written. The text and graphs are polished and it's easy to follow the idea and contributions. - The proposed modification of using an exponential scale over depth is simple and effective, which could be directly applicable to any UNet architecture. - The experiments are thorough and the results are strong. The proposed CS and LS algorithms are shown to be better than the vanilla UNet and heuristic $1/\sqrt(2)$ scaling rule in terms of stability and performance.

Weaknesses

- Explanation of the algorithm: Although the CS and LS algorithms are shown effective, an intuitive explanation on why they work is missing. - For CS, how is the direction of applying the exponentially decayed scale determined? For instance, can we do $\kappa_i=\kappa^{D-i-1}$? In the explanation L251-259, I don't see anything specific to the direction or order. - For LS, I'm assuming there is no additional supervision to the calibration network. In that case, is there any explanation of why it can discover a reasonable scaling rule (Fig.7)? Since both the weights and the scale contribute to the final output of each block, I think decoupling these would be challenging without additional regularization and supervision. Is there any explanation why LS discovers a decaying scaling curve similar to the CS heuristics, but not an increasing scale? - Experiments - For most of the experiments, the results are reported for "Org" (no scaling) and "CS-$1/\sqrt(2)$". However, Fig. 5 does not contain results for "CS-$1/\sqrt(2)$". How does the proposed method compare to it in terms of convergence? - Significance - Although the paper primarily focuses on training diffusion models, many of the analyses seem general enough to be applicable to other scenarios that use UNets, e.g., depth prediction and segmentations. I wonder if this is true and if there is a reason for only showing results on the generation tasks.

Questions

Minor: - L86: There should be a space between Unet and training. - Fig.3 (a): What does the unit m mean on the x-axis? I'm also missing the point this figure is trying to make, as referred to by L161-L163.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Yes

Reviewer 7ePs2023-08-16

I appreciate the authors' thoughtful response. I've carefully read other reviews and the authors' rebuttals, and I've decided to stick with my original score. Based on the feedbacks (particularly concerning typos, figures, and writing), I believe the work can be further polished.

Reviewer xZoY2023-08-17

Thank you for the detailed response. In particular, I am satisfied with your empirical validation of constant C1, and Table 2 comparisons to prior works. The charts in Figure 2 of the global author rebuttal are also welcome. The figure does indeed show the "oscillations" the author's hypothesize and try to fix by bounding. As such, I am raising my review score.

Reviewer VLJa2023-08-17

Response to the rebuttal

Thank you for the response. For the LC parameterization, my question was the latter: for some cases, the CS case adequately achieves good performance for learning $\kappa_i$, and is it beneficial to additively learn with some adaptive module? The rebuttal resolved most of the questions. As I expect this work to contribute to the architectural design of the well-performing diffusion model, I am raising the review score.

Reviewer Ehww2023-08-17

Rebuttal read

I confirm I have read the rebuttal and would like to keep my score.

Reviewer KfUJ2023-08-18

Thank you for the clear explanation. I'd like to keep my original rating of acceptance.

Area Chair hSdX2023-08-18

Reviewer Ttzr

Dear Reviewer Ttzr, Could you please comment on whether the rebuttal addresses your comments and concerns? Best, AC

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC