Unit 5.7 · about 6 h
86 / 324Stable log-sum-exp, online reductions and mixed precision
Long reductions can fail even when each individual input is representable. Stable identities and explicit accumulator dtypes give you independent checks beneath probability losses and mixed-precision training.
- derive
- implement
You can now
- Derive log-sum-exp's shift identity, max bounds and convexity without relying on a loss implementation.
- Implement naive, balanced pairwise and Kahan reductions and distinguish an error bound from a measured error.
- Prove online mean and centered-sum invariants and test low-precision storage with higher-precision accumulation.
Covers
- Stable softmax, log-sum-exp and online reductions
- Mixed precision and accumulation choices
Study, in this order
4 steps
- Log-sum-exp as a smooth convex maximumConvex Optimization
Section 3.1.5: Log-sum-exp bounds on p. 72 and the Hessian/Cauchy-Schwarz proof on p. 74
Depth · Read only the log-sum-exp passages. Derive the max-shift identity yourself; leave gradients of a classification loss to the linked later unit.
- Summation trees and compensationThe Accuracy of Floating Point Summation
Section 2 (pp. 784 to 786), especially equations 2.1 to 2.6; Section 3's general error expression and pairwise summation (pp. 787 to 788) and Kahan's compensated summation (pp. 790 to 791, equation 3.10 on p. 791)
Depth · Skip Section 2's sorting examples and Section 3's insertion and +/- methods. Derive the basic recursive bound, then inspect the balanced-tree and compensated algorithms; no probabilistic rounding model is assigned.
- Online mean and centered second momentAccurately computing running variance
Opening cancellation example and recurrence formulas for \(M_k\), \(S_k\) and sample variance \(S_k/(k-1)\)
Depth · Stop before the C++
RunningStatclass. Prove the invariant in exact arithmetic before translating it to NumPy. A variance can be zero, and overflow can still occur. - Precision budget for inputs and accumulatorsPyTorch numerical type information
torch.finfo:eps,maxandsmallest_normalforfloat16,bfloat16,float32andfloat64
Focus on
- For finite nonempty \(x\), \(\operatorname{LSE}(x)=m+\log(\sum\exp(x_i-m))\), \(m=\max x\). Every exponential is \(\le 1\) and at least one is 1, giving \(m\le\operatorname{LSE}\le m+\log(n)\). The Hessian \(\operatorname{diag}(p)-pp^\top\) is PSD by a weighted-variance argument.
- With unit roundoff \(u\), no overflow or underflow and \((n-1)u<1\), recursive summation has absolute error \(\le\gamma_{n-1}\sum|x_i|\), \(\gamma_k=ku/(1-ku)\). A balanced tree shortens rounding paths; it does not remove conditioning from cancellation.
- Online mean uses \(\delta=x_n-\mathit{mean}_{\text{old}}\) and \(\mathit{mean}_{\text{new}}=\mathit{mean}_{\text{old}}+\delta/n\). The centered sum updates \(M2_{\text{new}}=M2_{\text{old}}+\delta(x_n-\mathit{mean}_{\text{new}})\). Check \(M2_n=\sum(x_i-\mathit{mean}_n)^2\) with base case \(n=1\), and divide by \(n-1\) only for sample variance with \(n\ge 2\).
- Higher-precision accumulation can prevent repeated rounding of the running total. It cannot restore detail lost when input values were quantized; distinguish operand, product and accumulator precision.
Reproduce closed-book
[1,2,3]: mean=2, M2=2, sample variance=1. In NumPy float64 use x=[1000,1001,999]; the shifted formula gives 1001.4076059644, allowing 1e-10 absolute error for this displayed rounding. Compare the unrounded result with scipy.special.logsumexp on SciPy 1.18.1: measured absolute difference is 0; require \(\le 10^{-12}\), not exact equality. Implement a left-to-right sum with float64 accumulator zero, a balanced sum splitting each array at len(a)//2 recursively (empty sum zero, singleton returns its value), and Kahan with s=c=0; for each v use y=v-c, t=s+y, c=(t-s)-y, s=t. On np.concatenate((np.array([1e16],dtype=np.float64),np.ones(1024,dtype=np.float64),np.array([-1e16],dtype=np.float64))), the outputs are 0, 1022 and 1024. Against math.fsum's 1024 the measured absolute errors are 1024, 2 and 0; allow 1e-12 for these displayed checks. Run the online recurrence on [1e8+1,1e8+2,1e8+3] in that order in float64: sample variance is 1, while \((\operatorname{sum}(x^2)-3\operatorname{mean}(x)^2)/2\) gives 0. The absolute difference from np.var(x,ddof=1) for the online result is 0, allowed \(\le 10^{-12}\). Finally sum 4096 float16 ones left-to-right, explicitly casting every updated accumulator to np.float16: the total stalls at 2048. np.sum on those same inputs with dtype=np.float32 gives 4096; measured absolute errors against the integer count are 2048 and 0. No random draws are needed. These examples do not establish that Kahan or an online variance algorithm wins for every input.Traps
- Using
np.sumas a claimed naive left-to-right baseline. Its reduction strategy depends onaxisand layout; implement the order you claim to test. - Rearranging Kahan's compensation into algebraically equivalent operations. Reassociation changes the rounding errors the correction is meant to retain.
- Ignoring empty inputs, infinities or the \(n=1\) sample-variance denominator. The finite nonempty \(\operatorname{LSE}\) contract and \(n\ge 2\) variance contract need explicit checks.
Practice
The section mini-project covers this: derive and implement a finite-input log-partition primitive, three explicit summation orders and online mean/variance; test cancellation, large offsets, low-precision storage and accumulator choices against independent oracles.
Self-test
-
Why is the LSE Hessian positive semidefinite but not positive definite?
For finite \(x\), \(p_i=\exp(x_i)/\sum\exp(x_j)\) are positive and sum to one. \(v^\top(\operatorname{diag}(p)-pp^\top)v=\sum p_i(v_i-\sum p_jv_j)^2\ge 0\). Positivity makes this a weighted variance; constant \(v\) gives zero, so the all-ones direction is a null direction. -
Where does the online centered-sum proof use its assumptions?
At \(n=1\), \(\mathit{mean}=x_1\) and \(M2=0\). For \(n\ge 2\), the old centered deviations sum to zero, so recentering the first \(n-1\) values adds \((n-1)(\delta/n)^2\) with no cross term. The new deviation is \(\delta(n-1)/n\); adding its square gives \(\delta^2(n-1)/n=\delta(x_n-\mathit{mean}_{\text{new}})\). Division uses \(n>0\); the algebra assumes exact arithmetic, not an exact floating-point invariant. -
Why can a small absolute summation error be a large relative one?
Cancellation can make \(|\sum x_i|\) tiny while \(\sum|x_i|\) is large. The relative error bound contains their ratio, the summation condition number. At a zero exact sum use an absolute criterion instead. -
Does
float32accumulation makefloat16inputs equivalent to originalfloat32inputs?No. It improves the accumulation of already rounded inputs. Any distinction erased by conversion tofloat16is still gone, even if every later addition were exact.
Completion criterion
Stop when
Your reduction and online-state tests reproduce the stated absolute errors, and you can prove the \(\operatorname{LSE}\) bounds and centered-sum invariant while explaining which precision losses cannot be repaired.
Mark passed, then open 5.8 Only mark this when you can do it without notes.Leads to
- 1.2 Stable softmax, cross-entropy and loss reductions Apply the log-partition identity to stable probability losses.
- 8.4 Precision formats, autocast and gradient scaling Apply the distinction between storage precision and accumulation precision.
- 4.1 Likelihood and perplexity, and their limits Use stable log probabilities and controlled reductions when aggregating likelihoods.
Resources for this unit
- Convex Optimization Read Chapters 1-5 and 9 with the course; use Chapter 10 only for equality-constrained Newton steps.
- PyTorch numerical type information Read
torch.finfo'seps,maxandsmallest_normalfields and its subnormal note; inspectfloat64,float32,float16andbfloat16on CPU. - The Accuracy of Floating Point Summation Read Section 2's recursive-sum error bound and Section 3's pairwise and compensated summation algorithms. Leave the statistical error models and historical machine details aside.
- Accurately computing running variance Read the Welford mean and centered-sum recurrences. Prove the centered-sum invariant and translate the recurrence to NumPy; skip the C++ class.
Also in this section
- Stanford EE364a - Convex Optimization Follow lectures 1-9 and 15-16; use lecture 17 only for KKT systems. Leave full interior-point treatment for later.
- D2L - Optimization Algorithms Use the convexity, GD, SGD, minibatch, momentum, AdaGrad, RMSProp, Adam and scheduler sections for ML derivations and code.
- Fundamentals of Numerical Computation Study §§1.1-1.4, 2.7-2.8, 3.2-3.4 and 4.2-4.3 for floating point, conditioning, stability, least squares and iteration error, in the Python edition.
- CVXPY basic and DCP tutorials Use as the solver oracle after deriving a constrained problem and its KKT conditions by hand.
- Proximal Algorithms Read Sections 1.1, 4.2 and the elementwise soft-thresholding formula in 6.5.2, then the lasso formulation on p. 196. Skip 7.1.1's penalty-scaling and Gram-rewrite errors on p. 197; derive the update from 4.2 and 6.5.2.
- PyTorch optimizer and gradient-clipping contracts Use AdamW's algorithm and
clip_grad_norm_'s parameters to locate weight decay and clipping precisely. Inspectstate_dictfor optimizer restart; no network training is assigned. - SciPy minimize: BFGS Read the BFGS method description and
gtol,maxiter,c1andc2options. Supply an analytic gradient and compare residuals and evaluation counts, not just iterations.