Sculpting the Latent Space: Disentanglement with WAE-MMD and Programmable Priors

Quentin Fruytier; Akshay Malhotra; Shahab Hamidi-Rad; Aditya Sant; Aryan Mokhtari; Sujay Sanghavi

cs.LGOct 13, 2025

Sculpting the Latent Space: Disentanglement with WAE-MMD and Programmable Priors

Authors: Quentin Fruytier, Akshay Malhotra, Shahab Hamidi-Rad, Aditya Sant, Aryan Mokhtari, Sujay Sanghavi

Organizations: The University of Texas at Austin · InterDigital Communications, AI Lab

Abstract

Learning disentangled representations, where semantic features are independently captured in a latent space, is dominated by Variational Auto Encoders (VAE) and Adversarial Training. In this paper, we provide direct quantitative and visual evidence that these methods consistently fail to enforce the target factorized representation on the aggregate posterior, leading to entangled representations. We quantify this failure using the Latent Predictability Score (LPS), an unsupervised metric of mutual predictability that empirically lower bounds the level of entanglement in a latent space. To address these limitations, we leverage the non-parametric Maximum Mean Discrepancy (MMD) to directly sculpt the aggregate posterior enabling programmable priors. We demonstrate (1) state-of-the-art unsupervised statistical independence, with a 70% reduction in detected entanglement compared to leading baselines; (2) the flexibility to optionally inject semi-supervised information for explicit semantic alignment; and (3) precise shaping of the aggregate posterior distribution (validated through quantization-aware training) without sacrificing reconstruction quality.

Figures & tables

Figure 1: Visualizing the Sculpting Framework.

Panel of Figure 2 — Figure 2: Visual comparison of learned latent distributions on MNIST for key baselines. Each column corresponds to a different model, showing the covariance matrix between the latent space features over the dataset (top row) and the histogram plots for the marginal distribution of the latent space features over the whole dataset (bottom row). The VAE-based models (a, b) and the standard AE (e) fail to enforce either independence or the target Gaussian geometry. While a batch-wise KLD regularizer (d) achieves a diagonal covariance, its parametric nature prevents it from correctly shaping the marginal distributions leading to co-dependence (LPS equal to 0.97). Only the MMD-regularized model (c) successfully enforces both properties, producing a representation that is truly independent and matches the target prior.

Panel of Figure 3 — Figure 3: Example images (represented in an $8\times 8$ grid) from the synthetic datasets used in our experiments. From left to right: (a) dSprites, with factors like shape, scale, and orientation; (b-d) The XY family of datasets, with factors for position (XY), color (C), and shape (S); (c) The CelebA dataset consisting of faces with 40 different semantic factors.

	ImageNet ( $d=64$ )		CIFAR10 ( $d=64$ )		MNIST ( $d=12$ )		dsprites ( $d=5$ )		XY ( $d=2$ )		XYC ( $d=3$ )		XYCS ( $d=4$ )		celebA ( $d=32$ )
Model	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)	LPS	NMSE (dB)
SWAE-MMD (ours)	0.03	-11.70	0.04	-15.61	0.22	-11.92	0.04	-17.89	-0.05	-15.22	-0.01	-16.31	0.11	-16.76	0.08	-14.37
DGAE	0.77	-10.3	0.69	-15.53	0.63	-11.48	0.56	-14.61	0.25	-6.74	0.50	-10.94	0.59	-15.02	0.48	-14.40
NASHAE	0.56	-11.32	0.17	-15.02	0.22	-11.32	0.10	-13.48	-0.07	-13.08	0.14	-14.93	0.37	-13.15	0.36	-13.67
FactorVAE ( $\gamma$ $=$ 10.0)	0.44	-10.75	0.63	-12.49	0.27	-8.62	0.08	-14.56	-0.04	- 15.36	0.00	-15.92	0.17	-15.52	0.07	-14.39
$\beta$ -TCVAE ( $\beta$ $=$ 4.0)	0.79	-9.80	0.89	-10.68	0.84	-5.97	0.22	-8.83	-0.04	-14.68	-0.01	-13.72	0.17	-12.40	0.59	-13.10

Table 1: Main quantitative results comparing our unsupervised MMD-based Autoencoder against baselines (the latent space dimension is specified next to the dataset name). Each cell presents two key metrics: LPS-lgbm (

R^{2}

) | reconstruction NMSE (dB) where lower values are better (

\downarrow

). The best performance for each metric on each dataset is highlighted in bold . Averaged over all

