Naively computing Cross Entropy is Expensive

Cross Entropy is a ubiquitous loss function in deep learning, including the training of LLMs (next token prediction is a classification task). It's a classification over the entire vocabulary, and vocabularies are enormous.

Concretely: Gemma 2 has a vocabulary of 256,128 tokens. For a batch of 8,192 tokens, a single fp32 logit tensor is $8192 \times 256128 \times 32 \text{ bits} \approx 8.4$ GB. Each entry is one scalar logit $z_{k,i}$: the model's raw score for vocabulary class $k$ at token $i$. The tensor is enormous because it collects that score for every token-class pair.

A baseline PyTorch loss also materializes additional log-probability and intermediate buffers. In the measurements reported by Apple's CCE paper, the baseline loss peaks at 24 GB and the combined loss-plus-gradient computation peaks at 28 GB for Gemma 2 (2B). Cut Cross Entropy computes the loss using 1 MB.

By reducing the memory footprint of cross-entropy, we can significantly decrease the hardware requirements for training LLMs!

This is the source code that's referenced throughout this article. It might be helpful to have this open in another tab: Apple ML Cut Cross Entropy

Quick Recap: Cross Entropy

If you are unfamiliar with the math, check out ❋derivation of cross-entropy loss

Cross-entropy can be broken into two terms:

1. The negative dot product: $\htmlClass{math-var-tooltip math-var-tooltip-0}{- z_{y_i}}$
2. The log-sum-exp: $\log \sum_{k} e^{\htmlClass{math-var-tooltip math-var-tooltip-1}{z_k}}$

Here $\htmlClass{math-var-tooltip math-var-tooltip-1}{z_k}$ is shorthand for the score $z_{k,i}$ at one token $i$. If $e_i \in \mathbb{R}^{D}$ is that token's hidden-state vector and $c_k \in \mathbb{R}^{D}$ is the classifier vector for class $k$, then:

Stack the token vectors as the columns of $E \in \mathbb{R}^{D \times B}$ and the class vectors as the columns of $C \in \mathbb{R}^{D \times V}$. Collecting every $z_{k,i}$ gives the logit matrix:

This is the same logit tensor from the opening memory calculation. The implementation stores token vectors as rows, so its equivalent computation is $E^{\mathsf{T}} C = Z^{\mathsf{T}}$ with shape (B, V).

Keep the two-term decomposition in mind (the negative dot product and the log-sum-exp). The forward pass implements each term with its own dedicated kernel, and the backward pass is where the two terms' gradients merge back together.

The villain of the story is now concrete: materializing $Z$ means storing every one of those token-class scores at once. The whole "cut cross entropy" approach is about evaluating the two loss terms without ever writing the full projection to global memory.

Notation

Let's establish our vocabulary of terms (mirroring the paper, though I've taken some liberties with the shapes so it checks against $Z = C^{\mathsf{T}} E$):

Overview of Cut Cross Entropy

Cut Cross Entropy saves memory by doing two things:

1. Blockwise computation: the logits are computed in small blocks that live entirely in ✦on-chip SRAM
   , with each block immediately reduced into a running log-sum-exp and then discarded. The full $V \times B$ logit matrix simply never exists in global memory. This is where the memory savings of over 1000 times compared to traditional cross entropy implementations come from.
2. "Filtering" (although I think this is better described as "skipping" or "early stopping"): in the backward pass, chunks of near-zero gradient are skipped entirely. This is the "cutting", and it's what makes the backward pass's recomputation fast — it buys back compute and memory bandwidth, rather than memory capacity.

Breakdown: Forward Pass

Inputs:

• Token embeddings $E$
• Classifier embeddings $C$
• Bias vector $(\text{vocab} \times 1)$
• Targets $y_i$

Notably, the logits have not been computed yet.

The forward pass maps directly onto our two loss terms — one kernel each:

1. indexed_dot_forward_kernel → term 1, the negative dot product $\htmlClass{math-var-tooltip math-var-tooltip-0}{- z_{y_i}}$.
   • A triton kernel that gathers only the correct-class column $c_{y_i}$ and dots it with the corresponding embedding $e_i$. From the paper:

     Our kernel retrieves the value $x_i$, the $x_i$-th column from C, and the $i$-th column from E, and stores them in on-chip shared memory (SRAM). It then performs a dot product between $C_{x_i}$ and $E_{i}$ and writes the result into global memory. The kernel uses only on-chip SRAM throughout and does not allocate any GPU memory.

2. cce_lse_forward_kernel → term 2, the log-sum-exp.
   • A triton kernel that computes the log-sum-exp over blockwise logits. It also computes logit_avg, the mean logit per vocabulary entry, as a by-product — this looks useless right now, but it becomes the sorting key for the backward pass's filtering. Hold that thought.
3. Aggregate loss
   • The two terms are summed to compute the final loss.
   • Supports standard reductions (mean, sum).
4. Save tensors for backward pass
   • Embeddings, targets, the per-token log-sum-exp, and logit_avg. Note what's not on this list: the logits.

Forward Kernel

Between the kernels, the cce_lse_forward_kernel is more interesting, so let's explore its Triton implementation.

