Avoiding scipy softmax/logsumexp

[markkho]

msdm/msdm/algorithms/vectorizedvalueiteration.py

Line 1 in 8bb0f88

from scipy.special import softmax, logsumexp

We should avoid scipy's softmax/logsumexp according to Blanchard et al. (2020)

[cgc]

Here's an implementation of Algorithm 4.1, which they use to compute logsumexp (f) and softmax (g).

def logsumexp_and_softmax(x):
    '''
    Algorithm 4.1 from Blanchard et al. 2020
    '''
    a = np.max(x)
    w = np.exp(x-a)
    s = np.sum(w) - 1 # Subtracting 1 for the max item.
    f = a + np.log1p(s)
    g = w/(1+s)
    return f, g # logsumexp and softmax

Haven't yet been successful in finding inputs with a large difference, though here's an example with some small differences.

from scipy.special import logsumexp, softmax
z = np.array([-2, 10], dtype=np.float32)
l, s = logsumexp_and_softmax(z)
print(l, s, s.dtype)
sl, ss = logsumexp(z), softmax(z)
print(sl, ss, ss.dtype)
print()
print('logsumexp diff', l-sl)
print('softmax diff', s-ss)

with output

10.000006198863844 [6.144174e-06 9.999938e-01] float32
10.000006 [6.1441774e-06 9.9999428e-01] float32

logsumexp diff 4.7681794512755005e-07
softmax diff [-3.1832315e-12 -4.7683716e-07]