Численно устойчивый способ умножения логарифмических вероятностных матриц в numpy

Мне нужно взять матричное произведение двух матриц NumPy (или других 2d-массивов), содержащих логарифмические вероятности. Наивный способ np.log(np.dot(np.exp(a), np.exp(b))) не является предпочтительным по очевидным причинам.

Использование

from scipy.misc import logsumexp
res = np.zeros((a.shape[0], b.shape[1]))
for n in range(b.shape[1]):
    # broadcast b[:,n] over rows of a, sum columns
    res[:, n] = logsumexp(a + b[:, n].T, axis=1)

работает, но работает примерно в 100 раз медленнее, чем np.log(np.dot(np.exp(a), np.exp(b)))

Использование

logsumexp((tile(a, (b.shape[1],1)) + repeat(b.T, a.shape[0], axis=0)).reshape(b.shape[1],a.shape[0],a.shape[1]), 2).T

или другие комбинации плитки и перестройки также работают, но работают еще медленнее, чем цикл выше из-за непомерно больших объемов памяти, необходимых для матриц входного сигнала с реалистичным размером.

В настоящее время я рассматриваю возможность написания расширения NumPy в C, чтобы вычислить это, но, конечно, я бы предпочел избежать этого. Есть ли установленный способ сделать это, или кто-нибудь знает об менее интенсивном в памяти способе выполнения этого вычисления?

EDIT: Благодаря larsmans для этого решения (см. Ниже для вывода):

def logdot(a, b):
    max_a, max_b = np.max(a), np.max(b)
    exp_a, exp_b = a - max_a, b - max_b
    np.exp(exp_a, out=exp_a)
    np.exp(exp_b, out=exp_b)
    c = np.dot(exp_a, exp_b)
    np.log(c, out=c)
    c += max_a + max_b
    return c

Быстрое сравнение этого метода с вышеописанным методом (logdot_old) с использованием функции iPython magic %timeit дает следующее:

In  [1] a = np.log(np.random.rand(1000,2000))

In  [2] b = np.log(np.random.rand(2000,1500))

In  [3] x = logdot(a, b)

In  [4] y = logdot_old(a, b) # this takes a while

In  [5] np.any(np.abs(x-y) > 1e-14)
Out [5] False

In  [6] %timeit logdot_old(a, b)
1 loops, best of 3: 1min 18s per loop

In  [6] %timeit logdot(a, b)
1 loops, best of 3: 264 ms per loop

Очевидно, метод larsmans уничтожает мой!

Ответы

Ответ 1

logsumexp работает, оценивая правую часть уравнения

log(∑ exp[a]) = max(a) + log(∑ exp[a - max(a)])

I.e., он вытягивает max до начала суммирования, чтобы предотвратить переполнение в exp. То же самое можно применить и перед выполнением векторных точечных продуктов:

log(exp[a] ⋅ exp[b])
 = log(∑ exp[a] × exp[b])
 = log(∑ exp[a + b])
 = max(a + b) + log(∑ exp[a + b - max(a + b)])     { this is logsumexp(a + b) }

но, проведя другой поворот в выводе, получим

log(∑ exp[a] × exp[b])
 = max(a) + max(b) + log(∑ exp[a - max(a)] × exp[b - max(b)])
 = max(a) + max(b) + log(exp[a - max(a)] ⋅ exp[b - max(b)])

Конечная форма имеет векторный точечный продукт в его внутренностях. Он также легко распространяется на матричное умножение, поэтому мы получаем алгоритм

def logdotexp(A, B):
    max_A = np.max(A)
    max_B = np.max(B)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

Это создает два A временных интервала и два B -размерных, но один из них может быть удален

exp_A = A - max_A
np.exp(exp_A, out=exp_A)

и аналогично для B. (Если входные матрицы могут быть изменены функцией, все временные файлы могут быть устранены.)