Introduction to generative models

Generative models have become a cornerstone of modern machine learning, giving us the ability to create new data instances that resemble those found in training sets. In our previous discussions on VAEs, we’ve already touched upon the concept of generative models. VAEs allow us to sample new data points from a learned latent space, producing outputs like images that feel familiar yet are entirely new.

At their core, generative models aim to model the distribution of the data itself. Given some input data, the goal is to learn a probability distribution that can then be sampled to produce new, unseen examples. For instance, given a set of images, a generative model can learn the characteristics of these images and then generate new ones that could plausibly belong to the same set.

These models have a wide range of applications, from generating realistic images, music, and text, to solving more complex problems like filling in missing parts of data, upscaling images, or even designing new molecules in drug discovery.

As we dive deeper into this blog, we’ll explore a particularly exciting class of generative models — diffusion models — which are the core of the stable diffusion solution and uncover how they’re used to create impressive visual outputs

Enters diffusion models

Diffusion models have emerged as a groundbreaking innovation in generative AI, with applications ranging from image generation to audio synthesis. These models gained widespread attention in 2022, powering well-known tools like DALL·E 2 and Google’s Imagen, reshaping the landscape of creative AI.

At their core, diffusion models borrow inspiration from non-equilibrium thermodynamics, particularly the concept of diffusion. But don’t worry — we’re not here to dive into a physics lecture! However, understanding the analogy helps clarify how these models work. Think of it like this: when you drop a tea bag into a cup of hot water, the tea’s molecules diffuse from an area of high concentration (the tea bag) into the water, gradually spreading and changing the color of the water. This diffusion process eventually leads to an equilibrium state, where the tea is evenly distributed throughout the water.

Now, in the case of diffusion models, the idea is similar but applied to data like images or audio. The “information” at the start — think of a crisp image or a clean sound sample — undergoes a forward process where noise is gradually added, akin to the tea molecules spreading into the water. As this process continues, the input data becomes increasingly noisy, losing more and more of its recognizable features.

The magic of diffusion models lies in their ability to reverse this process. Just as it would be impossible to “undo” the tea diffusion in water, these models aim to reverse the diffusion process in data. By learning how to add and remove noise at each step, diffusion models can generate new, high-quality samples from a noisy starting point, whether it’s reconstructing an image or generating entirely new content.

In technical terms, diffusion models work by gradually adding noise to data in a series of steps, modeled by a process called a Markov chain. At each step, a small amount of noise is applied to the input, moving from a clear image, or data point, to a fully noisy version. This is known as the forward process and is represented as \(q(x_t∣x_{t−1})\), where each step \(t\) depends on the previous step \(t−1\).

The goal during training is to learn the reverse process, \(p_\theta(x_{t−1}∣x_t)\), which gradually removes the noise added in the forward steps.

Equations… A lot of Equations!

Now, let’s dive into the mathematics. Don’t worry if the equations seem complex at first—grasping the intuition behind them will make things much easier to follow.

First, we define the forward process, where noise is progressively added to the input data:

Now, the goal of the reverse process is to undo this destruction, starting from noise and gradually “denoising” the data to recover the original clean image:

The reverse process learns to predict \(x_{t−1}\) based on \(x_t\). Instead of explicitly specifying the mean and variance, the model is trained to predict the mean \(\mu_\theta(x_t)\), which helps recover the denoised version of the data at each step. The variance \(\Sigma_\theta(x_t)\) is typically fixed to simplify the computation, as it’s assumed to follow a Gaussian distribution.

Visualizing this pipeline, you might be wondering how the model generates an image when you only provide a text prompt, as seen in various applications. In this scenario, the process starts with an image—but not just any image. It’s pure noise, essentially a canvas of random pixels. The model then refines this noisy image step by step, guided entirely by the text prompt you provide. This technique, known as classifier-free guidance, ensures that the generated image aligns closely with the meaning and context of the text. We’ll dive deeper into how this guidance works shortly.

Going back now to the maths behind diffusion models. As seen previously, these models have two steps: a forward process and a backward process. The goal of training a diffusion model is to find the reverse Markov transitions that maximize the likelihood of the training data. In other words, the model learns to predict the noise that was added to \(x_{t-1}\) to turn it into \(x_t\). The objective is to make the predicted noise as close as possible to the real noise added. More generally, training is achieved by optimizing the variational bound on the negative log-likelihood (NLL).

This might look complex at first glance, but let’s break it down step-by-step. The model’s goal is to learn the parameters required to reverse the noise-adding Markov process. Ideally, we would maximize the likelihood of the true distribution over our dataset, \(\log p_\theta(x_0)\). However, directly maximizing this likelihood is practically infeasible due to the high computational cost involved. This is what we call an intractable problem.

To make this optimization feasible, we instead find a lower bound to this likelihood and maximize it. Maximizing this lower bound indirectly maximizes \(\log p_\theta(x_0)\), achieving our goal. This approach is known as the Evidence Lower Bound (ELBO), defined as:

Think of this with an analogy: imagine you own a company where your total revenue comes from multiple sources like sales, investments, and partnerships. Since your total revenue is the sum of all these channels, any individual channel, like sales, represents a lower bound to your total revenue. If you focus on maximizing sales, your total revenue will increase indirectly. Similarly, in training diffusion models, we maximize this lower bound (the ELBO), which indirectly optimizes the true data likelihood.

But how do we define this lower bound? A useful starting point is the data distribution from the forward process. By comparing this distribution to the joint distribution over all variables from \(x_0\) to \(x_T\), we can express the log-likelihood of our data distribution as,

Using Jensen’s Inequality, we derive a lower bound by moving the logarithm inside the expectation,

And voilà, we have our lower bound that we can maximize instead. This bound, the ELBO, is what drives our optimization process during training.

Practically, once we establish the ELBO, we can use it to define the model’s loss function, as shown in step 5 of the training algorithm. In simple terms, the model learns to predict the noise added at each timestep \(t\), given a noisy image of the form \(\sqrt{\bar{\alpha_t}}x_0 + \sqrt{1 - \bar{\alpha_t}}\). Here, \(x_0\) is the original image, and \(\epsilon\) is the added noise sampled from a Gaussian distribution. By performing gradient descent on this loss function, we optimize the model to minimize the difference between the predicted and true noise values. This, in turn, maximizes the ELBO, which indirectly maximizes the likelihood of our data distribution, helping the model learn the reverse process effectively.

Implementation Details

You might be thinking, “Okay… This might be really hard to implement.” But what if I told you that you can get this up in less than 20 lines of code? Let me show you how straightforward it is.

class Diffusion(keras.models.Model):
    def __init__(self):
        super().__init__()
        self.time_embedding = TimeEmbedding(320)
        self.unet = UNET()
        self.final = UNET_OutputLayer(320, 4)
        
    def call(self, inputs):
        # Inputs:
        # latent: (batch_size, 4, height/8, width/8) - Image latent vector
        # context: (batch_size, seq_len, dim) - Text prompt embedding
        # time: (1, 320) - Time embedding vector
        latent, context, time = inputs
        
        # Time embedding transformation
        # Maps (1, 320) to (1, 1280) for compatibility with UNET
        time = self.time_embedding(time)
        
        # Process latent with UNET
        # Input: latent (batch_size, 4, height/8, width/8)
        # Output: processed latent (batch_size, 320, height/8, width/8)
        output = self.unet(latent, context, time)        
        
        # Refine the output for final prediction
        # Maps back to the original latent format: (batch_size, 4, height/8, width/8)
        output = self.final(output)
        
        
        return output

That’s it! The heavy lifting here is done by the UNET architecture, which we’ll break down shortly. But before diving into that, let’s unpack what’s happening in this function and why these components are essential.

The diffusion model takes three inputs, each serving a crucial role:

• Latent Vector - This represents your input image after being processed by the VAE (Variational Autoencoder) encoder. In simpler terms, it’s a compressed representation of your image data.

• Context - This is the text prompt embedding obtained from the CLIP encoder. It links the text description to the image, guiding the diffusion process toward generating coherent outputs.

• Timestep - A scalar value embedded as a vector, representing the current step in the denoising process. It helps the model understand when it is in the sequence of removing noise.

Each of these inputs flows through a carefully designed pipeline. The time embedding enriches the model’s understanding of the timestep by transforming it into a higher-dimensional space. Then, the UNET does most of the work: it processes the latent image, guided by the text context and timestep, to refine it at each step of the diffusion process. Finally, the output layer ensures the result is mapped back into the latent space’s original dimensions, ready for further processing or decoding.

UNET

