<figure><img src="image_1.png" /></figure>

HAL
open science

# Disentanglement Learning via Topology

Nikita Balabin, Daria Voronkova, Evgeny Burnaev, Serguei Barannikov

► To cite this version:

Nikita Balabin, Daria Voronkova, Evgeny Burnaev, Serguei Barannikov. Disentanglement Learning via Topology. ICML, 2024, Vienne, Austria. hal-05399453

HAL Id: hal-05399453

https://hal.science/hal-05399453v1

Submitted on 4 Dec 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons Attribution - ShareAlike 4.0 International License
# Disentanglement Learning via Topology

Nikita Balabin¹, Daria Voronkova¹, Ilya Trofimov¹, Evgeny Burnaev¹,², Serguei Barannikov¹,³

We propose TopDis (Topological Disentanglement), a method for learning disentangled representations by adding a multi-scale topological loss term.

* TopDis utilizes shifts in the latent space preserving the Gaussian distribution;

* We improve the reconstruction quality with gradient orthogonalization;

* TopDis enhances disentanglement metrics w.r.t. benchmark models.

## Motivation

<figure><img src="image_3.png" /></figure>

The figure presents an example of various shifts in the latent space for the dSprites dataset with known factors of variations. In this figure, we demonstrate that transformations in disentangled directions (1-2 and 3-4 rows in the figure) have minimal topological dissimilarities as measured by RTD between two sets of samples. As illustrated by this example, minimizing RTD should favor the decomposition of latent space into disentangled directions.

## Related works

### Representation Topology Divergence
(RTD) is a topological tool comparing point clouds $X, \tilde{X}$ with one-to-one correspondence between points. The distances inside clouds $X, \tilde{X}$ define two weighted graphs $G^W, G^{\tilde{W}}$ with the same vertex set $X, w_{AB} = \text{dist}(A, B), \tilde{w}_{AB} = \text{dist}(\tilde{A}, B)$.

RTD tracks the differences in multi-scale topology between $G^{W \le a}, G^{\tilde{W} \le a}$ by comparing them with the graph $G^{\min(w, \tilde{w}) \le a}$. For connected components $C_1, C_2$ of $G^{W \le a}$, let $\alpha_1$ be the scale at which $C_1, C_2$ merge into one component in $G^{\min(w, \tilde{w}) \le a}$. Let at some $\alpha_2 > \alpha_1, C_1, C_2$ become connected in $G^{W \le a}$. R-Cross-Barcode$_1(\tilde{X}, \tilde{X})$ is the multi-set of intervals like $[\alpha_1, \alpha_2]$. By definition, $\text{RTD}(X, \tilde{X}) = 0.5 \cdot (\text{R-Cross-Barcode}_1(\tilde{X}, X) + \text{R-Cross-Barcode}_1(X, \tilde{X}))$

### The Variational Autoencoder (VAE),
A widely used generative model, is not able to achieve disentanglement alone. To address this limitation, researchers have proposed different modifications of VAE losses, most of them are focused on reweighting and decomposing KL term to achieve independence of latent codes.

$$ \mathcal{L}_{\text{VAE}} = \mathcal{L}_{\text{rec}} + \mathcal{L}_{\text{KL}} = -\frac{1}{N} \sum_{n=1}^{N} E_q \left[ \log p_{\theta}(x_n | z) \right] - \frac{1}{N} \sum_{n=1}^{N} \text{KL}(q_{\phi}(z | x_n) || p(z)) $$

<figure><img src="image_4.png" /></figure>

Given a batch $X = x_1, \dots, x_N$, we sample latent representations, $z_n \sim q_{\phi}(z | x_n)$, and the reconstructed samples, $\hat{x}_n \sim p_{\theta}(x | z_n)$.

To ensure that the shifts in a latent code preserve the prior Gaussian distribution, we propose the shifts defined by the equation:

$$ z_{\text{shifted}} = F^{-1}(F(z | \rho, \sigma^2) + C | \rho, \sigma^2) \quad F(z | \rho, \sigma^2) \text{ - the cumulative function of } \mathcal{N}(\rho, \sigma^2). $$

Parameters $\rho, \sigma^2$ are calculated as the empirical mean and variance of the latent code.

The TopDis loss is calculated using the RTD between the original and samples shifted in the latent space:

$$ \mathcal{L}_{\text{TD}} = \text{RTD}^{(P)}(\hat{X}_{\text{original}}, \hat{X}_{\text{shifted}}) \quad \hat{X}_{\text{original}} \sim p_{\theta}(x | z_{\text{original}}), \quad \hat{X}_{\text{shifted}} \sim p_{\theta}(x | z_{\text{shifted}}) $$

Final loss: $\mathcal{L} = \mathcal{L}_{\text{VAE-based}} + \gamma \mathcal{L}_{\text{TD}}$.