Skipping over pointer and offset initialization, we hit the first math operation:

accum = tl.zeros((BLOCK_B, BLOCK_V), dtype=tl.float32)
for d in range(0, tl.cdiv(D, BLOCK_D)):
   e = tl.load(e_ptrs, mask=e_mask, other=0.0)
   c = tl.load(c_ptrs, mask=c_mask, other=0.0)

   accum = tl.dot(e, c, accum, input_precision=DOT_PRECISION)

   e_ptrs += BLOCK_D * stride_ed
   c_ptrs += BLOCK_D * stride_cd

(I've stripped out the mask bookkeeping: e_mask/c_mask just handle the ragged edges when $D$ or $V$ don't divide evenly into blocks.)

Here we are computing accum. Each pass loads the next BLOCK_D features from the token and class vectors, multiplies them, and adds the result to the same accumulator. After the loop has covered all $D$ features, accum[i, k] holds $z_{k,i}$ for one block of tokens and vocabulary classes.

The important part is where accum lives. Each program instance owns one $(\text{BLOCK\_B}, \text{BLOCK\_V})$ tile of $Z^{\mathsf{T}}$, matching the code's (B, V) layout, and that tile exists only in on-chip SRAM.

The source then casts the accumulator, adds the optional bias, masks vocabulary positions beyond $V$ to $-\infty$, and applies optional soft-capping. The resulting logits variable is still the same SRAM-resident tile; it has not been written to global memory.

Now each token row in that tile can be reduced to a partial log-sum-exp and the tile can be discarded. Only the final (B,) log-sum-exp vector survives in global memory. That's the memory trick in one line: $O(B \cdot V)$ becomes $O(B)$.

this_mx = tl.max(logits, axis=1)e = tl.exp(logits - this_mx[:, None])this_lse = this_mx + tl.log(tl.sum(e, axis=1)) # ... this_locks = Locks + (pid_b // tl.cdiv(B, BLOCK_B * num_locks))while tl.atomic_cas(this_locks, 0, 1) == 1:   pass lse = tl.load(lse_ptrs, mask=o_mask, other=0.0, eviction_policy="evict_last")lse = tl_logaddexp(lse, this_lse)tl.store(lse_ptrs, lse, mask=o_mask, eviction_policy="evict_last") tl.debug_barrier()tl.atomic_xchg(this_locks, 0)

There are four distinct steps here:

1. Reduce the local tile. this_mx finds the largest logit in each token row. Subtracting that maximum before exponentiating avoids overflow, and adding it back after the sum gives one stable partial log-sum-exp per token.
2. Acquire the lock. Several vocabulary tiles contribute to the same token rows. tl.atomic_cas spins until this program has exclusive access to those rows of the running global result. Without the lock, two programs could load the same old value and overwrite one another's updates.
3. Merge and store. The global vector starts at -inf, the identity for log-add-exp: $\operatorname{logaddexp}(-\infty, x) = x$. The kernel loads the running value, merges the tile's partial result with tl_logaddexp, and stores the update. Repeating this across vocabulary tiles is equivalent to taking one log-sum-exp over the entire vocabulary.
4. Release and discard. After the store is complete, tl.atomic_xchg(this_locks, 0) releases the lock. The program is finished with its logit tile, so that tile disappears from SRAM and the next vocabulary tile can repeat the process.

In short: every vocabulary tile briefly computes logits, reduces them to one partial statistic per token, merges that statistic into the running log-sum-exp, and disappears. At no point does the full logit tensor need a home in global memory.

Breakdown: Backward Pass

The backward pass needs to produce two main gradients, plus an optional bias gradient:

• the gradient with respect to the query embeddings, $\nabla E$
• the gradient with respect to the class embeddings, $\nabla C$
• if present, the gradient with respect to bias, $\partial b$. We're gonna ignore this.

Here's where the two-term map pays off. Differentiate each term with respect to the logits:

• The log-sum-exp term gives the softmax: $\frac{\partial \operatorname{LSE}}{\partial \htmlClass{math-var-tooltip math-var-tooltip-1}{z_k}} = \frac{e^{\htmlClass{math-var-tooltip math-var-tooltip-1}{z_k}}}{\sum_j e^{z_j}} = S$.
• The $\htmlClass{math-var-tooltip math-var-tooltip-0}{- z_{y_i}}$ term gives $-1$ at exactly the correct class and $0$ everywhere else — that's $-Y$.

Sum them and you get the classic softmax-minus-onehot, scaled by the upstream gradient:

This $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}$ is the key intermediate product of the backward pass. One more application of chain rule through $Z = C^{\mathsf{T}} E$ turns it into both of our gradients:

Two things from the paper make this computable at all. First, $S$ needs no fresh normalization pass: $S = \text{softmax}(C^{\mathsf{T}} E) = \exp(C^{\mathsf{T}} E - \operatorname{LSE})$, reusing the per-token log-sum-exp we already paid for in the forward pass. Second, $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}$ is just as enormous as the logits — so the backward pass recomputes $C^{\mathsf{T}} E$ blockwise in SRAM, exactly like the forward pass did, rather than ever storing it.