Now that we’ve explored the training objective of diffusion models—learning to predict and remove noise—let’s look at the architecture that makes this possible. In Stable Diffusion, the model responsible for predicting the noise at each timestep is called UNet.

UNet was originally introduced as an architecture for biomedical image segmentation, designed to handle the challenge of locating fine details in medical images. Moreover, there is a large consent that deep neural networks require many thousand annotated training samples to be performing. UNET is able to perform really well using few annotated samples by leveraging its built in data augmentation, and is relatively fast compared to other methods. Created by Ronneberger et al. in 2015¹, UNet’s structure features a unique encoder-decoder design (not an actual autoencoder) that captures contextual information at multiple resolutions. The encoder-like part gradually down-samples the input to identify broad patterns, while the symmetric part up-samples the features, restoring spatial details lost in down-sampling.

Stable Diffusion leverages these qualities, repurposing UNet’s structure to predict noise in images. The UNET in stable diffusion takes as input the noisy latent image, produced by the encoder of the VAE, as well as a text prompt and predicts how much noise was added to the latent image, which is the difference between a less noisy image and the input image. The process is used for Denoising Diffusion Probabilistic Models (DDPM) type diffusion models, another approach using the gradient between two steps is used in a score-based diffusion model.

The diffusion process consists in taking a noisy latent image and pass it through the UNET several times. The process ends after a given number of steps, and the output image should represent a noiseless image similar to that of the training data. That is the model should learn to remove the noise and produce an almost exact image as the one of the training data.

We’ve discussed how UNet operates on a noisy image and its role within diffusion models. However, in Stable Diffusion, UNet takes in not only the noisy image but also two additional inputs: the text prompt embedding and the timestep. The text prompt embedding, derived from the CLIP model which we will see later, acts as guidance for the UNet, steering it to generate an image that aligns with the provided text description. Meanwhile, the timestep indicates the amount of noise present at each step, helping the model predict and remove noise more accurately as it moves through the denoising process. These additional inputs enhance UNet’s ability to create high-quality, text-aligned images by combining visual noise reduction with semantic guidance.

Implementation Details

The UNET architecture is the backbone of the diffusion process, responsible for progressively denoising the latent image. It operates in three stages: Input Blocks, a Middle Block, and Output Blocks, connected by skip connections that allow the encoder and decoder to share details. Let’s break it down:

class SpatialTransformer(keras.layers.Layer):
    def __init__(self, channels, n_heads, d_head):
        super().__init__()
        self.norm = keras.layers.GroupNormalization(epsilon=1e-5)
        assert channels == n_heads * d_head
        self.proj_in = PaddedConv2D(n_heads * d_head, kernel_size=1)
        self.transformer_blocks = [BasicTransformerBlock(channels, n_heads, d_head)]
        self.proj_out = PaddedConv2D(channels, kernel_size=1)

    def call(self, inputs):
        x, context = inputs
        b, h, w, c = x.shape
        x_in = x
        x = self.norm(x)
        x = self.proj_in(x)
        x = tf.reshape(x, (-1, h * w, c))
        for block in self.transformer_blocks:
            x = block([x, context])
        x = tf.reshape(x, (-1, h, w, c))
        return self.proj_out(x) + x_in


class Downsample(keras.layers.Layer):
    def __init__(self, channels):
        super().__init__()
        self.op = PaddedConv2D(channels, 3, strides=2, padding=1)

    def call(self, x):
        return self.op(x)


class Upsample(keras.layers.Layer):
    def __init__(self, channels):
        super().__init__()
        self.ups = keras.layers.UpSampling2D(size=(2, 2))
        self.conv = PaddedConv2D(channels, 3, padding=1)

    def call(self, x):
        x = self.ups(x)
        return self.conv(x)


class UNetModel(keras.models.Model):
    def __init__(self):
        super().__init__()
        self.time_embed = [
            keras.layers.Dense(1280),
            keras.activations.swish,
            keras.layers.Dense(1280),
        ]
        self.input_blocks = [
            [PaddedConv2D(320, kernel_size=3, padding=1)],
            [ResBlock(320, 320), SpatialTransformer(320, 8, 40)],
            [ResBlock(320, 320), SpatialTransformer(320, 8, 40)],
            [Downsample(320)],
            [ResBlock(320, 640), SpatialTransformer(640, 8, 80)],
            [ResBlock(640, 640), SpatialTransformer(640, 8, 80)],
            [Downsample(640)],
            [ResBlock(640, 1280), SpatialTransformer(1280, 8, 160)],
            [ResBlock(1280, 1280), SpatialTransformer(1280, 8, 160)],
            [Downsample(1280)],
            [ResBlock(1280, 1280)],
            [ResBlock(1280, 1280)],
        ]
        self.middle_block = [
            ResBlock(1280, 1280),
            SpatialTransformer(1280, 8, 160),
            ResBlock(1280, 1280),
        ]
        self.output_blocks = [
            [ResBlock(2560, 1280)],
            [ResBlock(2560, 1280)],
            [ResBlock(2560, 1280), Upsample(1280)],
            [ResBlock(2560, 1280), SpatialTransformer(1280, 8, 160)],
            [ResBlock(2560, 1280), SpatialTransformer(1280, 8, 160)],
            [
                ResBlock(1920, 1280),
                SpatialTransformer(1280, 8, 160),
                Upsample(1280),
            ],
            [ResBlock(1920, 640), SpatialTransformer(640, 8, 80)],  # 6
            [ResBlock(1280, 640), SpatialTransformer(640, 8, 80)],
            [
                ResBlock(960, 640),
                SpatialTransformer(640, 8, 80),
                Upsample(640),
            ],
            [ResBlock(960, 320), SpatialTransformer(320, 8, 40)],
            [ResBlock(640, 320), SpatialTransformer(320, 8, 40)],
            [ResBlock(640, 320), SpatialTransformer(320, 8, 40)],
        ]
        self.out = [
            keras.layers.GroupNormalization(epsilon=1e-5),
            keras.activations.swish,
            PaddedConv2D(4, kernel_size=3, padding=1),
        ]

    def call(self, inputs):
        x, t_emb, context = inputs
        emb = apply_seq(t_emb, self.time_embed)

        def apply(x, layer):
            if isinstance(layer, ResBlock):
                x = layer([x, emb])
            elif isinstance(layer, SpatialTransformer):
                x = layer([x, context])
            else:
                x = layer(x)
            return x

        saved_inputs = []
        for b in self.input_blocks:
            for layer in b:
                x = apply(x, layer)
            saved_inputs.append(x)

        for layer in self.middle_block:
            x = apply(x, layer)

        for b in self.output_blocks:
            x = tf.concat([x, saved_inputs.pop()], axis=-1)
            for layer in b:
                x = apply(x, layer)
        return apply_seq(x, self.out)

Now this might seem like complicated code, but let me walk you through it. UNET operates in three stages:

• Input Block - which encode the input latent features by progressively downsampling the spatial resolution of the image, where each block contains convolution layers PaddedConv2D for feature extraction, residual blocks ResBlock, SpatialTransformer to apply Cross-attention since we have different inputs (image, text) and finally Downsample to reduce feature map resolution.
• Middle Block - which is the bottleneck of the network, consisting of 2 ResBlock and a SpatialTransformer
• Output_blocks - which decode the features back to the original spatial resolution. There are two additional steps in the code:
• Time Embedding - The timestep embedding is transformed into a feature vector (1280 dimensions) through dense layers and a swish activation to condition the model on the denoising step.
• Final Output - The final block normalizes and refines the output feature map before projecting it back to the latent space dimensions (4 channels) using a PaddedConv2D.

Classifier-free guidance

Recall that the goal of a diffusion model is to learn \(p_\theta(x)\), the distribution over the training data. However, this function alone doesn’t incorporate additional information, such as a text prompt. In other words, while the model learns to generate images similar to the training data, it doesn’t understand the link between a text prompt, like “a cat” or “a mountain”, and the image generated. To address this, we need a way for the text prompt to serve as a conditioning signal, guiding image generation.

One straightforward solution could be to train the model to learn a joint distribution over the data and the prompt \(c\), i.e., \(p_\theta(x, c)\). However, this approach would require the model to heavily rely on the prompt context, which could risk losing its ability to generate diverse images and would require training a separate model for each conditioning signal, which is inefficient. So, the challenge becomes: how can we make sure the model learns both \(p_\theta(x)\) on its own and while conditioned on the prompt?

Classifier-free guidance offers an elegant solution. Instead of using separate models for conditioned and unconditioned data, we use a single model and occasionally drop the prompt signal. At each timestep, the model receives either the noisy image, the timestep, and the prompt to predict the noise to remove (conditioned case), or the noisy image and timestep without the prompt (unconditioned case). This way, the model learns both to generate images aligned with the prompt and without it.

