Technical Summary: Disentanglement Beyond Generative Models with Riemannian ICA

Problem Statement

There is a significant gap between the theoretical foundations of disentanglement and modern representation learning practices. Traditional theoretical frameworks, particularly Independent Component Analysis (ICA) and its nonlinear variants, rely on a global generative model where data is assumed to be generated by statistically independent latent variables. While these frameworks offer strong theoretical guarantees (identifiability), they require strict assumptions, most notably that the mapping between data and latents is globally invertible. This assumption is often unrealistic for high-dimensional data with non-trivial topology and excludes non-invertible neural network architectures (e.g., many pretrained encoders) from theoretical analysis.

Conversely, modern self-supervised learning models often learn features that exhibit disentangled properties without explicit generative assumptions. However, there is currently no general theory to interpret these features as independent factors of variation without forcing them into a global generative model that may be misspecified. The core challenge is to provide a theoretical basis for studying local disentanglement in pretrained representations without assuming a global invertible generative model.

Methodology: Riemannian ICA (RICA)

The paper introduces Riemannian ICA (RICA), a local geometric analogue of ICA that replaces the global generative model with a local geometric structure defined by a Riemannian metric and a data density.

Core Insights

1. Local vs. Global: Disentanglement is fundamentally a local question regarding how an individual data point varies with respect to latent variations. Global invertibility is not required; instead, the method relies on local invertibility around a chosen base point.
2. Geometric Interpretation: In ICA, factors of variation correspond to axis-aligned lines in latent space. RICA interprets these as radial geodesics emanating from a base point in the data space. A variation in a single latent component corresponds to moving along a specific geodesic curve defined by the metric.
3. Normal Coordinates: The method utilizes Riemann normal coordinates centered at a base point $x_0$. These coordinates serve as adaptive latent variables. The exponential map $\exp_{x_0}$ maps latent vectors (tangent space) to data points, and the logarithmic map $\log_{x_0}$ performs the inverse.

The Disentanglement Tensor

The central theoretical contribution is the disentanglement tensor ($D$), which encodes a second-order notion of disentanglement called pointwise disentanglement. Pointwise disentanglement is defined as the condition where the Hessian of the latent log-likelihood is diagonal at the base point.

The disentanglement tensor is defined as:
$$D = \nabla^2 \log \rho - \frac{1}{3} \text{Ric}$$
Where:

• $\nabla^2$ is the covariant Hessian under the Levi-Civita connection.
• $\rho$ is the Hausdorff density (an intrinsic density function defined with respect to the Riemannian volume element).
• $\text{Ric}$ is the Ricci curvature tensor.

This tensor captures the second-order statistical independence of the latent variables. The term $-\frac{1}{3}\text{Ric}$ acts as a correction to the standard change-of-variables formula required when operating on a curved manifold.

Algorithm

To achieve pointwise disentanglement, RICA solves a generalized eigenvalue problem. Given the disentanglement tensor $D_x(x_0)$ in data coordinates and the metric $g_x(x_0)$, the algorithm finds a $g$-orthogonal matrix $J^*_x$ such that:
$$J^{*\top}_x D_x(x_0) J^*_x = \Lambda$$
$$J^{*\top}_x g_x(x_0) J^*_x = I_n$$
where $\Lambda$ is diagonal. The resulting normal coordinate chart $s = h_{x_0, J^*_x}(x)$ yields latents that are pointwise disentangled at $x_0$.

Key Contributions

1. Riemannian ICA Framework: A theoretical framework that defines independent factors of variation locally using Riemannian geometry, removing the need for a globally invertible generative model.
2. Pointwise Disentanglement: A necessary condition for disentanglement that captures second-order statistical independence at a specific point, formalized via the Hessian of the latent log-likelihood.
3. The Disentanglement Tensor: A tensorial object ($D = \nabla^2 \log \rho - \frac{1}{3} \text{Ric}$) that encodes pointwise disentanglement. Its diagonalization identifies the orientation of the latent space relative to the data density and metric.
4. Coordinate Invariance: The method is constructed to be invariant to the coordinate system used to represent the observations, unlike standard ICA baselines which are sensitive to the chart.

Experimental Results

The authors validated the theory in a controlled source recovery setting where ground-truth sources, metrics, and densities were known. Experiments were conducted on four manifolds: the sphere ($S^d$), hyperbolic space ($H^d$), flat torus ($T^d$), and Symmetric Positive Definite (SPD) matrices.

• Performance: RICA successfully recovered ground-truth sources across all tested manifolds and coordinate charts, achieving a Mean Correlation Coefficient (MCC) of $1.00 \pm 0.00$.
• Baseline Comparison:
  • FastICA: Performed well when samples were near the base point (where the map is approximately linear) but showed sensitivity to the coordinate chart used.
  • Nonlinear ICA (Conformal): Performed poorly, particularly as the dimension increased or source scales became more anisotropic. Its performance degraded significantly in higher dimensions ($n \approx 64$) and was highly sensitive to the coordinate chart.
• Failure Modes: RICA fails when samples are generated outside the injectivity radius of the exponential map, as the logarithmic map cannot invert the exponential map in that region.

Significance and Claims

The paper claims to provide a theoretical basis for studying local disentanglement without assuming a global generative model. By recasting disentanglement as a local geometric property of a Riemannian metric relative to a data density, RICA allows for the interpretation of features in pretrained representations (such as those from self-supervised learning) as independent factors of variation.

The authors emphasize that this work is a step toward a theory for modern representations. They acknowledge practical limitations, noting that computing the disentanglement tensor requires third-order automatic differentiation and closed-form log-likelihoods, which are currently prohibitive for large-scale pretrained models. Consequently, the paper does not claim to solve disentanglement for all modern models immediately but rather offers a rigorous geometric framework for analyzing local structure where such metrics and densities can be defined. The ultimate goal is to enable theoretically grounded interpretability of pretrained representations by recovering local factors of variation.