Selecting Appropriate Loss Functions for Autoencoders

Once you've designed the basic architecture of your encoder and decoder and decided on the dimensionality of your latent space, the next important step is to choose an appropriate loss function. The loss function, also known as a cost function or objective function, quantifies how far the autoencoder's reconstruction is from the original input. During training, the autoencoder's weights are adjusted to minimize this loss, effectively learning to create better and better reconstructions, and consequently, more meaningful latent representations.

The choice of loss function is not arbitrary; it primarily depends on the type and characteristics of the data you are working with and the desired properties of the reconstruction. Let's look at the most common choices.

Loss Functions for Continuous Data

When your input data consists of continuous, real-valued numbers, such as pixel intensities in an image (after normalization) or measurements in a tabular dataset, the Mean Squared Error (MSE) is a very common and effective choice.

Mean Squared Error (MSE)

MSE measures the average squared difference between the original input values and the reconstructed values. For a single data instance with $D$ features (e.g., $D$ pixels or $D$ columns in tabular data), if $x_j$ is the $j$-th feature of the original input and $x'_j$ is the $j$-th feature of the reconstructed output, the MSE is calculated as:

$L_{MSE} = \frac{1}{D} \sum_{j=1}^{D} (x_j - x'_j)^2$

The total loss for a batch of $N$ data instances would then be the average of these individual MSE values.

Why MSE?

• Intuitive: It directly penalizes deviations from the original value.
• Differentiable: It's smooth and differentiable, which is essential for gradient-based optimization algorithms like backpropagation.
• Emphasis on Large Errors: Because the errors are squared, MSE penalizes larger errors more heavily than smaller ones. This can be beneficial if large deviations are particularly undesirable.

When using MSE, the output layer of your decoder typically uses a linear activation function if the target values are unbounded or normalized to a range like (-1, 1) (in which case, a tanh activation might also be used for the output layer if inputs are similarly scaled). If your continuous data is normalized to be within [0,1], a sigmoid activation can also be used with MSE, though Binary Cross-Entropy is often preferred in that specific [0,1] scenario for probabilistic interpretations.

Mean Absolute Error (MAE)

Another option for continuous data is the Mean Absolute Error (MAE). It measures the average absolute difference between the original and reconstructed values:

$L_{MAE} = \frac{1}{D} \sum_{j=1}^{D} |x_j - x'_j|$

MSE vs. MAE MAE is generally less sensitive to outliers compared to MSE. If your dataset contains significant outliers that you don't want to overly influence the training process, MAE might be a better choice. However, MSE's stronger penalization of large errors can sometimes lead to reconstructions that are, on average, closer to the original for well-behaved data. The choice often comes down to empirical testing or specific domain requirements.

The following chart illustrates how MSE applies a much larger penalty to larger errors compared to MAE.

This chart shows the loss value attributed by MSE and MAE as a function of the error between a single predicted value and its target (assumed to be 0 for simplicity). Notice the quadratic increase for MSE versus the linear increase for MAE.

Loss Functions for Binary or Probabilistic Data

When your input data is binary (e.g., black and white images where pixels are 0 or 1) or can be interpreted as probabilities (e.g., pixel intensities normalized strictly between 0 and 1), Binary Cross-Entropy (BCE) is typically the preferred loss function.

Binary Cross-Entropy (BCE)

Binary Cross-Entropy, also known as log loss, measures the dissimilarity between two probability distributions. In the context of autoencoders, it compares the distribution of the original binary/probabilistic input data with the distribution of the reconstructed output, which is usually passed through a sigmoid activation function to ensure its values are also between 0 and 1.

For a single data instance with $D$ features, where $x_j$ is the original value and $x'_j$ is the reconstructed probability for the $j$-th feature, the BCE loss is:

$L_{BCE} = - \frac{1}{D} \sum_{j=1}^{D} [x_j \log(x'_j) + (1-x_j) \log(1-x'_j)]$

Why BCE?

• Probabilistic Interpretation: It's well-suited when outputs are probabilities. If an original input $x_j$ is 1, the loss term becomes $-\log(x'_j)$. To minimize this, $x'_j$ must approach 1. Conversely, if $x_j$ is 0, the term becomes $-\log(1-x'_j)$, and $x'_j$ must approach 0 to minimize loss.
• Stronger Gradients: Compared to MSE for values near 0 or 1 with a sigmoid output, BCE often provides stronger gradients, which can lead to faster training.

Important Note for BCE:

• Output Activation: The decoder's output layer must use a sigmoid activation function. This squashes the output values to the range (0, 1), making them suitable for comparison with the input data (which should also be scaled to [0, 1]).
• Input Scaling: Ensure your input data is scaled to the [0, 1] range.
• Numerical Stability: Implementations of BCE often include mechanisms to prevent $\log(0)$ errors by clipping predicted values to a small epsilon away from 0 and 1.

Relationship with Output Layer Activation

As highlighted, the choice of loss function is tightly coupled with the activation function used in the output layer of the decoder:

• MSE:
  • Use a linear activation (i.e., no activation) if your target data can take any real value (after normalization, this might be, for example, data centered around zero with some standard deviation).
  • Use a tanh activation if your target data is normalized to the range [-1, 1]. tanh outputs values in this range.
  • You can use a sigmoid activation if your target data is normalized to [0, 1], but BCE is often preferred in this scenario.
• BCE:
  • Almost always paired with a sigmoid activation in the output layer, as BCE expects inputs (both target and prediction) to be in the [0, 1] range, representing probabilities.

Mismatched loss functions and output activations can lead to poor training performance or nonsensical results. For instance, using BCE with a linear output layer that produces values outside [0,1] will cause errors or lead the model to learn incorrectly.

Practical Implementation

Most deep learning frameworks make it easy to specify your chosen loss function.

• In TensorFlow/Keras: You can pass the string identifier or the class instance to the compile method of your model:

  # For MSE
  model.compile(optimizer='adam', loss='mean_squared_error')
  # or
  # model.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
  
  # For BCE
  model.compile(optimizer='adam', loss='binary_crossentropy')
  # or
  # model.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy())
  

• In PyTorch: You instantiate the loss function class and then call it during your training loop:

  # For MSE
  criterion = torch.nn.MSELoss()
  # ... in training loop:
  # loss = criterion(reconstructions, original_inputs)
  
  # For BCE
  criterion = torch.nn.BCELoss()
  # ... in training loop (ensure reconstructions are sigmoid outputs):
  # loss = criterion(reconstructions, original_inputs)
  

Summary for Choosing a Loss Function

1. Analyze your input data:

   • Is it continuous (e.g., image pixels normalized to [-1,1] or [0,1], sensor readings)?
   • Is it binary (0s and 1s)?
   • How is it normalized? (e.g., [0,1], [-1,1], z-score normalized).

2. Select based on data type and normalization:

   • Continuous data normalized to [0,1] or general continuous data: MSE is a strong default. If data is specifically in [0,1] and you want a probabilistic output from a sigmoid, BCE can also be an option.
   • Continuous data normalized to [-1,1]: MSE with a tanh output activation.
   • Binary data (0 or 1): BCE with a sigmoid output activation is the standard. Make sure your input data is also represented as 0s and 1s (or floats in that range).

3. Align with output layer activation:

   • MSE: Linear, tanh, or sometimes sigmoid.
   • BCE: Sigmoid.

Choosing the right loss function is a foundational step in building an effective autoencoder. It directly influences what the model learns and how well it performs its primary task of reconstruction, which in turn affects the quality of the features you'll extract from the bottleneck layer. In the upcoming hands-on session, you'll see one of these in action.