8

datasets, AE-MMD obtains a 70% reduction in detected latent space entanglement compared to Factor VAE, the next best baseline, all while maintaining a better image reconstruction quality.

Panel of Figure 4 — Figure 4: Aligning the Latent Space to Semantic Features Using The Programmable Prior On dSprites. The plots show the DCI Disentanglement Score (left), the Latent Predictability Score (LPS) using LGBM (middle), and the reconstruction NMSE in decibels (right) as the number of pseudo-labels used for alignment increases. The DCI scores were always evaluated with respect to all 5 semantic features (Shape, Scale, Orientation, X-Position, Y-Position).

Panel of Figure 6 — Figure 6: Visualization of Alignment Mechanism in Programmable Prior Framework. The figure shows the covariance matrix of the joint distribution $p(z,f(u))$ for the target (left) and the learned joint distribution (right) with $2$ and $4$ pseudo-labels respectively on the dsprites dataset.

Panel of Figure 8 — Figure 8: Compression with the Programmable Prior. We compare the reconstruction NMSE (dB) of our MMD-regularized models against baselines with a uniform (left) or a Gaussian (right) quantizer.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Number of Layers	Latent Dimension ( $d$ )	Training Epochs
XY, XYC, XYCS	6, 6, 6	2, 3, 4	500, 250, 125
dSprites	6	5	50
CelebA	6	32	75
MNIST	4	12	250
CIFAR-10	4	64	250
Tiny ImageNet	6	64	75

Appendix

Table 2: Dataset-specific hyperparameters. The "Number of Layers" refers to the depth of the encoder and decoder individually (e.g., "6" indicates 6 layers for the encoder and 6 for the decoder).

	L=64
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	NMSE_dB	NMSE_linear
AE-MMD-Gaussian (Ours)	54.386 $\pm$ 2.220	0.029 $\pm$ 0.002	0.001 $\pm$ 0.004	-11.704 $\pm$ 0.007	0.068 $\pm$ 0.000
FactorVAE ( $\gamma$ =10.0)	2.338 $\pm$ 0.286	0.437 $\pm$ 0.064	0.483 $\pm$ 0.065	-10.754 $\pm$ 0.239	0.084 $\pm$ 0.005
$\beta$ -TCVAE ( $\beta$ =4.0)	4.950 $\pm$ 0.264	0.794 $\pm$ 0.018	0.932 $\pm$ 0.015	-9.798 $\pm$ 0.027	0.105 $\pm$ 0.001
$\beta$ -VAE ( $\beta$ =4)	4.934 $\pm$ 0.104	0.726 $\pm$ 0.018	0.872 $\pm$ 0.035	-9.866 $\pm$ 0.009	0.103 $\pm$ 0.000
Standard AE	7.699 $\pm$ 0.207	0.455 $\pm$ 0.006	0.460 $\pm$ 0.007	-11.775 $\pm$ 0.005	0.066 $\pm$ 0.000

Appendix

Table 3: Summary for TINY IMAGENET

	L=64
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	NMSE_dB	NMSE_linear
AE-MMD-Gaussian (Ours)	65.917 $\pm$ 1.788	0.038 $\pm$ 0.004	0.025 $\pm$ 0.008	-15.615 $\pm$ 0.008	0.027 $\pm$ 0.000
FactorVAE ( $\gamma$ =10.0)	3.347 $\pm$ 0.075	0.628 $\pm$ 0.004	0.788 $\pm$ 0.006	-12.491 $\pm$ 0.012	0.056 $\pm$ 0.000
$\beta$ -TCVAE ( $\beta$ =4.0)	10.711 $\pm$ 0.350	0.890 $\pm$ 0.008	0.953 $\pm$ 0.006	-10.684 $\pm$ 0.005	0.085 $\pm$ 0.000
$\beta$ -VAE ( $\beta$ =4)	11.506 $\pm$ 0.238	0.846 $\pm$ 0.008	0.925 $\pm$ 0.006	-10.492 $\pm$ 0.006	0.089 $\pm$ 0.000
Standard AE	6.929 $\pm$ 0.192	0.508 $\pm$ 0.005	0.513 $\pm$ 0.005	-15.834 $\pm$ 0.007	0.026 $\pm$ 0.000

Appendix