<figure><img src="image_5.png" /></figure>

In order to preserve reconstruction quality and achieve state-of-the-art results while minimizing the topological regularization term, we apply the gradient orthogonalization between $\mathcal{L}_{\text{TD}}$ and the reconstruction loss term $\mathcal{L}_{\text{rec}}$.

$$ \nabla^{\text{ort}} \mathcal{L}_{\text{TD}} = \nabla \mathcal{L}_{\text{TD}} - \frac{\langle \nabla \mathcal{L}_{\text{TD}}, \nabla \mathcal{L}_{\text{rec}} \rangle}{\langle \nabla \mathcal{L}_{\text{rec}}, \nabla \mathcal{L}_{\text{rec}} \rangle} \nabla \mathcal{L}_{\text{rec}} $$

<table><thead><tr><th>Method</th><th>FactorVAE score</th><th>MIG</th><th>SAP</th><th>DCI, dis.</th></tr></thead><tbody><tr><td>β-VAE</td><td>0.807 ± 0.037</td><td>0.272 ± 0.101</td><td>0.065 ± 0.002</td><td>0.440 ± 0.102</td></tr><tr><td>β-VAE + TopDis (ours)</td><td>0.833 ± 0.016</td><td>0.348 ± 0.028</td><td>0.066 ± 0.015</td><td>0.506 ± 0.050</td></tr><tr><td>FactorVAE</td><td>0.819 ± 0.028</td><td>0.295 ± 0.049</td><td>0.053 ± 0.006</td><td>0.534 ± 0.029</td></tr><tr><td>FactorVAE + TopDis (ours)</td><td>0.824 ± 0.038</td><td>0.356 ± 0.025</td><td>0.082 ± 0.001</td><td>0.521 ± 0.044</td></tr><tr><td>β-TCVAE</td><td>0.810 ± 0.058</td><td>0.332 ± 0.029</td><td>0.045 ± 0.004</td><td>0.543 ± 0.049</td></tr><tr><td>β-TCVAE + TopDis (ours)</td><td>0.821 ± 0.034</td><td>0.341 ± 0.021</td><td>0.051 ± 0.004</td><td>0.556 ± 0.042</td></tr><tr><td>ControlVAE</td><td>0.806 ± 0.012</td><td>0.333 ± 0.037</td><td>0.056 ± 0.002</td><td>0.557 ± 0.009</td></tr><tr><td>ControlVAE + TopDis (ours)</td><td>0.810 ± 0.012</td><td>0.344 ± 0.029</td><td>0.059 ± 0.002</td><td>0.578 ± 0.007</td></tr><tr><td>DAVA</td><td>0.746 ± 0.099</td><td>0.253 ± 0.058</td><td>0.024 ± 0.015</td><td>0.395 ± 0.054</td></tr><tr><td>DAVA + TopDis (ours)</td><td>0.807 ± 0.010</td><td>0.344 ± 0.010</td><td>0.048 ± 0.012</td><td>0.551 ± 0.019</td></tr></tbody></table>

<table><thead><tr><th></th><th colspan="4">3D Shapes</th></tr></thead><tbody><tr><td>β-VAE</td><td>0.965 ± 0.060</td><td>0.740 ± 0.141</td><td>0.143 ± 0.071</td><td>0.913 ± 0.147</td></tr><tr><td>β-VAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.839 ± 0.077</td><td>0.195 ± 0.030</td><td>0.998 ± 0.004</td></tr><tr><td>FactorVAE</td><td>0.934 ± 0.058</td><td>0.698 ± 0.151</td><td>0.099 ± 0.064</td><td>0.848 ± 0.129</td></tr><tr><td>FactorVAE + TopDis (ours)</td><td>0.975 ± 0.044</td><td>0.779 ± 0.036</td><td>0.159 ± 0.032</td><td>0.940 ± 0.089</td></tr><tr><td>β-TCVAE</td><td>0.909 ± 0.079</td><td>0.693 ± 0.053</td><td>0.113 ± 0.070</td><td>0.877 ± 0.018</td></tr><tr><td>β-TCVAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.751 ± 0.051</td><td>0.147 ± 0.064</td><td>0.901 ± 0.014</td></tr><tr><td>ControlVAE</td><td>0.746 ± 0.094</td><td>0.433 ± 0.094</td><td>0.091 ± 0.068</td><td>0.633 ± 0.093</td></tr><tr><td>ControlVAE + TopDis (ours)</td><td>0.806 ± 0.046</td><td>0.591 ± 0.055</td><td>0.125 ± 0.02</td><td>0.795 ± 0.098</td></tr><tr><td>DAVA</td><td>0.800 ± 0.095</td><td>0.625 ± 0.061</td><td>0.099 ± 0.016</td><td>0.762 ± 0.088</td></tr><tr><td>DAVA + TopDis (ours)</td><td>0.847 ± 0.092</td><td>0.679 ± 0.112</td><td>0.101 ± 0.043</td><td>0.836 ± 0.074</td></tr></tbody></table>

