VAEs for Anomaly Detection and Out-of-Distribution Detection

Variational Autoencoders, by their very nature of learning a compressed representation and a generative model of the training data, offer a powerful framework for identifying samples that deviate from the norm. This capability makes them well-suited for anomaly detection and out-of-distribution (OOD) detection tasks. The core assumption is that a VAE trained on "normal" data will be proficient at reconstructing familiar samples while struggling with anomalous or OOD inputs.

Leveraging Reconstruction Error

The most straightforward approach to anomaly detection with VAEs is based on the reconstruction probability. A VAE is trained to minimize the reconstruction error for data points similar to its training distribution. The Evidence Lower Bound (ELBO), which VAEs maximize, includes a reconstruction term, typically a log-likelihood such as:

$$\mathcal{L}_{rec}(x, \hat{x}) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]$$

where $x$ is the input and $\hat{x}$ is its reconstruction through the VAE (i.e., $z \sim q_\phi(z|x)$, $\hat{x} \sim p_\theta(x|z)$).

For an input $x_{new}$, its reconstruction error (e.g., Mean Squared Error for continuous data, or negative log-likelihood of $p_\theta(x_{new}|z)$ for $z$ sampled from $q_\phi(z|x_{new})$) can serve as an anomaly score.

$$\text{AnomalyScore}(x_{new}) = ||x_{new} - \hat{x}_{new}||^2$$

or more generally, $-\log p_\theta(x_{new}|\hat{z})$, where $\hat{z}$ is the mean of $q_\phi(z|x_{new})$. Samples that are "normal" or in-distribution should have low reconstruction errors, while anomalous or OOD samples are expected to yield higher reconstruction errors because the VAE's decoder $p_\theta(x|z)$ hasn't been optimized to generate them accurately from any latent code $z$ produced by the encoder $q_\phi(z|x)$.

Reconstruction error scores for normal data tend to be clustered at lower values, while anomalous data points typically exhibit higher scores, allowing for separation.

Latent Space Characteristics for Anomaly Detection

Beyond reconstruction, the latent space $z$ learned by the VAE provides another avenue for identifying anomalies.

1. Distance from Prior: The encoder $q_\phi(z|x)$ maps inputs to distributions in the latent space. For in-distribution data, these encoded distributions are regularized by the KL divergence term $D_{KL}(q_\phi(z|x) || p_\theta(z))$ to be close to the prior $p_\theta(z)$ (often a standard Normal distribution $\mathcal{N}(0, I)$). Anomalous inputs might be mapped to regions of the latent space that are far from where $p_\theta(z)$ has high density. The KL divergence value itself, or the Mahalanobis distance of the mean of $q_\phi(z|x)$ from the prior's mean, can be used as an anomaly score.

2. Density in Latent Space: Normal data points are expected to map to denser regions of the learned latent manifold. Anomalies might fall into sparser, less populated areas. One could estimate the density around $z = \mu_\phi(x)$ using techniques like Kernel Density Estimation (KDE) on latent codes from a validation set of normal data, though this can be computationally intensive and sensitive to dimensionality.

Normal data (blue points) typically cluster within or near the region favored by the latent prior (dashed ellipse). Anomalous data (red points A1, A2, A3) may be encoded to locations far from this central mass or into sparser areas of the latent space.

Using the ELBO or its Components

The ELBO itself, $\mathcal{L}(x; \theta, \phi)$, is an approximation of the log marginal likelihood $\log p(x)$.

$$\mathcal{L}(x; \theta, \phi) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) || p_\theta(z))$$

A lower ELBO value could indicate an OOD sample. However, relying solely on ELBO can sometimes be misleading. Highly flexible decoders might assign high likelihood (and thus high ELBO) to simple OOD samples that are easy to reconstruct, even if they are semantically different from the training data. Some research suggests that the typical set for $p(x)$ (where VAEs perform well) and the high-density regions of $p(x)$ might not perfectly align for complex data distributions, which can complicate OOD detection using raw likelihoods.