To guide generation, we combine the two outputs (conditioned and unconditioned) with a weighting factor that controls the degree of conditioning. This weighting is similar to the temperature parameter you might encounter in generative models: the higher the weight given to the prompt, the closer the output will be to the prompt’s description, and the lower the weight, the less influence the prompt has. The final output is calculated as:

This approach allows for fine control over how strongly the prompt influences the generated image, making it a versatile and efficient way to incorporate conditioning into diffusion models.

CLIP

Contrastive Language–Image Pretraining (CLIP) marks a significant advance in multimodal learning. Developed by OpenAI in 2021, CLIP was designed to bridge visual and textual representations, making it possible for models to understand images by being instructed with natural language. Unlike earlier approaches, CLIP enables models to perform a wide range of classification tasks without direct optimization for each task—a concept similar to “zero-shot” learning. Before CLIP, models typically excelled at recognizing specific objects or features but struggled with open-ended prompts, such as understanding the subtle differences between a “sunset on the beach” and a “sunset in the mountains.”

During training, CLIP was exposed to paired images and captions, learning to associate them by embedding both into a shared vector space. Instead of predicting specific labels, CLIP’s training taught it to relate images with descriptive text based on content, allowing it to generalize across a vast range of prompts and contexts.

In Stable Diffusion, only the text encoder part of CLIP is used. When given a prompt like “a futuristic cityscape at sunset,” the CLIP encoder transforms this text into an embedding, a numerical representation of the prompt’s meaning. This text embedding serves as a conditioning signal that guides the UNet model in Stable Diffusion, aligning the generated image with the user’s description.

Implementation

As mentioned previously, The CLIP Text Encoder is the backbone of how Stable Diffusion understands and processes text prompts. At its core, it converts a sequence of words into a dense vector or embedding that captures the semantic meaning of the text. This embedding is then used to guide the image generation process.

class ClipTextEmbeddings(keras.layers.Layer):
    def __init__(self, n_vocab, n_embd, n_tokens):
        super().__init__()
        self.token_embedding = keras.layers.Embedding(n_vocab, n_embd)
        self.position_embedding = keras.layers.Embedding(n_tokens, n_embd)
    
    def call(self, inputs):
        # (batch_size, Seq_len) -> (batch_size, Seq_len, dim)
        input_ids, position_ids = inputs
        word_embeds = self.token_embedding(input_ids)
        pos_embeds = self.position_embedding(position_ids)
        return word_embeds + pos_embeds

class ClipEncoderLayer(keras.layers.Layer):
    def __init__(self, n_head, n_embd):
        super().__init__()
        self.layer_norm_1 = keras.layers.LayerNormalization(epsilon=1e-5)
        self.self_attn = CLIPAttention(n_head, n_embd)
        self.layer_norm_2 = keras.layers.LayerNormalization(epsilon=1e-5)
        self.fc1 = keras.layers.Dense(4 * n_embd)
        self.fc2 = keras.layers.Dense(n_embd)
        
    def call(self, inputs):
        #(batch_size, Seq_Len, dim)
        # Self attention
        z, causal_attention_mask = inputs
        residue = z
        z = self.layer_norm_1(z)
        z = self.self_attn([z, causal_attention_mask])
        z += residue
        
        # Feed Forward Layer
        residue = z
        z = self.layer_norm_2(z)
        z = self.fc1(z)
        z = z * tf.sigmoid(1.702 * z) # QuickGelu activation function
        z = self.fc2(z)
        z += residue
        return z

