Denoising Diffusion Generative Models Secretly Calculate Attentions

A diffusion system is trained to denoise noisy images in a recurssive manner. In the forward process, Gaussian noise is gradually added to the true image $x_{0}$ in hundreds or thousands of steps [13]:

for $\alpha<‌1$. Therefore, each image in the forward process can be written based on the true initial image as:

This eventually leads to a noise-only image at the final stage of the forward process since $\lim_{t\to+\infty}\alpha^{t}=0$. We train a neural network to recover the true image $x_{0}$ from each noisy image $x_{t}$. Theoretically, this results in a neural network that can produce a realistic image from a pure noise and based on the true distribution of the data which is represented by the dataset used for training.

In practice, we train the denoiser neural network on a finite dataset $\mathcal{D}=\{s_{1},\ldots,s_{N}\}$. This finite dataset training situation was first analyzed in [8]. Adding Gaussian noise incurs a mixture of Gaussian distribution on dataset points in signal space:

Then [8] uses a flow matching continuous-time model which is solved backward in time for signal generation:

in which $\gamma_{t}$ is a constant and $s_{t}(y):=\nabla_{y}\log p_{t}(y)$ is the score function. The score function of the mixture distribution in (2) is then [8]:

Then [8] continues the analysis to develop a theory for combinatorial creativity of diffusion models when the denoiser is a CNN network.

Our analysis starts from the same finite dataset training assumption and then proceeds in a different direction to find the relation of the diffusion model to transformers and autoencoders as current dominant models in machine learning research.

When we train the network on the finite dataset $\mathcal{D}$, the network learns to recover each $s_{i}$ from its noisy versions. We guess that when sampling the system for a new signal, although we use a random noise as the input, the system still outputs just an interpolation of the signals it had seen in the dataset since it has not seen anything else.

For signal $s_{i}$ and Gaussian noise $n$, DDPM minimizes Mean-Squared Error (MSE) loss function:

in which $x_{0}$ is the real signal, $\hat{x}_{0}$ its estimate, $t\in\{0,\ldots,T\}$ time index, and $\theta$ is the neural network’s parameters. If we infinitely train the network with fixed $s_{i}$ and $y$, it learns to convert input $y$ to output $s_{i}$.

But we do not train the network just for one sample image. Adding Gaussian noise to any signal $s_{i}$ for sufficiently many times, can reach any point $y$ in the signal space many times. A geometrical scheme of the problem is depicted in Fig. 1. Reaching any fixed point $y$ from various signals $s_{i}$ occurs in accordance with the Gaussian distribution of the noise term $n$. Therefore, if we fix $y$, the system is trained with an average loss function:

For fixed $y$, (7) is a Quadratic Programming (QP) convex optimization problem:

Where (9) is a conic criterion for the strictly convex optimization in (8). The optimal solution for the above optimization problem can be derived using gradients [14]:

Assume that $\sigma$ is sufficiently small. Then just those signals $s_{i}$ that are in the neighborhood of $y$ are effective in the convex combination in (11). Therefore, if $N$ is large enough and the space is sufficiently populated with data points, we can safely assume that $\|s_{i}\|$’s are approximately constant. Other situations also may hold the condition of approximately constant $\|s_{i}\|$ like in image dataset when total energy of most of the images are in the same range. In this case we will have:

in which $<\cdot\,,\cdot>$ denotes the inner vector product.

As we can see in (12), the optimum solution to the input signal $y$ in a diffusion system with fixed $\sigma$ is the attention vector. But this is just one step of the overall diffusion system since each step in the forward diffusion process has a distinct increasing value of $\sigma$. We will discuss this convergence issue in the next section II-D.

The derivation resulting in (12) is also a mathematical basis for the very attention system. It shows that attention is the optimum solution for recovering signals in Gaussian noise as it was shown earlier in [7]. Note that in the conventional attention system [1], the noise variance $\sigma^{2}=\sqrt{p}$, where $p$ is the dimension of the space.

This derivation also shows a basis for the softmax layer. The exponential form in (12) with $\sigma=1$, is the softmax output when $s_{i}$ is a standard one-hot label vector $\mathbb{I}_{i}$:

which shows the equivalence of softmax and attention layers:

Note that the result in (11) can be deduced using the results in [8], replacing (4) and (5) in (3). However, the authors in [8] did not proceed to deduce the attention mechanism from their closed-form score function solution of the diffusion system in (4).

In Sec. II-C, it was shown that each step of a diffusion system calculates attention for a predefined $\sigma$ value. The reverse generative process in DDPM starts with a random Gaussian $y$ and refines it with denoising network trained with noisy images from dataset. Theoretically, there should be hundreds of denoising networks each trained for denoising with a specific $\sigma$. But usually a shared network is trained and used for every values of $\sigma$ to reduce complexity. Therefore, output of each step is then recursively fed in the network to generate another less noisy output image. This is repeated hundreds of times to reach a high quality image.

As can be seen from (1), $\sigma^{2}_{t}=1-\alpha^{t}$ is increasing with $t$ in the forward process. This will result in a decreasing $\sigma_{t}$ in the backward generative process. We start with a Gaussian random $y$. In the early stages when $\sigma$ is large, differences in distances $\|y-s_{i}\|$ in (11) make negligible effect. Therefore, the output sequence converges to the mean value of the nearby vectors.

When $\sigma$ gets small, differences in distances are promoted in (11) and the results quickly converge to the nearest signal, with some contribution from nearby signals. This is a sparse representation, as the simulation of (11) for a simple setup shows in Fig. 2. In this simulation, $\sigma_{t}=0.96\,\sigma_{t+1}$ is a decreasing sequence which is decreased gradually from $1$ to $0.37$ in the reverse process. Although, as we will discuss later in Sec. II-E, it seems that $\sigma$ does not converge to zero because of practical limitations in the denoising neural network design and training.

The diffusion process initially converges to the mean point of the distribution. Then in the long term, it converges to the vicinity of some dense region of the dataset. This is in accordance with the fact that the diffusion process simulates the real data distribution as reflected in the dataset.

Knowledge about the convergence of the diffusion process enables us to circumvent the whole convergence process by directly using a convex combination of some nearby dataset signals as output of the generation process. This basically reduces not only the lengthy generation process, but also the very lengthy training process of denoising neural network.

In this manuscript, we have generally assumed infinite training of denoising network. This enables the training process to cover every $y$ point in the space. It also enables the network through many gradient descent updates, to achieve the ground-truth closed form solution of the training process optimization.

This is rarely the case in practice. Due to limited resources like silicon, time, and energy, we usually perform minimal training up to the point of acceptable network performance. Here we discuss the effect of non-asymptotic training on the network behavior.

First, the network does not experience every $y$ in the space during the training. It only sees finite adjacent inputs $\{y_{1},\ldots,y_{k}\}$. It learns the correct direction of move from these inputs to outputs $\{o_{1},\ldots,o_{k}\}$, respectively. Any new $y$ is sorrounded by multiple seen points. Since the network is a continouos system, we expect an average behavior in $y$ if the training set is sufficiently dense. This guides $y$ to regions of higher density which are experienced more in training and gradually we experience better refinements.

where $\{\alpha_{i}\}$ and $\{\beta_{i}\}$ are two convex coefficient sets, $\mathbb{E}$ is the encoder function and $\mathbb{D}$ is the decoder.

The second consequence of finite training is that the network does not learn the optimal solution of the optimization in (7) for each training input $y_{i}$. We can model this training error as an additive noise in the output of the network. The same is true for the unseen input error and we can model the error in $\mathbb{D}\left(\mathbb{E}\left(y_{i}\right)\right)$ also as an additive noise.

When noise is added to the output of the denoiser network in a diffusion system, the output will no longer be a sparse attention combination even when $\sigma$ of the diffusion system is very small. This is in fact a new additive noise in the backward process:

where $n^{\prime\prime}_{t}$ is the new noise term indicating errors emanating from finite training process.

Therefore, the overall effect of finite training can be modeled as a higher value of $\sigma$ which increases the number of intermittent signals. This means that although in the diffusion framework, theoretically $\sigma$ gradually converges to zero, practical limitations leads to a stable higher limiting value for $\sigma$. This prevents the optimization process to achieve a very sparse solution and the final output will be a convex combination of many signals.

As we showed in Sec. II, there is a deep relationship between diffusion generative systems and attention mechanism. Here, we also show a strong resemblance between AEs, attention systems, and diffusion systems. This will present a great unity between three of the most important artificial intelligence systems currently in use.

First of all, notice that the derivations in Sec. II-C is not limited to a diffusion system. In fact, it is a widespread setting in unsupervised learning systems. Assume we have an AE network which as usual use MSE loss function to reconstruct input images or other signals in an unsupervised manner. Regardless of the AE’s inner structure, we assume perfect training of the system. Most of the times, an augmentation technique of adding Gaussian noise to limited dataset signals is used for better training and performance. Now, we can see that the mathematical setup is identical to the optimization in (7). Therefore, the same solution as in (11) applies to the problem. This certifies that such an AE system is equivalent to an attention system in (12) with fixed $\sigma$ in the infinite training setup. Attention parameter $\sigma$ in (11) can be tuned according to the structural parameters of the AE for a satisfactory equivalence.

The equivalence between AE and attention is double-sided. It means that we can in principle replace current attention systems with AEs. A good example for such a replacement is the linear attention [15], which is known to be a finite-state attention mechanism with compressed rather than full state. This is reminiscent of recursive auto-encoder structure [16], especially when we consider the causal implementation of the linear attention.

Assume an auto-encoder with $m$-dimensional input-output and $k$-dimensional code layer as in Fig. 3. Manifold hypothesis assumes that data from the dataset is located on a high-dimensional manifold in input space [10]. It is well known that the autoencoder estimates the data manifold during the training phase [11]. We will show that AE projects the $m$-dimensional input signal on to a $k$-dimensional manifold which represents locus of data in the dataset and their meaningful augmentations in the space.

Assume that the AE is fully trained with the dataset and therefore, it can now regenerate signal $s$ in the output approximately. In the input space, we add a differential length vector $v_{1}$ to $s$. With high probability, no ReLU-like switch in the network is on its change position. Therefore, for any differential change in the input, the system is linear. Then, changing the input $s$ with $v_{1}$ results in a differential $c_{1}$ vector change in the code layer:

Continue these changes with orthogonal differential vectors $v_{i}$ that results each in code layer change of $c_{i}$. In the $k+1$’th step, the set of $\{c_{1},\ldots,c_{k+1}\}$ is a linearly dependent set since the code layer vector space is $k$-dimensional. Therefore, a linear combination of these vectors sums to zero:

The same set of $\alpha_{i}$ coefficients can be used to find a direction in the input space that keeps the code layer, and therefore the output constant:

Therefore, we showed that any $k+1$ vectors $v_{i}$ in the input space can form a direction with null effect in the output. This direction is:

The local null space of the encoder at $s$ is defined as the local Affine vector space of all such $w$ vectors:

where $\delta$ is a small constant. The null space $\mathbb{N}(s)$ is local to $s$ and is effective just in a $\delta$-neighborhood of $s$. At most $k$ input vectors $v_{i}$ can be chosen that does not lie in this local null space. Therefore:

Suppose $\mathbb{M}(s)$ is the $k$-dimensional local Affine space orthogonal to the local null space:

In the vicinity of $s$, any displacement along the null space $\mathbb{N}(s)$ does not change the values of the code layer and therefore, the output. But displacement along manifold $\mathbb{M}(s)$ translates to changes in code layer and output. Therefore, the AE will tune its null space $\mathbb{N}(s_{i})$ such that it direct farthest from the neighbors of $s_{i}$. This will set the estimated data subspace $\mathbb{M}(s_{i})$ in the nearest position that best cover the neighbors of $s_{i}$ and orthogonal to the local null space. This is in accordance with autoencoder estimating a manifold with good fit to neighboring data points [11] with extensive training on sufficiently dense datasets.

Here we assume that the dataset is dense enough that the continuity of the data manifold can be tracked by the AE as in Fig. 4 . We also assume that the AE does not have a specific structure (like convolutional or any other weight sharing scheme). This helps the AE to follow any direction in the data space without any bias. It is well-known that networks like convolutional networks better estimate the data manifold. This is because some of the data structure, like image content translation invariance is embedded in such networks. This is equivalent to more data points which helps the network better estimate the data manifold.

Any input $y$ outside the manifold defines a specific local null space $\mathbb{N}(y)$ which aggregates directions of moves that does not change the code layer and output layers values. Therefore, AE moves $y$ in a curved path along local null spaces to finally reach the AE estimated data manifold $\mathbb{F}$ in $h$. Since $\mathbb{N}(h)\perp\mathbb{M}(h)$ we reach to the AE manifold perpendicularly in the last touch as is shown in Fig. 4.

Here, we show that the attention mechanism also projects data on a low dimensional manifold. This is in-line with our core insight that basically these three structures:‌ AE, attention, and diffusion systems are similar.

Assume that $y$ is near the data manifold. Dataset satisfies conditions that we can use (11) in place of the better-known attention formula in (12). Also, assume that $\sigma$ is small enough that only $k$ nearest signals:

are effective in the attention sum in (11). Then, the attention coefficients are:

The adjacency set $A_{i}$ forms a $k-1$-dimensional Affine subspace. Assume that the orthogonal projection of $y$ onto this subspace is $h$ as in Fig. 5. Then:

Then $\|y-h\|^{2}$ is shared in numerator and denominator of (27), eliminating it:

Output of $h$ from the attention mechanism is a convex combination of $s\in A_{i}$ and therefore is on their Affine subspace. Also, the fact that the output of $y$ is the same as the output for $h$ shows that there is a local Affine null space orthogonal to the local affine space of nearby signals in $A_{i}$ that does not change the output of the attention.

The attention mechanism when act on a point $y$ near the data manifold $\mathbb{F}$, orthogonally projects it on the manifold at point $h$, and then transform $h$ to another point in the convex hull of signals in $A_{i}$ in the output. Overall, we can conclude that the attention mechanism projects any point $y$ near the data manifold to some point on the manifold too and this is what also auto encoders do.

The attention mechanism depends on the parameter $\sigma$, while the AE is characterized by its code layer dimension $k$. The equivalence between attention and AE requires fine tuning these parameters. As stated before, setting a small $\sigma$ is required in order that approximately $k+1$ nearby signals are effective in the attention. This will result in the estimation of a $k$-dimensional manifold for data similar to an AE with $k$-dimensional code layer. This is not an exact equivalence because the required $\sigma$ parameter may not be constant in different parts of the input space.

In a conventional diffusion system, we have an outer autoencoder which compresses data to a lower dimension. Then a recursive inner autoencoder denoises an input noise. We have shown that not only the high level diffusion framework is equivalent to the attention mechanism, but also both autoencoders in the network structure realization perform in a similar way to the attention mechanism. Therefore, we can alternate between these three structures based on the situation.

If attention and auto encoder structures are basically similar, what is their difference?‌ The difference lies in where the computational load is concentrated. In an auto encoder, we use vast amount of data in training phase while the test phase is as simple as one forward calculation of the network. The attention mechanism is reverse. Calculating attention does not require any training while all the computational load is concentrated in the test time where the forward pass of the network requires vast amount of product calculations with a long context of vectors. Therefore, we may be able to reduce the complex training phase of the auto encoder and diffusion systems by using some simpler form of the attention mechanism.