Vectorised computation of logsumexp

Question

In this related post there is an explanation of how you can add together two very small probabilities using the logsumexp function, and how this can be programmed into base R. How can this method be extended to get a general vectorised function that calculates:

$$\ell_+ = \ln \Bigg( \sum_{i=1}^n \exp(\ell_i) \Bigg),$$

where the values $\ell_1,...,\ell_n$ are log-probabilities, and the corresponding probability values $\exp(\ell_i)$ are too small to be distinguished from zero in the computational platform (e.g., in base R).

Answer 1

Before showing how this can be done, it is worth noting that the logsumexp function has already been programmed into R in the matrixStats package. This function calls an underlying function programmed in C, so it computes the result faster than forming an equivalent function using standard R commands. Nevertheless, for the purposes of showing the mathematical method of computing the logsumexp function, this answer will show how the function can be defined in terms of underlying functions in the base package.

The linked question shows how to do this for a sum of two argument values. To vectorise this to a larger set of log-probabilities we use a similar representation for sums of exponentials. To do this we let $\ell_{(1)} \leqslant ... \leqslant \ell_{(n)}$ and we define the partial sums:

$$S_k \equiv \ln \Bigg( \sum_{i=1}^k \exp(\ell_{(i)}) \Bigg).$$

For these partial sums we can establish the recursive relationship:

$$\begin{equation} \begin{aligned} \exp(S_{k+1}) &= \exp(\ell_{(k+1)}) + \exp(S_k) \\[6pt] &= \exp(\max(\ell_{(k+1)}, S_k)) + \exp(\min(\ell_{(k+1)}, S_k)) \\[6pt] &= \exp(\max(\ell_{(k+1)}, S_k)) (1 + \exp(\max(\ell_{(k+1)}, S_k) - \min(\ell_{(k+1)}, S_k)) \\[6pt] &= \exp(\max(\ell_{(k+1)}, S_k)) \Big( 1 + \exp( -|\ell_{(k+1)} - S_k|) \Big). \\[6pt] \end{aligned} \end{equation}$$

Hence, we have the recursive equation:

$$\begin{equation} \begin{aligned} S_{k+1} = \max(\ell_{(k+1)}, S_k) + \ln \Big( 1 + \exp( -|\ell_{(k+1)} - S_k|) \Big). \\[6pt] \end{aligned} \end{equation}$$

This recursive equation gives us a procedure to calculate $\ell_+ = S_n$ via use of the log1p function, without ever working directly with the small probability values. We have:

logsumexp <- function(l) { 
    n <- length(l)
    L <- sort(l, decreasing = TRUE)
    S <- rep(L[1], n)
    if (n > 1) {
    for (k in 1:(n-1)) {
        S[k+1] <- max(L[k+1], S[k]) + log1p(exp(-abs(L[k+1] - S[k]))) } }
    S[n] }

It can be confirmed that this method yields essentially the same result as the logSumExp function in the matrixStats package (with a small difference due to rounding in the intermediate steps), but the latter is programmed using an underlying C function so it is faster:

#Generate large number of small log-probabilities
set.seed(1)
n <- 10^6
l <- rnorm(n, -3000, 1)

#Calculate logsumexp using above function and package
library(matrixStats)
l1 <- logsumexp(l)
l2 <- matrixStats::logSumExp(l) 

print(l1, digits = 20)
[1] -2985.6845568559206
print(l2, digits = 20)
[1] -2985.6845568559038

#Test calculation speed of functions (specific to my PC - not reproducible)
TIME1 <- system.time(logsumexp(l))
TIME2 <- system.time(matrixStats::logSumExp(l))

TIME1[3]
elapsed 
    0.7

TIME2[3]
elapsed 
   0.05

The function programmed here takes about 14 times as long to calculate the desired log-probability as the matrixStats::logSumExp function.