<table><thead><tr><th></th><th colspan="4">3D Faces</th></tr></thead><tbody><tr><td>β-VAE</td><td>1.0 ± 0.0</td><td>0.561 ± 0.017</td><td>0.058 ± 0.008</td><td>0.873 ± 0.018</td></tr><tr><td>β-VAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.545 ± 0.005</td><td>0.052 ± 0.004</td><td>0.854 ± 0.013</td></tr><tr><td>FactorVAE</td><td>1.0 ± 0.0</td><td>0.593 ± 0.058</td><td>0.061 ± 0.014</td><td>0.848 ± 0.011</td></tr><tr><td>FactorVAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.626 ± 0.026</td><td>0.062 ± 0.013</td><td>0.867 ± 0.037</td></tr><tr><td>β-TCVAE</td><td>1.0 ± 0.0</td><td>0.568 ± 0.063</td><td>0.060 ± 0.017</td><td>0.822 ± 0.033</td></tr><tr><td>β-TCVAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.591 ± 0.058</td><td>0.062 ± 0.011</td><td>0.859 ± 0.031</td></tr><tr><td>ControlVAE</td><td>1.0 ± 0.0</td><td>0.447 ± 0.011</td><td>0.058 ± 0.008</td><td>0.713 ± 0.007</td></tr><tr><td>ControlVAE + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.477 ± 0.004</td><td>0.074 ± 0.007</td><td>0.760 ± 0.014</td></tr><tr><td>DAVA</td><td>1.0 ± 0.0</td><td>0.527 ± 0.002</td><td>0.047 ± 0.009</td><td>0.822 ± 0.006</td></tr><tr><td>DAVA + TopDis (ours)</td><td>1.0 ± 0.0</td><td>0.536 ± 0.012</td><td>0.052 ± 0.011</td><td>0.814 ± 0.008</td></tr></tbody></table>

<table><caption>Table 1. Evaluation on the benchmark datasets. **Bold** denotes the best variant in the pair with vs. without the TopDis loss. **Blue** denotes the best method for a dataset/metric. The best variant is almost always a variant with the TopDis loss.</caption><thead><tr><th></th><th colspan="4">MPI 3D</th></tr></thead><tbody><tr><td>β-VAE</td><td>0.428 ± 0.054</td><td>0.221 ± 0.087</td><td>0.092 ± 0.035</td><td>0.238 ± 0.049</td></tr><tr><td>β-VAE + TopDis (ours)</td><td>0.479 ± 0.040</td><td>0.335 ± 0.056</td><td>0.172 ± 0.032</td><td>0.337 ± 0.036</td></tr><tr><td>FactorVAE</td><td>0.589 ± 0.053</td><td>0.336 ± 0.056</td><td>0.179 ± 0.052</td><td>0.391 ± 0.056</td></tr><tr><td>FactorVAE + TopDis (ours)</td><td>0.665 ± 0.041</td><td>0.377 ± 0.053</td><td>0.238 ± 0.040</td><td>0.438 ± 0.065</td></tr><tr><td>β-TCVAE</td><td>0.377 ± 0.039</td><td>0.168 ± 0.021</td><td>0.084 ± 0.012</td><td>0.233 ± 0.059</td></tr><tr><td>β-TCVAE + TopDis (ours)</td><td>0.501 ± 0.023</td><td>0.287 ± 0.011</td><td>0.149 ± 0.006</td><td>0.356 ± 0.045</td></tr><tr><td>ControlVAE</td><td>0.391 ± 0.021</td><td>0.180 ± 0.048</td><td>0.107 ± 0.003</td><td>0.178 ± 0.037</td></tr><tr><td>ControlVAE + TopDis (ours)</td><td>0.554 ± 0.026</td><td>0.232 ± 0.016</td><td>0.154 ± 0.003</td><td>0.274 ± 0.028</td></tr><tr><td>DAVA</td><td>0.404 ± 0.080</td><td>0.234 ± 0.075</td><td>0.086 ± 0.043</td><td>0.268 ± 0.051</td></tr><tr><td>DAVA + TopDis (ours)</td><td>0.606 ± 0.036</td><td>0.337 ± 0.067</td><td>0.181 ± 0.041</td><td>0.401 ± 0.049</td></tr></tbody></table>

<figure><img src="image_6.png" /></figure>

Disentangled directions (azimuth, smile, hair color) found in a pre-trained StyleGAN within the latent space using principal components optimized by multi-scale topological difference after a shift along this axis:

$$ \text{RTD}(\tilde{X}_{\text{original}}, \tilde{X}_{\text{shifted}}) $$