class ClipEncoder(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.layers = [ClipEncoderLayer(12,768) for i in range(12)]

    def call(self, inputs):
        [hidden_states, causal_attention_mask] = inputs
        for l in self.layers:
            hidden_states = l([hidden_states, causal_attention_mask])
        return hidden_states 

class ClipTextTransformer(keras.models.Model):
    def __init__(self):
        super().__init__()
        self.embeddings = ClipTextEmbeddings(49408, 768, 77)
        self.encoder = ClipEncoder()
        self.final_layer_norm = keras.layers.LayerNormalization(epsilon=1e-5)
        self.causal_attention_mask = tf.constant(
            np.triu(np.ones((1, 1, 77, 77), dtype="float32") * -np.inf, k=1)
        )
        
    def call(self, inputs):
        # tokens = keras.ops.cast(tokens, tf.int64)
        # (batch_size, Seq_Len) -> (Batch_size, Seq_Len, Dim)
        input_ids, position_ids = inputs
        x = self.embeddings([input_ids, position_ids])
        x = self.encoder([x, self.causal_attention_mask])
        return self.final_layer_norm(x)

The main class, ClipTextTransformer, handles the process, calling several helper components along the way. Let’s break it down step by step.

• The first step is converting the input text or tokens into numerical embeddings. This is done using the class ClipTextEmbeddings, which combines token embedding and position embedding that adds information about the position of each token in the sequence.
• Once the embeddings are ready, they are passed through multiple layers of a transformer encoder ClipEncoder. This encoder has two main componenents; self-attention and feed-forward networks
• Finally, the output is normalized to ensure consistent embedding representations.

Wrap-Up: Bringing It All Together

Stable Diffusion is more than just an impressive tool for generating art—it’s a convergence of years of research into convolutional networks, attention mechanisms, and generative modeling. In this blog series, we’ve carefully unpacked the building blocks of Stable Diffusion, stage by stage, explaining key concepts like convolutions, attention mechanisms, VAEs, and the innovations behind diffusion models. Each section has been designed to give you a solid understanding of these foundations, bridging theory with practical insights.

Through the series, we’ve explored the individual stages of the Stable Diffusion pipeline in detail, from how text is encoded into meaningful representations, to how noisy latents are progressively refined into coherent, high-quality images. For those eager to see how all these pieces fit together into the complete pipeline, the full implementation can be explored through the provided code in the next section. These resources will allow you to follow along, experiment, and build on what you’ve learned.

What’s Next?

The full pipeline implementation in TensorFlow is waiting for you here and the PyTorch implementation here.

Stable Diffusion is a testament to how far AI has come—and where it’s headed. By mastering its foundations, you’re stepping into a world of infinite creativity and innovation. So, grab the code, start experimenting, and let your creativity flow.

Here’s to building the future—one diffusion step at a time!

The Need for Efficient Data Representations

Imagine I give you two number sequences and ask which one is easier to remember:

1. Sequence A - 3, 7, 12, 19, 24
2. Sequence B - 1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29

Although Sequence A has fewer numbers, each one is unique and independent. You’d have to remember each number individually. On the other hand, while Sequence B is longer, it follows a simple pattern: it’s the sequence of odd numbers from 1 to 29. This pattern makes it much easier to remember because instead of memorizing all 15 numbers, you only need to store the starting number (1), the ending number (29), and the rule (“odd numbers”).

This analogy highlights the concept of efficient data representations. Instead of storing or processing all the details of the data, we seek a compact and structured way to capture the essential features, reducing redundancy and preserving the key information. This is especially useful when dealing with high-dimensional data like images, audio, or text. Efficient representations enable models to focus on meaningful patterns and relationships within the data, leading to better performance in tasks such as classification, clustering, and generation.

What Are Autoencoders?

Autoencoders are a type of neural network architecture designed to learn efficient data representations in an unsupervised manner. Unlike supervised learning, where models learn to map inputs to specific outputs, autoencoders learn to encode the input data into a compressed form and then decode it back to reconstruct the original input. This process helps them identify and retain the most critical features while discarding irrelevant information.

The architecture of a basic autoencoder consists of two main parts:

• Encoder - The encoder compresses the input data into a lower-dimensional space known as the latent representation. This compression forces the network to learn the essential characteristics of the input data and discard the less significant details.
• Decoder - The decoder takes the latent representation and reconstructs it back to match the original input as closely as possible. This reconstruction step ensures that the encoder has effectively captured the underlying structure of the input data.

Visually, autoencoders resemble a typical neural network with an input layer, hidden layers, and an output layer. However, unlike other neural networks, the number of neurons in the output layer of an autoencoder must be equal to the number of inputs to ensure it can reconstruct the original data.

Since the latent representation has a lower dimensionality than the input data, the network is unable to simply copy the input features into the latent space. Instead, it’s forced to learn the most salient features and representations that summarize the input. This makes autoencoders highly efficient at dimensionality reduction and feature extraction.

Moreover, autoencoders can be used as generative models. By training on a given dataset, they learn to generate new data samples that resemble the original training data.

Convolutional Autoencoders

In our previous discussions on images, we highlighted how convolutional layers are more effective than dense layers for capturing spatial hierarchies and local features. When building autoencoders for images, it’s beneficial to use Convolutional Autoencoders (CAEs) instead of fully connected ones. This allows us to leverage the strengths of convolutional operations, such as detecting edges, textures, and shapes, while preserving spatial structure.

But rest assured, the core architecture—comprising the encoder, latent representation, and decoder—remains the same with some minor modifications. The key difference is how these components are constructed:

• Encoder - The encoder in a convolutional autoencoder is similar to that of a standard CNN. It consists of a series of convolutional and pooling layers (or downsampling layers). As with a typical CNN, the purpose of the encoder is to progressively reduce the height and width of the image while increasing its depth, meaning it transforms the original image into a set of lower-dimensional feature maps. These feature maps capture the key patterns and spatial structures of the image in a compact form.
• Decoder - The decoder reverses the transformations applied by the encoder. It uses upsampling layers (like transposed convolutions or interpolation layers) to gradually increase the height and width of the feature maps while reducing their depth, reconstructing the input image. The objective is to restore the image’s original dimensions and structure as accurately as possible, based on the compressed latent representation.

For now, we won’t delve into the implementation details until we cover the theoretical aspects more thoroughly. However, later in this chapter, we’ll show how convolutional autoencoders can be tailored for more complex tasks such as Stable Diffusion. In fact, Stable Diffusion uses convolutional autoencoders as part of its architecture. We’ll explore how and why these models are used in greater detail as we proceed.

Variational Autoencoders

Variational Autoencoders (VAEs) are one of the most widely used and popular types of autoencoders, introduced in the paper¹. What sets them apart from traditional autoencoders is their probabilistic nature and ability to generate new data points that resemble the training data.

The VAE plays a crucial role in diffusion model architecture. The encoder produces the latent representation of an input image, in the case of image-to-image, or a noisy input in the case of text-to-image, while the decoder is the final step that reconstructs and outputs the image you requested with a text prompt.

The key distinction between a VAE and a standard autoencoder is that instead of mapping the input to a fixed latent vector, VAEs aim to map it to a distribution. More precisely, the latent vector is replaced by two components: the mean \(\mu\) and the standard deviation \(\sigma\). From this Gaussian distribution, defined by \(\mu\) and \(\sigma\), we sample a latent vector, which is then passed to the decoder to reconstruct the input, just like in traditional autoencoders. This leads to two important characteristics of VAEs:

• Probabilistic Nature - The latent space of VAEs is probabilistic, meaning that each time we sample from the Gaussian distribution, we get slightly different latent representations. As a result, the output of the decoder is also probabilistic, introducing variability in the reconstructions.
• Generative Capabilities - Because of their stochastic nature, VAEs are generative models. They can generate new instances that look similar to the data they were trained on but are, in fact, unique. For example, if a VAE is trained on images of dogs, it could generate a new image of a dog that resembles the training data but includes subtle variations, like a dog with three ears. These new samples come from the same underlying distribution as the training data but introduce novel variations.

As illustrated in the diagram above, after the encoder processes the input, the final layer of the encoder outputs two vectors: \(\mu\) and \(\sigma\) of the latent distribution. Instead of directly sampling from these vectors, an additional step is introduced — the reparameterization trick.

The Reparameterization Trick

In the ideal scenario, we would sample directly from the latent distribution using the computed \(\mu\) and \(\sigma\). However, during training, we use backpropagation to compute the gradient of the loss with respect to every trainable parameter in the network. The stochastic nature of \(\mu\) and \(\sigma\) creates a problem here, as the gradients cannot propagate through the sampling step.

The solution is the reparameterization trick, which introduces a clever workaround: Instead of sampling directly from the Gaussian distribution defined by \(\mu\) and \(\sigma\), we transform the process into something more manageable for backpropagation. Here’s how it works:

1. We sample an auxiliary variable \(\epsilon\) from a standard normal distribution with \(\mu = 0\) and \(\sigma = 1\).
2. The latent vector is then computed as \(z=μ+σ\ast \epsilon\). This operation allows the mean and standard deviation to remain deterministic, which makes them suitable for gradient-based learning. The stochastic aspect is now captured by \(\epsilon\).

By using this trick, we move the randomness into a non-trainable node \(\epsilon\) while keeping \(\mu\) and \(\sigma\) as trainable parameters, allowing gradients to flow during backpropagation.

Why Not Train \(\epsilon\)?

A common question might arise: “Why not train \(\epsilon\)?” The answer is simple — \(\epsilon\) is sampled fresh in every forward pass, acting as a fixed noise source. Its role is to introduce controlled randomness, and it doesn’t need to be learned since it always comes from the standard Gaussian distribution (mean 0, standard deviation 1).

VAE Equation

Now that we’ve covered how VAEs differ from traditional autoencoders, let’s dive deeper into their loss function. The VAE loss combines two important components that work together to ensure the model can generate new data while keeping the latent space well-organized.

Below is the general formula for the VAE loss:

It consists of two parts:

• Reconstruction Loss - The first term helps the model recreate the input data as accurately as possible.
• Regularization Loss - The second term, also known as the Kullback-Leibler (KL) divergence, penalizes the model if the latent space deviates too much from the normal distribution.

Let’s break these components down further.

Reconstruction Loss

The primary goal of the VAE is to take an input (e.g., an image), encode it into a simpler latent representation, and then decode it to reconstruct the original input. The reconstruction loss measures how close the reconstructed output is to the original input.

To make this clearer: Suppose you give the model a picture of a cat. The model compresses the image into a latent vector and then reconstructs it. The reconstruction loss checks how similar the new image is to the original. The more accurate the reconstruction, the smaller this loss becomes.

Technically, this term is the likelihood of the original input \(x\) given the latent representation \(z\) from the encoder. We aim to maximize this likelihood so that the decoder, represented by \(p_\theta(x∣z)\), can generate data as close as possible to the original input, which was compressed by the encoder \(q_\phi(z∣x)\).

Regularization Loss

The second part of the VAE loss function is the regularization term, which is where the probabilistic nature of VAEs comes into play. Recall that the encoder doesn’t output a single deterministic latent vector but instead generates a mean \(\mu\) and standard deviation \(\sigma\), allowing us to sample from a Gaussian distribution to create the latent representation.

This stochastic property adds flexibility, allowing the VAE to generate smooth, continuous variations of new data. However, to ensure that the latent space is well-structured and meaningful, we need to regularize it. This is achieved through the KL divergence, which measures how much the learned latent distribution, produced by the encoder, deviates from a standard Gaussian distribution. The goal is to make the latent vectors follow this normal distribution so that similar inputs produce similar latent representations. If the latent space is poorly organized — i.e., if latent vectors are scattered randomly — generating new, coherent data points would be difficult. The KL divergence penalizes the model when the latent vectors stray too far from the Gaussian distribution, encouraging a well-organized and continuous latent space.

In summary, the regularization term ensures that the latent space remains structured, preventing it from becoming chaotic, and helping the model generate meaningful and smooth variations from the training data.

Residual Blocks

Residual blocks, or skip connections, were a game-changing innovation introduced in the ResNet paper². They addressed a rather surprising issue in deep learning: adding more layers to a neural network doesn’t always lead to better performance. In fact, as networks grew deeper, their performance often degraded. This ran contrary to the fundamental idea of deep learning—that deeper networks should be able to capture more complex features and perform better.

So, what was going wrong? As networks became deeper, they struggled to learn optimal weights. Information struggles to pass through all the layers, and the gradients during backpropagation became increasingly small, leading to what’s known as the vanishing gradient problem. This made it difficult for the network to update weights, especially in the earlier layers, causing deep models to underperform compared to shallower ones.

Residual blocks provided a clever solution to this problem. Rather than forcing the network to directly learn the full mapping \(f(x)\), residual blocks allow the network to learn the residual \(h(x)=f(x)-x\), which simplifies to \(f(x)=h(x)+x\). The key is an identity connection — a shortcut — that bypasses one or more layers, letting information skip through the network without obstruction. This allows gradients to flow more freely during backpropagation, addressing the vanishing gradient issue.

But the benefits of residual blocks don’t stop there. These identity connections also speed up training. Early in training, when weights are initialized near zero, approximating the identity function provides a helpful kickstart. In traditional networks, some layers may struggle to “wake up” because of small gradients or poor initialization. But with residual blocks, the model can start by learning the identity function \(f(x)\approx x\) , which allows it to make rapid initial progress before moving on to learn more complex mappings. This ensures that the model doesn’t get stuck early on, making learning faster and more efficient.

Now, you might wonder: why are residual blocks important in the context of stable diffusion? In architectures like Stable Diffusion, residual blocks play a critical role in maintaining a smooth flow of information through the deep layers of the network. Generating high-quality images depends on this effective information flow. By incorporating residual connections within the autoencoders, the model ensures that transformations in the latent space remain consistent and stable, even as it manipulates intricate details in images.

Implementation

In this section, we will implement the building blocks of a VAE for Stable Diffusion. Specifically, we’ll focus on the encoder, decoder, and the all-important residual block. By the end, you’ll not only understand how these components work but also how they come together to form a functional VAE.

Before diving into the encoder and decoder implementation, let’s first look at the residual block which is implemented as follows:

class ResBlock(keras.layers.Layer):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        
        # Layers to process the input features
        self.in_layers = [
            keras.layers.GroupNormalization(epsilon=1e-5),
            keras.activations.swish,
            PaddedConv2D(out_channels, kernel_size=3, padding=1)
        ]

        # Layers to process the time embedding
        self.emb_layers = [
            keras.activations.swish,
            keras.layers.Dense(out_channels)
        ]

        # Layers to further refine the merged features
        self.out_layers = [
            keras.layers.GroupNormalization(epsilon=1e-5),
            keras.activations.swish,
            PaddedConv2D(out_channels, kernel_size=3, padding=1)
        ]

        # Skip connection for residual learning
        self.skip_connection = (
            lambda x: x if in_channels == out_channels else PaddedConv2D(out_channels, kernel_size=1)
        )

    def call(self, inputs):
        # Unpack the inputs: feature maps and time embedding
        z, time = inputs
        residue = z  # Save the input for the skip connection

        # Apply the input layers
        z = apply_seq(z, self.in_layers)

        # Process the time embedding
        time = apply_seq(time, self.emb_layers)
        
        # Merge the feature maps with the time embedding
        merged = z + time[:, None, None]

        # Apply the output layers
        merged = apply_seq(merged, self.out_layers)

        # Add the skip connection and return the result
        return merged + self.skip_connection(residue)

The Residual Block uses a skip connection, which helps retain information from earlier layers. It also integrates time embeddings which are essential for the Diffusion model.

Now moving on to the encoder, which is the component responsible for transforming the input image into a latent representation. Below is the implementation in TensorFlow.

import tensorflow as tf
from tensorflow import keras

# Define the encoder for the Variational Autoencoder
class VAE_Encoder(keras.Sequential):
    def __init__(self):
        super().__init__([
            # Initial convolution to extract 128 features
            PaddedConv2D(128, kernel_size=3, padding=1),  # (batch_size, 128, height, width)

            # Stack of ResNet blocks to refine features
            ResnetBlock(128, 128),
            ResnetBlock(128, 128),  # (batch_size, 128, height, width)

            # Downsample using strided convolution
            PaddedConv2D(128, kernel_size=3, strides=2, padding=(0, 1)),  # (batch_size, 128, height/2, width/2)

            # Further feature extraction and dimension increase
            ResnetBlock(128, 256),  # (batch_size, 256, height/2, width/2)
            ResnetBlock(256, 256),

            # Downsample again
            PaddedConv2D(256, kernel_size=3, strides=2, padding=(0, 1)),  # (batch_size, 256, height/4, width/4)

            # Continue with higher-dimensional feature extraction
            ResnetBlock(256, 512),  # (batch_size, 512, height/4, width/4)
            ResnetBlock(512, 512),

            # Final downsampling to reduce spatial dimensions
            PaddedConv2D(512, kernel_size=3, strides=2, padding=(0, 1)),  # (batch_size, 512, height/8, width/8)

            # Deep feature extraction using multiple ResNet blocks
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),

            # Attention block for contextual feature aggregation
            AttentionBlock(512),

            # Additional refinement with ResNet block
            ResnetBlock(512, 512),

            # Normalize and activate features
            keras.layers.GroupNormalization(epsilon=1e-5),
            keras.layers.Activation('swish'),

            # Final convolution to reduce feature dimensions
            PaddedConv2D(8, kernel_size=3, padding=1),  # (batch_size, 8, height/8, width/8)
            PaddedConv2D(8, kernel_size=1),  # (batch_size, 8, height/8, width/8)

            # Scale latent representation
            keras.layers.Lambda(lambda x: x[..., :4] * 0.18215),
        ])

