Ищете более быстрый способ реализовать logSumExp в многомерном массиве

Question

У меня есть строка в каком-то R-коде, который я пишу, и она довольно медленная. Он применяет logSumExp к 4-мерному массиву с помощью команды apply. Мне интересно, есть ли способы его ускорить!

Реплекс: (это может занять 10 секунд или больше)

library(microbenchmark)
library(matrixStats)

array4d <- array( runif(5*500*50*5 ,-1,0),
                  dim = c(5, 500, 50, 5) )
microbenchmark(
    result <- apply(array4d, c(1,2,3), logSumExp)
)

Любые советы приветствуются!

Answer 1

rowSums - это менее общая версия apply, оптимизированная для скорости суммирования, поэтому ее можно использовать для ускорения вычислений. Обратите внимание на предостережение в файле помощи ?rowSums, если важно сохранить разницу в вычислениях между NA и NaN.

library(microbenchmark)
library(matrixStats)

array4d <- array( runif(5*500*50*5 ,-1,0),
                  dim = c(5, 500, 50, 5) )
microbenchmark(
  result <- apply(array4d, c(1,2,3), logSumExp),
  result2 <- log(rowSums(exp(array4d), dims=3))
)


# Unit: milliseconds
#                                            expr      min       lq      mean    median        uq      max neval
# result <- apply(array4d, c(1, 2, 3), logSumExp) 249.4757 274.8227 305.24680 297.30245 328.90610 405.5038   100
# result2 <- log(rowSums(exp(array4d), dims = 3))  31.8783  32.7493  35.20605  33.01965  33.45205 133.3257   100

all.equal(result, result2)

#TRUE

Это приводит к 9-кратному увеличению скорости на моем компьютере

Answer 2

В остальном отличное решение от @Miff заставляло мой код обрабатывать sh с определенными наборами данных, поскольку создавались бесконечности, что, как я в конечном итоге выяснил, было связано с проблемой недостаточного заполнения, которую можно избежать, используя трюк logSumExp: https://www.xarg.org/2016/06/the-log-sum-exp-trick-in-machine-learning/

Вдохновленный кодом @Miff и функцией R apply(), я создал новую функцию, которая дает более быстрые вычисления, избегая при этом проблемы недостаточного заполнения. Однако не так быстро, как решение @Miff. Размещение на случай, если это поможет другим

apply_logSumExp <- function (X) {
    MARGIN <- c(1, 2, 3) # fixing the margins as have not tested other dims
    dl <- length(dim(X)) # get length of dim
    d <- dim(X) # get dim
    dn <- dimnames(X) # get dimnames
    ds <- seq_len(dl) # makes sequences of length of dims
    d.call <- d[-MARGIN]    # gets index of dim not included in MARGIN
    d.ans <- d[MARGIN]  # define dim for answer array
    s.call <- ds[-MARGIN] # used to define permute
    s.ans <- ds[MARGIN]     # used to define permute
    d2 <- prod(d.ans)   # length of results object
    
    newX <- aperm(X, c(s.call, s.ans)) # permute X such that dims omitted from calc are first dim
    dim(newX) <- c(prod(d.call), d2) # voodoo. Preserves ommitted dim dimension but collapses the rest into 1
    
    maxes <- colMaxs(newX)
    ans <- maxes + log(colSums(exp( sweep(newX, 2, maxes, "-"))) )
    ans <- array(ans, d.ans)
    
    return(ans)
}

 > microbenchmark(
+     res1 <- apply(array4d, c(1,2,3), logSumExp),
+     res2 <- log(rowSums(exp(array4d), dims=3)),
+     res3 <- apply_logSumExp(array4d)
+ )
Unit: milliseconds
                                          expr        min         lq       mean    median        uq       max
 res1 <- apply(array4d, c(1, 2, 3), logSumExp) 176.286670 213.882443 247.420334 236.44593 267.81127 486.41072
  res2 <- log(rowSums(exp(array4d), dims = 3))   4.664907   5.821601   7.588448   5.97765   7.47814  30.58002
              res3 <- apply_logSumExp(array4d)  12.119875  14.673011  19.635265  15.20385  18.30471  90.59859
 neval cld
   100   c
   100 a  
   100  b