Numerically Stable Log-Sum-Exp and Softmax
Naive computation of $\log \sum_i e^{x_i}$ and $\text{softmax}(x)_j = e^{x_j} / \sum_i e^{x_i}$ blows up or underflows when entries of $x$ are large in magnitude. Implement numerically stable versions of both, then build cross-entropy loss and its gradient on top -- all bundled into a single function.
Implement stable_ops(logits, target) where logits is a list of n floats and target is the integer index of the correct class. Return ONE flat list of length 2 + 2*n:
[ lse ]-- $\log \sum_i e^{x_i}$, computed via the max-shift trick $m + \log \sum_i e^{x_i - m}$ with $m = \max(x)$.softmax(logits)-- $n$ values, $e^{x_j - m} / \sum_i e^{x_i - m}$.[ ce_loss ]-- the cross-entropy loss $-\log s_{\text{target}}$, computed stably aslse - logits[target].cross_entropy_gradient-- $n$ values, $\partial L / \partial x_j = s_j - \mathbf{1}[j = \text{target}]$.
All outputs must be free of NaN and Inf for valid inputs, even when entries are on the order of $\pm 10^3$.
Constraints: - Input is a 1-D list of floats with $1 \le n \le 10^5$. - Do not use any library softmax / logsumexp.
Example
stable_ops([1.0, 2.0, 3.0], 0) -> [3.4076059644443806, 0.09003057317038046, 0.24472847105479764, 0.6652409557748218, 2.4076059644443806, -0.9099694268296196, 0.24472847105479764, 0.6652409557748218]
Here $n=3$, so the returned list has length $2 + 2\cdot3 = 8$: the log-sum-exp (3.4076...), then the 3 softmax probabilities, then the cross-entropy loss for target=0 (2.4076...), then the 3 gradient components (note the target's gradient is softmax[0] - 1 = -0.9099...).
Hints
- The problem with naive exponentiation is that $e^{x}$ overflows for $x > 709$ in float64. Think about what algebraic identity lets you factor out a constant from inside $\log \sum e^{x_i}$.
- Use the identity $\log \sum e^{x_i} = m + \log \sum e^{x_i - m}$ with $m = \max(x)$. This ensures the largest exponent is exactly $e^0 = 1$.
- For cross-entropy, don't compute softmax and then take $\log$ -- instead compute $-x_{\text{target}} + \text{log\_sum\_exp}(x)$ directly. This avoids catastrophic cancellation when the target probability is close to 1.
Worked Solution
How to Think About It: The core issue is that $e^{x_i}$ overflows for $x_i > 709$ (in float64) and underflows to zero for very negative $x_i$. If every entry overflows, $\log(\texttt{inf})$ is still $\texttt{inf}$ -- and if every entry underflows, $\log(0)$ is $-\texttt{inf}$. But the actual answer is finite. The trick every production ML framework uses: factor out the maximum. Since $\log \sum e^{x_i} = m + \log \sum e^{x_i - m}$ for any constant $m$, choosing $m = \max_i x_i$ guarantees the largest exponent is $e^0 = 1$ and every other exponent is $\le 1$. No overflow, and the terms that underflow to zero are genuinely negligible.
Contract: The judge calls a single bundled function stable_ops(logits, target) and expects ONE flat list of length 2 + 2*n:
[ lse ]-- $\log \sum_i e^{x_i}$ (1 value)softmax(logits)-- $n$ values, $e^{x_j - m} / \sum_i e^{x_i - m}$[ ce_loss ]-- $-\log s_{\text{target}}$ (1 value)cross_entropy_gradient-- $n$ values, $s_j - \mathbf{1}[j = \text{target}]$
Key precision detail: Compute the cross-entropy loss with the stable *log-softmax* formulation ce_loss = lse - logits[target], NOT -log(softmax[target]). The latter loses precision for small probabilities and produces a slightly different last-ULP result. The gradient is simply softmax minus a one-hot vector at the target index.
Code:
```python import math
def stable_ops(logits, target): # Numerically stable log-sum-exp / softmax / cross-entropy bundle. n = len(logits) m = max(logits) exps = [math.exp(x - m) for x in logits] s = sum(exps) lse = m + math.log(s) softmax = [e / s for e in exps] ce_loss = lse - logits[target] grad = [softmax[j] - (1.0 if j == target else 0.0) for j in range(n)] return [lse] + softmax + [ce_loss] + grad ```
Intuition
This problem tests whether you understand the single most important trick in numerical computing with exponentials. Every neural network framework, every softmax classifier, every Boltzmann sampler uses the same idea: shift by the maximum before exponentiating. The reason it works is that $\log \sum e^{x_i}$ is translation-invariant up to an additive constant -- adding $c$ to every $x_i$ just adds $c$ to the result. By choosing $c = -\max(x)$, you pin the largest exponent to 1 and push everything else toward zero, which is exactly the safe zone for floating point.
The cross-entropy piece illustrates a subtler point: even after you have a stable softmax, chaining $-\log(\text{softmax})$ can lose precision when the probability is very close to 1, because $\log(1 - \epsilon) \approx -\epsilon$ requires full precision in $\epsilon$. The log-softmax formulation $-x_t + \text{LSE}(x)$ sidesteps this entirely. In production ML code, log_softmax is always a single fused operation for exactly this reason. If you see someone computing np.log(softmax(x)) in a training loop, that's a code smell.