Table 4: Summary for CIFAR10

	L=12
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	NMSE_dB	NMSE_linear
AE-MMD-Gaussian (Ours)	46.635 $\pm$ 2.442	0.220 $\pm$ 0.007	0.333 $\pm$ 0.006	-11.922 $\pm$ 0.024	0.064 $\pm$ 0.000
DGAE	6.191 $\pm$ 0.368	0.632 $\pm$ 0.002	0.671 $\pm$ 0.005	-11.478 $\pm$ 0.027	0.071 $\pm$ 0.000
NASHAE	47.835 $\pm$ 6.503	0.210 $\pm$ 0.014	0.407 $\pm$ 0.023	-11.389 $\pm$ 0.133	0.073 $\pm$ 0.002
FactorVAE ( $\gamma$ =10.0)	2.628 $\pm$ 0.094	0.271 $\pm$ 0.066	0.275 $\pm$ 0.098	-8.621 $\pm$ 0.023	0.137 $\pm$ 0.001
$\beta$ -TCVAE ( $\beta$ =4.0)	8.071 $\pm$ 2.238	0.838 $\pm$ 0.022	0.893 $\pm$ 0.015	-5.967 $\pm$ 0.008	0.253 $\pm$ 0.000

Appendix

Table 5: Summary for MNIST

	L=32
Baseline	NMSE_dB	NMSE_linear	completeness	disentanglement	feature_dependence	informativeness	lps_lgbm	lps_mlp	mig	sap
AE-MMD-Gaussian	-14.370 $\pm$ 0.012	0.037 $\pm$ 0.000	0.009 $\pm$ 0.000	0.008 $\pm$ 0.000	61.678 $\pm$ 3.671	0.868 $\pm$ 0.001	0.080 $\pm$ 0.003	0.227 $\pm$ 0.009	0.006 $\pm$ 0.002	0.000 $\pm$ 0.000
DGAE	-14.404 $\pm$ 0.016	0.036 $\pm$ 0.000	0.014 $\pm$ 0.001	0.013 $\pm$ 0.001	5.742 $\pm$ 0.318	0.873 $\pm$ 0.001	0.481 $\pm$ 0.011	0.529 $\pm$ 0.010	0.011 $\pm$ 0.003	0.000 $\pm$ 0.000
FactorVAE ( $\gamma$ =10.0)	-14.394 $\pm$ 0.150	0.036 $\pm$ 0.001	0.014 $\pm$ 0.001	0.013 $\pm$ 0.001	62.873 $\pm$ 2.077	0.871 $\pm$ 0.001	0.071 $\pm$ 0.007	0.106 $\pm$ 0.018	0.014 $\pm$ 0.003	0.000 $\pm$ 0.000
NASHAE	-13.672 $\pm$ 0.340	0.043 $\pm$ 0.003	0.010 $\pm$ 0.001	0.009 $\pm$ 0.001	18.684 $\pm$ 14.394	0.867 $\pm$ 0.001	0.362 $\pm$ 0.185	0.477 $\pm$ 0.181	0.007 $\pm$ 0.001	0.000 $\pm$ 0.000
Standard AE	-14.546 $\pm$ 0.012	0.035 $\pm$ 0.000	0.012 $\pm$ 0.001	0.012 $\pm$ 0.001	7.490 $\pm$ 0.169	0.871 $\pm$ 0.000	0.441 $\pm$ 0.004	0.479 $\pm$ 0.005	0.008 $\pm$ 0.001	0.000 $\pm$ 0.000

Appendix

Table 6: Summary for CELEBA (L=32)

	L=64
