Generative Diffusion Models Talk (24): Take Fewer Shortcuts, Arrive Faster

生成扩散模型漫谈（二十四）：少走捷径，更快到达 苏剑林 · 2024-04-23 · Machine Learning differential equation generative model diffusion

How to reduce the number of sampling steps while preserving generation quality is a key issue at the application level of diffusion models. Among the earliest attempts at this, Generative Diffusion Models Talk (4): DDIM = DDPM from a Higher Perspective introduced DDIM, which can be regarded as the first serious attempt at accelerating sampling. Later, works introduced in Generative Diffusion Models Talk (5): The General Framework, SDE Edition and Generative Diffusion Models Talk (5): The General Framework, ODE Edition connected diffusion models with SDEs and ODEs, so that the corresponding numerical integration techniques could be directly applied to accelerate diffusion model sampling. Among these, the relatively simple ODE-based acceleration techniques turned out to be especially abundant; we covered one such example in Generative Diffusion Models Talk (21): Accelerating ODE Sampling with the Mean Value Theorem.

In this article we introduce another particularly simple and effective acceleration trick—Skip Tuning—from the paper The Surprising Effectiveness of Skip-Tuning in Diffusion Sampling. To be precise, it is meant to be used in conjunction with existing acceleration techniques to further improve sampling quality. This means that, while keeping the same sampling quality, it can further compress the number of sampling steps, thereby achieving acceleration.more

Revisiting the Model

Everything has to start with U-Net, the mainstream architecture for current diffusion models. The later U-ViT retains roughly the same overall shape, except that the CNN-based ResBlocks are replaced with Attention-based ones.

U-Net comes from the paper U-Net: Convolutional Networks for Biomedical Image Segmentation, originally designed for image segmentation. Its defining feature is that the input and output have the same size, which happens to match exactly what diffusion models need, so it was naturally carried over into diffusion models. Structurally, U-Net looks a lot like a conventional AutoEncoder—both progressively downsample and then progressively upsample—but U-Net adds extra Skip Connections to address the information bottleneck of AutoEncoders:

Illustration from the U-Net paper

Different papers' implementations of U-Net may differ in detail, but they all share the same kind of Skip Connection: roughly speaking, the output of the first layer (block) has a "shortcut" leading directly to the second-to-last layer, the output of the second layer has a "shortcut" leading directly to the third-to-last layer, and so on. These "shortcuts" are the Skip Connections. Without Skip Connections, due to the bottleneck effect, the model's information flow would be constrained by the feature map with the smallest resolution, which would lead to blurry results for tasks that require full information, such as reconstruction or denoising.

Besides avoiding the information bottleneck, Skip Connections also serve as a form of linear regularization. Clearly, if a layer close to the output only used the Skip Connection as its input, then everything after it would essentially have been added in vain, and the model would move closer and closer to a shallow model, or even a linear model. So the presence of Skip Connections encourages the model to prefer as simple (i.e., as close to linear) a prediction logic as possible, resorting to more complex logic only when necessary—this is one of its inductive biases.

Just a Few Lines

Once we understand U-Net, Skip Tuning can actually be explained in just a couple of sentences. We know that sampling in a diffusion model is a multi-step recursive process going from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$, which together constitutes a complex nonlinear mapping from $\boldsymbol{x}_T$ to $\boldsymbol{x}_0$. For practical reasons, we always want to reduce the number of sampling steps used, and regardless of which specific acceleration technique is used, this inevitably reduces the nonlinear capacity of the overall sampling map.

Many algorithms, such as ReFlow, take the approach of adjusting the noise schedule so that the sampling process follows as "straight" a path as possible, making the sampling function itself as close to linear as possible, thereby reducing the quality degradation brought about by acceleration techniques. Skip Tuning instead thinks the other way around: since acceleration techniques cost us some nonlinear capacity, can we compensate for it somewhere else? The answer lies in the Skip Connections. As we just mentioned, their presence encourages the model to simplify its prediction logic—the heavier the Skip Connection, the closer the model gets to a simple linear model or even an identity model. Conversely, then, lowering the weight of the Skip Connection can increase the model's nonlinear capacity.

Of course, this is just one way of increasing the model's nonlinear capacity, and there's no guarantee that the nonlinear capacity gained this way exactly matches the nonlinear capacity lost to sampling acceleration. Yet the experimental results of Skip Tuning show that the two are indeed roughly equivalent! So, true to its name, tuning the weights of the Skip Connections can further improve sampling quality after acceleration, or alternatively reduce the number of sampling steps needed while preserving sampling quality. The tuning procedure is very simple: suppose there are $k + 1$ Skip Connections; we multiply the Skip Connection closest to the input layer by $\rho_{\text{top}}$, the one farthest from the input layer by $\rho_{\text{bottom}}$, and let the rest vary uniformly with depth in between. In most cases we set $\rho_{\text{top}}=1$, so essentially there's only one parameter, $\rho_{\text{bottom}}$, that needs tuning.

The experimental results of Skip Tuning are quite impressive as well. Below are two tables excerpted from the paper; for more experimental figures, readers can consult the original paper.

Skip Tuning results 1Skip Tuning results 2

Personal Reflections

This is probably the simplest article in the diffusion series so far—no lengthy exposition, no complicated formulas—and readers who go straight to the original paper would certainly find it easy to understand. Still, I wanted to introduce it here. Like the previous article, Generative Diffusion Models Talk (23): Signal-to-Noise Ratio and Large Image Generation (Part 2), this work reflects a kind of inventive imagination and observational insight on the part of the authors that I feel I myself sorely lack.

A paper closely related to Skip Tuning is FreeU: Free Lunch in Diffusion U-Net, which analyzes the roles played by different components of U-Net in diffusion models, finding that Skip Connections are mainly responsible for adding high-frequency detail, while the main trunk is mainly responsible for denoising. This gives us another way to understand Skip Tuning: Skip Tuning's experiments are mostly on ODE-based diffusion models, and this type of diffusion model tends to develop more noise as the number of sampling steps is reduced. So shrinking the Skip Connections effectively increases the relative weight of the trunk, strengthening its denoising ability—a case of "treating the symptom with the right medicine." Conversely, for SDE-based diffusion models, one might need to reduce the shrinkage ratio applied to the Skip Connections, or even increase their weight instead, since this type of diffusion model tends to produce overly smooth results when the number of sampling steps is reduced.

Skip Tuning adjusts the Skip Connections, so does that mean architectures like DiT, which have no Skip Connections, have no opportunity to apply this technique? Not necessarily. Although DiT has no Skip Connections, it still has residual connections, and the design of the identity branch is fundamentally the same inductive bias of linear regularization. So even without Skip Connections, tuning the residual connections might yield similar benefits.

Summary

This article introduced a technique that effectively improves the generation quality of diffusion models after accelerated sampling—by lowering the weight of U-Net's "shortcuts" (i.e., Skip Connections). The overall method is remarkably simple and intuitive, and well worth learning about.