LogCumsumExp

[agadetsky]

🚀 Feature

Add numerically stable cumulative logsumexp function. Also we have associated PR on cummax that is needed for numerically stable implementation (#20240).

Motivation

This is useful when computing sum of probabilities and have different applications.

Pitch

Torch has cumsum and cumprod so I suggest logcumsumexp to be added.

Alternatives

Current workaround for me and for people that need it can be written as follows (it is quite fast, only overhead for converting between numpy/torch):

import numpy as np
import torch

def cummax(x, dim):
    x_np = x.detach().cpu().numpy()
    ret = np.maximum.accumulate(x_np, axis=dim)
    return torch.from_numpy(ret).to(x)


def logcumsumexp(x, dim=-1):
    if (dim != -1) or (dim != x.ndimension() - 1):
        x = x.transpose(dim, -1)

    init_size = x.size()
    last_dim_size = init_size[-1]
    x_resized = x.contiguous().view(-1, last_dim_size)
    d1, d2 = x_resized.size()
    x_cummax = cummax(x_resized, -1).view(d1, d2, 1)
    x_expand = x_resized.unsqueeze(1).expand(d1, d2, last_dim_size)
    mask = torch.tril(torch.ones(last_dim_size, last_dim_size)).unsqueeze(0)
    ret = torch.log(torch.sum(torch.exp(x_expand - x_cummax) * mask, dim=-1)) + x_cummax.view(d1, d2)
    ret = ret.view(*init_size)

    if (dim != -1) or (dim != x.ndimension() - 1):
        ret = ret.transpose(-1, dim)

    return ret

UPDATE: code above is not always numerically stable. Still most numerically stable, but slow way to compute logcumsumexp is use for-loop:

def logcumsumexp(x, dim):
    # slow implementation, but ok for now
    if (dim != -1) or (dim != x.ndimension() - 1):
        x = x.transpose(dim, -1)

    out = []
    for i in range(1, x.size(-1) + 1):
        out.append(torch.logsumexp(x[..., :i], dim=-1, keepdim=True))
    out = torch.cat(out, dim=-1)

    if (dim != -1) or (dim != x.ndimension() - 1):
        out = out.transpose(-1, dim)
    return out

[opheliagame]

I would like to work on this.

[agadetsky]

@opheliagame see #20240 (comment) for cummax implementation

[michiboo]

HI @agadetsky , I created a PR #26832 for this issues. But I am not sure If I am doing it right in the implementation part, I stop just before the line

mask = torch.tril(torch.ones(last_dim_size, last_dim_size)).unsqueeze(0)

in python implementation.

It would be great if you can review THTensor_(logcumsumexp) in THTensormoreMath.cpp.

I want to make sure I am going in right direction.

Many thanks!

[agadetsky]

@michiboo hello! firstly I think we need to end this PR on cummax (#26637). Can you work on it?

[michiboo]

Hi, Yes I could try work on it.

[michiboo]

@agadetsky what do you use to loop over multi-dimensional tensor in C++, I am trying to compute the convert value of the tensor to 1s and 0s

[agadetsky]

@michiboo unfortunately I am not familiar with pytorch C++ implementation.

[arktrail]

@michiboo Hi, are you still working on this issue?

[pandeykartikey]

I will try to submit a PR for this by the weekend, if not needed urgently.

PS: @agadetsky I can't seem to assign myself the issue 😅

[agadetsky]

@pandeykartikey thank you, waiting for your PR

[agadetsky]

@pandeykartikey you can use implemented cummax for stable implementation of logcumsumexp

[pandeykartikey]

@agadetsky Wouldn't that create issues while building a feature because It can undergo drastic changes in implementation as it is being reviewed right now?

[agadetsky]

@pandeykartikey better to ask someone else, atm I don't know, but I think that without cummax it is harder to implement logcumsumexp properly

[agadetsky]

@pandeykartikey hi! do you have any progress?

[agadetsky]

@pandeykartikey cummax is merged now (#30881), is it possible to implement logcumsumexp for you using cummax?

[pandeykartikey]

Hi @agadetsky, I am almost done with the implementation. Only the backprop and tests are left. But I am having trouble figuring out the backprop algorithm. Can you help me with that?

[rebelmachina]

@pandeykartikey what's the status on this ? Still open ?

[agadetsky]

@whoiscris yes, maybe @anjali411 will finish it soon here #32876

[agadetsky]

up, #32876 still needs to be finished