The softmax function is an crucial part of modern deep learning. It is especially important for the attention mechanism used in Transformer models. In self-attention, softmax converts raw similarity scores into a probability distribution. This distribution tells the model how much focus to put on each token in a sequence. We usually define scaled dot product attention like this:

When we calculate the dot product of queries and keys, we get an $N \times N$ a sequence of length $N$. This creates a major problem because the time and memory needed grow at a quadratic rate $O(N^2)$:

For very long sequences, a single GPU cannot store the entire attention matrix. The Ring Attention paper highlights just how extreme this memory demand can be:

To put the memory demand in perspective, even when dealing with a batch size of 1, processing 100 million tokens requires over 1000 GB of memory for a modest model with a hidden size of 1024.

To solve this, we can split the attention matrix into smaller blocks and run calculations on those blocks individually :

Processing in blocks is a great idea, but the standard softmax function makes it difficult because it needs the sum of every element in the sequence to calculate the denominator.

Naive Softmax

The standard softmax formula is defined as:

We can write a simple version of this in PyTorch:

import torch

def naive_softmax(x: torch.Tensor) -> torch.Tensor:
    return x.exp() / x.exp().sum()

If we compare this naive function to the official PyTorch version, the results look identical for small numbers:

x         = torch.randn(8)
reference = torch.softmax(x, dim=-1)
naive     = naive_softmax(x)

print(f"Torch: {reference}")
print(f"Naive: {naive}")
print(f"allclose: {torch.allclose(reference, naive)}")

# Output
Torch: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
Naive: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
allclose: True

However, this naive version is not stable. When the input values are large, the exponential function grows too fast and causes a numerical overflow. This results in “nan” errors in naive softmax:

naive_softmax(x * 1000)

# Output
tensor([0., 0., nan, nan, 0., 0., 0., 0.])

Safe Softmax

Softmax has the shift invariance property. This means that if we add or subtract the same constant from every input, the output stays the same. We can use this to keep our numbers small and manageable. By subtracting the maximum value from the entire vector, we ensure that the largest exponent is exactly zero.

By doing this, every exponent is less than or equal to zero, which completely prevents overflow. Here is how we implement it:

def safe_softmax(x: torch.Tensor) -> torch.Tensor:
    m = x.max()
    return (x - m).exp() / (x - m).exp().sum()

reference = torch.softmax(x, dim=-1)
safe      = safe_softmax(x)

print(f"Torch: {reference}")
print(f"Safe: {safe}")
print(f"allclose: {torch.allclose(reference, safe)}")

# Output
Torch: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
Safe: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
allclose: True

This version is also stable for large inputs:

safe_softmax(x * 1000)

# Output
tensor([0., 0., 1., 0., 0., 0., 0., 0.])

Online Softmax

In the attention mechanism, we apply softmax row by row. To calculate the denominator for a single query, we need to access every single key vector in that row. This global dependency is exactly what makes it hard to process the matrix in chunks:

Online softmax allows us to calculate the result step by step. If we track the local maximum $m$ and the local sum $d$ for each small block, we can update them as we move through the sequence. When we receive a new element $x_S$, we update our statistics like this:

This logic allows us to merge several local results into one global distribution. If we have $B$ blocks, each with its own max value $m_i$ and sum $d_i$, we can combine them by rescaling the local probabilities. This approach perfectly matches the standard softmax output:

We first need a modified version of safe softmax that returns the local maximum and the denominator sum:

def safe_softmax2(x: torch.Tensor) -> torch.Tensor:
    m = x.max()
    a = (x - m).exp() # subtract maximum value
    d = a.sum() # normalization factor
    return a / d, m, d

Next, we create the online softmax function. It takes these local blocks and merges them by finding a global maximum and adjusting the sums accordingly:

def online_softmax(
    *blocks: Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
) -> torch.Tensor:
    p_blocks, m_blocks, d_blocks = zip(*blocks)
    m_max = torch.stack(m_blocks).max() # get global maximum
    d_total = sum( # compute global normalizer
        d * torch.exp(m - m_max)
        for d, m in zip(d_blocks, m_blocks)
    )
    return torch.cat(
        [
            p * d * torch.exp(m - m_max) / d_total
            for p, m, d in zip(p_blocks, m_blocks, d_blocks)
        ]
    )

Finally, we can verify that splitting the input into four chunks and processing them with our online softmax function gives the correct result:

blocks = []

for chunk in list(x.chunk(4)):
    blocks.append(safe_softmax2(chunk)) # compute p, m, d values per chunk

online = online_softmax(*blocks)

print(f"Torch: {reference}")
print(f"Online: {online}")
print(f"allclose: {torch.allclose(reference, online)}")

# Output
Torch: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
Online: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
allclose: True

Optimized Online Softmax

We can make this even more efficient . Instead of tracking the max value and the sum separately, we can combine them into a single value called the logarithm of the sum of exponentiated values or Log Sum Exp (LSE). We define this as $\text{lse}_i = m_i + \log(d_i)$. This simplifies our global formula significantly:

To use this approach, we slightly modify the safe_softmax function so that it also returns the lse value instead of m and d separately:

def safe_softmax3(x: torch.Tensor) -> torch.Tensor:
    m = x.max()
    a = (x - m).exp()
    b = a.sum()
    lse = m + torch.log(b)
    return a / b, lse # return local softmax and lse outputs

The new online softmax function now uses these LSE weights to rescale the local probabilities before joining them together:

def online_softmax2(
    p_blocks = List[torch.Tensor],
    lse_blocks = List[torch.Tensor]
) -> torch.Tensor:
    weights = torch.exp(torch.stack(lse_blocks))         
    return torch.cat([
            p * w
            for p, w in zip(p_blocks, weights)
        ]) / weights.sum()

This version still gives us the exact same output but with a cleaner mathematical structure:

p_blocks = []
lse_blocks = []

for chunk in list(x.chunk(4)):
    p, lse = safe_softmax3(chunk) # compute p, lse values per chunk
    p_blocks.append(p)
    lse_blocks.append(lse)

online2 = online_softmax2(p_blocks, lse_blocks)
print(f"Torch: {reference}")
print(f"Online2: {online2}")
print(f"allclose: {torch.allclose(reference, online2)}")

# Output
Torch: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
Online2: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
allclose: True

Practical Implementation

Although the previous generalized version is useful for computing global softmax from multiple local softmax outputs, Ring/Striped attention uses the next key/value pairs to update the softmax output at each step. As the data moves between GPUs, the hardware calculates attention for the local block and then passes it along:

Because we are only dealing with two blocks, we can simplify the formula even further. Using the property of $\frac{1}{1 + B/A} = \frac{A}{A + B}$: we can reduce the number of times we use the exponential function:

Since our goal is to iteratively compute updated softmax value for given new $p$ and $lse$, we also need to compute $lse_{new}$. Since $lse_{i} = m_i + \log(d_i)$, we can find $lse_{new}$ as follows:

The implementation of this iterative update looks like this:

def online_softmax3(
    p0: torch.Tensor,
    lse0: torch.Tensor,
    p1: torch.Tensor,
    lse1: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    out0 = p0 / (1 + torch.exp(lse1 - lse0))
    out1 = p1 / (1 + torch.exp(lse0 - lse1))
    new_lse = lse0 + torch.log(1 + torch.exp(lse1 - lse0))
    return torch.cat([out0, out1]), new_lse

We can now compute the global softmax result by looping through the chunks one by one:

p, lse = p_blocks[0], lse_blocks[0]
for p_new, lse_new in zip(p_blocks[1:], lse_blocks[1:]):
    p, lse = online_softmax3(p, lse, p_new, lse_new)

print(f"Torch: {reference}")
print(f"Online3: {p}")
print(f"allclose: {torch.allclose(reference, p)}")

# Output
Torch: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
Online3: tensor([0.1085, 0.0730, 0.3312, 0.2468, 0.1182, 0.0637, 0.0182, 0.0404])
allclose: True

If you look at the source code for Ring Flash Attention in PyTorch, you will see a very similar pattern . They use a function to update the output and the LSE values iteratively, which allows them to handle massive sequences across multiple GPUs:

def _update_out_and_lse(
    out: torch.Tensor,
    lse: torch.Tensor,
    block_out: torch.Tensor,
    block_lse: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    block_out = block_out.to(torch.float32)
    block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
    new_lse = lse + torch.log(1 + torch.exp(block_lse - lse))
    out = torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out
    lse = new_lse
    return out, lse

Personally, although I managed to understand this concept while preparing for deep learning exams, I usually lose my intuition a few months later. Reviewing it again requires extra effort to bring everything back together. That’s why I decided to write this technical blog post, both as a reminder for myself and as a guide for anyone trying to understand backpropagation in matrix form.

Forward Propagation

When training a deep learning model, the first step is to compute a linear combination of input features. Given an input vector $x \in \mathbb{R}^M$ with $M$ features and a corresponding weight vector $w \in \mathbb{R}^M$, we calculate a weighted sum and add a bias term $b \in \mathbb{R}^1$ to produce a single output $z$:

This is simply the dot product between $w$ and $x$ with an added bias term:

To visualize this computation, we can represent it in matrix form:

or in image representation:

We can simplify our notation by incorporating the bias term directly into the weight vector. This involves adding the bias $b$ as the first element of the weight vector and including a constant input $x_0 = 1$:

For notational clarity, we’ll denote the bias term as $w_0 = b$ and the constant input as $x_0 = 1$. Under this convention, both vectors have dimension $M+1$: $w \in \mathbb{R}^{M+1}$ and $x \in \mathbb{R}^{M+1}$. Our equation becomes:

or in matrix form:

or in image representation:

Adding Non-Linearity

The linear combination alone is insufficient for learning complex patterns. To add non-linearity, we apply an activation function $\sigma(\cdot)$ to the output $z$.

Although $\sigma$ often denotes the Sigmoid function, here it represents a general activation function. For this explanation, we’ll use the ReLU activation function in hidden layers due to its simplicity, computational efficiency, resistance to vanishing gradients, and widespread popularity:

we can show this with visual representation:

Multi-class Classification

In practice, deep learning models often solve multi-class problems where we need to predict one of $C$ possible classes. This requires the output $a$ to be $C$-dimensional vector, with each dimension representing the score for a particular class. The predicted class corresponds to the entry with the highest activation value.

To generate $C$ outputs, we need $C$ distinct weight vectors, each of dimension $(M+1)$. To achieve this, we need a separate weight vector for each class. Stacking them gives a weight matrix $W \in \mathbb{R}^{C \times (M+1)}$:

or in expanded form:

or in visual form:

Each figure represents the same network, but highlights different output paths from the same inputs, and each path is computed by its own set of output weights.

Batch Processing

The computations described so far process only a single input example (batch size = 1). In practice, we process multiple inputs in parallel to improve computational efficiency.

If we process batch size of $N$ inputs simultaneously, then our input becomes $X \in \mathbb{R}^{N \times (M+1)}$ (note the uppercase $X$ since we now have a matrix rather than a vector). To handle $N$ outputs while maintaining proper matrix dimensions, we transpose our weight matrix to $W \in \mathbb{R}^{(M+1) \times C}$:

or in matrix form:

When dealing with batches, it can sometimes be hard to picture how forward propagation works. We can visualize the process like this:

Note that in this illustration, the non-linearity (activation function) is not explicitly shown. The final output $A$ actually corresponds to the activated values.

Multi-layer Networks

The calculations presented so far describe a single layer. Deep learning models stack multiple layers to learn increasingly complex representations. We can denote each layer with a superscript $[l]$ to represent the layer number:

Here $H_l$ represents the number of hidden units in layer $l$. The input to layer $l$ activated output from the previous layer $l-1$. For simplicity, we usually define $A^{[0]} = X$ (the input layer).

Softmax for Classification

For classification problems, we apply the softmax activation function to the final layer’s output to convert the raw scores into a probability distribution. The softmax function ensures that all outputs sum to 1, allowing us to interpret them as class probabilities.

The softmax function is applied row-wise (for each input example) across the $C$ classes:

Intuitively, for input example $i$, this computes the probability of class $j$ by normalizing the exponential of its score by the sum of exponentials across all $C$ possible classes.

The matrix representation shows this row-wise operation:

If we denote the last layer as $L$, then the model’s output for $N$ inputs and $C$ classes can be represented as follows:

Cross-Entropy Loss Function

After obtaining the predicted probabilities $\hat{Y}$, we measure the model’s performance by comparing these predictions with the ground truth labels $Y \in \mathbb{R}^{N \times C}$ using a loss function. The ground truth is represented as one-hot encoded vectors, where each input example has exactly one correct class.

For multi-class classification, we typically use Cross-Entropy (CE) loss:

This formula computes the element-wise product between the true labels $Y$ and the logarithm of predicted probabilities $\log\hat{Y}$ then averages across all examples to produce a single scalar loss value.

We can express this using matrix operations with the element-wise (Hadamard) product $\odot$:

Expanding this matrix operation:

We can simplify this by recognizing that multiplying by vectors of ones simply sums all elements in the matrix. Therefore:

We can show this summation with following image (excluding scaling factor $-\frac{1}{N}$):

Although this is a general representation, we denote the output cross-entropy loss as $-\log\hat{Y}$. Since the ground truth $Y$ is one-hot vector, it is 1 for true class and 0 for others in each input. Therefore, the assumption in this case can be:

This completes our mathematical framework for forward propagation in deep learning models, from single predictions to batch processing with multi-class classification and loss computation.

Simple Computation Graph

Before diving lots of matrix gradient calculation, I would like to show a computation graph for forward and backward propagation so we can see the cases we need to be careful when we compute gradients. In this case, I want to only show scalar calculations first instead of thinking about matrices so we can adapt this into matrix backpropagation.

where $[l]$ and $[L]$ represents arbitrary layer $l$ and last layer $L$, respectively. $\sigma(\cdot)$ is activation function (ReLU in this case). In this representation, everything is scalar except $w^{[L]}$ because the output must be multi-class for softmax so I just played in the last layer to show a meaningful example.

We can visualize the computation graph of this network as follows:

Chain Rule

Backpropagation relies on the chain rule of calculus to compute gradients efficiently. For a composite function $f(g(h(x)))$, the chain rule states:

In neural networks, we apply this principle to decompose the loss gradient into a product of simpler derivatives based on any parameter.

Scalar Backward Propagation

After computing loss, we can start computing gradients with respect to it. Since there are $C$ classes, we need to calculate derivative of each prediction $\hat{y}_i$:

We can visualize this gradient flow like this:

Since the predicted output $\hat{y}_i$ is computed using $z^{[L]}_i$ in both the numerator and the denominator

the derivative $\frac{\partial \mathcal{L}}{\partial z^{[L]}_k}$ depends not only on the numerator but also on the denominator:

It might be easier to understand when you see visual representation:

Since we already know $\frac{\partial \mathcal{L}}{\partial \hat{y}_i}$, the gradient expression simplifies to:

The most confusing part begins here, because we now need to carefully compute the local gradient $\frac{\partial \hat{y}_i}{\partial z^{[L]}_k}$. To do this, we apply the quotient rule:

For this derivative, there are two distinct cases to consider:

Case 1: $i = k$

The derivative of numerator:

The derivative of the denominator is:

Substituting these into the quotient rule:

Factorizing into:

Shortly, this is equivalent to:

Case 2: $i \neq k$

The derivative of the numerator is zero, because $e^{z_i}$ does not depend on $z_k$. The derivative of the denominator is:

Thus, we have:

Simplifying:

Shortly, this is equivalent to:

Finally, we can summarize these two cases as

In literature, it is pretty common to use the following form:

where $\delta_{ik}$ is 1 if $i = k$, 0 otherwise. When we combine the gradient of cross-entropy and local gradient, we have:

where $\hat{y}_i$ cancels each other:

Since ground truth value $y_i$ is one-hot vector (1 for true class and 0 for others), equation simplifies to:

This compact form is what makes softmax combined with cross-entropy loss so useful. Instead of dealing with complicated fractions, we can directly use this result to propagate gradients to the next layers.

From this point, gradient calculations across layers will start to follow a recurring pattern. Each layer essentially repeats the same process: we compute gradients with respect to its inputs and its weights.

For the layer output $z^{[L]}$, we have two key variables to differentiate with respect to:

• The weights $w^{[L]}$: their gradients are crucial because they are the trainable parameters we want to optimize during learning.

• The activations $a^{[L-1]}$: while not parameters themselves, their gradients are equally important since they ensure the flow of gradients backward through the network, enabling earlier layers to update as well.

Thus, even though only the weight gradients directly influence optimization, the activation gradients play a important role in keeping backpropagation alive throughout the entire network:

• If we calculate derivatives with respect to weights $w^{[l]}$:

• If we calculate derivatives with respect to weights $a^{[l-1]}$:

Lastly, we need to calculate the derivative of $a^{[l]}$ with respect to $z^{[l]}$ in $\sigma(z^{[l]})$, assuming $\sigma(z^{[l]}) = \text{ReLU}(z^{[l]})$:

where $1(z^{[l]} > 0)$ outputs 1 if $(z^{[l]} > 0)$, 0 otherwise as a indicator function.

Therefore, we can summarize backpropagation for scalar values with following pattern:

Backward Propagation

After computing the loss through forward propagation, we need to update the model’s weights to minimize this loss. Backpropagation is the algorithm that computes the gradients of the loss function with respect to each parameter in the network. We’ll derive these gradients step by step, working backwards from the loss to the input layer.

Gradient of Cross-Entropy with Softmax

Starting with loss function, we calculate relative gradients and go back step by step to calculate further gradients. Since the model has predictions with Softmax activation and the loss is calculated with cross-entropy, the combination of these methods has a nice property which simplifies the gradient:

or in matrix representation:

or in visual representation:

Gradient of Arbitrary Layer $l$

At this stage, we can begin calculating the gradients of the weights in the corresponding layer. Since the forward propagation is computed as

There are two possible local gradients we might consider:

You may wonder why we would need the gradient with respect to $A^{[l-1]}$ since our goal is to update the weight matrix $W^{[l]}$. The reason is that the derivative with respect to $A^{[l-1]}$ becomes necessary for updating the weights of the previous layer, because

When we want to calculate the gradients of weight matrix $W^{[l]}$ with respect to the loss function $\mathcal{L}$, we use chain rule as follows:

The local gradient gives us:

Since there are $H_{l-1} \times H_l$ values in weight matrix $W^{[l]}$, we need to find each individual gradient so the shape of $\frac{\partial \mathcal{L}}{\partial W^{[l]}}$ should be $H_{l-1} \times H_l$. However, the dimensions of matrices does not match when we try to calculate matrix multiplication:

If we take transpose of $A^{[l-1]}$ and place it to the left side of upstream gradient, we have the correct dimension matching:

At first glance, this transformation may seem arbitrary. Why do we transpose? How can we be sure this result is correct? To clarify, let’s check the matrix representation:

If we look closely only one gradient calculation, such as $\frac{\partial \mathcal{L}}{\partial W^{[l]}_{11}}$:

Since single value of weight matrix is fixed for each input value, the gradient of this weight is just summation of all $N$ input derivatives:

In the image below, you can see the visualization of this calculation for 4 weights in the network:

At the same time, we need to calculate the derivative of $A^{[l-1]}$, $\frac{\partial \mathcal{L}}{\partial A^{[l-1]}}$, so upstream gradients can flow to previous layer and the gradient of weight matrix can be calculated in that layer:

The local gradient is:

Once again, the dimensions of matrices do not match:

We can transpose the weight matrix to align matrices:

We can see the reasoning more easily if we check matrix representation again:

By focusing on only first value, \(\frac{\partial \mathcal{L}}{\partial A^{[l-1]}_{11}}\), the gradient of \(A^{l}_{11}\) depends on the values that it affected during forward propagation (the same logic as \(\frac{\partial \mathcal{L}}{\partial W^{[l]}_{11}}\)):

Now we want to compute the gradient $\frac{\partial \mathcal{L}}{\partial Z^{[l-1]}}$. Recall that $A^{[l-1]}$ is obtained by applying an activation function to $Z^{[l-1]}$:

In our examples, we use the ReLU activation function. Its derivative with respect to each element is given by:

For $N \times H_{l-1}$ inputs, we must compute a gradient for each element, producing a matrix of 0s and 1s. The derivative of ReLU applied element-wise can be written as the indicator (mask) of positive entries:

where $1(\cdot)$ is the element-wise indicator function (1 when the condition is true, 0 otherwise). In matrix form:

Because the derivative is applied element-wise, the upstream gradient $\frac{\partial \mathcal{L}}{\partial A^{[l-1]}}$ is combined with this mask element-wise, not by matrix multiplication:

or

The Pattern

At this point, we completed the backpropagation using the matrix multiplication step. From here on, the process is essentially a repetition: each layer applies the same logic and propagates the gradients backwards through the weights (also biases implicitly) and activations.

The only real exception was the output layer, where we combined the softmax function with cross-entropy loss. This required a more detailed derivation, but the model is consistent for all hidden layers and is repeated across the entire network. You can see the pattern as follows:

I know it looks ugly but if you pay some attention, you will see the pattern for gradient calculation. There are just some rules that you need to follow for each block.