Inversely, the decoder takes the latent representation and reconstructs the image by progressively upsampling and refining the features.

# Define the decoder for the Variational Autoencoder
class VAE_Decoder(keras.Sequential):
    def __init__(self):
        super().__init__([
            # Rescale the latent input
            keras.layers.Lambda(lambda x: 1 / 0.18215 * x),

            # Initial convolution to expand features
            PaddedConv2D(4, kernel_size=1),
            PaddedConv2D(512, kernel_size=3, padding=1),

            # Stack of ResNet and Attention blocks to refine features
            ResnetBlock(512, 512),
            AttentionBlock(512),
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),

            # Upsample spatial dimensions
            keras.layers.UpSampling2D(size=(2, 2)),
            PaddedConv2D(512, kernel_size=3, padding=1),

            # Further refinement with ResNet blocks
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),
            ResnetBlock(512, 512),

            # Upsample and refine again
            keras.layers.UpSampling2D(size=(2, 2)),
            PaddedConv2D(512, kernel_size=3, padding=1),
            ResnetBlock(512, 256),
            ResnetBlock(256, 256),
            ResnetBlock(256, 256),

            # Final upsampling to original image dimensions
            keras.layers.UpSampling2D(size=(2, 2)),
            PaddedConv2D(256, kernel_size=3, padding=1),
            ResnetBlock(256, 128),
            ResnetBlock(128, 128),
            ResnetBlock(128, 128),

            # Final normalization and activation
            keras.layers.GroupNormalization(32),
            keras.layers.Activation('swish'),

            # Final convolution to map back to RGB channels
            PaddedConv2D(3, kernel_size=3, padding=1),  # (batch_size, 3, height, width)
        ])

Introduction

In 1950, Alan Turing proposed the famous Turing Test, challenging machines to exhibit human-like intelligence — particularly in mastering language. He envisioned a chatbot based on hardcoded rules, designed to fool a human interlocutor into believing they were conversing with another human. Today, we stand at a point where Large Language Models (LLMs), like GPT-4o, Llama3, and Mistral, have progressed far beyond Turing’s vision, often outperforming humans in tasks like summarization and translation. But how did we reach this point?

Before LLMs became dominant, earlier architectures, such as Recurrent Neural Networks (RNNs), laid the groundwork. Despite different approaches, these models all shared a core principle: predicting the next word based on previous ones. The more “previous words” or context a model could handle, the better its predictions — something we’ll delve into shortly.

But how can machines even begin to understand human language? The key lies in Embeddings — a technique for transforming words into numerical representations that machines can process.

Tokenization

Before discussing embeddings, it’s important to grasp tokenization—the process of breaking down a sentence into smaller units called tokens, which are then converted into embeddings. Tokenization can vary across models, but there are three common methods:

• Word Tokenization - The text is split into individual words.
• Character Tokenization - The text is split into individual characters.
• Subword Tokenization - The text is split into partial words or character sets. This method is commonly used in OpenAI’s GPT models via Byte-Pair Encoding (BPE).

Word embeddings

Once a sentence is tokenized, each token is transformed into a vector of floating-point numbers — known as an embedding. These embeddings allow the machine to capture the meaning and relationships between words. For example, the words “elementary,” “dear,” and “Watson” can be represented as vectors in a 3D space:

• “elementary” → [-2, 0, 1]
• “dear” → [-1, 0, -3]
• “Watson” → [2, 0, 2]

This 3D representation is a simplified example. In reality, embeddings used in models like GPT-3 or GPT-4 have much higher dimensions. GPT-3’s embeddings, for instance, span 12,288 dimensions. Such high-dimensional embeddings can’t be visualized, but their purpose remains the same: to represent semantic relationships between words. These embeddings exist within what’s called the embedding space or vector space. In this space, words with similar meanings or relationships are positioned closer together, while words with different meanings are farther apart. For example, the words “king” and “queen” may be positioned near each other due to shared attributes, while “apple” would lie farther away, representing a completely different concept.

As the model trains, it refines these word embeddings to best capture the meanings it learns. Different models have their own unique ways of learning embeddings, but in the case of the transformer architecture, the core architecture of LLMs, embeddings do more than just encode the meaning of individual words. They also include richer information, such as the word’s position in the sentence and its surrounding context.

