Stabilizing the Kumaraswamy Distribution

Max Wasserman, Gonzalo Mateos

Abstract: Large-scale latent variable models require expressive continuous distributions that support efficient sampling and low-variance differentiation, achievable through the reparameterization trick. The Kumaraswamy (KS) distribution is both expressive and supports the reparameterization trick with a simple closed-form inverse CDF. Yet, its adoption remains limited. We identify and resolve numerical instabilities in the log-pdf, CDF, and inverse CDF, exposing issues in libraries like PyTorch and TensorFlow. We then introduce simple and scalable latent variable models to address exploration-exploitation trade-offs in contextual multi-armed bandits and facilitate uncertainty quantification for link prediction with graph neural networks. We find these models to be most performant when paired with the stable KS. Our results support the stabilized KS distribution as a core component in scalable variational models for bounded latent variables.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: Camera-ready version: de-anon, remove blue color indicating revisions

Code: https://github.com/maxwass/stabilizing-the-kumaraswamy-distribution

Assigned Action Editor: ~Diana_Cai1

Submission Number: 3913