Stabilizing the Kumaraswamy Distribution
Authors: Max Wasserman
Gonzalo Mateos
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (03/2025)
Stabilizing the Kumaraswamy Distribution
Max Wasserman
mwwsserff@cs.rrchsster.om
Deportment of Computer Science University of Rochester
Gonzalo Mateos
cannteosb@eccrroocheee edu
Department of Electrical and Computer Engineering University of Rochester
Reviewed on OpenReview: https: // openreview. net/ metteeettuuueeeeeeee forum? id- XXXX
Abstract
Large-scale latent variable models require expressive continuous distributions that support efficient sampling and low-variance differentiation, achievable through the reparameteri- zation trick. The Kumaraswamy (KS) distribution is both expressive and supports the reparameterization trick with 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 Py'Torch and TensorFlow. We then introduce simple and scalable latent variable models to address exploration-cxploitatee multi-armed bandits and facilitate uncertainty quantification for link prediction with graph neural networks. We find these models be most performant when paired with the stable KS. Our results support the stabilized KS distribution core component in scalable variational models for bounded latent variables.
tradc-offs in contextual
1 Introduction
Probabilistic models use probability distributions as building blocks to model complex joint distributions between random variables. Such distributions can model unobserved 'latent' variables z, or observed 'data' variables æ. Bounded interval-supported latent variables are central to many key applications, such as unobserved probabilities (e.g., user clicks in recommendation systems or links between network nodes), missing measurements in control systems (e.g., joint angles in [0, 2T]),, and stochastic policies over bounded actions in reinforcement learning (c.g., motor torque in [-10, 10)).
To meet the demands of large-scale latent variable models, distributions supported on bounded intervals must satisfy the following criteria: (i) support the reparameterization trick through an explicit reparameterization function, such as a closod-form inverse CDF, enabling low-variance gradient estimation and efficiont sampling: (ii) provide sufficient expressiveness capture complex latent spaces; and (iii) offer simple distribution--rlalel functions (log-pdf, explicit reparameterization function, and gradients) that allow fast and accurate evaluation. In Section 2 argue that the Kumaraswamy (KS) distribution uniquely meets these criteria, yet remains surprisingly underused. In Section g we demonstrate that the KS distribation-relateed functions exhibit numerical instabilities concealed by standard parameterizations and exacerbated in large-scale latent variable models.
In this paper, we make the following technical contributions:
We introduce an unconstrained logarithmic parameterization of the KS's log-pdf, CDF, inverse CDF, and gradients, which isolate the dominant numerical instabilities, allowing application recently developed stabilization techniques (Section 3)
We illustrate how the KS can be used as an integral component of latent variable models and demonstrate its benefits alternative distributions. In addition to the workhorse variational autoencoder (VAF),

---