Combining metrics, such as a weighted sum of the negative reconstruction error and the KL divergence, often yields more robust anomaly scores:

$$\text{AnomalyScore}(x) = \alpha \cdot (-\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]) + (1-\alpha) \cdot D_{KL}(q_\phi(z|x) || p_\theta(z))$$

The hyperparameter $\alpha$ balances the two terms and usually requires tuning.

Setting Detection Thresholds

Once an anomaly scoring mechanism is in place (e.g., reconstruction error, latent KL divergence, or a combined score), a threshold must be determined. Any sample with a score exceeding this threshold is flagged as an anomaly. Common strategies for threshold selection include:

• Empirical Percentile: Calculate anomaly scores for a held-out validation set of normal data. The threshold can be set to a high percentile (e.g., 95th or 99th) of these scores. This approach assumes that a small fraction of false positives on normal data is acceptable.
• Statistical Modeling: Fit a distribution (e.g., Gaussian or Extreme Value Theory distribution) to the anomaly scores from the normal validation data and set the threshold based on a desired p-value or false positive rate.
• Supervised Thresholding (if available): If a small labeled dataset of both normal and anomalous samples is available for validation, the threshold can be chosen to optimize a specific metric like F1-score or to meet a specific precision/recall target on this validation set.

Practical Considerations and Challenges

While VAEs offer a principled way for anomaly detection, several factors can influence their performance:

• Model Architecture: The capacity and architecture of the encoder and decoder networks are important. A VAE that is too expressive might learn to reconstruct some types_of anomalies too well (overgeneralization), while one with insufficient capacity might not model the normal data distribution accurately, leading to high reconstruction errors even for normal samples.
• Latent Dimensionality: The size of the latent space $d_z$ can impact performance. Too small a $d_z$ might overly compress information, losing subtle details that distinguish normal from anomalous. Too large a $d_z$ can make the KL divergence term difficult to optimize or lead to the "posterior collapse" issue where $q_\phi(z|x)$ becomes uninformative and similar to $p_\theta(z)$ for all $x$.
• Nature of Training Data: The VAE's definition of "normal" is entirely dictated by its training data. If the training set inadvertently contains anomalies or does not comprehensively cover all variations of normal behavior, the VAE's performance will suffer.
• Type of Anomaly: VAEs are generally better at detecting anomalies that are structurally different from the training data or lie off the learned data manifold. They might struggle with anomalies that are subtle variations within the manifold or share many low-level features with normal data.
• Posterior Collapse: If the KL divergence term in the ELBO dominates training, leading to $q_\phi(z|x) \approx p_\theta(z)$ (posterior collapse), the latent representations $z$ become less informative about $x$. This can weaken anomaly detection methods that rely on latent space characteristics. Techniques discussed in earlier chapters to mitigate posterior collapse (e.g., $\beta$-VAE, free bits) can be beneficial here.

Extensions and Sophistications

More advanced VAE-based anomaly detection might involve:

• Ensembles of VAEs: Training multiple VAEs with different initializations or architectures and combining their anomaly scores can improve robustness.
• Iterative Refinement/Detection: Some methods use an iterative process where high-confidence anomalies are removed from the training set (or down-weighted), and the VAE is retrained.
• Contrastive Approaches: Training VAEs to explicitly distinguish between normal data and artificially generated "borderline" anomalies.
• VAEs with Normalizing Flows: Using normalizing flows for more expressive priors $p_\theta(z)$ or posteriors $q_\phi(z|x)$ can lead to better density estimation and potentially improved OOD detection, although these models can be more complex to train.

In summary, VAEs provide a versatile toolkit for anomaly and OOD detection by learning a model of normal data. The choice of anomaly score, whether based on reconstruction error, latent space properties, or the ELBO itself, along with careful model design and threshold selection, are significant for achieving effective performance in practice. Understanding these nuances allows engineers and researchers to adapt VAEs to a wide array of detection scenarios.