Latent Space Representations

A variational autoencoder, or VAE, takes an image through an encoder and a decoder, but it does not pass through a single fixed code in the middle. The middle of the model is a latent space, and the way a VAE fills that space is the subject of this article. Instead of asking the encoder for one exact point that stands for an input, a VAE asks the encoder to describe a region. A point is then drawn from that region, and the decoder turns the drawn point back into an image.

Keep that picture in mind as you read. The rest of the article traces one input from the moment it enters the encoder to the moment a reconstructed output comes out of the decoder, paying close attention to what is stored at each stage and what the sampling step does with it.

A VAE has three stages: encoding, sampling, and decoding. In the encoding stage, the input image goes through the encoder. In the sampling stage, the encoder's output is used to draw a latent point z. In the decoding stage, the decoder receives z and produces an output that is meant to resemble the original input. The sampled point z, not the original image and not the distribution parameters alone, is the value that gets decoded.

Read the diagram as a chain of hand-offs. The encoder does not hand the decoder an image. It hands the sampling step a pair of vectors. The sampling step does not hand the decoder those vectors either. It hands over one point drawn from the distribution they describe. Each stage changes what the data represents, and the decoder only ever sees the last of those representations.

The encoder is a convolutional network that maps an input image to two vectors. In the implementation described by the source, convolutional layers process the image, the result is flattened, and a dense layer produces an intermediate representation. Two further dense layers then produce z_mean and z_log_var. Together these outputs describe the distribution used for sampling in latent space.

A common first reading is that z_mean and z_log_var are two alternative images or two competing reconstructions. They are neither. They are numbers that describe a statistical distribution. z_mean locates the center of the region the encoder has chosen for this input, and z_log_var is a variance-related parameter that describes how far samples are allowed to spread around that center.

The encoder's output becomes useful only after the sampling step. Before sampling, z_mean and z_log_var are just two sets of numbers. The sampling step is what converts them into a single latent point the decoder can use.

The VAE combines z_mean and z_log_variance with epsilon, which is a random tensor of small values. The result is z, the latent-space point that is passed to the decoder. The standard way to write this combination is shown below. Here $\epsilon$ is the random draw, and the exponential of half of z_log_variance gives the spread used to scale that draw.

Trace the flow in the diagram. z_mean supplies the center. z_log_variance, after being halved and exponentiated, supplies the spread. epsilon supplies the random direction and size of the draw, and it is multiplied by the spread before being added to the center. Because epsilon is random, two passes through the same input can produce two different values of z, even though the encoder outputs stayed the same.

In Keras, code that is not already part of a built-in layer can be placed inside a Lambda layer. This is how the sampling operation is incorporated into the network, so the random draw becomes one more step in the model rather than something handled outside it.

The comparison below sets a single fixed code against the distribution parameters a VAE uses. In the fixed case, the encoder's output for an input is one point, and that same point is passed to the decoder every time. In the VAE case, the encoder's output is a mean and a variance-related parameter, and the decoder receives a fresh sample from the region they describe.

Randomness in a VAE is not there merely to make the model unpredictable. Because samples are drawn around the encoded location, nearby latent points are encouraged to decode into similar images. This gives the latent space continuity: moving to a nearby point should not abruptly change the decoded result. Along with low dimensionality, this continuity helps make directions in latent space meaningful for changing aspects of the data. This is the basis for the source's description of the approach as useful for image generation and for image editing through concept vectors.

Think of the encoder's output as marking a neighborhood rather than a dot. Points sampled from the same neighborhood should decode to outputs that look alike, and the space between neighborhoods is filled with smoothly varying results rather than sudden jumps.

A VAE is trained with two loss functions. Reconstruction loss encourages the decoded sample to match the original input. Regularization loss encourages a well-formed latent space and helps reduce overfitting to the training data. The two pressures act on the same model but favor different things, and both are needed.

Reconstruction loss, on its own, would push the model toward reproducing each training input as accurately as possible. Regularization loss, on its own, would push the latent space toward a well-organized structure. Reconstruction alone focuses on reproducing inputs, while regularization supports the organized latent structure that makes sampling and manipulation useful. Training balances the two.