3.5 Concise Implementation of Softmax Regression

Just as high-level deep learning frameworks made it easier to implement linear regression (see Section 2.5), they are similarly convenient here.

from d2l import torch as d2l
import torch
from torch import nn
from torch.nn import functional as F

from d2l import tensorflow as d2l
import tensorflow as tf

from d2l import jax as d2l
from flax import nnx
import optax

from d2l import mxnet as d2l
from mxnet import gluon, init, npx
from mxnet.gluon import nn
npx.set_np()

3.5.1 Defining the Model

As in Section 2.5, we construct our fully connected layer using the built-in layer. The built-in __call__ method then invokes forward whenever we need to apply the network to some input.

We use a Flatten layer to convert the fourth-order tensor X to second order by keeping the dimensionality along the first axis unchanged.

We use a Flatten layer to convert the fourth-order tensor X by keeping the dimension along the first axis unchanged.

NNX layers are ordinary stateful Python objects. Since Fashion-MNIST images have \(28\times28=784\) features, we declare that input width when constructing the linear layer and flatten each minibatch in forward.

Even though the input X is a fourth-order tensor, the built-in Dense layer will automatically convert X into a second-order tensor by keeping the dimensionality along the first axis unchanged.

class SoftmaxRegression(d2l.Classifier):
    """The softmax regression model."""
    def __init__(self, num_outputs, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential(nn.Flatten(),
                                 nn.LazyLinear(num_outputs))

    def forward(self, X):
        return self.net(X)

class SoftmaxRegression(d2l.Classifier):
    """The softmax regression model."""
    def __init__(self, num_outputs, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = tf.keras.models.Sequential()
        self.net.add(tf.keras.layers.Flatten())
        self.net.add(tf.keras.layers.Dense(num_outputs))

    def forward(self, X):
        return self.net(X)

class SoftmaxRegression(d2l.Classifier):
    def __init__(self, num_outputs, lr, num_inputs=784, rngs=None):
        super().__init__()
        self.save_hyperparameters(ignore=['rngs'])
        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
        self.net = nnx.Linear(num_inputs, num_outputs, rngs=rngs)

    def forward(self, X):
        X = X.reshape((X.shape[0], -1))  # Flatten
        return self.net(X)

class SoftmaxRegression(d2l.Classifier):
    """The softmax regression model."""
    def __init__(self, num_outputs, lr):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Dense(num_outputs)
        self.net.initialize()
    def forward(self, X):
        return self.net(X)

3.5.2 Softmax Revisited

In Section 3.4 we computed the softmax explicitly and then took its logarithm inside the cross-entropy loss. To keep that version usable we had to clamp the probabilities away from zero, a band-aid that prevents \(\log 0\) but still forms the overflow-prone softmax first and silently kills the gradient on any clamped entry. Here we remove the problem at its source rather than patch its symptom.

Recall that the softmax function computes probabilities via \(\hat y_j = \frac{\exp(o_j)}{\sum_k \exp(o_k)}\). If some of the \(o_k\) are very large, i.e., very positive, then \(\exp(o_k)\) might be larger than the largest number we can have for certain data types. This is called overflow. Conversely, a strongly negative \(o_k\) makes \(\exp(o_k)\) underflow to \(0\). Single-precision floats span roughly \(10^{-38}\) to \(10^{38}\), so \(\exp\) overflows once an argument exceeds about \(88\) and gradually underflows past about \(-88\) (entering the subnormal range), reaching exactly \(0\) only near \(-104\). A single large positive logit therefore overflows the numerator, while strongly negative logits underflow individual terms to \(0\): harmless in the sum, but fatal once we take a logarithm. A way round this problem is to subtract \(\bar{o} \stackrel{\textrm{def}}{=} \max_k o_k\) from all entries:

\[ \hat y_j = \frac{\exp o_j}{\sum_k \exp o_k} = \frac{\exp(o_j - \bar{o}) \exp \bar{o}}{\sum_k \exp (o_k - \bar{o}) \exp \bar{o}} = \frac{\exp(o_j - \bar{o})}{\sum_k \exp (o_k - \bar{o})}. \]

By construction we know that \(o_j - \bar{o} \leq 0\) for all \(j\). As such, for a \(q\)-class classification problem, the denominator is contained in the interval \([1, q]\). Moreover, the numerator never exceeds \(1\), thus preventing numerical overflow. Numerical underflow only occurs when \(\exp(o_j - \bar{o})\) numerically evaluates as \(0\). Nonetheless, a few steps down the road we might find ourselves in trouble when we want to compute \(\log \hat{y}_j\) as \(\log 0\). In particular, in backpropagation, we might find ourselves faced with a screenful of the dreaded NaN (Not a Number) results.

Fortunately, we are saved by the fact that even though we are computing exponential functions, we ultimately intend to take their log (when calculating the cross-entropy loss). By combining softmax and cross-entropy, we can escape the numerical stability issues altogether. We have:

\[ \log \hat{y}_j = \log \frac{\exp(o_j - \bar{o})}{\sum_k \exp (o_k - \bar{o})} = o_j - \bar{o} - \log \sum_k \exp (o_k - \bar{o}). \]

This avoids both overflow and underflow. We are not quite done, though, because the object we actually need is not \(\log \hat y_j\) but the loss. For an example with true class \(y\), the cross-entropy loss is \(\ell = -\log \hat y_y\). Substituting the stabilized expression above turns the loss into a function of the logits alone:

\[ \ell(y, \mathbf{o}) = \log \sum_{k} \exp(o_k) - o_y = \underbrace{\bar{o} + \log \sum_{k} \exp(o_k - \bar{o})}_{\textrm{numerically stable}} - o_y, \qquad \bar{o} = \max_k o_k. \]

The first term, \(\log \sum_k \exp(o_k)\), is the log-sum-exp function, a smooth upper bound on \(\max_k o_k\) (you proved this, including the fact that the gap never exceeds \(\log q\), in Section 3.1, exercise 6); the second equality is the only safe way to evaluate it, since every exponent \(o_k - \bar{o} \leq 0\). This is precisely what the built-in cross-entropy loss computes when handed raw logits: it never forms the softmax probabilities at all, so neither \(\exp\) of a large number nor \(\log\) of a zero ever occurs. Because the fused loss differentiates this exact expression, its gradient is exactly \(\partial_{o_j}\ell = \mathrm{softmax}(\mathbf{o})_j - y_j\) derived in Section 3.1.3.2, with no clamp to perturb it. We keep the explicit softmax of Section 3.4 only for reading off predicted probabilities at inference time; for the loss we pass logits and let the fused operation do the rest.

For two classes with logits \((x, 0)\) the loss’s first term is \(\mathrm{lse}(x, 0) = \log(1 + e^x)\), which hugs \(\max(x, 0)\) from above, with the gap largest at the tie \(x = 0\), where it equals \(\log 2 \approx 0.69\), our bound \(\log q\) for \(q = 2\):

x = d2l.arange(-4.0, 4.0, 0.01)
lse, mx = d2l.log(1 + d2l.exp(x)), (x + d2l.abs(x)) / 2
d2l.plot(x, [lse, mx, lse - mx], 'x', legend=['lse(x, 0)', 'max(x, 0)', 'gap'])

@d2l.add_to_class(d2l.Classifier)
def loss(self, Y_hat, Y, averaged=True):
    Y_hat = d2l.reshape(Y_hat, (-1, Y_hat.shape[-1]))
    Y = d2l.reshape(Y, (-1,))
    return F.cross_entropy(
        Y_hat, Y, reduction='mean' if averaged else 'none')

@d2l.add_to_class(d2l.Classifier)
def loss(self, Y_hat, Y, averaged=True):
    Y_hat = d2l.reshape(Y_hat, (-1, Y_hat.shape[-1]))
    Y = d2l.reshape(Y, (-1,))
    reduction = (tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE
                 if averaged else tf.keras.losses.Reduction.NONE)
    fn = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction=reduction)
    return fn(Y, Y_hat)

@d2l.add_to_class(d2l.Classifier)
def loss(self, Y_hat, Y, averaged=True):
    Y_hat = d2l.reshape(Y_hat, (-1, Y_hat.shape[-1]))
    Y = d2l.reshape(Y, (-1,))
    fn = optax.softmax_cross_entropy_with_integer_labels
    return fn(Y_hat, Y).mean() if averaged else fn(Y_hat, Y)

@d2l.add_to_class(d2l.Classifier)
def loss(self, Y_hat, Y, averaged=True):
    Y_hat = d2l.reshape(Y_hat, (-1, Y_hat.shape[-1]))
    Y = d2l.reshape(Y, (-1,))
    fn = gluon.loss.SoftmaxCrossEntropyLoss()
    l = fn(Y_hat, Y)
    return l.mean() if averaged else l

The built-in fused loss (named differently in each library) takes logits, not probabilities: passing softmax outputs would apply the softmax twice. Correspondingly, the model’s forward returns raw logits and contains no softmax: the loss owns that step. We defined this loss on the base Classifier (note the #@save), so every classifier in the rest of the book inherits the numerically stable version.

3.5.3 Training

Next we train our model. We use Fashion-MNIST images, flattened to 784-dimensional feature vectors.

data = d2l.FashionMNIST(batch_size=256)
model = SoftmaxRegression(num_outputs=10, lr=0.1)
trainer = d2l.Trainer(max_epochs=10)
trainer.fit(model, data)

As before, training converges to about 83–84% validation accuracy, the same solution as the from-scratch version of Section 3.4 (read it off the validation curve), now in far fewer lines of code.

3.5.4 Summary

High-level APIs are very convenient at hiding from their user potentially dangerous aspects, such as numerical stability. Moreover, they allow users to design models concisely with very few lines of code. This is both a blessing and a curse. The obvious benefit is that it makes things highly accessible, even to engineers who never took a single class of statistics in their life (in fact, they are part of the target audience of the book). But hiding the sharp edges also comes with a price: a disincentive to add new and different components on your own, since there is little muscle memory for doing it. Moreover, it makes it more difficult to fix things whenever the protective padding of a framework fails to cover all the corner cases entirely. Again, this is due to lack of familiarity.

As such, we strongly urge you to review both the bare bones and the concise versions of many of the implementations that follow. While we emphasize ease of understanding, the implementations are nonetheless usually quite performant (convolutions are the big exception here). It is our intention to allow you to build on these when you invent something new that no framework can give you.

The fused cross-entropy loss in this section is the implementation to reach for: it computes the log-sum-exp directly and never materializes an unstable softmax.

3.5.5 Exercises

Deep learning uses many different number formats, including FP64 double precision (used extremely rarely), FP32 single precision, BFLOAT16 (good for compressed representations), FP16 (very unstable), TF32 (a new format from NVIDIA), and INT8. Compute the smallest and largest argument of the exponential function for which the result does not lead to numerical underflow or overflow.
INT8 is a very limited format consisting of integers in \([-128, 127]\) (or \([0, 255]\) for the unsigned variant). How could you extend its dynamic range without using more bits? Do standard multiplication and addition still work?
Take the from-scratch softmax of Section 3.4 and feed it the logits \(\mathbf{o} = (1000, 0, 0)\). What do you get, and why? Now compute the loss for the same logits with the framework’s cross_entropy, passing the logits directly. Why is it finite? Verify that on benign logits, e.g., \(\mathbf{o} = (2, 1, 0)\), the two routes agree to floating-point precision.
Show, using the identity \(\ell = \log\sum_k \exp(o_k) - o_y\), that adding the same constant \(c\) to every logit leaves the loss unchanged. Why does this make subtracting \(\bar{o} = \max_k o_k\) a free and safe choice?

Discussions