Baseline	NMSE_dB	NMSE_linear	completeness	disentanglement	feature_dependence	informativeness	lps_lgbm	lps_mlp	mig	sap
AE-MMD-Gaussian	-16.087 $\pm$ 0.011	0.025 $\pm$ 0.000	0.011 $\pm$ 0.001	0.012 $\pm$ 0.001	63.822 $\pm$ 2.724	0.866 $\pm$ 0.001	0.097 $\pm$ 0.003	0.309 $\pm$ 0.007	0.004 $\pm$ 0.001	0.000 $\pm$ 0.000
DGAE	-15.894 $\pm$ 0.030	0.026 $\pm$ 0.000	0.019 $\pm$ 0.002	0.021 $\pm$ 0.003	7.992 $\pm$ 0.522	0.874 $\pm$ 0.001	0.498 $\pm$ 0.008	0.839 $\pm$ 0.079	0.012 $\pm$ 0.005	0.000 $\pm$ 0.000
FactorVAE ( $\gamma$ =10.0)	-14.885 $\pm$ 0.088	0.032 $\pm$ 0.001	0.018 $\pm$ 0.002	0.020 $\pm$ 0.002	89.942 $\pm$ 7.703	0.871 $\pm$ 0.001	0.566 $\pm$ 0.007	0.829 $\pm$ 0.016	0.013 $\pm$ 0.002	0.000 $\pm$ 0.000
NASHAE	-13.557 $\pm$ 1.355	0.046 $\pm$ 0.014	0.011 $\pm$ 0.001	0.013 $\pm$ 0.001	8.334 $\pm$ 3.465	0.866 $\pm$ 0.000	0.759 $\pm$ 0.185	0.930 $\pm$ 0.067	0.008 $\pm$ 0.003	0.000 $\pm$ 0.000
Standard AE	-16.251 $\pm$ 0.016	0.024 $\pm$ 0.000	0.015 $\pm$ 0.001	0.017 $\pm$ 0.001	8.984 $\pm$ 0.336	0.870 $\pm$ 0.000	0.516 $\pm$ 0.007	0.566 $\pm$ 0.007	0.007 $\pm$ 0.001	0.000 $\pm$ 0.000

Appendix

Table 7: Summary for CELEBA (L=64)

	L=5
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian (Ours)	102.265 $\pm$ 21.202	0.039 $\pm$ 0.002	0.033 $\pm$ 0.006	0.058 $\pm$ 0.024	-17.888 $\pm$ 0.281	0.016 $\pm$ 0.001	0.000 $\pm$ 0.000	0.026 $\pm$ 0.018	0.025 $\pm$ 0.017	0.950 $\pm$ 0.002
AE-MMD-hybrid-sampler (Ours)	263.218 $\pm$ 87.577	0.105 $\pm$ 0.008	0.087 $\pm$ 0.012	0.247 $\pm$ 0.030	-16.311 $\pm$ 0.261	0.023 $\pm$ 0.001	0.246 $\pm$ 0.027	0.200 $\pm$ 0.038	0.196 $\pm$ 0.038	0.957 $\pm$ 0.001
DGAE	4.065 $\pm$ 1.233	0.557 $\pm$ 0.112	0.573 $\pm$ 0.147	0.261 $\pm$ 0.198	-14.614 $\pm$ 2.173	0.040 $\pm$ 0.026	0.000 $\pm$ 0.000	0.226 $\pm$ 0.104	0.219 $\pm$ 0.101	0.962 $\pm$ 0.004
FactorVAE ( $\gamma$ =10.0)	8.477 $\pm$ 1.004	0.082 $\pm$ 0.010	0.053 $\pm$ 0.017	0.226 $\pm$ 0.079	-14.564 $\pm$ 0.110	0.035 $\pm$ 0.001	0.000 $\pm$ 0.000	0.081 $\pm$ 0.033	0.079 $\pm$ 0.032	0.964 $\pm$ 0.004
$\beta$ -TCVAE ( $\beta$ =4.0)	6.469 $\pm$ 4.162	0.223 $\pm$ 0.165	0.180 $\pm$ 0.182	0.408 $\pm$ 0.215	-8.831 $\pm$ 0.085	0.131 $\pm$ 0.003	0.000 $\pm$ 0.000	0.212 $\pm$ 0.119	0.195 $\pm$ 0.105	0.960 $\pm$ 0.005

Appendix

Table 8: Summary for DSPRITES

	L=2
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian (Ours)	240.970 $\pm$ 243.891	-0.053 $\pm$ 0.007	-0.004 $\pm$ 0.003	0.751 $\pm$ 0.554	-15.222 $\pm$ 0.403	0.030 $\pm$ 0.003	0.423 $\pm$ 0.316	0.199 $\pm$ 0.244	0.199 $\pm$ 0.243	1.000 $\pm$ 0.000
AE-MMD-hybrid-sampler (Ours)	380.827 $\pm$ 408.352	-0.065 $\pm$ 0.020	-0.003 $\pm$ 0.005	1.047 $\pm$ 0.414	-15.825 $\pm$ 0.154	0.026 $\pm$ 0.001	0.713 $\pm$ 0.139	0.471 $\pm$ 0.230	0.469 $\pm$ 0.233	1.000 $\pm$ 0.000
DGAE	16.631 $\pm$ 19.417	0.248 $\pm$ 0.380	0.241 $\pm$ 0.385	1.580 $\pm$ 1.414	-6.742 $\pm$ 2.522	0.262 $\pm$ 0.206	0.000 $\pm$ 0.000	0.521 $\pm$ 0.435	0.521 $\pm$ 0.435	0.999 $\pm$ 0.001
NASHAE	52.672 $\pm$ 22.762	-0.066 $\pm$ 0.012	0.007 $\pm$ 0.004	0.124 $\pm$ 0.059	-13.082 $\pm$ 0.849	0.050 $\pm$ 0.010	0.325 $\pm$ 0.186	0.134 $\pm$ 0.144	0.131 $\pm$ 0.142	0.999 $\pm$ 0.001
FactorVAE ( $\gamma$ =10.0)	2.022 $\pm$ 1.392	-0.040 $\pm$ 0.011	0.015 $\pm$ 0.013	0.693 $\pm$ 0.545	-15.364 $\pm$ 0.359	0.029 $\pm$ 0.002	0.537 $\pm$ 0.268	0.337 $\pm$ 0.255	0.336 $\pm$ 0.255	1.000 $\pm$ 0.000

Appendix

Table 9: Summary for XY

	L=3
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian (Ours)	49.341 $\pm$ 12.168	-0.014 $\pm$ 0.016	0.016 $\pm$ 0.007	0.256 $\pm$ 0.118	-16.311 $\pm$ 0.245	0.023 $\pm$ 0.001	0.237 $\pm$ 0.077	0.115 $\pm$ 0.039	0.113 $\pm$ 0.038	0.999 $\pm$ 0.000
AE-MMD-hybrid-sampler (Ours)	324.644 $\pm$ 274.142	0.168 $\pm$ 0.030	0.117 $\pm$ 0.035	0.465 $\pm$ 0.336	-15.962 $\pm$ 0.205	0.025 $\pm$ 0.001	0.266 $\pm$ 0.240	0.228 $\pm$ 0.224	0.227 $\pm$ 0.224	1.000 $\pm$ 0.000
DGAE	4.802 $\pm$ 2.926	0.498 $\pm$ 0.417	0.503 $\pm$ 0.414	0.340 $\pm$ 0.181	-10.936 $\pm$ 3.421	0.110 $\pm$ 0.083	0.000 $\pm$ 0.000	0.184 $\pm$ 0.105	0.175 $\pm$ 0.098	0.996 $\pm$ 0.005
NASHAE	16.476 $\pm$ 5.414	0.145 $\pm$ 0.091	0.138 $\pm$ 0.084	0.048 $\pm$ 0.004	-14.933 $\pm$ 0.701	0.033 $\pm$ 0.005	0.217 $\pm$ 0.080	0.095 $\pm$ 0.035	0.094 $\pm$ 0.034	0.999 $\pm$ 0.000
FactorVAE ( $\gamma$ =10.0)	1.581 $\pm$ 0.076	-0.000 $\pm$ 0.006	0.032 $\pm$ 0.005	0.140 $\pm$ 0.059	-15.919 $\pm$ 0.219	0.026 $\pm$ 0.001	0.183 $\pm$ 0.157	0.079 $\pm$ 0.069	0.076 $\pm$ 0.065	1.000 $\pm$ 0.000

Appendix

Table 10: Summary for XYC

	L=4
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian (Ours)	42.117 $\pm$ 2.445	0.111 $\pm$ 0.011	0.106 $\pm$ 0.020	0.092 $\pm$ 0.036	-16.757 $\pm$ 0.143	0.021 $\pm$ 0.001	0.110 $\pm$ 0.032	0.048 $\pm$ 0.016	0.048 $\pm$ 0.016	0.997 $\pm$ 0.000
AE-MMD-hybrid-sampler (Ours)	133.148 $\pm$ 52.617	0.203 $\pm$ 0.021	0.151 $\pm$ 0.025	0.235 $\pm$ 0.070	-16.101 $\pm$ 0.283	0.025 $\pm$ 0.002	0.278 $\pm$ 0.060	0.149 $\pm$ 0.051	0.146 $\pm$ 0.050	0.998 $\pm$ 0.000
DGAE	7.087 $\pm$ 3.365	0.586 $\pm$ 0.235	0.583 $\pm$ 0.233	0.275 $\pm$ 0.159	-15.023 $\pm$ 0.701	0.032 $\pm$ 0.005	0.000 $\pm$ 0.000	0.190 $\pm$ 0.076	0.187 $\pm$ 0.074	0.997 $\pm$ 0.001
NASHAE	8.577 $\pm$ 4.533	0.368 $\pm$ 0.026	0.383 $\pm$ 0.023	0.040 $\pm$ 0.009	-13.152 $\pm$ 0.681	0.049 $\pm$ 0.007	0.168 $\pm$ 0.007	0.145 $\pm$ 0.043	0.143 $\pm$ 0.044	0.997 $\pm$ 0.001
FactorVAE ( $\gamma$ =10.0)	1.609 $\pm$ 0.127	0.167 $\pm$ 0.020	0.165 $\pm$ 0.032	0.176 $\pm$ 0.114	-15.517 $\pm$ 0.366	0.028 $\pm$ 0.002	0.192 $\pm$ 0.090	0.149 $\pm$ 0.098	0.148 $\pm$ 0.097	0.997 $\pm$ 0.001

