Add numerically stable log1mexp = log(1 - exp(-|x|)) function

[simonepri]

🚀 Feature

Add the implementation for a numerically stable log(1 - exp(-|x|)) function.

See https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf

Ref https://www.rdocumentation.org/packages/VGAM/versions/1.1-3/topics/log1mexp

Motivation

The function can be implemented using other pytorch functions but it might be tricky to get it right.

Some implementations found on GitHub:

@wouterkool/estimating-gradients-without-replacement/blob/9d8bf8b/bernoulli/gumbel.py#L7-L11

@visinf/n3net/blob/5d5883a/src_denoising/models/non_local.py#L94-L108

cc: @wouterkool @visinf

[mruberry]

Hey @simonepri, thanks for the suggestion. Would you be interested in submitting a PR for it?

Follow-up question, how numerically stable are PyTorch's log1p and expm1 implementations?

[xiaomengy]

If no one took this, I think I can help add them at my spare time.

[mruberry]

@BIT-silence We already have log1p and expm1, which can be used to write a reasonable stable version of this function (so long as they themselves are stable). I'd start by checking the numerical stability of those ops.

[cmpute]

Thanks for the links provided! This function is also useful in my work!

[Balandat]

If anyone still wants to add this (courtesy of @SebastianAment):

def log1mexp(x: Tensor) -> Tensor:
    """Numerically accurate evaluation of log(1 - exp(x)) for x < 0.
    See [Maechler2012accurate]_ for details.
    """
    mask = -math.log(2) < x  # x < 0
    return torch.where(
        mask,
        (-x.expm1()).log(),
        (-x.exp()).log1p(),
    )

[Maechler2012accurate] M. Mächler. Accurately Computing log (1 - exp (-| a|)). Assessed by the Rmpfr package. Technical report, 2012. https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf