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). The catch is that 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 4 \text{ bytes} \approx 8.4$ GB. 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: GitHub - Apple ML Cut Cross Entropy

Quick Recap: Cross Entropy

If you are unfamiliar with the math, check out derivation of cross-entropy loss. But here is the structure in brief:

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 exponent (LSE): $\log \sum_{k} e^{\htmlClass{math-var-tooltip math-var-tooltip-1}{z_k}}$

Keep this decomposition in mind — it's the map for everything that follows. 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 the logit matrix $Z = C^{\mathsf{T}} E$: the linear classifier $C^{\mathsf{T}}$ projects the embeddings $E$ onto an extremely large vocabulary space, and the whole "cut cross entropy" approach is about never materializing that projection (which we'd otherwise still have to compute further terms of, like softmax).

Notation

Let's establish our vocabulary of terms (mirroring the paper, though I've fixed the shapes so everything type-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 statistic (the LSE) 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 $\htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}}$.
   • 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 LSE, 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 lets explore its triton implementation.

Skipping over the initialisation (of the pointer and offset) 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: the dot product of the query and class embeddings (represented generally by e and c). This is done in chunks of BLOCK_D.

The important part is where accum lives. Each program instance owns one $(\text{BLOCK\_B}, \text{BLOCK\_V})$ tile of the logit matrix, and that tile exists only in on-chip SRAM. It gets reduced into a partial LSE (next snippet) and then thrown away. The full $V \times B$ logit matrix never touches global memory — only the final $(B,)$ LSE vector does. That's the memory trick in one line: $O(B \cdot V)$ becomes $O(B)$.

After that, comes computation of the linear-log-sum-exponent (LSE).

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 3 main steps here:

1. Compute the "partial LSE": the local LSE for each batch, using the (BLOCK_B, BLOCK_V) chunk of logits
2. Use a spin lock to atomically update the global LSE
   • tl.atomic_cas is used to implement a spin lock.
   • The lock index, as defined by (pid_b // tl.cdiv(B, BLOCK_B * num_locks)), causes multiple pid_b blocks to share the same lock. Not shown in all this code is that pid_v blocks that correspond to the same pid_b (same block, but different vocabulary slices) will share the same pid_b lock.
   • Yeah, this might require writing it out or staring at the code to catch.
   • tl;dr: we are only allowing one block/"vocabulary chunk" to update the global LSE at a time.
3. Store the global LSE back to global memory
   • lse = tl.load(lse_ptrs...) loads the current global value of LSE. This is intialized as -inf (because $e^{-\infty} = 0$).
   • lse = tl_logaddexp(lse, this_lse) is the core aggregation step. By repeatedly computing $\log(\exp(\text{lse}) + \exp(\text{this\_lse}))$, different blocks/vocabulary chunks will combine their partial LSEs into a global LSE.
   • tl.store(lse_ptrs, lse, mask=o_mask, eviction_policy="evict_last") stores the updated global LSE back to global memory.
   • tl.atomic_xchg(this_locks, 0) releases the lock, allowing other threads to update the global LSE.

In short: to ensure synchronicity, we use a spin lock so that only one vocabulary chunk can update the global LSE at a time. Each block computes a partial LSE, which is then merged into the global one. This is an aggregation step to get our final result, the $\htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}}$ term.

Breakdown: Backward Pass

The backward pass needs to produce two gradients:

• 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 $\htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}}$ term gives the softmax: $\frac{\partial \htmlClass{math-var-tooltip math-var-tooltip-2}{\text{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 - \htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}})$, reusing the LSE 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 LSE, 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 - \htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}})$, the softmax probabilities.
4. Compute $\htmlClass{math-var-tooltip math-var-tooltip-3}{\hat{S}} = (S - Y) \cdot \nabla \htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}}$, using the upstream gradient $\nabla \htmlClass{math-var-tooltip math-var-tooltip-2}{\text{LSE}}$ (or 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 legal? 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 LSE) to avoid renormalization, and skip work that provably can't matter. Any loss over a huge output space is a candidate — I've since been applying the same treatment to other losses, and future posts will dig into those.

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.