Contextual word embeddings

Language is contextual — words can change their meaning depending on the surrounding words. Transformer models take this into account by adjusting the embedding of a word based on the context in which it appears. The larger the context window, the better the model is at capturing meaning. For example, a model with a 32,000-token context window can consider 32,000 preceding words when generating embeddings for the current word. Naturally, a model with a larger context window performs better than one with a smaller window, as it can “remember” more of the conversation or text.

One way to test this is to compare two models with different context windows: a model with a smaller window will quickly lose track of a long conversation, while a model with a larger window will retain much more context, leading to more coherent responses.

So far, we’ve discussed how embeddings capture the meaning of words in context. But not all words in a context window are equally important. In fact, certain words contribute much more to understanding the meaning of a specific word than others. This is where the Attention mechanism comes into play, a fundamental concept in transformer models. Attention allows the model to focus on the most relevant words in the context, giving more weight to those that are crucial for understanding the current word or token.

How does Attention mechanism work?

The attention mechanism was originally introduced by Dzmitry Bahdanau and colleagues¹ to enhance sequence-to-sequence models. Its purpose was to allow the decoder to focus on specific parts of the input sequence—encoded by the encoder—at each step. This proved particularly useful in tasks like machine translation, where the decoder must generate output based on relevant sections of the source sentence.

Over time, several variations of attention have been proposed², but we will focus on the Scaled Dot-Product Attention introduced in the transformer architecture³. This mechanism is central to modern models like GPT and BERT, providing a way to weigh different parts of the input sequence based on their importance to the current task.

To build a deeper intuition for how transformers work, I highly recommend watching Grant Sanderson’s explanation on his YouTube channel, 3Blue1Brown. His visual breakdowns are incredibly insightful.

Now, let’s develop a clear intuition for the attention mechanism before breaking down its underlying mathematics.

Imagine you’re working on a machine translation model, using an encoder-decoder architecture. The encoder processes the input sentence, “I drink milk,” and generates embeddings that represent the sentence’s meaning. The decoder then takes these embeddings and generates the translation, “Je bois du lait.”

During this process, the encoder not only encodes the individual words but also captures the grammatical structure — like recognizing that “I” is the subject and “drink” is the verb. Now, let’s say the decoder has already translated the subject and needs to move on to translating the verb. How does it know which word in the input corresponds to the verb?

It’s helpful to think of the encoder as having created an implicit dictionary like this:

• “subject” → “I”
• “verb” → “drink”
• …

When the decoder is ready to translate the verb, it essentially needs to “look up” the corresponding verb from this dictionary. However, instead of having explicit keys like “subject” or “verb,” the model relies on the embeddings generated by the encoder.

Here’s where attention comes into play. The decoder generates a query, essentially asking, “What part of the input sentence corresponds to the verb I need to translate?” The query is then compared to a set of keys derived from the encoder’s output embeddings. The comparison is done using a similarity measure, specifically the dot-product. The key that most closely matches the query tells the decoder which part of the input to focus on.

Once the attention mechanism identifies the most relevant key, the decoder retrieves the corresponding value, also derived from the encoder’s embeddings, to help generate the next word in the translation.

What we have described so far is how the attention mechanism operates in the context of an encoder-decoder architecture. In such models, the encoder processes the entire input sentence and the decoder focuses on specific parts of the encoded representation during translation. However, there’s another type of attention, called self-attention, which is crucial in transformer architectures.

Self-attention operates similarly to the encoder-decoder attention, but with one key difference: instead of focusing on an external sequence, each word in a sentence attends to every other word in the same sentence. This allows the model to build rich contextual representations of each word, accounting for its relationships with all other words in the sequence.

For example, in the sentence “The cat sat on the mat,” the word “cat” would attend to every other word to understand its role in the sentence. This includes attending to “sat” to infer an action associated with the subject and attending to “mat” to identify where the action took place. By doing so, self-attention captures context and dependencies within the sentence more effectively.

Understanding the Scaled Dot-Product Attention

Now that we have an intuition behind the attention mechanism, let’s dive deeper into the Scaled Dot-Product Attention used in transformers. This form of attention is defined mathematically as follows,

Each matrix is derived by linearly projecting the original token embeddings through learned weight matrices. Let’s break down each component step-by-step:

• Query Matrix \(Q\) - The query matrix is computed by multiplying each token embedding \(E_i\) (at position \(i\)) by a learned weight matrix \(W_Q\). This transformation maps the token embeddings into a new space specifically designed for querying relevant information. The shape of the query matrix is \((n_{queries},d)\),where \(n_{queries}\) is the number of queries (typically equal to the number of tokens in the input) and \(d\) is the dimension of each query vector. For example, in GPT-3, $d$ is 128.
• Key Matrix \(K\) - The key matrix is obtained by multiplying the same token embeddings \(E_i\) by another learned weight matrix \(W_K\). This transformation projects the token embeddings into a new space that is specific to “keys.” Its shape is \((n_{keys},d)\), where \(n_{keys}\) is number of keys and matches the number of queries \(n_{queries}\) in the context of self-attention.
• Dot-product \(QK^T\) - The dot-product between the query and key matrices results in a matrix of shape \((n_{queries},n_{keys})\). Each entry in this matrix represents the similarity score between a query vector and a key vector. Intuitively, if a query closely matches a key, the dot-product will yield a high, positive value. Conversely, if they don’t match, the resulting score will be close to zero or even negative.
• Softmax Normalization - To convert these similarity scores into weights, we apply the softmax function. The softmax operation normalizes the scores such that they sum to 1 along each row, transforming the scores into a probability distribution. This helps the model focus on the most relevant parts of the input while suppressing less important information.
• Scaling by \(\sqrt{d}\) - The division by \(\sqrt{d}\) is introduced to prevent the dot-product values from growing too large, which can lead to gradients that are either too small or too large during training. This scaling ensures more stable optimization.
• Computing the Attention Weights - The result of applying softmax to \(\frac{QK^T}{\sqrt{d}}\) gives us the attention weights. These weights indicate the importance of each word in the input sentence with respect to the query. In other words, the weights tell us how much each word should contribute to the meaning of the word currently being updated.
• Value Matrix \(V\) - Finally, we use a third weight matrix \(W_V\) to project the token embeddings into the value space, forming the value matrix \(V\). This matrix holds the information that we want to pass on, based on the computed attention weights.
• Weighted Sum of Values - The final step is to compute the weighted sum of the value vectors, using the attention weights as coefficients. This step results in a new representation for each word that incorporates contextual information from all other words in the sentence.

Intuition with an example

Let’s return to our example: the sentence “You are a wizard, Harry”. Suppose we want to understand the word “Harry” better by considering its surrounding context. Initially, the model assigns a random embedding to the word “Harry.” However, when we compute the attention weights, we might find that the word “wizard” has a high relevance score with “Harry.” This means the information from the value vector of “wizard” will be heavily weighted when updating the embedding for “Harry.”

As a result, the new embedding for “Harry” will be influenced mostly by its relationship to “wizard,” enabling the model to understand that “Harry” refers to the fictional character, Harry Potter. This context-aware embedding is what gives attention-based models their power in understanding and generating natural language.

Masking

In some cases, we don’t want the model to have access to all the words in a sentence simultaneously. Consider the sentence “Building AI models is cool”. If we want the model to predict the word “AI”, it would be unfair if it could attend to the future words “…models is cool” — essentially giving it access to information it should not yet know. This would be like giving away the answer and would prevent the model from learning proper word dependencies and sequences.

To prevent this, we introduce a technique called masking. Masking allows us to limit the model’s access to certain tokens based on the current context, ensuring that each word prediction only considers the relevant parts of the input.

For example, when predicting the word “AI”, the model should only consider the previous tokens “Building” and not see any of the future words. Similarly, when predicting the next word “models”, it should have access to “Building AI”, but not to “is cool”. This constraint helps the model learn causal relationships and generate text in a step-by-step manner.

How Does it work? The simplest way to implement masking is by modifying the attention weights of future tokens. Specifically, we set these weights to zero, so that the model does not pay any attention to these tokens. This is often achieved by adding a large negative number, typically -\(\infty\) (infinity) to the attention scores of the masked positions before applying the softmax function.

Multi-head Attention

We’re almost at the end of this chapter on attention mechanisms, but there’s one more important concept to cover: multi-head attention.

Multi-head attention is an extension of the scaled dot-product attention mechanism we’ve already discussed. Rather than having just one set of attention weights and outputs, multi-head attention allows the model to use multiple attention mechanisms (or “heads”) in parallel. Each head operates independently and captures different aspects of the input sequence.

