Entropy-Based Characterisation of the Polarised Regime in Latent Variable Models

Technical Summary: Entropy-Based Characterisation of the Polarised Regime in Latent Variable Models

Problem Statement
Variational Autoencoders (VAEs) and related latent variable models often exhibit a "polarised regime," where latent dimensions segregate into three distinct subsets: active (encoding maximum data variation), passive (collapsing to the prior and encoding negligible information), and mixed (switching between active and passive states depending on the input). Identifying these subsets is crucial for understanding model behavior, particularly regarding "posterior collapse," where the model fails to utilize the latent space effectively.

Existing methods for distinguishing active from passive dimensions rely heavily on specific priors, typically Gaussian. Common approaches include setting thresholds on the Kullback-Leibler (KL) divergence or analyzing the statistics of the mean and variance representations ($\mu$ and $\sigma$) relative to a standard Gaussian prior. These methods are limited because they are not model-agnostic; they struggle to apply to deterministic architectures (like standard autoencoders) or models with non-Gaussian priors (such as Identifiable VAEs). Furthermore, recent findings suggest that "passive" variables may retain subtle correlations with other variables, implying that discarding them based on strict prior-based criteria might remove useful information. There is a need for a principled, prior-agnostic criterion to characterize latent activity.

Methodology
The authors propose an information-theoretic framework based on the entropy of the mean representation ($H(\mu_i)$) of the latent dimensions.

1. Entropy Criterion: Instead of measuring deviation from a prior (via KL divergence), the method quantifies the disorder or unpredictability of the empirical marginal distribution of each latent dimension's mean output ($\mu_i$) across the dataset.

   • Active: $H(\mu_i) > \tau$ (High entropy, indicating significant variation across data points).
   • Passive: $H(\mu_i) \le \tau$ (Low entropy, indicating the dimension is nearly constant).
   • Mixed: Dimensions exhibiting high entropy for some data subsets and low for others.

2. Theoretical Analysis: The paper establishes a theoretical link between this entropy-based criterion and existing definitions:

   • Entropy vs. KL Divergence: Using entropy-variance inequalities (specifically $h(X) \le \frac{1}{2}\log(2\pi e \text{Var}(X))$), the authors show that high entropy in the mean representation implies a lower bound on the KL divergence. Thus, high entropy correlates with the "active" state defined by high KL contributions.
   • Entropy vs. Bonheme's Criteria: The authors relate their approach to Bonheme and Grzes's definition, which relies on the variance of $\mu$ and $\sigma$. They demonstrate that $H(\mu)$ captures the same statistical variation as $\text{Var}(\mu)$ but in an information-theoretic manner.
   • Limitations: The analysis clarifies that entropy of the mean alone cannot reliably distinguish active from mixed dimensions without additional signals from the variance representation, as both can exhibit high uncertainty in the mean.

3. Estimation: For empirical evaluation, the authors utilize the matrix-based Rényi $\alpha$-entropy estimator. This approach avoids explicit density estimation, making it suitable for high-dimensional latent spaces and capable of handling continuous variables directly via Gram matrices.

Experimental Evaluation
The authors evaluated the entropy criterion across four distinct model classes using standard benchmarks (MNIST, smallNORB, d-Sprites):

• $\beta$-VAEs: Models with varying $\beta$ hyperparameters to control the strength of the KL penalty.
• Identifiable VAEs (iVAEs): Models using auxiliary information to enforce identifiability.
• Least-Volume Autoencoders (LV-AEs): Deterministic models regularized by the volume of the latent space.
• L2-Regularised Autoencoders (L2-AEs): Deterministic baselines with L2 penalties on latent codes.

Key Results

1. Recovery of Polarised Regime: The entropy criterion consistently recovered the polarised regime (separation into high-entropy active and near-zero-entropy passive dimensions) across all tested architectures, including deterministic models where KL-based definitions are inapplicable.
2. Coherence with Other Measures: Entropy, variance, and mutual information ($I(X; \mu)$) were found to evolve coherently during training. Entropy was selected as the primary metric due to its robustness and lack of computational overhead compared to mutual information estimation.
3. Generalizability:
   • In LV-AEs, the entropy criterion successfully identified the polarised regime induced by volume regularization, confirming that polarization is a general phenomenon of regularized representation learning, not just a variational artifact.
   • In iVAEs, while total posterior collapse was avoided (consistent with identifiability theory), the entropy criterion detected selective collapse (polarisation). Crucially, "collapsed" units in iVAEs retained small but non-zero entropy, reconciling the empirical observation of suppression with the theoretical impossibility of total collapse.
4. Downstream Utility: When latent codes were used for downstream logistic regression tasks, models trained with raw latent vectors showed accuracy plateaus once active variables were included. However, when latent codes were normalized, including "passive" variables yielded small but consistent improvements in accuracy. This suggests that passive dimensions are not entirely devoid of information but are suppressed in scale rather than completely removed.

Significance and Claims
The paper claims that the proposed entropy-based criterion offers a model-agnostic and prior-agnostic tool for characterizing latent activity. Its significance lies in:

• Universality: It applies to both variational and deterministic architectures, overcoming the limitations of KL-based heuristics tied to Gaussian priors.
• Theoretical Grounding: It provides a principled link between information theory (entropy) and the mechanisms driving the polarised regime (KL minimization and variance constraints).
• Nuanced Understanding of Collapse: The findings challenge the binary view of posterior collapse. The authors argue that collapse is often a matter of scale rather than absolute information removal, as "passive" dimensions retain detectable entropy and downstream utility when appropriately normalized.

Limitations
The authors modestly acknowledge several limitations:

• Mixed Variables: The entropy of the mean alone cannot cleanly distinguish between active and mixed variables; additional variance signals are required.
• Threshold Heuristic: The threshold $\tau$ for classifying activity is heuristic and lacks a universally optimal, principled selection method across different datasets and architectures.
• Scope: Empirical validation is currently restricted to image datasets and specific autoencoder architectures. The behavior in sequential models or diffusion-based architectures remains an open question.
• Practical Utility: While passive variables offer small gains in downstream tasks, the improvements are modest, suggesting the criterion is primarily valuable for diagnostic analysis and understanding latent structure rather than as a direct performance optimization tool.