37. Custom Autograd Function: Numerically Stable LogSumExp

Problem

logsumexp(x) = log(sum(exp(x))) is the literal denominator inside softmax and cross-entropy, so it needs to be both correct and stable at the magnitudes real logits reach. Computed literally, exp(x) overflows to inf in float32 once x exceeds roughly 88, and log(inf) is inf instead of the true finite answer. The standard fix is the max-subtraction trick: logsumexp(x) = max(x) + log(sum(exp(x - max(x)))), mathematically identical but never exponentiates a large positive number.

Autograd can differentiate straightforward code like this automatically, so writing it as a torch.autograd.Function is not strictly necessary here — but it's exactly this kind of primitive (small, numerically delicate, reused everywhere) where teams hand-write one to pin down a cheaper or more stable backward than what autograd would trace through the forward, or to make the stability trick (the .detach() on the max) explicit rather than relying on autograd not to differentiate through it by accident. The gradient has a closed form worth deriving by hand: d(logsumexp)/dx_i = softmax(x)_i.

Source: src/37_custom_autograd_logsumexp.py

class StableLogSumExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor, dim: int) -> torch.Tensor: ...
    @staticmethod
    def backward(ctx, grad_output: torch.Tensor): ...

def stable_logsumexp(x: torch.Tensor, dim: int = -1) -> torch.Tensor: ...

>>> x = torch.tensor([-1.0, 0.0, 1.0])
>>> stable_logsumexp(x)
tensor(1.4076)

>>> x = torch.tensor([500.0, 501.0])  # naive log(sum(exp(x))) overflows
>>> stable_logsumexp(x)
tensor(501.3133)

Step-by-Step Approach

1. In forward, compute x_max = x.amax(dim=dim, keepdim=True) and detach it — it's a numerical-stability shift, not a differentiable part of the function.
2. Compute shifted = x - x_max, then sumexp = shifted.exp().sum(dim=dim, keepdim=True); every exponentiated value is now ≤ 1, so it can't overflow.
3. Return sumexp.log() + x_max, squeezed on dim — adding x_max back recovers the exact true value.
4. Before returning, save softmax_x = shifted.exp() / sumexp via ctx.save_for_backward — it's exactly the gradient backward needs, and you already have the pieces from the forward pass.
5. In backward, retrieve the saved softmax and multiply it by the (unsqueezed) upstream gradient: grad_output.unsqueeze(dim) * softmax_x.
6. Return one gradient per forward input, in order — None for dim, since it's an int, not a differentiable tensor.

The key insight is that the forward pass already computes everything backward needs (the softmax weights), so the efficient move is to save that intermediate in ctx during forward rather than recomputing exp() a second time in backward.

Reference solution

import torch


class StableLogSumExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor, dim: int) -> torch.Tensor:
        # max-subtraction trick: exp(x - max) never overflows, and adding
        # max back after the log recovers the exact true value
        x_max = x.amax(dim=dim, keepdim=True)
        # detach: the max is only a stability shift, not differentiable --
        # letting gradient flow through it would double-count
        shifted = x - x_max.detach()
        sumexp = shifted.exp().sum(dim=dim, keepdim=True)
        result = sumexp.log() + x_max
        # softmax(x) = exp(x - max) / sumexp is exactly what backward needs
        ctx.save_for_backward(shifted.exp() / sumexp)
        ctx.dim = dim
        return result.squeeze(dim)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        (softmax_x,) = ctx.saved_tensors
        # d(logsumexp)/dx_i = softmax(x)_i
        grad_input = grad_output.unsqueeze(ctx.dim) * softmax_x
        return grad_input, None  # None: no gradient w.r.t. the int `dim` arg


def stable_logsumexp(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    return StableLogSumExp.apply(x, dim)

Key Functions & Tricks

• torch.autograd.Function with forward/backward static methods — the standard way to hand-write a custom differentiable op.
• ctx.save_for_backward(...) — the correct way to stash tensors between forward and backward (vs. plain attribute assignment, which autograd doesn't track for memory/versioning purposes).
• tensor.amax(dim=dim, keepdim=True) — reduction max that keeps the dim for clean broadcasting against the un-reduced tensor.
• .detach() on the max — explicitly cuts the max out of the autograd graph so backward only sees the intended closed-form gradient.
• tensor.squeeze(dim) / grad_output.unsqueeze(dim) — forward reduces a dimension, so backward must reintroduce it before broadcasting the gradient back to the input's shape.
• Function.apply(...) — the only correct way to invoke a custom Function; calling forward directly bypasses autograd bookkeeping entirely.

How to Recognize This Pattern