Appendix

Table 11: Summary for XYCS

Figure 9: The Paradox of Independence vs. Alignment. Each point represents a model. The ideal model would be in the bottom-right (high alignment, high independence). We observe that models like DGAE occupy the unfavorable top-left quadrant, achieving a high DCI score only by sacrificing independence. Our method (with a Gaussian prior) occupies the bottom-left, prioritizing true independence over potentially flawed alignment metrics.

	dsprites		xy		xyc		xycs
Baseline	Disent. ( $\uparrow$ )	LPS ( $\downarrow$ )	Disent. ( $\uparrow$ )	LPS ( $\downarrow$ )	Disent. ( $\uparrow$ )	LPS ( $\downarrow$ )	Disent. ( $\uparrow$ )	LPS ( $\downarrow$ )
AE-MMD-Gaussian	0.025	0.039	0.199	-0.053	0.113	-0.014	0.048	0.111
AE-MMD-hybrid	0.196	0.105	0.469	-0.065	0.227	0.168	0.146	0.203
DGAE	0.219	0.557	0.521	0.248	0.175	0.498	0.187	0.596
FactorVAE	0.079	0.082	0.336	-0.040	0.076	-0.000	0.148	0.187
$\beta$ -TCVAE	0.195	0.222	0.368	-0.042	0.090	-0.010	0.245	0.171

Appendix

Table 12: Comparison of Disentanglement (DCI) and Latent Predictability Score (LPS) across models and synthetic datasets. For Disentanglement, higher is better (

\uparrow

), and for LPS, lower is better (

\downarrow

Panel of Figure 10 — Figure 10: Comparison of metric stability across 5 random seeds. The error bars represent one standard deviation. (a) The DCI scores exhibit extremely high variance, with the standard deviation often being as large as the mean itself, highlighting the metric’s unreliability. (b) In stark contrast, our proposed LPS metric shows minimal variance, demonstrating its stability and reliability as an intrinsic measure of representation quality.

Panel of Figure 11 — Figure 11: Latent Independence vs. Latent Space Dimension on CIFAR-10. The plots show the reconstruction NMSE in decibels (left), and the Latent Predictability Score (LPS) using MLP (middle) and LGBM (right) regression as the size of the latent bottleneck increases. Lower scores are better for all metrics. Our method (AE-MMD-Gauss, green line) maintains a high degree of independence across all dimensions, while most baselines learn increasingly entangled representations. Moreover, we achieve this stability at no significant cost to reconstruction error, which is surpassed only by the unregularized standard AE model.

Panel of Figure 12 — Figure 12: Ablation Study on CIFAR-10: Regularization Coefficient and Batch Size. We report the LPS-lgbm score and the test reconstruction NMSE (dB) against regularization coefficients ranging from $0.1$ to $5$ (left), and training batch sizes ranging from $64$ to $2048$ (right).

Figure 13: Ablation Study of Kernel Type on CIFAR-10. We report a scatter plot of the LPS-lgbm score versus the test reconstruction NMSE (dB) for exponential, quadratic, and Gaussian (RBF) kernel types. We also compare configurations with different numbers of bandwidths for kernel averaging according to the median heuristic.

Panel of Figure 14 — Figure 14: Qualitative Comparison on the XY Dataset. Each column compares the learned latent marginals (top) with the corresponding latent traversals (bottom). (a) Our MMD-trained model successfully learns the true uniform distribution, which results in clean, disentangled traversals. (b) In contrast, the standard AE defaults to a uni-modal distribution, leading to entangled traversals.

Prior	NMSE (dB)	LPS-lgbm	LPS-mlp	DCI-RF	MIG
AE-Gaussian	-15.2	-0.053	-0.004	0.199	0.751
AE-Uniform	-15.8	-0.064	-0.003	0.469	1.047
AE-GMM (2 component)	-14.9	-0.051	0.000	0.251	0.863

Appendix

Table 13: Results for our MMD-based Autoencoder on the XY dataset with different engineered prior distributions. We find selecting a Uniform target prior significantly improve alignment metrics.

Figure 15: Learned marginal distributions for the AE-MMD-hybrid-sampler on dSprites. The histograms provide direct visual proof that the model successfully learned to match the specified target distributions for each latent dimension.

	L=5
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian	102.265 $\pm$ 21.202	0.039 $\pm$ 0.002	0.033 $\pm$ 0.006	0.058 $\pm$ 0.024	-17.888 $\pm$ 0.281	0.016 $\pm$ 0.001	0.000 $\pm$ 0.000	0.026 $\pm$ 0.018	0.025 $\pm$ 0.017	0.950 $\pm$ 0.002
AE-MMD-hybrid-sampler	263.218 $\pm$ 87.577	0.105 $\pm$ 0.008	0.087 $\pm$ 0.012	0.247 $\pm$ 0.030	-16.311 $\pm$ 0.261	0.023 $\pm$ 0.001	0.246 $\pm$ 0.027	0.200 $\pm$ 0.038	0.196 $\pm$ 0.038	0.957 $\pm$ 0.001

Appendix

Table 14: Summary for DSPRITES

	L=2
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian	240.970 $\pm$ 243.891	-0.053 $\pm$ 0.007	-0.004 $\pm$ 0.003	0.751 $\pm$ 0.554	-15.222 $\pm$ 0.403	0.030 $\pm$ 0.003	0.423 $\pm$ 0.316	0.199 $\pm$ 0.244	0.199 $\pm$ 0.243	1.000 $\pm$ 0.000
AE-MMD-hybrid-sampler	380.827 $\pm$ 408.352	-0.065 $\pm$ 0.020	-0.003 $\pm$ 0.005	1.047 $\pm$ 0.414	-15.825 $\pm$ 0.154	0.026 $\pm$ 0.001	0.713 $\pm$ 0.139	0.471 $\pm$ 0.230	0.469 $\pm$ 0.233	1.000 $\pm$ 0.000

Appendix

Table 15: Summary for XY

Figure 16: Learned marginal distributions for the AE-MMD-hybrid-sampler on XYC.

	L=3
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian	49.341 $\pm$ 12.168	-0.014 $\pm$ 0.016	0.016 $\pm$ 0.007	0.256 $\pm$ 0.118	-16.311 $\pm$ 0.245	0.023 $\pm$ 0.001	0.237 $\pm$ 0.077	0.115 $\pm$ 0.039	0.113 $\pm$ 0.038	0.999 $\pm$ 0.000
AE-MMD-hybrid-sampler	324.644 $\pm$ 274.142	0.168 $\pm$ 0.030	0.117 $\pm$ 0.035	0.465 $\pm$ 0.336	-15.962 $\pm$ 0.205	0.025 $\pm$ 0.001	0.266 $\pm$ 0.240	0.228 $\pm$ 0.224	0.227 $\pm$ 0.224	1.000 $\pm$ 0.000

Appendix

Table 16: Summary for XYC

Figure 17: Learned marginal distributions for the AE-MMD-hybrid-sampler on XYCS.

	L=4
Baseline	Covariance Ratio	LPS_LGBM	LPS_MLP	MIG_mig_score	NMSE_dB	NMSE_linear	SAP_sap_score	completeness	disentanglement	informativeness
AE-MMD-Gaussian	42.117 $\pm$ 2.445	0.111 $\pm$ 0.011	0.106 $\pm$ 0.020	0.092 $\pm$ 0.036	-16.757 $\pm$ 0.143	0.021 $\pm$ 0.001	0.110 $\pm$ 0.032	0.048 $\pm$ 0.016	0.048 $\pm$ 0.016	0.997 $\pm$ 0.000
AE-MMD-hybrid-sampler	133.148 $\pm$ 52.617	0.203 $\pm$ 0.021	0.151 $\pm$ 0.025	0.235 $\pm$ 0.070	-16.101 $\pm$ 0.283	0.025 $\pm$ 0.002	0.278 $\pm$ 0.060	0.149 $\pm$ 0.051	0.146 $\pm$ 0.050	0.998 $\pm$ 0.000

Appendix

Table 17: Summary for XYCS

Panel of Figure 18 — Figure 18: Visualizing the Latent Space Copying experiment on MNIST. (b) First, a standard Autoencoder was trained, and its complex, entangled latent distribution was saved to serve as an empirical prior. (a) Then, a new "student" model was trained using our MMD-based sculpting regularizer, tasked with replicating this non-analytic prior. The figure displays the covariance matrix between the latent space features over the whole dataset and the histogram plots of the marginal distribution of the latent space features over the whole dataset. The near-identical covariance matrices and marginal distributions visually confirm our method’s ability to enforce a complex, co-dependent latent geometry that would be impossible to specify analytically.

Figure 19: Failure of Knowledge Transfer via Aggregate Posterior Matching. (Left) Our MMD regularizer successfully forces a "student" model to replicate the latent distribution of a "teacher" model. (Right) Despite this, latent traversals show that the semantic meaning of the latent dimensions is completely different. For the same inputs (connected by lines), traversing a latent feature produces different visual changes, demonstrating that the underlying representations are not the same.

Disentangled Representation Learning

Disentangled representation learning separates distinct factors of variation so models can interpret, control, and generalize their predictions.

See all 231 papers

A Unified Latent Space Disentanglement VAE Framework with Robust Disentanglement Effectiveness Evaluation
Mar 11, 2026Xiaoan Lang, Md Mostafizer Rahman, Fang LiuVariational Autoencoders Disentangled Representation Learning
Department of Applied and Computational Mathematics and Statistics University of Notre Dame, Notre Dame, IN, USA · Lucy Family Institute for Data & Society University of Notre Dame, Notre Dame, IN, USA
XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision
Jan 29, 2026Alexandre Myara, Nicolas Bourriez, Thomas Boyer +3Variational Autoencoders Disentangled Representation Learning
ENS Paris
The Trade-off Between Covariate Dependence and Latent Structure in Representation Learning
Aug 17, 2026Małgorzata Łazęcka, Ewa SzczurekDisentangled Representation Learning Latent Space Alignment
Faculty of Mathematics, Informatics, and Mechanics, University of Warsaw Warsaw, Poland · Institute of AI for Health, Helmholtz Munich Munich, Germany

Mar 11, 2026J/K move · Enter open · S save

A Unified Latent Space Disentanglement VAE Framework with Robust Disentanglement Effectiveness Evaluation

Xiaoan Lang, Md Mostafizer Rahman, Fang Liu

Department of Applied and Computational Mathematics and Statistics University of Notre Dame, Notre Dame, IN, USA · Lucy Family Institute for Data & Society University of Notre Dame, Notre Dame, IN, USA

Variational Autoencoders Disentangled Representation Learning Representation Disentanglement Neural Network Interpretability

Open paper AskPDF arXiv

Evaluating and interpreting latent representations, such as variational autoencoders (VAEs), remains a significant challenge for diverse data types, especially when ground-truth generative factors are unknown. To address this, we unify several state-of-the-art disentangled VAE approaches for latent space disentanglement into one framework -- bfVAE. To assess the effectiveness of a disentangled VAE model and enhance latent space interpretability, we propose Feature Variance Heterogeneity via Latent Traversal (FVH-LT) and Dirty Block Sparse Regression in Latent Space (DBSR-LS). To ensure robust interpretability of learned latent space, we develop a greedy alignment strategy (GAS) that mitigates label switching and aligns latent dimensions across runs to set the foundation of result aggregation. We also introduce a convenient scalar latent space separation index (LSSI) based on the GAS-aligned outputs of FVH-LT and DBSR-LS to summarize the overall latent structural separation without knowledge of the ground-truth generative factors. We compare bfVAE to five VAE models and validate the effectiveness FVH-LT, DBSR-LS, and LSSI in on seven tabular and image datasets. Under our examined experimental settings, bfVAE provides a more flexible disentanglement framework achieves more favorable overall trade-off between disentanglement and reconstruction than the benchmark VAE models; FVH-LT and DBSR-LS reliably uncover semantically meaningful and domain-relevant latent structures and generally yield consistent results; and LSSI makes an effective quantitative summary of latent structural separation.

Abstract

Disentangled Representation Learning

Explore similar work

A Unified Latent Space Disentanglement VAE Framework with Robust Disentanglement Effectiveness Evaluation

XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision

The Trade-off Between Covariate Dependence and Latent Structure in Representation Learning