Published in Transactions on Machine Learning Research (03/2025)
Sampling (s)
CB
tanh
Beta
KS
Gradient (le
Time
CB
tanhy
Beta
KS
0
00
0
10
1
Figure Comparison of relevant bounded interval-supported distributions. Left: Time for sampling and differentiating through samples. The Beta lacks explicit reparameterization, and has slower sampling and gradients. Right: Expressiveness in terms of attainable prototypical shapes.
experiments showcase: (i) the Variational Bandit Encoder (VBE), addressing exploration-cxploitation trade-offs in contextual Bernoulli multi-armed bandits (Section 4.2g and (ii) the Variational Edge Encoder (VEE), facilitating uncertainty quantification in link prediction with graph neural networks (Section 4.3
The VBE and VEE scalable latent variable models with bounded interval-supported latent variables. Unlike traditional methods which tend to model global latent variables (Section 5), such the parameters of shared neural network (NN), the VBE and VEE define local latent variables per bandit arm or network link. This allows the models to incorporate prior knowledge precisely where domain expertise tends to reside-at at a granular level, such as the expected reward of a specific arm or the probability of a particular link's existence. Our numerical experiments demonstrate that both models tend to perform best when paired with the stabilized KS distribution in their variational posterior, reinforcing its role as core component large-scale bounded latent variable modeling. Our experiments in Sections 4.2 4.3 provide evidence toward the benefits of adopting the KS in timely application domains enabled by our stable parameterization without making claims improvements over state-of-the-arl models, which future work may investigate.
2 Background
The KS distribution JJones 2009
Kumaraswamy 1980 has pdf f(x)) abx xa)b-1 CDF F(x)
x²)², (1-x*),, and inverse CDF F ( (1 d all defined for x, (0, 1) and parameterized by
a, b 0. The differential entropy of KS with parameters a, is
H(KS) - [( f(z))og f(r)dr
=1 1 b (1 a) (σ!0) (b-1+1)++) (b)0((((((+++111)
loog log b,
where ())) the digamma function and 7 R 0.577 is the Fuler-Mascheroni constant. The digamma function and its gradient, the trigamma function (x), can be represented as infinite scries which converge rapidly and thus can be used effectively in numerical applications. They are included as standard functions in common difteeiiffeentiiation frameworks.
Continuous distributions with bounded interval support. Among distributions with bounded interval support, the KS uniquely satisfies desiderata (i) (iii in Section [ It supports the reparamederization trick through its formm differentiable inverse CDF, providing efficient sampling and variance gradients. The KS supports four distinct prototypical shapes bell, U, increasing, and decreasing (Figure right) - providing expressivity for diverse modeling tasks. Its log-pdf, CDF, and inverse CDF, along with their gradients, composed only of affine transformations, exponentials, and logarithms, and can be
2

---

Published in Transactions on Machine Learning Research (03/2025)
Property / Distribs.
CB
tanhy
Beta
KS
Relative Expressiveness
low explicit
high explicit
high
high
Gradient Reparam.
implicit
explicit
Contains Uniform
X
Closed-form CDF
x
x
Closed- form inverse CDF
x high
x low
Numerical Issues
mild
low
Complex Functions
tanh R /
log(1- tanh
B, I
None
Parameterization
R²
R.
R²
A Analytical Moments
x
/
/
Closed-form KL
Exp. Family
tanhy x
Exp. Family
Beta
Entropy H
Table 1: Comparison of bounded interval-supported distribution families.
parameterizod directly terms of unconstrained logarithmic values; see Scction 3 This enables straightforward implementation with minimal dependencies and keeps most computation in log-space, enhancing stability and accuracy. The racond logarithmic parameterization makes it well-suited for NNs, eliminating the need for positivity-enforeing link functions. Additionally, the KS has differentiable, closed-form expressions for moments, median, differential entropy M(SS), and the Kullback-Leibler (KL) divergence to the Beta distribution, facilitating efficient incorporation of prior information.
We briefly introduce workhorse bounded-interval supported distribution families, namely the the Continuous Bernoulli, the Beta, and the lanb-squashed--Gausee
The Continuous Bernoulli (CB) [Loaiza-Ganee & Cunningham 2019) arises in deep learning for modeling continuous [0, 1]-valued pixel intensities in natural images. It provides normalized probabilistic counterpart to the commonly used binary cross-cntropy loss, with density p(=; A) = (1 - 2)"-=, x E [o, 1], A E (0, 1), where C(A) = (2 if X = = else 2tanh 2tanh-"(1-2M) is the normalizing constant. The Beta distribution is a flexible two- iwo-parameter family, widely used for modcling probabilities and proportions. Its density, parameterized by a, b > 0, is given by: p(z; b) B(a,b) (oa.b)-2---10 -)0-1,a, x E (0, 1), where B(a, b) the Beta function. The saab--sqasheed-CGe -aausiinn (tanhy) maps Gaussian samples through the tanh function to produce outputs in [-1, y y= tanh(=), NN o²). It is widely used in reinforcement learning over continuous bounded action spaces (Haarnoja et al. 2018) due
1-2A
to its support for the reparameterization trick.
Table D comparcs these boundod-intorval supported distribution familics across important propertics for latent variable modeling. Relative eapressiveness measures the variety of prototypical shapes a distribution can represent. All distributions except CB exhibit four prototypical shapes; CB is limited to two. Admittedly, none of these offer support for multimodal distributions over bounded intervals. Appendix H discusses other flexible alternatives (which have their own limitations), c.g., copulas (Tran et al. 2015) normalising flows (Rezende & Mohamed 2015). Contains wniform refers to the ability to represent the uniform distribution, critical for modeling complete uncertainty. All distributions except tanhw can express the uniform. Closed-form CDF indicates whether a closed-form CDF is available, which could be a desirable feature for use in copula-based modeling, simplifying the representation of complex dependencies (Tran et al. 2015). Only CB and KS provide such expressions. Similarly, closed-form inverse CDF indicates the availability of closod-form inverse CDF, with only CB and KS satisfying this criterion. Numerical issues capture challenges in stable evaluation. For example, the CB log-pdf requires Taylor expansion to handle singularities as X + 0.5. The tanhy distribution requires log-pdf clipping and parameter regularization to maintain stability, in various implementations (Haarnoja et al.) 2018). Complez functions highlight reliance on non-affine, non-logarithmic, or non-exponential operations. The tanhy log-pdf involves computing log (1- (x)) which is numerically unstable (Björck et al. 2021). The Beta distribution relies on the Beta function and the regularized incomplete Beta function in its log-pdf and CDF, respectively, both requiring numerical approximations. In contrast, our novel parameterization of the KS distribution (Section 3 avoids complex functions in the log-pdf, CDF, and inverse CDF; note a1 1 is computed via exp(- avoiding division.
appears
3

---

Published in Transactions on Machine Learning Research (03/2025)
The KS differential entropy contains the digamma function 600) and the trigamma function e(), but because there exist rapidly converging polynomial expansions, they can be evaluated with very high accuracy in fast, efficient, and numerically stable manner. Perumeterization examines whether a distribution can be effectively expressed with unconstrained parameters. Both CB (via log A E R) and tanhy (via (H, log Ø) E RR2 support racoedt parameterization. We introduce the first reconsttrained parameterization for KS in Section 3 using (log a, log b) E RR2. The Beta distribution, due to its dependence on the Beta function, resists effective unconstrained parameterization. Closed-forn KL functions refer to analytical KL divergence expressions. The CB and Beta distributions, members of the exponential family, admit closed-form KL expressions with other exponential family members. The KS also has closed-form KL, expressions with Beta family members, while tanhy is restricted to closed-form KL expressions within its own family. Entropy considers the availability of closed-form expressions for differential entropy. This property is present for all distributions except tanhy'
Latent variable modeling with stochastic variational inference (SVI). The primary method for fitting large-scale latent. variable models is SVI (Hoffman et. al. 2013). Consider model po(ar) = po( >))(( ())(z) where E RM is the observation, E R" is vector valued latent variable, pa(x x) is the likelihood function with parameters 0, and p(z) is the prior distribution. Except for a few special cases, maximum likelihood learning in such models is intractable because of the difficulyy of the integrals involved. Variational inference (Jaakkola & Jordan 2000) provides a tractable alternative by introducing a variational posterior distribution (z) and maximizing a lower bound on the marginal log-likelihodk called the ELBO:
C(x, 0, Φ)- = Ex)(() [logro(> (0 x)] - DKL (90(x) (90(z)]lp(z)) [l((z) I log))(x)).
(1)
Training models with modern SVT Kingma & Welling, 2014 Rezende et al., 2014) involves gradient-based optimization of this bound w.r.t. both the model paramcters O and the variational parameters Φ. The first term in 1 encourages the model to assign high likelihood to the data, but its exact evaluation and gradients are typically intractable and so the expectation is often approximated with samples from (z)... The KL divergence term incorporates prior information by penalizing deviations of the variational posterior from the prior p(z). Closed-form expressions of DKL. (go (=) |/ (z)) allow efficient encoding of prior information; otherwise, sample-based approximations are required. In the common setting of i.i.d. data with per-datapoint latent variables, amortized inference introduces a shared NN, parameterized by 'inference parameters' Φ, to map observations to variational paramcters, approximating their individual postcriors as 94 æ)) Modifying the ELBO by scaling the KL term with a parameter fK1. > is often necessary to balance the trade-off between data likelihood and prior regularization (Alemi et al. 2018 We denote the sample-based approximation of this modified ELBO as LOKL.
Gradient reparameterization explicit and implicit. A distribution (0(z) is said to be explicilly reparameterizable, or amenable to the 'reparameterization trick', if it can be expressed as a deterministic, differentiable transformation z = g(e, Φ) of a base distribution e ~ p(e). This base distribution is typically simple, such as Uniform or standard Normal, enabling fast sample generation by first sampling from the base and then applying g. This enables the use of backpropagation to estimate gradients of the form [cf. 1
V= = Epre VV +))] = Epte) Va.ellVV V Vøg(e,
(2)
an expectation with form encompassing the ELBO. Explicit reparameterization is compatible with distributions in the location-scale family (e.g., Gaussian, Laplace, Cauchy), distributions with tractable inverse CDFs (e.g., exponential, KS, CB), or those expressible as deterministic transformations of such distributions (e.g.,
tanhw). When explicit reparametterization not available, implicit reparameterization (Figurnov et al. 2018 is commonly used for distributions with numerically tractable CDFs, such as truncated, mixture, Gamma, Beta, Dirichlet, von Mises distributions. This method expresses the parameter gradient through the sample Vaz as a function only of the CDF gradients, not its inverse. Such CDF gradients are commonly found using numerical methods, e.g., forward mode auto-differentiatin on truncaled iterations to estimate the CDF, as in the Gamma and Beta distributions. Such numerical methods introduce approximation error, and thus potentially higher variance and numerical stability issues in the estimation of (2) (Mohamed et al.] 2020), and tend to introduce significantly more complexity than explicit reparameterization which only require implementing o(c,)).
4

---

Published in Transactions on Machine Learning Research (03/2025)
Relative error in computing log(1 exp(x))
2-5
log(1 exp(x))
loglmexp(x)
2
2-25
--35 35
-149
-225 - -100
log2 |x|
-50
-25
Naive computation of log (1 - exp (x)) (red) becomes unstable - 0 due to catastrophic cancellation, while lxg1mexp(x) (blue) ensures accurate computation. Relative error is the absolute difference between single-precision and 1024-decimal-plar ovaluations, normalized by the 1024-docimal-plac result.
Figure 2:
3 Stabilizing the Kumaraswamy
The KS distribution's utility relies on stable computation of its log-pdf, CDF, inverse CDF, and their gradients. In the standard parameterization, these functions contain instabilities from hidden log(1 - exp(x)) terms. We address this by introducing an unconstrained logarithmic paramcterization that isolates these unstable terms, enabling their straightforward replacement with the stable loglmexp function. Finally, show why naive stabilization techniques, such parameter clipping, fail in high-dimensional applications.
Identifying the instability: log (1 exp (z)). Naive computation of log (1 -xxp (x)) for < 0 leads to significant numerical errors as approaches 0 (Figure 2
red). These errors grow so large that they can cause numerical instabilily, i.e., an irrecoverable error such as -inf. These errors result from catastrophic cancellation, which occurs when subtracting nearly equal numbers here, - exp(x). As x + 0, exp(x) & 1, exp results the cancellation of leading significant bits, leaving only a few less significant, less accurate bits to represent the result. This causes large relative errors in 1 exp which when input to the logarithm its magnitude grows sharply
amplified zero. If the cancellation is complete,
exp (x) underflows to 0 and the logarithm returns -inf, seen in Figure (red) when log2 | < 242.
Numerical building blocks for accurate log (1 exp (z)) computation. When |æ| < 1, both log(1 + æ) and exp(x) - can suffer from severe cancollation: the former between and x, the latter between cxp(x) and In both cases, a simple solution for accurate computation to use a few terms of the Taylor series, as
x²
3
log1p(x) == log(1 + x)
for |a| 1,
2
2
3
expm1 (x) == cxp(x) - 1 x
for | < 1,
2!
31
where n! denotes the factorial. Standard single-procistee logip and expm1 implementations typically use 4 terms in their Taylor series, which we found sulficient for our needs. These functions form the basis for two common methods to compute log exp(x)): 1og(-expm1 (x)) and logip(-exp(x)). Mächler (2012) showed that neither method alone provides sufficient accuracy across the domain, but each is accurate in complementary regions. To address this, Mächler 2012hh introduced
log1mexp(x) = 1ogg(eexp11 (x)) - llog 2 < x <
(3)
1og1p(-exp(x))
log
which computes log (1 - exp (x)) accurately throughout single precision, shown Figure 2 (blue).
A stable Kumaraswamy. The direct implementation of the KS's odg CDF, and inverse CDF
found in all core auto-differentiation libraries produces numerical instabilities. Here, introduce novel
5

---

Published in Transactions on Machine Learning Research (03/2025)
parameterization in terms of unconstrained logarithmic parameter values, which isolates and makes explicit the unstable terms
Wb 1 (u) = log(1 - ub ") = log(1 - exp(b=lloguu) Wa (z) log(1 -xx x²) log(1 exp(alogg x));
eliminates the need for positivity-enforcing link functions, and whose expressions involve only affine, ex- ponential, and logarithmic transformations. This allows the log-pdf and its gradients to be expressed
as
log f() log a + logb (a- 1) log x + (6 (x)
(4)
V al log (x) (a 1) (b 1) exp(a w. (2) + log a)
(5)
V Vloga log [(x + a log x (1 (b 1) exp(a wa (2)))
(6)
V log log f() + 6. (.)))
(7)
Likewise for the CDF
F(x) - (1- - wa (x))
(8)
V (x)) = exp(log log + (a - 1). log x + (b 1). wx(x))
(9)
V logg F(x) exp(log a log b- + a log x + (b (1-1). 1) ww (x)) log
(10)
V log 6 F(x) exp(log b wa(x)) - (- u. (x)),
(11)
and the inverse CDF
F (1) = (1- - jaa= = exx((e 1ooa--uw-100 (b))
(12)
V log a
F (u) explp- log a + a m (-wh-1 (u))
(13)
log 8
(u) exp(- log a log b + b log u + (a - 1) (u)) log
(14)
This parameterization's algebraic form allows direct replacement of the dominant unstable terms, substituting (p))(u) with loglmexp(1- log u) and wa (2) with logimexp(a l) Access to log a and log avoids errors from unnecessary transitions in-and-out of log-space. We also avoid the error prone expressions produced in backpropogation's direct application of the chain rule, c.g.,
1 a
(1 log 1-exp ( log (oogs))))) -11-(((-0pp (I-cxp ( (1-Cx((H logu))
exp ( (ogu) log u
-1 5
Ving 8 F-
exp
b
and (14] are equivalent expressions for Voog but their computed values can differ greatly for extreme parameter values. Desirable KS distributions can obtain such problematic extreme parameter values, e.g.,
the KS distributions in Figure g have b ~ 106, See Section A for further discussion on how instability in the unmodified KS can worsen with increasing evidence.
Both PyTorch and TensorFlow implement the KS distribution similarly and fail to address the loglmexp instabilities we highlight. Here, we use PyTorch as representative example, though our results and mothods apply equally to TensorFlow. Figure 3 compares the PDF, inverse CDF, and histograms of reparameterized samples for KS distributions which are typical to real-world modeling scenarios. The PyTorch implementation (top row) shows jaggedness in both the PDF and inverse CDF, caused by catastrophic cancellation in the unstable terms wa (2) and Wp-1(ul (u). Additionally, the PyTorrc inverse CDF underflows beyond u ~ 1 - 39.3: here, (u) = -0, and [-) exp(a-1 -00) = 0. This underflow results in a point mass at x = 0 (a point outside of the KS support) with probability z 39.3 in each of the reparameterized sampling distributions, and produces infinite gradicnts via F- 0o Jcf. [13 This infinite gradient triggers a cascade: infinite
parameter values after the optimizer step and NaN activations in the next forward pass, which is what users ultimately observe when training fails.

---

Published in Transactions on Machine Learning Research (03/2025)
PDF
Inverse CDF (u)
Reparameterized Samples
PyTorch Ks
39.3% samples underflow
Stable KS
lag
X
X
Stabilizing log(1 - exp(x)) terms climinates numerical instabilitics in the KS log-pdf and inverse CDF. We compare the unstable Py PyTorc KS implementation (lop row) and our stable KS (boltom row) for realistic KS distributions (log, b = 24, varying a). Catastrophic cancellation in the log(1 - exp(x)) terms in the PyTorch KS causes jagged curves and inverse CDF underflow beyond 1 39.3, resulting in a point mass of ze 39.3 I 0 in the sampling distribution. Our stable KS removes the instability by using logimexp.
Figure 3:
3.1 The inadequacy of parameter clipping in large-scale settings
Numerical instability in the KS is inherently stochastic, and in high-dimensional settings, the compounded probability of failure across multiple variables makes program failure almost certain. As the program goes unstable if any single KS goes unstable, the overall instability probability can be modeled as the probability of at least one failure in D independent Bernoulli trials: - p(KS stable) for D KS latent variables. practical large-scale settings, e.g., recommendation systems with 107 users and D = 10 recommendation items, the probability of instability approaches across all reasonable parameter clipping values, rendering clipping ineffective stabilization strategy.
Quantitative illustration. Consider the stochastic instability arising from the term log(1 exp(b-1 log u)), where catastrophic cancellation occurs if 1 exp(b- log u) becomes too small. To avoid logarithric domain errors in single precision, we enforce -6-1 log > 2 2- (Figure 2 We aim to select a 6max to satisfy this constraint: a larger bmax cxpands the variational family allowing improved posterior approximation, but worsens stability. Consider the moderate eutropy KS distribuutions Figure 3 which use 224, Using 6max -224, only u 0.6321 satisfies the stability condition, i.e., p(KS stable) z 0.6321 per sample. With D = 107, the overall probability of instability becomes - 0.6321 10" = 1. Now consider aggressively restricting bmax = 21, as done (Nalisnick et al. 2016). Now и 0.9999 sat isfies the stability condition. Even then, introducing D = 107 variables, we still have the overall probability of instability is 0.999910" Re 1. Thus, even extreme clipping fails to stabilize KS distributions at scale. Further, this analysis considers only a single posterior sample. In practice, training with SVI requires T ~ 10 optimization steps, each requiring posterior samples for gradient estimation. This compounds the instability probability to 1 P(KSS stable)TD, making clipping ineffective in realistic large-scale scenarios.
4
Experiments
Using the well established VAE framework on MNIST and CIFAR-10 datasets, show that the stabilized KS enables reliable training as both a variational posterior [Equs. 12 14h and likelihood function Eqqs. 4 7- We then introduce two new deep variational architectures that leverage bounded interval- supported latent variables: the Variational Bandit Encoder (VBE) for addressing exploration-eexloitaalee trade-offs in contextual multi-armed bandits (Section 4.2), and the Variational Edge Encoder (VEE) for facilitating uncertainty quantification in fink-prrdiction with graph neural networks (Section 4.3 Across the experimental domains, our stable KS is performant and often casier to use than alternative variational distributions supported bounded intervals. For instance, tanhy models require log-pdf clipping for stability,
7

---

Published in Transactions on Machine Learning Research (03/2025)
Table 2: VAE on MNIST and CIFAR-10.
Table
MNIST test digit VAE reconstructions.
Inputs
0
1
2
3
4
5
6
7
8
Prior go(2)z) po
MNIST
CIFAR-10
N-CB
0
1
2
8
4
5
1
7
2
9
LL
Acc.
LL
Acc.
KS-CB 0
/
2
8
4 4
5
6
7
8 3
4
N
CB
1860
97.3 1191
37.9
KS
CB
1855
97.4
1194
41.5
Beta CB
0
2
3
5
6
7
A
((a,))
Beta
CB
1856
97.5
1189
40.3
A-Beta 0
8
7
2
Mo.
N
Beta
4209
92.1
3648
48.5
KS-Beta O
2
7
U(0
KS
Beta
4160
91.3
3579
50.1
U(,
Beta
Beta
4198
90.1
N/A
N/A
Beta-Beta
O
2.2.
8
7
Mo.1
N
KS
3393
96.4
1875
47.1
W-KS 0 1
8 4
5
7
3
9
U
KS
KS
3401
96.8
1939
48.8
U(0.1)
Beta
KS
3405
97.1
N/A
N/A
KS-KS
0
/
2
3
15670
7
3
9
Beta-KS
0
/
2
8
6
7
8
9
while Beta models show significant performance variability based on the chosen positivity-enforcingg link function and often fail to converge, e.g., on CIFAR-10 in Section 4.1 Finally, our new variational models are fast: the VBEs in Section 4.2 are 22x faster than the sfate-of-the-art bascline.
Remark. Across all three experimental settings, models using the wnstable KS produce NaN errors in training and are therefore excluded. Prior work using the KS in low-dinensional latent variable models (Nalisnick eet al. 2016 Nalisnick & Smyth 2017 Stirn et al.) 2019) similarly find NaN errors, and rely parameter clipping to avoid instability. See Section 3.1 for why this is approach does not work in large-scale settings. Our stabilization approach directly resolves these numerical issues, enabling stable training at scale.
4.1 Image variational
The VAE (Kingma & Welling 2014) is generative latent variable model trained using amortized variational inference. Both the variational posterior q (2(2))) and the conditional likelihood po(z x)) are parameterized using NNs, known as the encoder ex( R + RRP and decoder do(z): RD - RM, respectively. VA typically use the standard Normal distribution as the prior and a factorized Normal as the variational posterior. The use of alternative variational distributions allows incorporating different prior assumptions about the latent factors of the data, such as bounded support or periodicity (igurnov et al. 2018).
Experimental setup and metrics. Inspired by (Loaiza- Ganem & Cunningham 2019), we train VA with fully factorized priors and variational posteriors MNIST and CIFAR-10 without pixel binarization, using unmodified ELBO (3KL= 1). We adopt the most effective likelihoods from their work (Beta and CB), identical latent dimension D (MNIST: D = 20, CIFAR-10: D 50), and the same standard NN architectures, which are detailed in Appendix B along with the training hyperparameters. For each variational posterior factor, we choose the canonical prior: Nio.1) for N, and 10,,) for KS and Beta. We evaluate the models with an importance weighted estimator (Burda et al. 2016) of the marginal Log Likelihood (L.I.) using 200 samples. To assess usefulness of the Icarned latent representations, we encode test data In, compute the mean of q0 (2, w.); and classify the test labels using a 15-nearest neighbor classifier; the classifier accuracy (%) is denoted "Acc." For subjective evaluation, display the mean decoded likelihood of single sample from the encoded posterior of random test digits in Figure 3
Discussion of results. Rather than introducing a more performant VAE architecture, the sole purpose of this experiment is to provide evidence toward the stabilization of the KS. Notably, stable KS VAEs maintain numerical stability while all VAEs with the unstable KS produce unstable training. VAEs with Beta-distributed variational posteriors often do not converge; indeed, (Figurnov et al. 2018) reported strong performance on binarized MNIST using softplus link function, but did not present. results on CIFAR-10,
8

---

Published in Transactions on Machine Learning Research (03/2025)
nor could we find other works that did. We suspect this is due to similar instability issues, with the higher gradient variance of the Beta's implicit reparameterization a likely explanation. In an attempt to overcome this instability in Beta VAEs we report the best metries across softplus or exp link functions in Table 2 When neither converges, we report N/A. The results in Table 2 show that across datasets, VAEs with KS-distributed variational posteriors consistently produce useful latent spaces, evidenced by strong latent nearest neighbor classifier accuracy, and yield reconstructions with competitive LLs and visual quality.
When paired with any variational posterior, a KS likelihood yields stronger MNIST reconstructions than Beta likelihoods: compare rows +-Beta to =S in Table 3 As in (Loaiza- Ganem & Cunningham 2019). we find CB likelihoods produce the most subjectively performant VAEs on MNIST, unsurprising as CB was introduced specifically for the approxinately binary MNIST pixel data.
4.2 Contextual Bernoulli multi-armed bandits
The contextual Bernoulli multi-armed bandit (MAB) problem involves a decision maker who, at each time step 1= 1, ., T, selects one arm from a finite set of K options. Fach arm has an associated context Zy E IRd, and pulling an arm yields a binary reward F3 ~ Bernoull((ux), where μ E [0, 1] is the unknown mean reward. MABs originate by analogy to casino slot machines, where each machine (arm) has a different payout rate, and the challenge lies in deciding which arms to pull in order to maximize total winnings while learning about their payout rates, a situation called the exploration-oxploitatai
dilemma. MABs have found applications in modern recommendation systems (Li et al.. 2010), clinical trials design (Villar et al. 2015), and mobile health (Tewari & Murphy, 2017). Thompson Sampling (TS) is a simple, empirically effective (Chapelle & Li. 2011). and scalable (Jun et al.] 2017)) arm selection heuristic. Tt selects the arm corresponding to the highest value drawn from the posterior distributions over the latent 2 's. This approach naturally balances exploration and exploitation: the uncertainty in the posteriors promole exploration, while concentration of probability mass on large mean rewards drive exploitation. See Section 5 for further discussion on TS-based Bernoulli MAB approaches.
Variational Bandit. Encoder (VBE). Consider K arms, each with context an E Rd and an associated latent variable 2 E [0, 1] representing its mean reward. Define X= = æx]. Let Ty be the set of time indicos at which arm & is chosen, resulting in binary rewards fru,cheTT- The joint distribution factorizes as
ro X) II (2(2 xx) II 1ET
Bernoullirra =x)]
(15)
where po(# mx)) is a conditional prior on the latent mean reward zu; see Appendix F for strategies to incorporate prior knowledge into this distribution. In the following, we take po(= æx) to be uniform, and so 0 = 0. Each reward VR,a E 1J is drawn from Bernoull((ux), so if an arm is never pulled (Tx 0), it simply yields no observed rewards.
With regards to variational approximations and lcarning, VBE's posit a fully factorized KS variational posterior II (=(=) mk). Similar to VAFs, we employ amortized inference using a shared NN encoder eg (x(x), ), which defines a reparameterizable variational distribution gø(2x xx). However, unlike VAEs, VBEs omit the decoder; samples εs (p(2) xx) directly parameterize the likelihood function. The arm selection at step [² We then draw reward r ~ Bernoulli(p.) and record it in the replay buller D + DU {( r)). We construct a sample approximation of the modified ELBO over the subset of arms KG. C (l..... K} that have been pulled by time t as
follows TS: a = argmax, {
(
Σ (E.D
logp(r 20) + 3KL
Σ Hlaa(z xx)], keK,
(16)
see Appendix D for the derivation. The second term promotes exploration by penalizing overconfidence with the exploration effect proportional to fKI. We maximize (66 w.r.l. & via gradient ascent, enabled by the reparameterizable KS. VBE execution is summarized in Algorithm
VBE advantages. VBEs provide four primary advantages over alternative TS-based Bernoulli MAB approaches, discussed in Section 5
9

---

Published in Transactions on Machine Learning Research (03/2025)
Algorithm Variational Bandit Encoder
10a
Regret
Require: { {u 9KL
1: Variation posterior q < KS
Replay buffer +-
for t 1 do
Reward
Encode: (aR, bx) esp (2tx) Sample: & ~ q(zki aR by)
10
6: TS: a argmax, (
7: Reward: r r Bernoulli(ua)
DU(((E 8: + -
((x,: D a,r
0
500
Step
1500
2000
Construct Lika as 16
Time (s)
153
10: Φ + Φ + + &L.K
11: end for
VBE-KS
VBE-Beta
VBE-tanhy
LMC-TS
Figure 4: Synthetic bandit performance over runs. VBE- KS best handles exploration-exploiiatiion trade-offs.
Scalability and Compotibility. VBE training consists of forward pass through NN, sampling explicitly reparameterized distribution, and a backward pass for gradient-based updates. This process is scalable and fully compatible with existing gradient-based infrastructure.
Prior Knowledge Incorporation. When prior knowledge exists on an arm we can encode it as po(# g = Beta(ag,b)) replacing H[go(z mx)) with -Dxt. (nø(2x ax)) po(#* mx));; see Appendix
Interpretability and Independence. Encoding æg with es produces KS distribution parameters, fully encapsulating the model's beliefs about μ4k- This is independent of other arms and past data.
Simplicity. VBEs lack numerous hyperparameters and complex architectural components.
A Alternative methods lack some or these properties because they do not directly model the mean rewards nor differentiate through reparameterized samples from variational mean reward posterior; instead, they tend to model parameters of context-to-mean reward function.
Experimental setup. We construct synthetic data with == 10 arms by first sampling weight vector w and features {= from (0,, We then compute [wT J and apply min-max normalization to produce probabilities (referred as "Original probabilities"
Original probabilities
Power (5) transformed probabilities
Figure 5) To introduce non- linearity, we raise these probabilities to the power 5 (shown
"Power (5) transformed probabilities' in Figure 5] Expo- nentiating the probabilities not only makes the mapping from features to mean rewards more challenging to learn, but it also significantly reduces the number of arms with high probabil- ities, forcing the agent to explore more. For instance, when raising the probabilities bilities the power of 5, the number of arms with large probabilitics drops from 167 just We consider T 2 3 steps.
Number of arms with reward probability 0.8: 1
1.0
Arm Reward Probability
We evaluate VBEs with either KS (VBE-KS), Beta (VBE- Beta), or tanhw (VBE-tanhy) all using 8KL = |K which makes the second term in (16) a mean. Gaussian variational posteriors are not used as their support K does not match the posterior support 1). VBE-tanhy's performance
Figure 5: High arm reward probabilities are reduced via a power 5 exponentiation, encouraging exploration.
sensitive to the number of samples used in its entropy estimate: we found degraded performance beyond 10 samples. The learning rate is set to 7 10-2 As baseline, we LMC-TS,, which employs Langevin Monte
10

---

Published in Transactions on Machine Learning Research (03/2025)
Carlo (LMC) to sample posterior parameters of a NN, known for state-of-the-art performance across various tasks [uu et al. 2022 All models use an MLP with 3 hidden layers of width 32. LMC-TS hyperparameters (inverse temperature, LMC steps, weight decay) are set or tuned based on the authors' code, We repeat experiments times on an Apple M2 CPU and report the mean and standard deviation across these runs in Figure 4
Metrics and evaluation. The optimal policy always selects the arm with the highest mean reward r*. Our objective is to minimize regret, defined the cumulative difference between the expected reward of the chosen action and the optimal action (accessible in the synthetic setting), i.e. Ei=(r" ra,). VBE-KS achieves lower regret and higher cumulative reward than all basclines. VBE-Beta performs significantly worse than VBE-KS and VBE-tanhy, highlighting the importance of explicit reparameterization. LMC-TS is performant but 8-22x slower than VBEs: VBEs avoid the computational overhead of LMC.
4.3 Variational link prediction with Graph Neural Networks
Graph Neural Networks (GNNs) have become a powerful tool for learning from graph-sttructrred data, with applications in critical areas like drug discovery (Zhang et al. 2022 and finance (Wang et al. 2022). A key task is link prediction, where the goal is to infer unobserved edges between nodes. However, real-world deployment of graph learning models is often hindered by lack of reliable uncertainty estimates and limited capacity to incorporate prior knowledge (Wasserman & Matcos 2024). To address these challenges, we propose variational approach where the GNN encodes a. fully factorized (across the edge set) KS variational posterior to model the unobserved probabilities of each network link's existence, enabling uncertainty quantification and prior knowledge integration with minimal computational overhead.
In a typical link prediction setup, the GNN has access to the features X E Vxd of all nodes, but only a subset of positive edges in the training D,, and validation D sets. Edge (u, takes value 4,,0 when present, or Lu, = when absent. The GNN generates edge embeddings through message passing and neighborhood aggregation, outputting probabilities Zu,u (0, 1) that parameterize Bernoulli likelihood. The seminal work of (Kipf & Welling, 2016) proposed Variational Graph Auto-encoders (VGAEs), which posits a Gaussian variational posterior over the final node embeddings. When used for link prediction it samples final node cmboddings from the variational postcrior and decodes them to produce edge probabilities. In contrast, our approach directly models the probability of an edge using the KS. An advantage of directly modeling edge probabilities is interpretability; deep nodal embeddings are often difficult to interpret, and priors are typically selected for computational tractability rather than their ability to incorporate meaningful prior information. However, the probability of edge (u, v) existing between two nodes is an interpretable quantity that can often be informed by domain expertise, and incorporated into the conditional edge prior po(, | x. D.r). For example, in gene regulatory networks, epidemiological networks, and social networks experts often have prior knowledge about the likelihood of specific interactions, transmissions, or friendships, respectively. We believe the limited exploration of variational modeling for edge probabilities is due to the previous lack of an expressive, stable, explicitly reparameterizable bounded-interval distributions.
Variational Edge Encoder (VEE). Let &ar be the edges included in the training set, which consist of positive edges Der and set of sampled negative edges D The joint distribution factorizes as
po x) II (o æx Eir) Bernoulli(I... (u,,)) (a.
(17)
Here, we take po(2u. x; E1r) to be uniform, and so = 0. Variational approzimations and learning. We propose fully factorized KS variational posterior II(u.w)E.. 90 X,,,). The GNN encoder ep parameterizes a KS distribution for each possible edge (u, E Exr. The remaining structure is highly similar to VBEs: a single sample Eu,r ~ qp(zu, X, Eir) directly paramotorizes the Bernoulli likelihood, and we maximize a sample approximation of the modified ELBO
L (Xx
Σ log Eu,w) BKG.. Σ Hlgss X ExE+))]- (u,D)EE
(18)
(a,,)EE.r
From their similarity with VBEs, VEEs inherit the same advantages outlined in Section 4.2
11

---

Published in Transactions on Machine Learning Research (03/2025)
Cora
Citeseer
Pubmed
VEE-KS
VEE-Beta
VEE-tanhy
Base GNN
AUC
P
Brier Score
AUC
Brier Score
AUC
p
Brier Score
Figure 6: VEE-KS produces high predictive accuracy (1 AUC), informative uncertainty estimates (f p), and improved calibration (l Brier Score) across graph datasets.
Models, metrics, and datasets. All models use 2-layer GNN with Graph Convolutional Network (GCN) layers and hidden/outpun nodal dimension of 32. In Base-GNN an MLP decodes the final nodal embeddings into link probabilities. In TIEE-KSBeeta/tanee variational distributions; all take 3KI. _05/Drr1-1. We use 10 samples in tanhy's entropy estimate; more did not produce significant performance differences. We train for 300 epochs, with a learning rate of ,01, averaging results over 5 runs with different seeds. The posterior predictive distribution over binary links p(A Er) = I p(A p(A(=))0(( distribution, parameterizing each edge Bernoulli distribution with such samples, followed by sampling binary edges. For Base-Gii GNNN we directly sample binary edges from the likelihood. Using 30 posterior predictive samples, compute the edge-wise posterior predictive mean (pred. mean) and standard deviation (pred. stdv.). We measure predictive accuracy using the area under the ROC curve (AUC) (higher is better), treating pred. mcan the predictor. To assess uncertainty estimation quality, we report the Pearson correlation (higher is better) between the predictive uncertainty (pred. stdv.) and error (the é difference between pred. mean and the true label). A strong positive correlation indicates that higher uncertainty corresponds to greater predictive error, signifying informative uncertainty estimates. Finally, we evaluate model calibration using the Brier Score (lower is better), defined as the squared difference between the predicted probability and the true label. Figure G presents results across three standard citation networks: Cora, Citeseer, and Pubmed.
MLP parameterizes the KS/Beta, JS//eta/tanhy,
X, z is estimated by using a single sample from each
Discussion of results. On all datasets and all metries, VEE-KS outperforms or is comparable to the most performant baselines, often providing higher predictive accuracy, better uncertainty estimation, and improved calibration. Similar to Section 4.2 we find Beta distributed variational posteriors perform significantly worse than those using KS or tanhy,, further underlining the importance of explicit roparamcterization. Moreover, models using explicitly reparameterizable latents are faster: on the largest dataset (Pubmed), the average time (ms) per epoch for VEE-KS, VEE-tanhy, and VEE-Beta was 381 + 61, 301 + 26, and 447 + 86 respectively, an Apple M2 CPU.
5 Related Work
TS-based Bernoulli MAB approaches. Existing basea approaches for Bernoulli MABs assume prior over model parameters p(Ø), which map contexts to rewards through cø- At each round, parameters are sampled from the posterior, o, p(c D), and used to compute mean reward posterior samples {eg, K= However, the Bernoulli likelihood often leads to intractable posteriors, making parameter sampling difficult. Common methods use either variational approximations (Chapelle & 2011 Urteaga & Wiggins 2018: Clavier et al.) 2024), primarily Laplace, or MCMC approaches like Gibbs sampling (Dumitrascu et al. 2018) or LMC [uu al. 2022). These approaches face several limitations. First, incorporating prior knowledge is challenging since the relationship between a parameter's value and its effect on rewards is often unclear, except in the simplest models. Second, scalability is a concern: Taplace
VBEs in context:
12

---

Published in Transactions on Machine Learning Research (03/2025)
approximations become inefficient with large context dimensions or model sizes, while MCMC-based methods are compute and memory intensive, requiring long burn-in periods (typically 102 iterations) and large machine memory to store the buffer D. Third, interpreting model beliefs over mean rewards requires drawing numerous posterior samples, adding further computational cost. Finally, these methods often introduce significant complexity through intricate algorithms, architectures, optimization steps, and hyperparameters, particularly MCMC parameters (e.g., burn-in iterations, chain length, LMC inverse temperature/ /weigt decay and their respective schedules). By directly modeling mean rewards with KS, instead of the parameters Φ, VBEs offer a simple, scalable, and interpretable approach to Bernoulli MABs.
Kumaraswany as a Beta surrogate. A simple approach to overcome the Beta distribution's lack of explicit reparameterization is to use the KS as a surrogate. This surrogate approach is feasible due to their significant similarities when defined by the same two parameters and the availability of a high-fidelity closed-form approximation of the KI. divergence between Beta and KS distributions. (Nalisnick et al. 2016 Nalisnick & Smyth, 2017) use KSs as surrogates for Betas in the Dirichlet Process stick-breaking construction to allow for stochastic latent. dimensionality in a VAE. However, both require parameter clipping for numerical
stability. In their published code (Nalisnick et al., 2016) constrains KS parameters log a, log 6 E [-2.3, 2.9],
significantly limiting the expressiveness of latent KS distributions. Also, (Nalisnick & Smyth (017)) comments
under Computational Issues section that 'If NaNs are encountered...clppinn the parameters of the variational
Kumaraswanys usually solve the problem.' (Stirn et al. 2019) improved upon (Nalismick et al. 2016 by
resolving the order-dependence issue in approximating a Beta with a KS. Similarly, (Singh et al. 2017) followed
a comparable process using an Indian Buffet Process. Both works maintained numerical stability by restricting the uniform base distribution's 's support from the unit interval to a narrower interval, before passing the
samples through the inverse CDF producing distortion of the reparameterized sampling distribution. This work eliminates the need for such distortions, enabling more accurale Beta approximations and simplifying the use of the KS distribution by ensuring numerical stability without additional interventions.
6 Conclusion, Limitations, and Future Work
We identified and resolved key numerical instabilities in the KS distribution, a uniquely attractive option in variational models for bounded latent variables. Our work demonstrates that the stabilized KS can tackle a wide range of large-scale machine learning challenges by powering simple deep variational models. We introduce the Variational Bandit Encoder, which enhances exploration-cxplotttiinn trade-offs in contextual Bernoulli MABs, and the Variational Edge Encoder, which improves uncertainty quantification in link prediction using GNNs. Our empirical results show these models are both performant and fast, achieving their best performance with the KS while avoiding the instability and complexity seen in alternatives like the Beta or tanhy distributions. These models open avenues for addressing other large-scale challenges, including in recommendation systems, reinforcement learning with continuous bounded action spaces, network analysis, and uncertainty quantification in deep learning, such as modeling per-parameter dropout probabilities using a Concrete distribution (Gal et al. 2017).
KS generalizations (Usman & ul Haq 2020) inherit - log(1-expzz) instabilities, which future work can resolve by building on our stabilization technique. A limitation of the current models is their inability to capture multimodal posteriors. Future work could explore KS mixtures or hierarchical latent spaces to bridge this gap. Further, optimizing the BKI. parameter with techniques like warm-up up schedules could yield further performance gains (Alemi et al.) 2018).. Applications of our stable KS distribution to non- non-parametric models like the Dirichlet Processes follows directly from prior work (Nalisnick & Smyth 2017 Stirn et al.] 2019) Lastly, a theoretical analysis of the VBE, particularly in proving regret bounds and asymplotic results, could extend its applicability to critical areas like clinical trials, where robust pocision-making under uncertainty is essential.
References
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing a broken ELBOO In ICML, 2018.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Johan Björck, Xiangyu Chen, Christopher De Sa, Carla P Gomes, and Kilian Weinberger. Low-precision reinforcement learning: running soft actor-critic in half precision. In ICML, 2021.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. ICLR, 2016.
Olivier Chapelle and Lihong Li. An empirical evaluation of Thompson sampling. Tn NewrIPS, 2011.
Pierre Clavier, Tom Huix, and Alain Durmus. VTTS: Variational inference Thomson sampling for contextual bandits. In ICML, 2024.
Bianca Dumitrascu, Karen Feng, and Barbara Engelhardt. PG-TS: Improved Thompson sampling for logistic contextual bandits. In NeurIPS, 2018.
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. NeurIPS, 2018.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. NeurIPS, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML, 2018.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. J. Mach. Learn. Res., 2013.
Tommi S Jaakkola and Michael Jordan. Bayesian parameter estimation via variational methods. Stat. Comput., 2000.
M.C. Jones. Kumaraswany's distribution: beta-type distribution with some tractability advantages, Stat. Methodol,, 2009.
Kwang-Sung Jun, Aniruddha Bhargava, Robert Nowak, and Rebecca Willett. Scalable generalized linear bandits: Online computation and hashing. In NewrIPS, 2017.
Diederik P Kingma and Max Welling, Auto-encoding variational Bayes. In ICLR, 2014.
Thomas Kipf and Max Welling. Variational graph auto-encoders. In NewrIPS Worksh. on Bayes. Deep Learn., 2016.
Ponnambalam Kumaraswany. A generalized probability density function for double-bounded random processes. J. Hydrol., 1980.
Lihong Li, Wei Chu, John Langford, and Robert Schapire. A contextual-bandit approach to personalized news article recommendation. In WWW, 2010.
Gabriel Loaiza-Ganom and John Cunningham. The continuous Bernoulli: Fixing a pervasive error variational autoencoders. In NewrIPS, 2019.
Martin Mächler. Accurately computing log(1 exp(-lal)) assessed by the Rmpfr package. CRAN, 2012.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. JMLR, 2020.
Eric Nalisnick and Padhraic Smyth. Stick-breaking variational autoencoders. In ICLR, 2017.
Eric Nalisnick, Lars Hertel, and Padhraic Smyth. Approximate inference for deep latent Gaussian mixtures.
In NewrIPS Worksh. on Bayes Deep Learn., 2016.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. Tn JCML, 2014.
14

---

Published in Transactions on Machine Learning Research (03/2025)
Rachit Singh, Jeffrey Ling, and Finale Doshi-Veler Structured variational autoencoders for the Reta-Bernonlii process. In NewrIPS Workshop Adv. Approz. Bayes. Infer., 2017.
Andrew Stirm, Tony Jebara, and David Knowles. new distribution on the simplex with auto-encoding applications. In NewrIPS, 2019.
Ambuj Tewari and Susan A Murphy. From ads to interventions: Contextual bandits in mobile health. Mob. Health: Sens, Analyt. Methods Appl. 2017.
Dustin Tran, David Blei, and Edo M Airoldi. Copula variational inference. NewrIPS, 2015.
Tnigo Urteaga and Chris Wiggins. Variational inference for the multi-armed contextual bandit. AISTATS, 2018.
Rana Muhammad Usman and Muhammad Ahsan ul Haq. The marshall-olkin extended inverted kumaraswamy distribution: Theory and applications. J. King Saud Univ. Sci., 2020.
Sofia S Villar, Jack Bowden, and James Wason. Multi-armed bandit models for the optimal design of clinical trials: benefits and challenges. Stat. Sci., 2015.
Jianian Wang, Sheng Zhang, Yanghua Xiao, and Rui Song. A review on graph neural network methods in financial applications. J. Data Sci., 2022.
Max Wasserman and Gonzalo Mateos. Graph structure learning with interpretable Bayesian neural networks.
Trans. Mach. Learn. Res., 2024.
Pan Xu, Hongkai Zheng, Eric V Mazumdar, Kamyar Azizzadenesheli, and Animashree Anandkumar. Langevin Monte Carlo for contextual bandits. In JCML, 2022.
Zehong Zhang, Lifan Chen, Feisheng Zhong, Dingyan Wang, Jiaxin Jiang, Sulin Zhang, Hualiang Jiang, Mingyue Zheng, and Xutong Li. Graph neural network approaches for interactions. Curr. Opin. Struct. Biol., 2022.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Differential Entropy (log)
Mean
30
30
Decr.
Bell
Decr.
Bell
25
1.0
0.8
15
15
0.6
109
5
0.2
Incr.
Incr.
-10
-10
100
-10
10 log
log
Figure 7: Differential entropy and mean of Kumaraswany distributions across a wide range of parameter values. Low-entropy distributions are concentrated near the origin, where log a = log b = corresponds to the uniform distribution. Distributions with a log differential entropy of and mean 0.5 are marked in black. Small changes in parameter magnitudes rapidly yield extremely high-entropy distributions essentially delta functions 0 or except in the narrow region around the black curve representing distributions with mean 0.5.
A Counter Intuitive Stability Properties of the Unstable Kumaraswamy
When using the unstable KS to model latent variables with SVI, the instability of the KS distribution can poradorically worsen as evidence increases. Here, evidence refers to observed data that sharpens the posterior distribution and reduces uncertainty. Representing sharper, high-entropy bell-shaped KS distributions indicative of reduced posterior uncertainty requires extremely large values. Figure illustrates this: bell-shaped KS distribution with mean 0.5 and differential entropy H Rs exp(2) necessitates logb = 24, and thus 6 e24( in the unstable KS implementation which lacks logarithmic parameterization; see Figure 3 for examples of such moderate entropy distributions with log 24. SVI will leverage the inverse CDF and its gradient expressions 12 14h, which critically depend on b through the term (b) = log(1 - exp(b-1 log u)). Large b values will act to worsen instability by driving exp(b-1 log u) closer to 1, increasing the risk catastrophic cancellation. We believe this counter-intuitive behavior likely frustrated modelers, but is no longer an issue in the stabilized KS.
As illustrative example, consider modeling the latent probability of heads in Bernoulli coin-flipping experiment using a KS distribution the variational posterior, where the true probability of heads is 0.5. With a uniform prior and a small number of observed flips, the posterior is well-approximated by a mild-entropy. bell-shaped KS distribution, characterized by low-magnitude parameters a, > 1. In this regime, 5-1 remains sufficiently far from zero, minimizing the risk of catastrophic cancellation in the term exppp-pppoopppoeee log u); as exp(b-1 log u) slays safely away from However, as the number of observed flips increases, the posterior sharpens to reflect reduced uncertainty, demanding larger values of represent the corresponding higher entropy KS distribution. This drives exp(b log u) closer to 1, increasing the risk of catastrophic cancellation and numerical instability.
VAE Architectural and Training Choices
The following almost identical to that used (Loaiza-Ganem Cunningham 2019), but provided here for completeness. For both experiments (MNIST and CIFAR-10) we use learning rate of 0.001, batch size of 500, and optimize with Adam for 200 epochs.
16

---

Published in Transactions on Machine Learning Research (03/2025)
Table VAE on MNIST and CIFAR-10.
Prior (=(=))) pa (2)
MNIST
CIFAR-10
LL
Acc.
LL
Acc.
No.
N
CB
1860 + 15
97.3
1191 907
37.9
(0
KS
CB
1855 + 14
97.4
+
41.5
Beta
CB
1856 + 15
97.5
1189
914
40.3
No.1
N
Beta
4209 = 323
92.1
3648
+ 1210
48.5
(,
KS
Beta
4160 331
91.3
3579
1154
50.1
U (,,1)
Beta
Beta
4198 126
90.1
N/A
N/A
Mo.1
N
KS
3393 + 44
96.4
1875 + 890
47.1
((o
KS
KS
3401
32
96.8
1939
889
48.8
((1,1)
Beta
KS
3405
40
97.1
N/A
N/A
Enforcing positive variational parameters.
Gaussian. When the variational posterior is Normal, the output layer of the encoder uses a softplus nonlincarity for the positive standard deviation.
KS. As we parameterize the KS by unconstrained log values, any required exponentiation occurs internally, so we require no nonlinearity on the output of the encoder.
Beta. The core software libraries do not implement the Beta distribution's reparameterized sampling with unconstrained log parameter values, so we use an exponential nonlincarity on the output of the encoder to enforce positivity. softplus nonlinearity was attempted which was found to less stable likely due to the model seeing very large latent parameter values, which is more stably accessible via an exp.
Enforcing positive likelihood parameters.
CB. When the likelihood is CB, the output of the decoder has sigmoid non-linearity enforce its parameter X E 1).
KS. As parameterize the KS by unconstrained log values, any required exponentiation occurs internally, require no further transformation the output of the decoder.
Beta. The core software librarics do not implement the Beta distribution's log-pdf with unconstrained log parameter values, so use softplus nonlinearity on the output of the decoder enforce positivity. An exponential nonlinearity attempted which was found to be less stable.
Data augmentation for (0, 1) likelihood functions. The CB has support [0, 1] and handles data on the support boundaries without issue. When the likelihood function is a Beta or KS, which have support (0, 1), we clamp pixel intensities to 12x255 2x255 to prevent non-finite gradient values.
For all our MNIST experiments we use a latent dimension D = 20, an encoder with two hidden layers with 500 units each, with leaky-ReLU non-linearities, followed by dropout layer (with parameter 0.9). The decoder also has two hidden layers with 500 units, leaky-ReLU non-linearities and dropout. For all use latent dimension of D = 40, an encoder with four convolutional layers, followed by two fully connected ones. The convolutions have respectively, 32, 32 and 32 features, kernel size 2, 2, 3 and 3, strides 1, 2, 1, and are followed by lcaky-ReLU non-lincaritios. The fully connected hidden layer has 128 units and a leaLy-ReL.U non linearity. The decoder has an analogous "reversed" architecture.
our CIFAR-10 experiments
Expanded experimental results.
Table 4 includes data from the image VAE experiments in Table from Section 4.1 but now including standard deviations of L.L values across test. samples.
17

---

Published in Transactions on Machine Learning Research (03/2025)
C
Kumaraswamy-Beta KL Divergence
The KL divergence between the Kumaraswamy distribution q(v) with paramcters b and the Beta distribution p(v) with parameters a; 3 is given by
Bg log ((v) p(v)
a
-7 (b) 5) + log ab log B(a,, 3) b- B
a
b
(3 - 1)b m + ab
where is Euler's constant, (-(-) is the Digamma function, and (..) is the Beta function. The infinite sum in the KL divergence arises from the Taylor expansion required to represent [oy[1 -a ]]; it is generally well approximated by the first few terms.
Kumaraswamy Moments
The Kumaraswamy distribution's n-th moment is expressed as:
bT + n/a)T(b)
mh T(1 + b bbb//// n/a)
6B (B(1 (1/aa
Here, B() represents the Beta function, and r(-) denotes the Gamma function. Using these raw moments, one can compute the variance, skewness, and excess kurtosis. As an example, the variance is given by 0² m2 - m m
E VBE Modified ELBO Derivation
Let X = [....... ØK] be matrix where the k-th column corresponds to the context feature æg. Assuming independence between arms and within-arm rewards, the data likelihood can be factorized as p(D 2)
I. p(r za). We adopt a fully factorized variational posterior of the form g X) = -I ( mg). We set the conditional prior uniform distribution: po(2 X) (0,1))- Recall that K., C (1,..., K} represents the subset of arms that have been pulled, and thus for which we have reward data. The modified ELBO is derived as follows:
K
t (D, = (lo =)] - 3Kr. DKL. (1 X) po(z X))
| (()
x)] + 8KL (90 X)), pe X)) U.
(by design)
H (90(za wa)! [ (( =)] + 3KG. Σ
kek.
(x((a)
Σ
log p(r 2a)
+ BKL Σ H (za xw)] keK.
logp(r (aa) 3KI.
Σ kEK.
H (z (a)), go(Ea x.)
where in the final step, we use a single sample approximation of the expectation.
F VBE Prior Knowledge Incorporation
A key advantage of the VBE framework is its capacity to integrate prior knowledge about the latent variable z directly into the model. Traditional approaches often impose priors on the weights or parameters of neural networks or logistic regression models parameters that are difficult interpret and for which prior information is scarce. In contrast, VBFs explieitly model the latent mean rewards, making it natural to
18

---

Published in Transactions on Machine Learning Research (03/2025)
incorporate prior information where it exists. In many multi-armed bandit applications, such personalized vecommendation systems or content curation, practitioners frequently have prior information on the probability of user engagement (e.g., the likelihood of click or purchase). Instead of embedding this knowledge indirectly in network weights, VBEs allow one to encode directly through the prior pe(2x x&). For example, suppose historical data or domain expertise indicates that a particular arm generally yields low rewards. In such a case, one can specify a Beta prior, [= m Beta (ax; b&), with parameters chosen to concentrate the probability mass on lower values. This approach is particularly useful for large-scale commercial platforms like Amazon.com or Walmart. Users can be segmented by demographics such as age, gender, and geographic region, while items can be segmented into categories such as baby care products (e-g., diapers, wipes), power tools (e.g,, circular saws, drills, screw guns), and gardening supplies. The Cartesian product of these user and item segments yields a set of user item segments. We may not hold strong prior beliefs over the majority of the user-item segments, and so assigning the user-item segment a a uniform prior is appropriate. Conversely, for segments where robust prior information cxists for example, females aged 35-45 in New York City purchasing baby care products
Beta prior with parameters biased loward higher probabilities is justified.
When ample historical data is available, the prior can be learned by including 0 in the KL divergence term of the ELBO, rather than the uniform distribution used in Scction 4.2 By parameterizing po x4a) using neural network or conditional normalizing flow for enhanced expressiveness, gradients can effectively propagate to 0.
G Additional Results for Variational Link Prediction with GNNs
For clarity, we present Table 5 which provides the exact numorical values corresponding to the results shown in Figure
Table 5: Numerical results corresponding to Figure G The values represent mcan + standard deviation.
Higher values are better for AUC and ps while lower values are better for Brier Score. The best value for each dataset-metric pair is highlighted in bold.
Dataset
Model
AUC f
Brier Score 4
P #t
VEE-KS
0.897 + 0.015
0.082 + 0.010
0.440 0.045
VEE-Beta
0.857 + 0.023
0.117 = 0.014
0.340 + 0.021
Cora
VEE-tanhy
0.853 + 0.066
0.115 0.045
0.440 = 0.043
Base-GNN
0.870 + 0.012
0.115 + 0.045
0.302 + 0.047
VEE-KS
0.851 + 0.010
0.104 + 0.008
0.304 0.035
VEE-Beta
0.789 + 0.021
0.140 = 0.018
0.124 = 0.045
Citeseer
VE.E-tany
0.850 + 0.015
0.111 0.004
0.288 + 0.053
Base-GNN
0.815 + 0.006
0.110 + 0.004
0.132 = 0.025
VEE-KS
0.970 + 0.004
0.061 + 0.004
0.654 + 0.015
VEE-Beta
0.966 + 0.002
0.069 0.002
0.590 0.014
Pubmed
VEE-tanhy
0.935 + 0.030
0.120 0.042
0.623 0.015
Dasc-GNN
0.970 + 0.003
0.121 0.042
0.590 0.019
H Trade-offs in Alternative Variational Approximations
The KS distribution is a compelling choice for modeling bounded latent variables, as it salisfies the key desiderata outlined in Section [ including support for the reparameterization trick, efficient sampling, expressiveness, and simple distribatiion-rclated functions. However, when a KS distribution is insufficiently flexible to approximate the true posterior, alternatives such as normalizing flows and copula-based variational inference (Copula VT) can offer greater expressivity, but with significant trade-offs.
19

---

Published in Transactions on Machine Learning Research (03/2025)
transform simple base distribution through sequence of invertible mappings, enabling rich, multimodal approxinations. However, they introduce additional computational overhead, can be sensitive to hyperparameters, and may suffer from numerical instabilities. Copula VI (Tan et al. 2015) instead models dependencies between latent variables while maintaining flexible marginal distributions, but this added structure increases computational cost and can lead higher-variance gradient estimates.
Normalizing flows (Rezende & Mohamed 2015
20


Reviews:
[{"review_id": "ztaIV07nPL", "reviewer": "Reviewer_M9y9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: This is a well written paper. The numerical issues of the Kumaraswamy distribution are well justified and illustrated. The proposed approach is very simple and could lead to applications beyond what is considered in the paper. The utility of the Kumaraswamy distribution is illustrated on three problems, namely, variational autoencoders, contextual multi-armed bandits and graph neural networks, where the need of a bounded distribution and issues with existing approaches are demonstrated with experimental results. The experimental results show that models using the proposed stable Kumaraswamy distribution are competitive, relatively computationally efficient and may open new avenues of research using such a distribution or its underlying approximation in problems involving variational inference with bounded distributions.\n\nWeaknesses: The novelty of the proposed approach reduces to leveraging the using the log1mexp(x) approximation of Machler (2012) in the calculation of the distribution and their gradients. The experiments though illustrate the capabilities of the proposed approximation are somewhat underwhelming. For instance:\n- The experiment on MNIST favors CB as discussed by the authors and though the results on CIFAR-10 show better performance with KS-based posteriors, they seem comparable to those from a simpler Normal-Beta model.\n- The multi-armed bandit experiment show more clear gains for the KS-based model, however, results are presented on a single artificial dataset, which makes the conclusions of the experiment less convincing.\n- The AUC results on the link prediction problem do not show substantial gains of the proposed KS_based model relative to the simpler base-GNN. However, the correlation values do. If the goal is to quantify the quality of the predictive uncertainty, why not consider more appropriate metrics for calibration. Moreover, the authors do not provide a justification of why it is expected that their approach is better calibrated than the alternative approaches.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors address the problem of alleviating numerical instabilities of the log-PDF, CDF and inverse CDF of the Kumaraswamy distribution, which is of importance in variational models with bounded latent variables. Then as illustration, they introduce a latent variable model for multi-armed bandits and graph neural networks. For the former, the goal is to improve the exploration-exploitation trade-ff while for the former the goal is to improve uncertainty quantification in link prediction tasks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- In the paragraph of Page 3 explaining Table 1 it is important to explain why the digamma and trigamma functions are not considered complex, and remove the redundant line about a^-1 as exp(-log a).\n- Describe how the error is calculated in Figure 2 and how many terms are necessary for the log1mexp(x) approximation.\n- Include error bars in Table 2.\n- Include additional datasets for the multi-armed bandit experiment.\n- Present other calibration metrics for the link prediction problem, e.g., Brier scores and calibration plots. The latter could be a good way to illustrate the regimes in which the KS-model perform better than the alternatives.\n- The authors mention TensorFlow in the abstract, however, no other mention or examples are presented in the paper (unlike Figure 3 for PyTorch).\n- In the multi-armed bandit experiment, what is the rationale for not considering a Gaussian variational distribution?", "broader_impact_concerns": "There are no potential ethical concerns that need a broader impact statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "S3R8m8cfH9", "reviewer": "Reviewer_V8yS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The main message of the paper is extremely clear, and the paper has tutorial-like value in explaining how KS should be used in practice. The recommendation is well supported by empirical evidence, in three concrete scenarios. In particular the first one, a clear demonstration in a maximally simple standard setting, is easy to follow and verifies the main claim. The technical development in Section 3 is easy to read and provides enough details for e.g. all software libraries to implement KS correctly.\n\nThe main weakness is the unclear role of Sections 4.2 and 4.3 that introduce specific variants of multi-armed bandits and graph NNs building on KS. If they are intended as concrete proposals of new contributions for these highly active research areas, they should be substantially stronger in terms of related work and experimentation and there should be clear claims. On the other hand, if they are intended as quick demonstrations, similar to Section 4.1, the authors should not emphasise them specifically as contributions and perhaps should also avoid even introducing the acronyms. While there are novel elements in the solutions, they still largely appear to be minor variants of previous methods that use the KS distribution in a natural way, roughly as we would expect anyone familiar with the tools to do. To clarify, I stress that the experiments and their results are fine -- clearly showing it is better to use KS compared to beta or tanh, especially in Section 4.2 -- and my criticism is only about trying to position the specific models as contributions for the MAB and graph-NN literature.\n\nSummary of strengths:\n - Well executed paper that provides highly practical tools for supporting use of KS in diverse applications\n - Detailed explanation of how KS should be implemented, likely to influence all key software libraries\n - Clear empirical evidence and broad illustration on how KS can be used in different kinds of applications\n\nSummary of weaknesses:\n - VEB and VEE are introduced as novel contributions, but there is nowhere near enough space to properly evaluate them or present the related literature in detail. Consequently, the (implicit) claims made regarding these newly proposed methods is not supported properly", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper provides detailed discussion on the KS distribution, advocating it as a flexible and efficient candidate for modelling random variables with bounded supports. Previously numerical instability has limited the use of KS and in this work the issues are resolved, by explaining a new parameterisation that makes the necessary computational elements stable. Use of KS, instead of other alternatives, is demonstrated in three concrete examples: VAE, multi-armed bandits and graph VAEs.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The list below includes both minor and major comments, in somewhat random order. My overall evaluation saying the claims are not supported by evidence relates only to points 3 and 4 (and the clarifying remark in 5), whereas the other points are more just recommendations on how to improve the paper. If points 3 and 4 are properly addressed, I would definitely support publishing the paper.\n\n1) I recommend providing the expressions for entropy, KL, moments etc for KS in Appendix, to increase the value of this paper as a stand-alone reference.\n\n2) I think the claim of expressiveness of KS is a bit exaggerated. While it is more expressive than some of the candidates, is is still limited to the four canonical types, with no support for e.g. multimodal distributions over bounded intervals. Brief discussion on still more flexible alternatives, e.g. copulas or normalising flows (that naturally have their own clear limitations), would be useful.\n\n3) You emphasize difference between explicit and implicit reparameterization on page 4, but do not actually show any empirical or theoretical evidence on what is the practical effect. You cite Kingma and Welling (2014) and Jang et al. (2017) with a general statement about positive qualities of explicit reparameterization, which is quite odd given that they predate Figurnov et al. (2018) that presented the basic idea of implicit reparameterization. I would strongly advice being more explicit about what is lost when resorting to implicit reparameterization (in terms of stability, computational efficiency, or whatever). Otherwise the claims related to that are not properly supported. In brief: Does it really matter if we need to use implicit reparameterization? If yes, how?\n\n4) The role of VBE and VEE needs to be clarified. Either make them properly clear contributions of this paper or reduce their role.\nFor the first choice, I would start by mentioning them already in the title and abstract, but additionally you would need to extend the literature and empirical experimentation so that they are properly contrasted with the literature. For the second choice, which I strongly recommend, re-writing selected parts is sufficient. All three experiments could be communicated in similar manner, as examples of how KS can be used in interesting models and as demonstrations of how it is better than using the alternatives (beta etc) as parts of the same model. You can, and should, still mention that the specific model variant is new, but you would not emphasize their originality and would refrain making claims on the performance relative to the current state-of-the-art models.\n5) Concrete examples of how VBE (Section 4.2) falls short of properly communicating a novel method, to clarify the above remark. Similar arguments, though to a lesser extent, could be raised about Section 4.3:\n    - No references to any multi-armed bandit work published after 2017 in the leading paragraph of Section 4.2.\n    - Only one data set and one comparison method (Xu et al. (2012)) and the data set is a specific construction with quite arbitrary elements (\"raise to power of 5\"). If you want to make claims on improved performance in MAB problems you need to follow standard experimentation protocols, including the data, the evaluation metrics and the baselines.\n    - The VBE advantages (prior knowledge incorporation, interpretability) are general for any Bayesian model for this, not related to KS specifically, and the last point (simplicity) is questionable -- you still require specification of the NN encoder\n    - The novelty of the proposed method is questionable; to me it looks like exactly standard MAB that just uses a variational approximation in Thompson sampling\n\n5) Is log1mexp as such a contribution of this work or from Mächler (2012)? Clarify the writing.\n\n6) $\\mathcal{K}(\\phi)$ is rather weird notation for accuracy -- why not just \"Acc.\" that would fit well in the table?\n\n7) Fig 6 is difficult to read -- please provide also the numbers in tabular format, probably in Appendix.", "broader_impact_concerns": "None. The paper presents purely technical contribution for improving stability of basic learning algorithms.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "A1on4Lq0rb", "reviewer": "Reviewer_Xp6w", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I think the paper has a nice contribution of proposing a method to use the KS distribution in a numerical stable way. The proposal is shown to perform significantly better than, for instance, PyTorch’s implementation.\n\nThey show across a range of small scale experiments that the proposed approach works better than existing alternatives for the KS distribution, and that it achieves competitive performance across a range of small-scale tasks. I’d consider this last point to be a weakness in this work; the empirical evaluation, while covering a range of tasks, is all on fairly simple/small models.\n\nA few comments / questions.\n\n(1) Why is being able to evaluate the CDF in closed form emphasized as one of the main motivations? In the context of probabilistic inference my impression is that as long as you can evaluate the PDF, and reparameterize (low variance gradients), it should be fine?\nRelated, closed form entropy and KL divergences against other families are also nice to have, in the context of inference this is not required (acknowledged in this work), and using empirical estimates instead of closed form can lead to better results. See, for instance, [1].\n\n(2) VAE experiments in 4.1 (CIFAR and MNIST).\n\n(2.1) The section states “we evaluate using a single sample approximation of the test ELBO”. Why was this chose made? The ELBO is an expectation, would it not make sense to use more samples to get better estimates? The ELBO is later on used to compare methods.\n\n(2.2) While table 3 shows better results for the KS likelihood (vs Beta likelihood), table 2 shows that the Beta likelihood yields better ELBOs. One way to compare methods, and that’s ultimately the goal, is the marginal likelihood of the data (the ELBO is a lower bound of that, but if we could optimize the marginal likelihood we would). Could such higher ELBOs be interpreted as the Beta likelihood fitting the data better? Ideally such a comparison should be done by estimating marginal likelihoods. I think it would be a nice addition to the results in Section 4.1.\n\n(2.3) Also, how meaningful is comparing ELBOs for different models? For instance, when the prior changes from a uniform to a Gaussian. (Table 2, for each choice for the likelihood function, compares models with different priors.) As mentioned above, I believe a potentially more meaningful way to do this involves the marginal likelihood.\n\nClarification suggestion for Section 4.2. I’m not very familiar with contextual bandits. While I don’t expect a deep introduction to the topic, a brief one is useful. While the work has a short intro, I think it can be improved a bit. A few suggestions (mostly restructuring a few things):\n\n(a) The first paragraph introduces the posterior over latent variables z, but the model was not introduced? So far the discussion only mentioned x_k, r_k, v_k. \n\n(b) How are v_k and x_k connected? Could it be assumed v_k is a function of x_k? This is actually explained later on, but a few words earlier could help.\n\n(c) Related to (a), the second paragraph describes the latents z_k and the distribution q(z|x), and then the likelihood p(r | v). So at this point it is unclear how the latents z fit in into the model. I think starting with the full model, and the posterior being modeled, would be useful. Again, this became clear as I kept on reading, but restructuring could help avoid these questions in the first few paragraphs.\n\nPossible typo? Section 4.1 states “ we encode the test data x+n, and compute the mean E[q(z | x)], and …” I think the expectation should be removed, as you care about the mean of the distribution q.\n\n[1] “Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference”. By Geoffrey Roeder, Yuhuai Wu, David Duvenaud.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper works on the Kumaraswamy (KS henceforth) distribution, a distribution with support bounded between 0 and 1, in the context of probabilistic modeling and inference. The main contributions involve introducing an unconstrained parameterization that enables the use of the KS distribution on a range of applications. While it had been used before, numerical tricks, such as clipping, had been used for numerical stability, which sometimes could result in deteriorated performance. The parameterization used in this work appears to address these issues.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would be interested in seeing marginal likelihood estimates for the experiments in section 4.1.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: During the initial round of reviewers, reviewers agreed the paper was well written, they appreciated the paper's story, and felt that the approach was well supported by the empirical evidence in the paper. On the other hand, reviewers pointed out some weaknesses in the empirical evaluation and novelty of the approach. 

The revision of the paper addressed most comments from the reviewers. Some reviewers felt the empirical evaluation could be still stronger, but overall agreed the paper should be accepted and that it is of interest to the TMLR community.  All reviewers agreed that the paper provides an interesting practical approach and that it addresses an important gap in the literature that could be useful for future research and applications.
Decision: Accept (poster)