On the Normalizing Constant of
the Continuous Categorical Distribution

Algorithm 1 is easy to code up and adds little computational overhead to most models. However, the summation in Step 3 involves positive and negative numbers, which can result in catastrophic cancellation (Goldberg, 1991). We stress that the log-sum-exp trick, while useful for preventing overflow, cannot address catastrophic cancellation (see Section 2.4). For example, consider the case $K=5$ with $\boldsymbol{\eta}=(1,2,3,4)$. In single-precision floating-point, the summation in Eq. 9 evaluates to:

Note that the output is an order of magnitude smaller than (at least one of) the summands, and as a result we have lost one digit in precision. As we increase the dimension of the CC, the cancellation becomes more severe. For example, in the case $K=10$ with $\boldsymbol{\eta}=(1,2,3,4,5,6,7,8,9)$, the same summation becomes:

We have now lost 3 digits, since the leading summand is 3 orders of magnitude greater than the output. If we continue increasing $K$ in the same way, by $K=20$ we will have lost 6 digits, and by $K=25$ we are no longer able to compute $C(\boldsymbol{\eta})$ to even a single significant digit. If we were to use double-precision floating-point instead, by $K=40$ we will have lost 13 digits, and by $K=50$ we can no longer compute $C(\boldsymbol{\eta})$ to a single significant digit.

We can summarize the problem as follows: when $C(\boldsymbol{\eta})$ is of a much lower order of magnitude than the leading term in the summation of Eq. 9, numerical computation fails due to catastrophic cancellation. To complicate things further, the relationship between the two orders of magnitude (that of $C(\boldsymbol{\eta})$ and that of the leading summand) is nontrivial. This relationship depends on the relative size of the exponential terms $e^{\eta_{k}}$ and the products of the differences $\eta_{k}-\eta_{i}$, and is not straightforward to analyze. However, if two elements of $\boldsymbol{\eta}$ are particularly close to one another, meaning that $|\eta_{j_{1}}-\eta_{j_{2}}|$ is close to 0 for some $j_{1}$ and $j_{2}$, it follows that the corresponding summands:

are (a) large in magnitude (due to the term $\eta_{j_{1}}-\eta_{j_{2}}$ in the denominator), (b) of opposing sign (due to the difference in $\eta_{j_{1}}-\eta_{j_{2}}$ versus $\eta_{j_{2}}-\eta_{j_{1}}$), and (c) similar in absolute value (since the terms in the product are approximately equal). Thus, it is likely that the terms in Eq. 12 are of leading order and catastrophically cancel each other out. Note that, as the dimensionality $K$ increases, it becomes more likely that some pair of the components are close to one another, and therefore computation becomes harder.

Another helpful intuition can be obtained by reasoning from the integral:

As $K$ increases, the Lebesgue measure of the simplex decays like $1/K!$.³³3As can be seen, for example, by taking Eq. 6 with $\alpha_{i}=0$ for all $i$. Therefore, assuming the components of $\boldsymbol{\eta}$ are $O(1)$, we have that $e^{\boldsymbol{\eta}^{\top}\mathbb{x}}=O(1)$ also, and therefore $C(\boldsymbol{\eta})=O(1/K!)$. Under this assumption, we also have that $\eta_{k}-\eta_{i}=O(1)$, and therefore the summands in Eq. 9 cannot decay factorially fast (they may, but need not, decay at most exponentially due to the product of $K-1$ terms of constant order in the denominator). Thus, such a regime implies catastrophic cancellation is inevitable for large enough $K$. On the other hand, when all the components of $\boldsymbol{\eta}$ are far from one another, we are spared of such cancellation and Algorithm 1 succeeds, including in high dimensions. We demonstrate these behaviors empirically in the following experiments (see Figures 1 and 2).

In this Section, we show that $C(\boldsymbol{\eta})$ can be written as the inverse Laplace transform of a function that does not suffer from catastrophic cancellation. In particular, said function can be passed to Laplace inversion methods (Davies and Martin, 1979) in order to evaluate $C(\boldsymbol{\eta})$ in the regime where Algorithm 1 fails.

Proposition: For a function $f:\mathbb{R}^{+}\to\mathbb{R}^{+}$, let $\mathcal{L}[f](s)=\int_{0}^{\infty}f(t)e^{-st}dt$ denote the Laplace transform. Define the function $c:\mathbb{R}^{+}\to\mathbb{R}^{+}$ by:

Then the Laplace transform of $c$ is equal to:

Remark: The function $c$ includes the normalizing constant of the CC as a special case $C(\boldsymbol{\eta})=c(1)$. More generally, we have that:

as can be seen from letting $\tilde{x}_{i}=x_{i}/t$ in the integral 14.

Proof: Let $f_{\eta}(x)=e^{\eta x}$, so that:

Next we apply the transformation from Wolpert and Wolf (1995), defined as $w_{1}=t$, $w_{k}=w_{k-1}-x_{k-1}$, or equivalently, $w_{k}=t-\sum_{i=1}^{k-1}x_{i}$. We can then write our integral as a convolution:

Next, since the Laplace transform of a convolution equals the product of the Laplace transforms, we have that:

but the univariate case is simply:

and the result follows. ∎

Corollary: The normalizing constant of the continuous categorical distribution can be written as the following inverse Laplace transform:

Proof: Take the inverse Laplace transform in Eq. 15 to find:

Taking $t=1$ gives the desired result. ∎

Unlike Eq. 9, the product in Eq. 15 does not suffer from catastrophic cancellation, nor does it diverge whenever $\eta_{j_{1}}\approx\eta_{j_{2}}$ for some $j_{1}\neq j_{2}$. The corresponding Laplace inversion, i.e., Eq. 21, provides an alternative method to compute our normalizing constant $C(\boldsymbol{\eta})$. Numerous numerical algorithms exist for inverting the Laplace transform; see (Cohen, 2007) for a survey. We note, however, that inverting the Laplace transform is generally a hard problem (Epstein and Schotland, 2008).

Empirically, we found some modest empiricasuccess in computing Eq. 21 numerically. We tested three inversion algorithms, due to Talbot (Talbot, 1979), Stehfest (Stehfest, 1970), and De Hoog (De Hoog et al., 1982). The experimental setup was identical to that of Section 2.1.1, and the results are incorporated into Figure 2. We found De Hoog’s method to be the most effective on our problem, whereas Talbot’s always failed and is omitted from the Figure. Importantly, De Hoog’s method showed some success in the regime where Algorithm 1 failed, meaning it could be used as a complementary technique. However, no inversion method was able to scale beyond $K=30$.

In this section, we provide an alternative algorithm to compute our normalizing constant. This algorithm will be based on the following recursive property of $C(\boldsymbol{\eta})$, which was implicitly used as part of the proof of Eq. 9 (Gordon-Rodriguez et al., 2020a).

Proposition: Define the subvector notation $\boldsymbol{\eta}_{:k}=(\eta_{1},\dots,\eta_{k-1})$, and make the dependence on $K$ explicit by writing $C_{K}(\boldsymbol{\eta})=C(\eta_{1},\dots,\eta_{K-1})$. We also use the notation $\boldsymbol{\eta}-\eta_{k}=(\eta_{1}-\eta_{k},\dots,\eta_{K-1}-\eta_{k})$. Then:

Proof: We start from the integral definition of the left hand side:

For the innermost integral, we have:

and the result follows by linearity of the integral. ∎

Remark: The base case $K=2$ corresponds to the univariate continuous Bernoulli distribution, which admits a straightforward Taylor expansion that is useful around the unstable region $\eta\approx 0$ (Loaiza-Ganem and Cunningham, 2019):

In words, Eq. 23 is stating that we can compute $C(\boldsymbol{\eta})$ for the full parameter vector $\boldsymbol{\eta}=(\eta_{1},\dots,\eta_{K-1})$ by first computing it for the lower-dimensional parameter vectors:

Substituting these back into Eq. 23, we have that:

Note that we are now left with not 4, but 3 new parameter vectors to recurse on: $\boldsymbol{\eta}_{1:K-3}$, $\boldsymbol{\eta}_{1:K-3}-\eta_{K-1}$, and $\boldsymbol{\eta}_{1:K-3}-\eta_{K-2}$. Repeating the argument $K-2$ times and working backwards we obtain Algorithm 2 for computing the normalizing constant.

We find the numerical properties of Algorithm 2 to perform similarly to Algorithm 1, suffering from the same cancellation issues in high dimensions. Nevertheless, we hope this alternative scheme may help to inspire further numerical improvements. For example, since the floating-point behavior of Algorithm 2 depends on the ordering of the elements of $\boldsymbol{\eta}$, it may be possible do design a reordering scheme that improves the overall precision of the algorithm; such reorderings have been explored in the context of sampling algorithms for the CC (Gordon-Rodriguez et al., 2020a). Other possibilities include Kahan summation (Kahan, 1965) or the compensated Horner algorithm (Langlois and Louvet, 2007); we leave their study to future work.

We conclude this section with a small remark on numerical overflow in the context of Algorithm 1. In high dimensions and with high $\eta$ values, overflow can occur for the terms in the summand, due to a large $e^{\eta_{j}}$ term or a small denominator. This can be addressed by re-writing the normalizing constant in a form that allows us to take advantage of the log-sum-exp trick:

First, we illustrate the main idea of our argument using an example with $K=3$ and $\eta_{1}=\eta_{2}\neq\eta_{3}=0$.⁴⁴4Note that the case $\eta_{3}\neq 0$ and the case $\eta_{1}\neq\eta_{2}=\eta_{3}$ are mathematically equivalent (albeit more algebraically cumbersome), since we can permute the elements of $\boldsymbol{\eta}$ and shift by a constant without loss of generality. By definition, the normalizing constant is then:

We apply the change of variables $u=x_{1}+x_{2}$, $v=x_{1}$ to obtain:

Note that the last integral corresponds to the expectation of a univariate CC random variable, an idea we now generalize.

We start by proving the following lemma, which relates $C(\boldsymbol{\eta})$, where some elements of $\boldsymbol{\eta}$ are repeated (potentially many times), to $C(\boldsymbol{\eta}^{\prime})$, where $\boldsymbol{\eta}^{\prime}$ collapses the repeated elements of $\boldsymbol{\eta}$ onto a single coordinate. We again use subvector notation $\mathbb{x}_{k:}=(x_{k},x_{k+1},\dots,x_{K-1})$.

Lemma: Let $\mathbb{x}\sim\mathcal{CC}_{K}(\boldsymbol{\eta})$ with $\eta_{1}=\eta_{2}=\dots=\eta_{k}$, for $1\leq k\leq K-1$, and let $f:\mathbb{R}^{K-k-1}\to\mathbb{R}$. Then:

Remark: We are not assuming that the last $K-k$ coordinates of $\boldsymbol{\eta}$ are all different, we are simply assuming that the first $k$ are identical. Note also that the positions of the coordinates could be arbitrary and need not be the first $k$ ones; we can always use this lemma to collapse repeated parameter values provided that the function $f$ does not depend on the corresponding coordinates (we can simply relabel the coordinates by a suitable permutation without loss of generality).

Proof: By definition of the left hand side:

Consider the following change of variable (note this is just like Section 3.1, but with more bookkeeping):

This change of variable amounts to an invertible linear transformation with the property that the absolute value of the determinant of its Jacobian is $1$, so that we have:

Note that the change of variables is such that the integrand does not depend on $v_{1},v_{2},\dots,v_{k-1}$. Therefore:

where:

But this is simply the Lebesgue measure of a simplex inscribed in the hypercube $[0,u_{1}]^{K-1}$, so that $g(u_{1})=u_{1}^{k-1}\mu(\mathbb{S}^{K})=u_{1}^{k-1}/(k-1)!$ (this can also be seen by changing variables to $\tilde{v}_{i}=v_{i}/u_{i}$, or by applying Eq. 16). Multiplying and dividing by $C_{K-k+1}(\boldsymbol{\eta}_{k:})$ gives the desired result. ∎

We can now derive a formula for the normalizing constant for an arbitrary parameter vector $\boldsymbol{\eta}$.

Corollary: Let $\boldsymbol{\eta}\in\mathbb{R}^{K}$ contain $D\leq K$ unique elements. Assume, without loss of generality, that $\boldsymbol{\eta}=(\eta_{1},\dots,\eta_{1},\eta_{2},\dots,\eta_{2},\dots,\eta_{D},\dots,\eta_{D})$, where each coordinate is repeated $1\leq r_{i}\leq K$ times, with $\sum_{i=1}^{D}r_{i}=K$. Then:

Proof: The result follows by applying the lemma $D$ times. First, we apply the lemma with $f(\cdot)=1$:

where $\boldsymbol{\eta}^{\prime}=\boldsymbol{\eta}_{r_{1}:}=(\eta_{1},\eta_{2},\dots,\eta_{2},\dots,\eta_{D},\dots,\eta_{D})$, i.e., we have collapsed the first parameter value onto a single coordinate. Next, we apply the lemma a second time on the new expectation term to collapse the $\eta_{2}$ values, this time using $f(u_{1})=u_{1}^{r_{1}-1}/(r_{1}-1)!$, which does not depend on the $\eta_{2}$ coordinates:

where $\boldsymbol{\eta}^{\prime\prime}=(\eta_{1},\eta_{2},\eta_{3},\dots,\eta_{3},\dots,\eta_{D},\dots,\eta_{D})$. Repeating $D$ times yields the desired result. ∎

Note that Eq. 35 can be computed using known results. The term $C_{D}(\eta_{1},\dots,\eta_{D})$ can be evaluated with Eq. 9, since all the parameter values are distinct. The expectation term can be computed by differentiating the moment generating function of $\mathbb{u}\sim\mathcal{CC}_{D}(\eta_{1},\eta_{2},\dots,\eta_{D})$, as discussed by Gordon-Rodriguez et al. (2020a). Note that in real data applications exact equality between parameter values may never occur, and it is unclear how close the elements of $\boldsymbol{\eta}$ should be in order to warrant applying Eq. 35. Nevertheless, Eq. 35 is important for theoretical completeness.