13. Numerically Stable Log-Space Cumulative Decay

Problem

Linear-recurrence and SSM layers compute a running product of per-step decay factors — the cumulative "how much of the original signal survives by timestep t" term that shows up in a closed-form unroll of a linear recurrence. Each decay factor sits in (0, 1), and over a long sequence (hundreds to thousands of steps, exactly the range a real-time voice session runs for) the literal product shrinks geometrically and underflows to exactly 0.0 in float32 long before the sequence ends — at which point any downstream computation that depends on it silently breaks.

The fix is the same log-sum-exp-family trick used for numerically stable softmax: never form the product directly. Take the log of each decay factor, accumulate with a running sum instead, and only exponentiate back out of log-space when a value is actually needed.

Source: src/13_log_space_cumulative_decay.py

def log_cumulative_decay(a: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: ...

>>> a = torch.full((1, 600), 0.8)
>>> torch.cumprod(a, dim=-1)[:, -1].item()   # naive product: underflows
0.0
>>> log_cumulative_decay(a)[:, -1].item()    # log-space: stays finite
-133.88...

Step-by-Step Approach

1. Clamp the input before taking a log: a.clamp_min(eps) — an exact 0.0 decay is a legal input (a fully closed gate), but log(0) = -inf would poison every later position once summed.
2. Take the elementwise log: log_a = torch.log(a.clamp_min(eps)).
3. Recognize the identity log(prod(a)) == sum(log(a)) — a cumulative product in linear space is exactly a cumulative sum in log space.
4. Apply torch.cumsum(log_a, dim=-1) to get the running log-cumulative-product at every position in one vectorized call.
5. Leave the result in log-space rather than exponentiating inside this function — callers that need the literal product back can exponentiate on demand, and a caller that only ever needs differences of cumulative log-values never needs to exponentiate at all.
6. Verify the stability claim directly: confirm torch.cumprod on the same input actually underflows to 0.0 for a long enough sequence, while this function's output stays finite.

The key insight is that summation degrades far more gracefully than repeated multiplication when the running value would otherwise leave the representable float range — the log-space value can go arbitrarily negative and remain a perfectly ordinary finite float, whereas the linear-space product hits a hard floor at the smallest representable positive number and clips to exactly zero.

Reference solution

import torch


def log_cumulative_decay(a: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    # clamp before log: an exact 0.0 decay is a legal input (fully closed
    # gate) but log(0) = -inf would poison every later position via cumsum
    log_a = torch.log(a.clamp_min(eps))
    # sum in log-space == product in linear space, and summing never
    # underflows the way repeated multiplication of small floats does
    return torch.cumsum(log_a, dim=-1)

Key Functions & Tricks

• tensor.clamp_min(eps) — floors an exact 0.0 before log, trading a tiny, harmless bias for avoiding -inf/nan propagation.
• torch.log — the entry point into log-space; everything downstream should stay additive rather than multiplicative.
• torch.cumsum(..., dim=-1) — the numerically stable stand-in for torch.cumprod once the input has been log-transformed.
• torch.cumprod (used only to demonstrate the failure mode) — underflows silently to exactly 0.0, with no warning or exception, once the true product falls below float32's smallest representable positive value.
• Staying in log-space as long as possible before a final, isolated torch.exp — the same principle behind why softmax is implemented via max-subtraction and log-sum-exp rather than naive exp then normalize.

How to Recognize This Pattern