Closed-form Expressions for Maximum Mean Discrepancy with Applications to Wasserstein Auto-Encoders

The Maximum Mean Discrepancy (MMD) has found numerous applications in\nstatistics and machine learning, most recently as a penalty in the Wasserstein\nAuto-Encoder (WAE). In this paper we compute closed-form expressions for\nestimating the Gaussian kernel based MMD between a given distribution and the\nstandard multivariate normal distribution. This formula reveals a connection to\nthe Baringhaus-Henze-Epps-Pulley (BHEP) statistic of the Henze-Zirkler test and\nprovides further insights about the MMD. We introduce the standardized version\nof MMD as a penalty for the WAE training objective, allowing for a better\ninterpretability of MMD values and more compatibility across different\nhyperparameter settings. Next, we propose using a version of batch\nnormalization at the code layer; this has the benefits of making the kernel\nwidth selection easier, reducing the training effort, and preventing outliers\nin the aggregate code distribution. Our experiments on synthetic and real data\nshow that the analytic formulation improves over the commonly used stochastic\napproximation of the MMD, and demonstrate that code normalization provides\nsignificant benefits when training WAEs.\n