27. Custom Autograd Function: Numerically Stable Softplus

Problem

Cartesia's models train on long audio/sequence data where activations can hit extreme magnitudes, so a naive nonlinearity can silently overflow. The softplus activation, softplus(x) = log(1 + exp(x)), is a smooth relative of ReLU used in gating and probabilistic-output layers. Computed literally, exp(x) overflows to inf for x much above ~88 in float32, and log(1 + inf) then evaluates to inf instead of the true answer (which is just ~x for large x).

PyTorch's autograd can usually differentiate through whatever forward code you write, but sometimes you need a custom torch.autograd.Function anyway: to hand-derive a backward pass that's cheaper or more stable than what autograd would build automatically, or to wrap a forward that isn't itself differentiable end to end. This problem is that exercise in miniature — implement both forward and backward by hand.

Source: src/27_custom_autograd_softplus.py

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

def stable_softplus(x: torch.Tensor) -> torch.Tensor: ...

>>> x = torch.tensor([-1.0, 0.0, 1.0])
>>> stable_softplus(x)
tensor([0.3133, 0.6931, 1.3133])

>>> x = torch.tensor([100.0])  # naive log(1+exp(x)) would overflow to inf
>>> stable_softplus(x)
tensor([100.])

Step-by-Step Approach

1. Recall the numerically stable softplus identity: softplus(x) = max(x, 0) + log1p(exp(-|x|)). Since exp(-|x|) is always in (0, 1] no matter how large |x| gets, this form never overflows.
2. In forward, compute the output with that identity instead of the literal log(1+exp(x)) formula, and save whatever tensor(s) backward will need via ctx.save_for_backward.
3. Work out the derivative by hand: d/dx softplus(x) = sigmoid(x). This is a standard calculus fact (differentiate log(1+exp(x)) and simplify) worth having memorized going in.
4. In backward, retrieve the saved input via ctx.saved_tensors and return grad_output * torch.sigmoid(x) — the chain rule applied to the upstream gradient.
5. Notice that sigmoid itself is already numerically stable for all inputs (it saturates toward 0 or 1 rather than overflowing), so unlike forward, backward needs no special-casing.
6. Wrap the autograd.Function in a plain function (stable_softplus) that calls .apply(x) — Function subclasses are invoked via .apply, never by calling forward directly.

The key insight is that forward-pass stability and backward-pass stability are separate problems with separate fixes: the exponential in the forward formula needed rewriting to avoid overflow, but its derivative (sigmoid) was already safe, so the custom backward ends up simpler than the custom forward.

Reference solution

import torch


class StableSoftplus(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
        # softplus(x) = log(1 + exp(x))
        # Stable identity: max(x, 0) + log1p(exp(-|x|))
        #   - for large positive x: max(x,0)=x dominates, exp(-|x|)->0, log1p(~0)=~0 -> ~x (correct)
        #   - for large negative x: max(x,0)=0, exp(-|x|)=exp(x)->0 -> ~0 (correct)
        # exp(-|x|) is always in (0, 1], so it never overflows regardless of x's magnitude.
        ctx.save_for_backward(x)
        return torch.clamp(x, min=0) + torch.log1p(torch.exp(-torch.abs(x)))

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
        # d/dx softplus(x) = sigmoid(x); sigmoid itself is numerically stable
        # for all x (it saturates to 0/1 rather than overflowing), so no
        # special-casing is needed here even though forward needed it.
        (x,) = ctx.saved_tensors
        return grad_output * torch.sigmoid(x)


def stable_softplus(x: torch.Tensor) -> torch.Tensor:
    return StableSoftplus.apply(x)


TEST_CASES = [
    {
        "name": "matches F.softplus on moderate values",
        "x": torch.tensor([-2.0, -0.5, 0.0, 0.5, 2.0]),
    },
    {
        "name": "large positive input stays finite (naive formula would overflow)",
        "x": torch.tensor([50.0, 88.0, 100.0, 500.0]),
    },
    {
        "name": "large negative input stays finite and near zero",
        "x": torch.tensor([-50.0, -100.0, -500.0]),
    },
    {
        "name": "gradient matches sigmoid(x)",
        "x": torch.tensor([-3.0, -0.2, 0.0, 0.2, 3.0]),
    },
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES[:3]):
        x = case["x"]
        print(f"Test {i}: {case['name']}, x.shape={tuple(x.shape)}")
        y = stable_softplus(x)
        print(f"  forward output: {y}")
        assert torch.isfinite(y).all(), f"non-finite output: {y}"
        expected = F.softplus(x.double()).float()
        torch.testing.assert_close(y, expected, atol=1e-4, rtol=1e-4)
        print("PASSED")

    case = TEST_CASES[3]
    x = case["x"].clone().requires_grad_(True)
    print(f"Test 3: {case['name']}, x.shape={tuple(x.shape)}")
    y = stable_softplus(x)
    y.sum().backward()
    print(f"  grad: {x.grad}")
    expected_grad = torch.sigmoid(x.detach())
    torch.testing.assert_close(x.grad, expected_grad, atol=1e-5, rtol=1e-5)
    print("PASSED")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• torch.autograd.Function — base class for defining custom forward/backward pairs; subclass it and implement forward and backward as @staticmethods.
• ctx.save_for_backward(x) — stashes tensors during forward so backward can retrieve them later without keeping the whole computation graph around.
• ctx.saved_tensors — retrieves whatever was stashed with save_for_backward, as a tuple.
• .apply(x) — the correct way to invoke a custom Function — calling StableSoftplus.forward(x) directly would skip autograd bookkeeping entirely.
• torch.log1p(y) — computes log(1+y) directly, more accurate than torch.log(1+y) when y is close to zero.
• torch.clamp(x, min=0) — equivalent to max(x, 0) elementwise, used as the stable identity's first term.
• torch.sigmoid(x) — the closed-form derivative of softplus; also happens to be its own numerically-stable building block.

How to Recognize This Pattern