Why use multiple attention heads? In the example of using scaled dot-product attention to analyze a sentence, we mentioned that attention can help capture certain grammatical relationships, such as identifying which word is the subject and which is the verb. However, a single head might be limited in what it can represent. It may capture grammatical roles well but miss other nuances like verb tense or semantic relationships between words.

Multi-head attention addresses this limitation by using multiple attention mechanisms in parallel, each focusing on different aspects of the input:

• One head might specialize in capturing grammatical relationships.
• Another head could learn to identify verb tense.
• Yet another might focus on identifying relationships between named entities.

Each head provides a unique “view” of the sentence, allowing the model to understand a richer set of features.

How Does Multi-Head Attention Work? Here’s a step-by-step breakdown of the multi-head attention process: For each token embedding in the input sequence, we apply multiple linear transformations using different sets of learned weight matrices. This results in multiple sets of query, key, and value vectors—one for each head. Each head computes its own scaled dot-product attention using its unique set of query, key, and value vectors. The outputs from all heads are concatenated into a single matrix, combining the unique features learned by each head. Finally, the concatenated output is then passed through a final linear transformation to project it back into the original embedding space.

Implementation

Now, let’s dive into implementing the multi-head attention mechanism. While TensorFlow provides a built-in implementation, we’ll build it from scratch to gain a deeper understanding of its inner workings.

import tensorflow as tf
import numpy as np
import keras
import math

def make_triu(x, dtype=tf.bool):
    """Creates an upper triangular mask for causal attention."""
    ones_like = tf.ones_like(x)
    triu_matrix = tf.linalg.band_part(ones_like, 0, -1)
    result = triu_matrix - tf.linalg.diag(tf.linalg.diag_part(triu_matrix))
    return tf.cast(result, dtype=dtype)

class SelfAttention(keras.layers.Layer):
    def __init__(self, n_heads, d_embd, in_proj_bias=True, out_proj_bias=True):
        super().__init__()
        # Linear layer to project input embeddings into query, key, and value spaces
        self.in_proj = keras.layers.Dense(3 * d_embd, use_bias=in_proj_bias)
        # Linear layer for output projection after attention computation
        self.out_proj = keras.layers.Dense(d_embd, use_bias=out_proj_bias)
        self.n_heads = n_heads  # Number of attention heads
        self.d_head = d_embd // n_heads  # Dimension of each head

    def call(self, inputs, causal_mask=False):
        # Retrieve the input shape: batch size, sequence length, embedding dimension
        input_shape = inputs.shape
        b, s, d_embd = input_shape

        # Reshape intermediate tensors to split dimensions for multi-head attention
        interm_shape = [b, s, self.n_heads, self.d_head]

        # Project input to obtain queries (Q), keys (K), and values (V)
        q, k, v = tf.split(self.in_proj(inputs), 3, axis=-1)  # 3 tensors of shape: (batch_size, seq_len, d_embd)

        # Reshape and transpose Q, K, and V for multi-head attention
        q = tf.transpose(tf.reshape(q, interm_shape), perm=[0, 2, 1, 3])  # (batch_size, n_heads, seq_len, d_head)
        k = tf.transpose(tf.reshape(k, interm_shape), perm=[0, 2, 1, 3])
        v = tf.transpose(tf.reshape(v, interm_shape), perm=[0, 2, 1, 3])

        # Compute attention weights by taking the dot product of Q and K^T
        weight = q @ tf.transpose(k, perm=[0, 1, 3, 2])  # (batch_size, n_heads, seq_len, seq_len)

        # Apply a causal mask if required (used for autoregressive models)
        if causal_mask:
            mask = make_triu(weight, dtype=tf.bool)
            weight = tf.where(mask, -np.inf, weight)  # Mask future tokens by setting scores to -inf

        # Scale weights by the square root of the head dimension for stability
        weight /= math.sqrt(self.d_head)

        # Apply softmax to normalize the weights along the last axis
        weight = tf.nn.softmax(weight, axis=-1)  # Attention probabilities

        # Compute the weighted sum of values (V) based on attention weights
        output = weight @ v  # (batch_size, n_heads, seq_len, d_head)

        # Rearrange and combine multi-head outputs into the original shape
        output = tf.transpose(output, perm=[0, 2, 1, 3])  # (batch_size, seq_len, n_heads, d_head)
        output = tf.reshape(output, input_shape)  # Combine the multi-head dimension

        # Apply the output projection layer
        output = self.out_proj(output)

        return output

The multi-head attention code is divided into 6 main sections:

• Input projection self.in_proj - This layer projects the input embeddings in a vector space then splits, along the last dimension, the resulting matrix into three separate matrices for queries, keys, and values, required for the attention mechanism.
• Multi-head splitting - The queries, keys, and values are reshaped into separate heads. Each head works independently on a smaller dimension d_head, which improves performance and allows for capturing diverse patterns.
• Scaled dot-product attention - The attention score is calculated using the dot product of queries \(Q\) and transposed keys \(K^T\). The result is scaled by the square root of the head dimension to maintain stability during training.
• Causal masking - For autoregressive tasks, this ensures that a token at position \(t\) can only attend to tokens at positions \(\leq t\).
• Attention Application - The softmax function normalizes the scores into probabilities. These probabilities weight the values \(V\) to compute the attention output.
• Output projection self.out_proj - After combining multi-head outputs into the original embedding space, a dense layer transforms the results back into the desired embedding dimension.

And voilà! Hopefully, this chapter has equipped you with the knowledge to implement and debug an attention mechanism from scratch.

A brief history and why we need it

In the early days of computer vision, researchers faced a significant challenge: how to enable machines to understand and process images as humans do. Traditional neural networks, struggled when applied to images. The main issue was the sheer complexity of visual data—images are high-dimensional, For example a single \(100\times100\)-pixel image has 10,000 pixels, and let us suppose the first layer of the traditional neural network has 512 neurons, that is a total of more than 5 million parameters, this combined with lack of sufficient computational power made the task of computer vision difficult.

Enter convolutional neural networks (CNNs), a groundbreaking approach that revolutionized how we process visual data. The concept of convolution itself isn’t new; it’s a mathematical operation that has been used for decades. However, its application in neural networks was a game changer.

It was in 1998, that LeCun et al. introduced the famous LeNet-5 architecture¹, widely used for handwritten digit recognition, which combined fully connected layers, sigmoid activation functions and introduced convolution and pooling layers.

What “Convolution” stands for in Convolution Neural Nets

Imagine you’re given two lists of numbers and asked to produce a third list based on these inputs. There are many operations you might consider—element-wise multiplication, addition, or subtraction. Another powerful operation is convolution. But what exactly does that mean? Let’s start by showing the convolution formula and break it down.

At first glance, this might look intimidating, but don’t worry—math is just a language for describing concepts. In simple terms, to convolve two inputs, f and g, you flip the second input, slide it across the first input, and accumulate the interactions between them at each position.

So how does this relate to our example of two lists of numbers? Well, \(f\) and \(g\) can represent those lists. Still not clear? Let’s visualize it with an animation showing the convolution between two lists, \(a\) and \(b\).

In a convolutional neural network, the input image of shape (height, width) is divided into smaller regions, each of shape (region height,region width). A matrix, called a filter or kernel (we will talk about it later), slides across these regions and applies the convolution operation we just discussed, producing a new layer called a convolutional layer. Instead of each neuron in this layer being connected to every pixel in the input image, it’s connected only to the pixels within its specific region—known as its receptive field.

As you stack more convolutional layers, neurons in each layer connect to small regions of the previous layer, creating a hierarchical structure. This hierarchy allows CNNs to efficiently capture important features in an image, making them incredibly effective for tasks like image recognition.

Technically, we say that a neuron in row \(i\), column \(j\) of a given layer is connected to the outputs of the convolution of neurons in the previous layer located in rows \(i\) to \(i+f_h-1\) and columns \(j\) to \(j+f_w-1\). Where \(f_w\) and \(f_h\) are the width and height of the filter.

You might notice that with each convolution, the layers tend to shrink in size. If we want to maintain the same height and width as the previous layer, we add zeros around the input—this is known as zero padding.

On the other hand, if we want the output to be much smaller, we can space out the positions where we apply the filter. The distance between the current position and the next position of the filter is called the stride.

Filters

Now, we’ve been talking about filters without introducing them. Simply put, Filters in CNNs play a role similar to that of weights in traditional Artificial Neural Networks (ANN). In an ANN, weights are learned parameters that define the influence of each input feature on the output. Similarly, in a CNN, filters are learned during training and define how different features, such as edges, textures, or shapes, are detected and emphasized in the input image. These filters function much like the receptive fields in the human brain, where each one is responsible for identifying specific features within an image.