Overall, this is what the backward pass does:

1. Retrieve saved tensors: embeddings, targets, the per-token log-sum-exp, and logit averages.
2. Recompute the logits ($C^{\mathsf{T}} E$) in chunks (like in the forward pass).
3. Compute $S = \exp(C^{\mathsf{T}} E - \operatorname{LSE})$, the softmax probabilities.
4. Compute $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}} = (S - Y) \odot g$, using the upstream gradient $g$ (d_out in the code) and the targets $Y$.
5. Compute $\nabla E = C \htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}$, the gradient of the query embeddings.
6. Compute $\nabla C = E \htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}^{\mathsf{T}}$, the gradient of the class embeddings.

Steps 3–4 look like this in code:

# Compute S
d_accum = tl.exp(accum - lse[:, None])
d_accum = tl.where(offs_v[None, :] < V, d_accum, 0.0)

# Compute S - Y
if HAS_TARGETS:
   targets = tl.load(Targets + target_offs_b, mask=target_offs_b < BMax, other=V + 1)
   is_target = targets[:, None] == offs_v[None, :]
   d_accum += tl.where(is_target, -1.0, 0.0) # <-- subtracting Y from S
else:
   is_target = None

# (gradient filtering happens here — covered in the next section)

d_out = grad_scale * d_out

# The final product we're looking for: S_hat = (S - Y) * upstream_grad
d_accum = d_accum * d_out

If this looks like a wall of mathy text, the main point is that steps 1–4 build the intermediate term $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}$, which steps 5–6 then consume. Steps 5–6 are matrix multiplications — expensive operations that we'd like to avoid if we can — and that brings us to the main optimization in the backward pass: gradient filtering.

Gradient Filtering

Before we dive into the gradient filtering, it helps to recall that $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}$ in the triton kernel is really $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}_{\text{block}}$, the gradient signal for a single block/vocabulary chunk. The same applies for $S \to S_{\text{block}}$ and the other terms.

Then let's zoom in on the gradient filtering segment of the code:

should_skip = False
if (FILTER_E_GRAD and COMPUTE_DE) and (FILTER_C_GRAD and COMPUTE_DC):
   if _block_is_filtered(tl.abs(d_accum), filter_eps):
      return
elif (FILTER_E_GRAD and COMPUTE_DE) or (FILTER_C_GRAD and COMPUTE_DC):
   should_skip = _block_is_filtered(tl.abs(d_accum), filter_eps)

# ...

if COMPUTE_DE:
   if FILTER_E_GRAD:
      should_skip_e = should_skip
   # ...

if COMPUTE_DC:
   if FILTER_C_GRAD:
      should_skip_c = should_skip
   # ...

# elsewhere:

@triton.jit
def _block_is_filtered(check_val: tl.tensor, filter_eps: tl.tensor) -> tl.tensor:
    return tl.reduce(check_val < filter_eps, None, tl_and_reduce_fn)

_block_is_filtered checks if all values within d_accum (which is $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}_{\text{block}}$) are below our "near-zero" threshold. If so, we terminate the kernel early, or skip over some of the remaining matrix multiplications involved in calculating the gradient outputs $\nabla E$ and $\nabla C$.

Why is this a reasonable approximation? After the softmax, the overwhelming majority of $S$ is astronomically small — the probability mass concentrates on a handful of plausible tokens. Contributions below filter_eps are around the noise floor of the accumulation's numerical precision anyway; skipping them trades an invisible amount of exactness for a lot of saved work.

Vocabulary Sorting

There's a catch with filtering at block granularity: a block is only skippable if every value in it is near zero. One confident token in an otherwise-dead block ruins the whole block.

This is what logit_avg from the forward pass was for. Before launching the backward kernels, the vocabulary is sorted by average logit, so the consistently-plausible tokens cluster together into a few dense blocks — and everything else collapses into large, contiguous, entirely-skippable regions. It's a heuristic, but it directly increases the number of blocks that _block_is_filtered gets to throw away.

As a whole, the gradient filtering process is:

1. Vocab sorting (host-side, before kernel launch): sort vocabulary entries by their average logit to make "significant" blocks as dense and contiguous as possible.
2. Gradient filtering (in-kernel): whenever $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}}_{\text{block}}$ is entirely below filter_eps, skip the expensive matrix multiplications in steps 5–6 and zero out those blocks' contributions.

Takeaways: Usage + Design

Cut cross entropy is likely going to become the go-to implementation of cross-entropy going forward. The memory savings it presents are too good to pass up.

One caveat worth knowing: gradient filtering makes CCE technically lossy. Gradients below filter_eps are genuinely discarded, not deferred. The paper shows this has no measurable effect on training loss curves, but it is a knob you can turn — tightening filter_eps trades speed back for exactness.

As for its design, the recipe here generalizes well beyond cross entropy: recompute blockwise in SRAM instead of materializing, reuse a cheap global statistic (the log-sum-exp) to avoid renormalization, and skip contributions below a chosen numerical threshold. Any loss over a huge output space could use this too.

I'll definitely be writing more about kernel optimization and other similar breakdowns by the way, so if you found this interesting, consider following me on any of my socials.