Idea
Autoencoders are a type of neural network used for unsupervised learning. Their architecture consists of three main components:
- Encoder. The first part of the neural network that takes input data and gradually reduces its dimensionality through its layers, ultimately reaching the bottleneck.
- Bottleneck. The network layer with the smallest dimensionality that contains the latent representation of the input data.
- Decoder. The part of the network connected to the bottleneck output that gradually expands the data’s dimensionality. As a result, at its last layer, it returns data of the same size as was initially passed to the encoder.
For images, the encoder and decoder are usually presented as convolutional neural networks.
In autoencoders, our ultimate goal during training is to make the network transform the input data into a more compressed representation in the bottleneck without losing too much information. During inference, we can pass the data to the encoder, extract the resulting embedding from the bottleneck, and then use it for our own purposes.
Let’s understand how training works in autoencoders.
Training
A great thing about autoencoders is that they don’t require any labeled data! Let’s see how they work.
As mentioned before, an input image is passed to the network, where it is compressed to a smaller size and then reconstructed to the original dimension. The question we should ask ourselves is what we want the decoder to output.
As you could have guessed, the decoder can simply try to reconstruct the original image from the compressed representation in the bottleneck. Why do so?
The idea behind this is simple:
- If the bottleneck’s compressed representation captures the main features of the encoder’s input well, then it should be relatively easy for the decoder to use that information to reconstruct the original image.
- If the bottleneck fails to capture the main features, the decoder won’t be able to reliably reconstruct the original image. Thus, the model will be penalized for a poor compressed representation.
This way, by asking the decoder to reconstruct the original image, we implicitly force the encoder to produce a rich yet compressed latent representation, helping the decoder efficiently achieve its task.
The space to which the input data is projected in the bottleneck is called the latent space.
Reconstruction loss
Given the original image and the reconstructed image from the decoder, what is the simplest way to compare the generated quality? The obvious answer is to compare the two images pixel-wise using the MSE loss, which, in the context of autoencoders, is called the reconstruction loss.
The calculated loss value is then used to perform backpropagation to update the model’s weights.
Latent space dimension
The latent space dimension is an important hyperparameter that directly affects the decoder’s performance.
On the one hand, the latent space dimension should be sufficient to efficiently encode the key input features. On the other hand, it should not be too large to maintain a high compression rate.
One well-known example is Stable Diffusion. It uses an autoencoder to transform the input image, 512 x 512 x 3, containing 786,432 values, into a 64 x 64 x 4 image with 16,384 values, resulting in a compression ratio of 48x.
Other autoencoder applications
One trick for training autoencoders is to have them learn to remove noise from images. The idea is simple: since autoencoders are good at reconstructing original images, we could add slight noise to the input images and then ask them to reconstruct the original images.
A great thing about this method is that for training, it is sufficient to have only the original images, to which you would then apply noise.
Another cool application of autoencoders is image inpainting, which involves passing images with masked patches to a model so it can unmask and fill in the missing image parts.
Similarly, autoencoders can be trained to remove specific objects from images. This is particularly useful for removing watermarks.
Blueriness problem
In reality, despite its simplicity, the MSE loss is not perfect for autoencoders. One common problem with using it is a tendency for the decoder to generate images with blurry pixels.
For example, we could imagine an image of size 512 x 512 with two vertical, non-overlapping black-and-white regions. We then take a horizontal row of that image whose pixels look like this:
[… 0 0 255 255 255 …]
The model doesn’t have any knowledge of the image structure; it only tries to minimize the MSE loss. Even if, for that image, a model made a prediction like [… 0 0 0 255 255 …], which is still very good because the region is shifted by only one pixel, the MSE loss in this case would be higher than in the case below, which a model might prefer:
[… 0 0 127 255 255 …]
In the latter scenario, despite the lower MSE, the middle pixel represents a blurry edge, which is visually unappealing.
This problem is addressed in more advanced autoencoder variations with adjusted loss functions.
Conclusion
As we can see, autoencoders are a simple yet powerful concept. By training the decoder to reconstruct the original image from compressed data, we gradually adjust the encoder to produce more informative features that can then be extracted for downstream tasks.
In addition to data compression, we saw that autoencoders have other applications, such as denoising images, performing image inpainting, and removing objects from images.
All images unless otherwise noted are by the author