A common task in machine learning is to estimate or optimize “log-sum-exp” functions with (potentially continuously) many terms such as $$ log Big( int_{\mathcal{X}} e^{v(x)} dq(x) \Big),$$ where \(v: \mathcal{X} \to \mathbb{R}\) is some potential function, and \(q\) is a probability distribution on the set \(\mathcal{X}\). This has many applications throughout data science, often through the normalization of probabilistic models, but also as a smooth approximation to the maximum, in transformers through its derivatives, or in reinforcement learning when using entropy regularization [19]. Sometimes the set \(\mathcal{X}\) is finite (potentially big) and the integral can be done by explicit summing, but often an exact computation is infeasible, and sampling from the probability distribution \(q\) is used instead.

The key difficulty comes from the variance of such estimates, in particular when \(v\) takes large values. In the simplest example, for \(z_1,\dots,z_n \in \mathbb{R}\) independent and normally distributed with mean \(\mu\) and variance \(\sigma^2\), the relative squared error for estimating \(\mathbb{E}[e^z]\) is $$\frac{ {\rm var}\big( \frac{1}{n} \sum_{i=1}^n e^{z_i} \big) }{( \mathbb{E}[ e^{z} ])^2} = \frac{1}{n} \frac{ {\rm var}(e^z) }{( \mathbb{E}[ e^{z} ])^2} = \frac{ e^{\sigma^2}-1}{n}.$$ It converges to zero when \(n\) grows (as can be expected from the law of large numbers), but explodes exponentially when \(\sigma\) grows. Even taking the logarithm does not change the exploding variance, that is, ${\rm var}\big( \log \big( \frac{1}{n} \sum_{i=1}^n e^{z_i} \big)\big)$ also can be shown to grow asymptotically similarly in \frac{ e^{\sigma^2}-1}{n} (when \(n\) is large, as can be obtained from the delta method).

While difficult to estimate, the log-sum-exp function comes with many nice properties (and that’s why people love it); I particularly like the fact that (1) it is a smooth approximation to the maximum (see, e.g., this earlier post), and (2) it is a way to normalize probabilistic models that is adapted to maximum likelihood estimation, in particular in hierarchical probabilistic models, where (conditional) independence assumptions lead to separability of associated loss functions (as thoroughly used in probabilistic graphical models). The main question I try to answer in this post is: Can we keep the advantages of optimizing log-sum-exp functions while being less exposed to their computational / statistical disadvantages? At the other end of the spectrum sits least-squares regression, with essentially the exact opposite features: On the positive side, we obtain computational and statistical simplicity in various forms, e.g., it leads to closed-form estimation for linear models through linear algebra, it is based on computing moments with fixed controlled variance, and it leads to sharp analyses in various setups (acceleration, stochastic gradient descent, etc.). See, e.g., this post on acceleration, and this one on averaging.

On the negative side, using least-squares regression for all prediction problems, particularly with discrete outputs, creates some artefacts. The traditional example is classification with Gaussian class-conditional data (with identical covariance matrices), where least-squares on the one-hot encoded outputs has problems, such as “masking” (see [13, Section 2.4] and the example below), or high approximation error compared with using multinomial logistic regression (a.k.a. softmax regression), because then the log conditional probabilities are affine. Can we reconcile them? In other words, is least-squares really all I need? (my colleagues sometimes mock me for my love of least-squares).

Note that there is another (classic) attempt at seeing the world through least-squares: doing it in series through Newton’s method, leading in this context to iteratively reweighted least-squares, but this is for computations only, with no statistical improvement. What we are aiming at is stron...