What’s particularly powerful about CNNs is that these filters don’t need to be manually defined. Instead, they are automatically learned during the training process through backpropagation and optimization. This means the network discovers the most effective filters for the task at hand, allowing it to efficiently extract meaningful patterns from images.

When a filter slides over an input image, the result is called a feature map. So far, we’ve considered examples with just one filter for simplicity, but in practice, a convolutional layer contains multiple filters, each producing its own feature map.

What sets CNNs apart from traditional ANNs is their ability to detect features regardless of their position in the input. In a CNN, once a feature — like a specific edge or texture — is learned, it can be recognized anywhere in the image. On the other hand, an ANN can only recognize a pattern if it appears in the exact location where it was learned during training. This position-invariance makes CNNs far more efficient, not only in detecting patterns but also in significantly reducing the number of parameters needed, leading to more manageable and effective models.

Convolution layer equation

All the previous explanations would allow us now to easily understand the general equation of a neuron in convolution layer. It is written as follows:

The equation shown in the image represents the computation of the output value \(𝑧_{𝑖,𝑗,𝑘}\) for a single neuron in a convolutional layer of a neural network. Here’s a step-by-step explanation:

• Output \(𝑧_{𝑖,𝑗,𝑘}\) - This is the value of the neuron located in row \(i\) column \(j\), and feature map \(k\) in the current convolutional layer. It represents the activation of this neuron after the convolution operation.
• Input \(x_{𝑖',𝑗',𝑘'}\) - This is the output from the previous layer (or the input image if this is the first layer), specifically from row \(i'\), column \(j'\) and feature map \(k'\).
• Weights \(w_{u,v,k,𝑘'}\) - These are the filter weights applied during convolution. Each weight connects a specific input from a receptive field to the current neuron. The indices \(u\) and \(v\) refer to the position within the receptive field.
• bias \(b_k\) - Each feature map \(k\) has an associated bias term. It adjusts the output of the entire feature map.

The equation sums over \(u\) and \(v\) which correspond to all positions in the receptive field of the input, as well as \(k'\), corresponding to all feature maps.

Implementation

Now that we’ve covered the theory behind convolution, let’s dive into how to implement a convolutional layer in TensorFlow. Fortunately, we don’t need to manually define or compute the convolution operation. TensorFlow provides a convenient Conv2D class that does the heavy lifting for us.

Let’s consider an example input image with the shape \((1, 64, 64, 3)\). Here’s what this shape means:

• The first dimension, \(1\), is the batch size, indicating that we have a single image.
• The next two dimensions, \(64, 64\), represent the height and width of the image.
• The final dimension, \(3\), represents the number of color channels. For an RGB image, there are three channels: red, green, and blue. For grayscale images, this value would be 1.

With that in mind, here’s how we define a convolutional layer in TensorFlow:

  import tensorflow as tf

  # Define a convolutional layer
  conv_layer = tf.keras.layers.Conv2D(
    filters=32,                   # Number of filters
    kernel_size=(3, 3),           # Size of each filter
    strides=(1, 1),               # Stride of the sliding window
    padding='same',               # Padding strategy
    data_format='channels_last'   # Input data format
  )

  # Example input: a batch of images with shape 
  # (batch_size, height, width, channels)
  input_tensor = tf.random.normal([1, 64, 64, 3]) # A single 64x64 RGB image

  # Apply the convolution
  output = conv_layer(input_tensor)
  print(output.shape)

This code snippet creates a convolutional layer, applies it to a random RGB image, and prints the shape of the output feature map. The output shape will depend on the filter size, stride, and padding we’ve defined. Now Let us understand each of the Conv2D parameters:

• Filters - This specifies the number of filters (or receptive fields) the convolutional layer will apply to the input. Each filter detects specific patterns or features in the image, such as edges, corners, or textures. As mentioned earlier, we can stack as many filters as needed, and each one will generate its own feature map.
• Kernel_size - This defines the dimensions of each filter. It can either be a single integer or a tuple. In our case, the filters are \(3\times 3\) matrices, which is a common choice for convolutional layers.
• Strides - The stride determines how far the filter moves across the image at each step. It can be a single integer or a tuple that specifies the stride for the height and width. A stride of \((1, 1)\) means the filter will move one pixel at a time both horizontally and vertically. A higher stride value would cause the filter to move in larger steps, reducing the size of the resulting feature map.

• Padding - There are two options:

  • valid: No padding is added, which means the output feature map will be smaller than the input.

  • same: Zero-padding is applied evenly around the input to maintain the same height and width in the output feature map. This option is useful when you want the output dimensions to match the input dimensions.

• Data_format - This determines the order of the dimensions in the input data. By default, TensorFlow uses channels_last, where the channels (RGB) come last in the shape (i.e., (batch_size, height, width, channels)). If your input has the channels dimension first, as in (batch_size, channels, height, width), you’d use channels_first to avoid any shape mismatch errors.

Welcome

Welcome to this blog series on implementing Stable Diffusion from scratch using TensorFlow!

In the fast-paced world of artificial intelligence, it’s easy to feel overwhelmed by the sheer volume of new developments, especially in the realm of generative models. Every month seems to bring a new breakthrough, a novel architecture, or a cutting-edge application that pushes the boundaries of what machines can create. Keeping up with these advancements can be daunting, and to make matters more challenging, finding comprehensive, easy-to-understand resources that break down these concepts often feels like searching for a needle in a haystack.

This blog series aims to be your guide through one of the most exciting areas of AI today: image generation, and more specifically, Stable Diffusion. Whether you’re an aspiring AI enthusiast or a seasoned machine learning practitioner, you’re about to embark on a journey that will not only deepen your understanding of generative models but also equip you with the skills to implement them from scratch using TensorFlow.

Why Stable Diffusion

Generative models have taken the machine learning world by storm, and for good reason. These models can create hyper-realistic images, generate human-like text, and even compose music, pushing the envelope of creativity in AI. However, with great power comes great complexity, and understanding the inner workings of these models can be challenging.

Among the various types of generative models, diffusion models—particularly Stable Diffusion—have emerged as powerful tools for image generation. What sets Stable Diffusion apart is its innovative approach: it generates data by gradually adding and then removing noise from an image. Imagine starting with a foggy, blurry picture and gradually clearing it up until a sharp, vivid image emerges. This core principle also enables Stable Diffusion to excel at tasks like text-to-image synthesis and image inpainting. The magic of this reverse journey through noise is powered by a symphony of components like convolutional layers, attention mechanisms, autoencoders, and more.

Note: Our goal in this series is to demystify these components, showing you how they work together to produce stunning results. The good news is that many generative models share similar building blocks, so once you grasp concepts like autoencoders, attention mechanisms, and transformers, you’ll find it easier to understand other models as well.

What to expect

This series is divided into multiple parts, each building on the last. By the end, you’ll have a thorough understanding of Stable Diffusion’s architecture and the know-how to implement it in Python. Here’s a sneak peek at what we’ll cover:

1. Convolution and Attention Mechanisms: These key building blocks are crucial for image processing and help the model focus on important features.
2. Variational Autoencoders (VAEs): VAEs are core components in generative models, helping create meaningful latent spaces.
3. Generative Models and Stable Diffusion: We’ll dive into generative models, zooming in on diffusion models to explain how they work and what makes them special.
4. Classifier-Free Guidance: Learn how to guide your diffusion model to generate specific images without requiring a classifier.
5. UNet Architecture: This architecture is central to many image-to-image tasks, including Stable Diffusion.
6. Contrastive Language-Image Pretraining (CLIP): We’ll explore how text and images are linked in the model, enabling powerful text-to-image synthesis.

Who is this for?

This series is for anyone with a basic understanding of machine learning and a desire to delve deeper into the world of generative models. Whether you’re looking to grasp the theory behind these models or you’re eager to get your hands dirty with code, there’s something here for you.

If you’ve ever been intrigued by how AI can generate art or create realistic images from scratch, this series will peel back the layers of complexity and show you the nuts and bolts of one of the most exciting developments in AI today.

How to follow along

Each part of the series will blend mathematical foundations, theoretical explanations, and practical coding examples to give you a well-rounded understanding.

To get the most out of this series, we recommend setting up your development environment with Python and TensorFlow. We’ll provide code snippets so you can experiment with the concepts as you learn.

So, are you ready to dive in? Let’s start with the basics and build our way up to a fully functional Stable Diffusion model!