ArXiv: 2411.09009

🎯 Pitch

Standard cross-entropy loss silently consumes up to 90% of GPU memory during LLM training because it materializes a giant logit matrix over the whole vocabulary. This paper eliminates that bottleneck by computing the loss on the fly and skipping negligible gradient updates, shrinking the memory for the loss layer from 24 GB to just 1 MB on Gemma 2B while matching training speed and convergence.

1. Executive Summary

This paper introduces Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss and its gradient without materializing the full logit matrix—the dominant memory consumer in large-vocabulary LLM training—into global GPU memory, instead computing logits as needed in on-chip SRAM via custom CUDA kernels and leveraging the inherent sparsity of softmax to skip negligible gradient contributions (gradient filtering with a threshold at the bfloat16 truncation boundary). Evaluated on models including Gemma 2 (2B), Phi 3.5 Mini, Qwen 2.5 (7B), and Mistral NeMo, CCE reduces the memory footprint of the loss computation from 24 GB to 1 MB for Gemma 2 (2B) and the total classifier-head memory from 28 GB to 1 GB, enabling batch size increases of 1.5× to 10× across frontier models without sacrificing training speed or convergence. The paper establishes that CCE’s gradient filtering requires no numerical-precision compromise for fine-tuning, while a variant with Kahan summation and selective filtering (CCE-Kahan-FullC) matches torch.compile’s validation perplexity during pretraining, demonstrating that the dramatic memory reduction applies to both regimes only when the appropriate numerical stability safeguards are employed.

2. Context and Motivation

The Core Problem: The Cross-Entropy Layer Has Become the Memory Bottleneck in LLM Training

This paper addresses a specific and growing imbalance in the memory consumption of large-language-model training: the cross-entropy loss computation at the classifier head consumes a disproportionate and increasingly unsustainable fraction of total GPU memory. While the research community has devoted enormous attention to reducing the memory footprint of attention mechanisms and transformer layers, the memory demands of the final loss layer—which computes log-probabilities over the entire vocabulary for every token in the batch—have quietly grown to dominate training memory.

The problem is fundamentally architectural. At training time, given a batch of NN tokens and a vocabulary of size ∣V∣|V|, the standard cross-entropy implementation materializes a logit matrix C⊤E∈R∣V∣×NC^\top E \in \mathbb{R}^{|V| \times N} into global GPU memory. The memory cost of this single matrix grows as O(N×∣V∣)O(N \times |V|), and with contemporary vocabularies reaching 128K–256K tokens and batch sequences frequently exceeding tens of thousands of tokens, this one layer can consume orders of magnitude more memory than the rest of the model combined.

Consider the examples the paper provides in Figure 1a. For Gemma 2 (2B) with a 256,128-token vocabulary, the log-probabilities materialized by the cross-entropy layer account for 89% of total training memory on a 16-GPU data-parallel setup. For Llama 3 (8B) with a 128K vocabulary, that figure is 65%. For Phi 3.5 Mini (32K vocabulary), it is 40%. These are not marginal overheads—the cross-entropy layer has become the single most expensive component of LLM training by a wide margin. For Gemma 2 (2B) specifically, processing a single sequence of length 80,000 consumes the entire available memory of an 80 GB H100 GPU just for the logits.

This trend is not accidental; it is the direct result of a deliberate design choice that has been reinforced by complementary infrastructure improvements. As the paper notes in Section 1, vocabulary sizes have grown because larger vocabularies are genuinely beneficial: they reduce sequence lengths by enabling individual tokens to represent multi-character spans, which compresses documents into shorter context windows, improves model comprehension, and reduces the computational cost of the transformer backbone. Scaling laws research (Tao et al., 2024) further suggests that even the largest contemporary vocabularies may benefit from expansion. Simultaneously, the success of other memory optimizations has shifted the relative burden toward the loss layer: FlashAttention reduced attention's memory from O(N2)O(N^2) to O(N)O(N), ZeRO sharded optimizer states and gradients, and activation checkpointing cut the intermediate activation footprint—leaving the cross-entropy logit matrix as the last remaining O(N×∣V∣)O(N \times |V|) memory hog that no prior technique addresses.

Why This Problem Matters: It Constrains Batch Size, Which Constrains Training

The practical consequence of this memory imbalance is not merely that training consumes more GPU memory; it is that this one layer determines the maximum feasible batch size for training. With data-parallel training, the batch size per GPU is bounded by the memory available after model parameters, optimizer states, gradients, and activation checkpoints are stored. Since the logit matrix scales with both vocabulary size and sequence length (the number of tokens per batch), it becomes the binding constraint that forces a choice between two undesirable outcomes: reducing batch size (potentially harming training stability and convergence) or consuming more GPUs (increasing cost and communication overhead).

Figure 1 quantifies this concretely. For a 16-GPU fully-sharded data-parallel setup, the standard cross-entropy implementation limits maximum batch size to roughly 1.1 million tokens for Gemma 2 (2B) and 740K tokens for Gemma 2 (27B). These are surprisingly small numbers for such capable models, and they represent a hard ceiling that no amount of parameter sharding or activation checkpointing can raise. The paper's CCE method increases these ceilings to 10.6 million and 2.5 million tokens respectively—increases of 9.5× and 3.4×—demonstrating just how much training capacity is currently being wasted by the loss layer's memory consumption.

This constraint also creates a pipeline-parallelism imbalance. When very large models are trained with pipeline parallelism, each stage of the pipeline should ideally have a similar memory-to-computation ratio to avoid bottlenecks. The classification head is currently an outlier with a disproportionately high memory-to-computation ratio (Section 6), which forces pipeline planners to either allocate excessive resources to that stage or accept underutilization elsewhere. As models continue to grow, this imbalance will only worsen.

Prior Approaches and Their Shortcomings

The paper identifies four categories of prior work that partially address vocabulary-related memory costs, but each has fundamental limitations that leave the core problem unsolved.

Chunked cross-entropy implementations. Torch Tune (Torch Tune Team, 2024) and Liger Kernels (Hsu et al., 2024) both reduce memory by dividing the logit computation into chunks—processing a subset of the vocabulary at a time rather than all at once. This approach has an inherent tradeoff: memory footprint is minimized when the number of chunks is high, but latency increases because each chunk incurs kernel launch overhead and requires repeated loads of the embedding matrix. The paper's Table 1 demonstrates this tension concretely. For a Gemma 2 (2B)-sized configuration with a batch of 8,192 tokens and a 256K vocabulary, Torch Tune (with 8 chunks) uses 9,631 MB and takes 169 ms for the combined loss+gradient computation. Liger Kernels uses only 1,474 MB but takes 304 ms—more than double the time of the fastest method. Both approaches still allocate memory proportional to O(N×D)O(N \times D) or O(N×∣V∣/chunks)O(N \times |V| / \text{chunks}), meaning they remain fundamentally dependent on vocabulary size and batch size. They trade memory for latency but cannot eliminate the dependence entirely.

Liger Kernels' fused forward-backward design. Liger Kernels achieves memory efficiency partly by computing the loss and gradient simultaneously rather than in separate forward and backward passes. This fusion avoids storing intermediate logit values between passes, but it comes with a practical cost: any transformation applied to the loss (such as masking for ignored tokens, loss weighting, or logit softcapping as used in Gemma 2) must be implemented inside the kernel itself, not in the user's training code. This reduces flexibility and composability. CCE, by contrast, maintains separate forward and backward stages, so user-defined loss transformations work without kernel modification—a design choice that the paper explicitly contrasts with Liger Kernels (Section 2).

Vocabulary reduction techniques. Grave et al. (2017) proposed hierarchical vocabulary structures where only a subset of the vocabulary is active at any given time, reducing the effective ∣V∣|V| in the logit computation. This approach modifies the model architecture itself rather than the loss computation, potentially affecting model quality. Yu et al. (2023) explored tokenization-free byte-level models that operate on dramatically smaller vocabularies (256 possible byte values versus 256K tokens). While this eliminates the large-vocabulary memory problem entirely, it forces the model to process much longer sequences since each token represents less information—trading vocabulary memory for sequence-length computation and potentially degrading quality on tasks where higher-level token representations are beneficial. Both approaches solve the memory problem by avoiding large vocabularies rather than supporting them efficiently, which is orthogonal to the growing consensus that large vocabularies are genuinely useful.

Sequence and model parallelism. Techniques like sequence parallelism (Jacobs et al., 2023; Li et al., 2023) and model parallelism (Huang et al., 2019; Shoeybi et al., 2019) distribute the vocabulary dimension across multiple GPUs, so each GPU only stores a fraction of the logit matrix. These approaches can train models with arbitrarily large vocabularies but at the cost of additional communication and GPU count. The memory problem is "solved" by buying more hardware and accepting the associated communication overhead—a solution that works for organizations with large GPU clusters but does nothing for researchers or practitioners operating at smaller scales. CCE reduces the per-GPU memory footprint without requiring additional GPUs or communication.

The crucial gap: no method eliminates the O(N×∣V∣)O(N \times |V|) materialization. Despite these efforts, every prior approach either (a) still materializes a chunk of the logit matrix into global memory, maintaining a dependence on ∣V∣|V| and NN, or (b) changes the model architecture or training setup to avoid the problem rather than solving it within the standard training loop. The paper's key observation is that the full logit matrix never needs to be materialized at all—the loss and gradient can be computed from the log-sum-exp reduction and the single ground-truth logit, both of which have O(N)O(N) or O(1)O(1) memory footprints. This observation is the foundation of CCE.

How the Paper Positions Itself

The paper positions CCE not as a model-architecture modification, nor as a parallelism scheme, nor as a chunking heuristic, but as a drop-in replacement for the standard cross-entropy loss computation that uses the same mathematical formulation, produces identical gradients (to within the chosen numerical precision), and requires zero changes to the model architecture, training hyperparameters, or convergence properties. This is emphasized throughout: CCE is a purely implementation-level contribution that changes how the cross-entropy is computed, not what it computes.

The intellectual lineage is explicitly drawn to FlashAttention (Dao et al., 2022). Just as FlashAttention eliminated the O(N2)O(N^2) attention matrix from global memory by computing softmax attention in blocks within on-chip SRAM and accumulating the normalization constant on the fly, CCE eliminates the O(N×∣V∣)O(N \times |V|) logit matrix from global memory by computing the log-sum-exp over the vocabulary in blocks within SRAM and accumulating the normalization constant via a thread-safe log-add-exp reduction. The analogy is not merely inspirational—the technical mechanisms are structurally parallel: both use online normalization (online softmax for FlashAttention, online log-sum-exp for CCE), both recompute intermediate values during the backward pass to avoid storing them during the forward pass, and both achieve asymptotic memory reductions that are independent of the problematic dimension (sequence length for FlashAttention, vocabulary size for CCE).

However, CCE introduces a key additional insight that FlashAttention did not need: gradient sparsity. The softmax function's output over a large vocabulary is inherently extremely sparse—for any given token prediction, only a few dozen vocabulary entries have non-negligible probabilities, with the probability of even the 50th most likely token typically falling below bfloat16 numerical precision (2−122^{-12}, or approximately 2.4×10−42.4 \times 10^{-4}). The paper quantifies this in Figure 3: less than 0.02% of softmax elements are non-zero in practice, and the sparsity increases with vocabulary size. This means that the expensive outer-product gradient computations ∇E=S^⋅C\nabla E = \hat{S} \cdot C and ∇C=S^⊤⋅E\nabla C = \hat{S}^\top \cdot E can skip the vast majority of operations without affecting the result. CCE exploits this by checking whether an entire block of the softmax matrix contains only values below the bfloat16 truncation threshold and, if so, skipping the corresponding gradient update entirely—achieving a 3.5× speedup over an unfiltered version without any measurable loss of precision.

The paper also positions itself carefully with respect to the fine-tuning versus pretraining distinction. CCE's basic gradient filtering threshold (ε=2−12\varepsilon = 2^{-12}) is set at the bfloat16 truncation boundary, meaning it only skips gradient contributions that would be numerically zero after floating-point addition anyway. For fine-tuning, this threshold is provably lossless: any value below 2−122^{-12} cannot affect a bfloat16 gradient accumulator when added to a sum that already contains values up to 1. For pretraining, however, the paper finds that two additional sources of numerical error matter: (1) gradient filtering applied to ∇C\nabla C can prevent gradient flow to rare tokens that see little training data, and (2) bfloat16 summation in global memory loses precision on the large, high-dimensional reductions. The paper addresses these with CCE-Kahan-FullC, a variant that uses Kahan summation for the global-memory reductions and disables gradient filtering specifically for the ∇C\nabla C computation, trading some of the memory and speed savings for numerical exactness during pretraining. This two-tier approach (basic CCE for fine-tuning, CCE-Kahan-FullC for pretraining) is a practically important design choice that acknowledges the different numerical sensitivity of the two training regimes.

Finally, the paper positions its contribution within a broader trend toward IO-aware algorithm design for deep learning. The fundamental insight—that global GPU memory (HBM) bandwidth is the primary bottleneck, not computation, and that algorithms should be redesigned to minimize reads and writes to global memory by performing as much work as possible in on-chip SRAM—has been applied to attention (FlashAttention), matrix multiplication (various tiled implementations), and now, with CCE, to the classification loss. The paper's use of custom Triton kernels to implement blockwise matrix multiplication fused with an online log-sum-exp reduction is a direct application of this design philosophy to a layer that had previously been treated as a simple combination of primitive operations (index, matmul, log-softmax) rather than as an algorithm that could be redesigned to exploit GPU memory hierarchy.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

This paper is a systems implementation paper that presents a new algorithm and corresponding CUDA/Triton kernels for computing the cross-entropy loss during language model training. The system solves the problem of the cross-entropy layer consuming up to 90% of training GPU memory by never materializing the full vocabulary-size logit matrix into global memory—instead computing only the ground-truth logit and the log-sum-exp normalization constant via blockwise operations in on-chip SRAM, and exploiting softmax sparsity to skip negligible gradient computations.

3.2 Big-picture architecture (diagram in words)

CCE decomposes the cross-entropy loss computation into three major kernel operations, each designed to operate within on-chip SRAM without allocating large buffers in global GPU memory:

  1. Indexed Matrix Multiplication (forward): Given the batch of embeddings $E \in \mathbb{R}^{D \times N}$, the classifier matrix $C \in \mathbb{R}^{D \times |V|}$, and the ground-truth token indices $x \in \mathbb{R}^N$, this kernel directly computes the vector $(C^\top E)_x \in \mathbb{R}^N$ containing only the logit of the correct token for each position, without computing the logits for all other vocabulary entries. It does this by indexing into $C$ using $x$ and performing the dot product with the corresponding embedding column entirely in SRAM.

  2. Linear-Log-Sum-Exp, Forward Pass: This kernel computes $\text{LSE} = \log\sum_j \exp(C_j^\top E) \in \mathbb{R}^N$, the log-sum-exp normalization constant for each token position. It processes the matrix multiplication $C^\top E$ in blocks—loading sub-matrices of $C$ and $E$ into SRAM, computing the local logits, performing an online log-sum-exp reduction within each block, and then atomically updating a shared $\text{LSE}$ vector in global memory using a thread-safe $\log(\exp(a) + \exp(b))$ operation. The logit matrix $C^\top E$ is never written to global memory as an intermediate.

  3. Linear-Log-Sum-Exp + Indexed MatMul, Backward Pass (combined): This single kernel computes both $\nabla E$ and $\nabla C$ by recomputing blocks of the logit matrix in SRAM, reconstructing the corresponding softmax block $S_{nv} = \exp(A_{nv} - \text{LSE}_n)$, checking whether all elements in the block fall below the bfloat16 truncation threshold $\varepsilon = 2^{-12}$, and—if any are significant—performing the double outer-product gradient accumulation $\nabla E^\top_{n,d} \;+\!\!= (\hat{S}_{nv} \cdot \nabla\text{LSE}_n) C_{v,d}$ and $\nabla C^\top_{v,d} \;+\!\!= (\hat{S}_{nv} \cdot \nabla\text{LSE}_n)^\top E_{n,d}$ using atomic additions into the gradient buffers stored in global memory. The indexed matrix multiplication gradient (the term for the correct-token logit) is fused into this same kernel since it shares the same access pattern.

Information flows as follows: embeddings $E$ and ground-truth indices $x$ enter the system → the indexed matmul kernel produces $(C^\top E)_x$ → the forward LSE kernel produces $\text{LSE}$ → the backward pass receives $\nabla \text{LSE}$ from the autograd system, recomputes $C^\top E$ blocks in SRAM, filters blocks by softmax magnitude, and accumulates $\nabla E$ and $\nabla C$ into pre-allocated gradient buffers. The loss $\ell = (C^\top E)_x - \text{LSE}$ is computed in a trivial element-wise subtraction that consumes negligible memory.

Auxiliary components—vocabulary sorting and gradient filtering—are integrated into the backward pass kernel to maximize the number of blocks that can be skipped without affecting the gradient result.

3.3 Roadmap for the deep dive

  • First, the mathematical decomposition of cross-entropy into the indexed term and the log-sum-exp term (Equation 4), because this reformulation is what makes memory-efficient computation possible—it separates the only vocabulary-dependent operation (the log-sum-exp) from the token-specific operation (the correct-token logit).
  • Second, the indexed matrix multiplication kernel (Algorithm 1), since it establishes the pattern of blockwise SRAM-resident computation with zero global-memory allocation that all other kernels follow.
  • Third, the forward pass of the linear-log-sum-exp kernel (Algorithm 2), because it introduces the online log-sum-exp reduction and the thread-safe atomic update mechanism that the backward pass depends on.
  • Fourth, the backward pass of the linear-log-sum-exp kernel (Algorithm 3 and the combined Algorithm 4), since this is where the bulk of the computation and the key innovations—gradient filtering and vocabulary sorting—reside.
  • Fifth, gradient filtering and vocabulary sorting, because they are the techniques that make the backward pass fast enough to be competitive with baseline implementations despite the added cost of recomputing the logits. This includes the rationale for the threshold $\varepsilon = 2^{-12}$ and the empirical sparsity measurements from Figure 3.
  • Sixth, the numerical stability variants (CCE-Kahan, CCE-Kahan-FullC), because they represent the practical tradeoffs necessary to achieve bitwise-equivalent training in pretraining regimes where bfloat16 accumulation loses precision.

3.4 Detailed, sentence-based technical breakdown

This is primarily a systems implementation paper whose core idea is that the cross-entropy loss over a large vocabulary can be computed without materializing the full $|V| \times N$ logit matrix by decomposing the computation into an indexed dot-product and a blockwise log-sum-exp reduction, both performed entirely in on-chip SRAM, with the backward pass exploiting softmax sparsity to achieve competitive speed despite recomputation.

The Mathematical Decomposition

The key enabling observation is that the cross-entropy loss separates into two terms with fundamentally different memory requirements:

ℓ=(C⊤E)x−log⁡∑jexp⁡(Cj⊤E)\ell = \left(C^\top E\right)_x - \log\sum_j \exp\left(C^\top_j E\right)

where $E = [E_1 \ldots E_N] \in \mathbb{R}^{D \times N}$ is the batch of output embeddings from the backbone network, $C \in \mathbb{R}^{D \times |V|}$ is the classifier weight matrix, $x \in \mathbb{R}^N$ is the vector of ground-truth token indices (so $x_i \in \{1, \ldots, |V|\}$ is the vocabulary index of the correct next token for position $i$), and $\left(C^\top E\right)_x = \left[C_{x_1}^\top E_1 \ldots C_{x_N}^\top E_N\right]^\top \in \mathbb{R}^N$ is the vector containing only the logit of the correct token for each sequence position.

What it computes: the standard autoregressive cross-entropy loss over a batch of $N$ token predictions. The first term extracts, for each position $i$, the logit (unnormalized log-probability) of the correct next token $x_i$ from the full matrix of logits over all vocabulary entries. The second term computes, for each position $i$, the log of the sum of exponentiated logits over all $|V|$ vocabulary entries—this is the log-normalization constant that converts logits into log-probabilities. Subtracting the second term from the first yields the log-probability of the correct token, $\log P(x_i \mid x_1 \ldots x_{i-1})$. The total loss is the sum (or average) over the $N$ positions.

Why this form: the decomposition exposes that the first term is vocabulary-sparse—it only needs the $|V|$ logit values corresponding to the specific $N$ ground-truth tokens, not the entire $|V| \times N$ matrix. This can be computed via an index-and-dot-product operation with $O(N)$ memory. The second term is vocabulary-dense—it needs the sum over all $|V|$ entries for each position—but produces only an $N$-dimensional output vector. The standard implementation computes the full $|V| \times N$ logit matrix explicitly, stores it in global memory, then applies log-softmax. CCE recognizes that this intermediate matrix is never needed in its entirety: the indexed term only accesses $N$ specific entries, and the log-sum-exp term can be computed incrementally by processing blocks of vocabulary entries in SRAM and accumulating the running log-sum-exp in a single $N$-dimensional vector. This reformulation is what makes the memory reduction from $O(N \times |V|)$ to $O(N)$ possible without changing the mathematical result.

Memory-Efficient Indexed Matrix Multiplication (Algorithm 1)

The indexed matrix multiplication computes $o = (C^\top E)_x \in \mathbb{R}^N$, where $o_i = C_{x_i}^\top E_i$ is the dot product between the embedding of the $i$-th token and the classifier vector for its ground-truth label.

Algorithm 1 structure. The computation proceeds in two nested loops over blocks:

  1. Outer loop over batch blocks: divide the embedding matrix $E$ and the index vector $x$ into blocks $E_n$ of size $D \times N_B$ and index blocks $x_n$ of size $N_B$. For each block, allocate a zero vector $o_n$ of size $N_B$ in on-chip SRAM.

  2. Inner loop over the hidden dimension: further subdivide $E_n$ into $E_{n,d}$ of size $D_B \times N_B$ along the feature dimension. For each sub-block, perform an indexed load $c = C_{x_n, d}$—this retrieves, from the classifier matrix $C$, only the rows corresponding to the $D_B$ feature dimensions and only the columns corresponding to the ground-truth indices in $x_n$. This indexed load brings a $D_B \times N_B$ sub-matrix of $C$ into SRAM, matching the dimensions of $E_{n,d}$.

  3. Column-wise dot product accumulation: for each column $i$ in the block, compute the partial dot product between the $i$-th column of $E_{n,d}$ and the $i$-th column of the indexed classifier sub-matrix $c$, and accumulate into $o_n[i]$.

  4. Write output: after all $D_B$-sized blocks along the feature dimension are processed (the inner loop completes), write the completed $o_n$ from SRAM to global GPU memory.

What this achieves: the kernel produces the $N$-dimensional vector of correct-token logits while allocating zero global GPU memory beyond the input and output buffers. All intermediate computation—the indexed classifier columns, the embedding sub-blocks, and the running dot-product accumulators—resides entirely in on-chip SRAM. The total global memory footprint for this operation is $O(N)$ for the output, plus the existing $O(ND)$ and $O(D|V|)$ buffers for the inputs, which are allocated by the rest of the training pipeline regardless.

Why this design: the naive alternative would compute all $|V| \times N$ logits and then index-select the correct ones, costing $O(N|V|)$ memory for the logit matrix. A slightly more memory-conscious alternative would index into $C$ first to construct $C_x \in \mathbb{R}^{D \times N}$ and then perform a column-wise dot product, but this still requires materializing the $D \times N$ matrix $C_x$ in global memory. CCE's fused kernel avoids this intermediate allocation entirely by performing the indexing and dot product in a single pass through SRAM. The choice of block sizes $N_B$ and $D_B$ is determined by the available SRAM capacity and the GPU's cache line structure—larger blocks amortize the cost of the indexed load but consume more SRAM, while smaller blocks may underutilize the memory bandwidth. The paper does not specify exact block sizes, as these are likely auto-tuned by the Triton compiler.

Memory-Efficient Linear-Log-Sum-Exp, Forward Pass (Algorithm 2)

The forward pass computes $\text{LSE} = \log\sum_j \exp(C_j^\top E) \in \mathbb{R}^N$—for each token position, the log of the sum of exponentiated logits over the entire vocabulary. This is the normalization constant needed to convert the indexed logit $(C^\top E)_x$ into a log-probability.

Algorithm 2 structure. The key insight is that a log-sum-exp can be computed incrementally (online) over blocks of the vocabulary using the identity:

LSE([a;b])=log⁡(exp⁡(LSE(a))+exp⁡(LSE(b)))\text{LSE}([a; b]) = \log(\exp(\text{LSE}(a)) + \exp(\text{LSE}(b)))

where $[a; b]$ denotes concatenation. This means the vocabulary can be processed in blocks, with each block computing its local matrix multiplication and local log-sum-exp, then merging the result into a running global log-sum-exp vector.

Algorithm 2 proceeds as follows:

  1. Initialization: allocate $\text{LSE} = -\infty^N$, a vector of $N$ negative infinities in global GPU memory. This serves as the identity element for the log-sum-exp reduction ($\log(\exp(-\infty) + \exp(x)) = x$).

  2. Double loop over blocks: iterate over all pairs of embedding blocks $E_n$ (of size $D \times N_B$) and classifier blocks $C_v$ (of size $D \times V_B$). Each $(n, v)$ pair corresponds to a sub-block of the full logit matrix covering $V_B$ vocabulary entries and $N_B$ batch positions.

  3. Blockwise matrix multiplication: for each $(n, v)$ pair, perform a tiled matrix multiplication $A_{nv} = C_v^\top E_n$ entirely in SRAM, producing a $V_B \times N_B$ matrix of local logits. This uses an inner loop over $D_B$-sized chunks along the feature dimension, loading $C_{v,d}$ (size $D_B \times V_B$) and $E_{n,d}$ (size $D_B \times N_B$) and accumulating $A_{nv} \;+\!\!= C_{v,d}^\top \cdot E_{n,d}$.

  4. Local log-sum-exp: compute $\text{LSE}_{nv} = \log\sum \exp(A_{nv}^\top)$ along the vocabulary dimension. This is a numerically stable implementation that first computes the per-column maximum $m_{nv} = \max(A_{nv})$, subtracts it, exponentiates, sums, takes the log, and adds the maximum back: $\text{LSE}_{nv} = m_{nv} + \log\sum\exp(A_{nv} - m_{nv})$.

  5. Atomic global log-add-exp: update the global $\text{LSE}_n$ vector using $\text{LSE}_n = \log(\exp(\text{LSE}_n) + \exp(\text{LSE}_{nv}))$. This is the thread-safe merge of the local result into the running global reduction. The paper uses a spin-lock on an atomic operation in global memory to synchronize this update across different CUDA blocks that may be writing to the same position $n$ (because multiple vocabulary blocks $v$ all contribute to the same batch position $n$). The paper notes that this is "simple to implement in Triton" and "incurs little overhead," though alternative methods like an atomic compare-and-swap loop may perform better in a direct CUDA implementation.

What this achieves: the output is an $N$-dimensional vector $\text{LSE}$ in global memory. At no point is the full $|V| \times N$ logit matrix written to global memory. Each $V_B \times N_B$ block of logits is computed in SRAM, immediately reduced via log-sum-exp, and then discarded. The global memory footprint is $O(N)$ for the LSE vector (plus a spin-lock integer per position, negligible).

Why this design: a naive implementation of log-sum-exp over a matrix requires the entire matrix to be stored—first to find the per-column maximum (for numerical stability), then to compute the sum of exponentiated differences. By processing the vocabulary in blocks and using the online update $\log(\exp(a) + \exp(b))$, CCE only needs to store the running LSE vector, not the individual logits. The atomic update is necessary because different CUDA blocks processing different vocabulary ranges $v$ but the same batch range $n$ will attempt to update $\text{LSE}_n$ concurrently. The spin-lock ensures that only one block updates $\text{LSE}_n$ at a time, preventing race conditions where two blocks read the same old value, each compute $\log(\exp(\text{old}) + \exp(\text{new}))$, and one's update overwrites the other's. This synchronization cost is amortized over the large $V_B \times N_B$ matrix multiplication that precedes each update.

Parallelization strategy. The work is parallelized over the $(n, v)$ block pairs. With typical training batch sizes in the thousands of tokens and vocabularies in the hundreds of thousands, the number of blocks is $(N / N_B) \times (|V| / V_B)$. For $N = 8192$, $|V| = 256000$, and typical block sizes $N_B, V_B \approx 64$–$128$, this yields thousands to tens of thousands of CUDA blocks—more than enough to saturate a modern GPU. This is the same parallelization strategy used by standard tiled matrix multiplication, which is why CCE can achieve comparable throughput: it exposes the same amount of parallelism as a matrix multiply, just with a different reduction at the end.

Memory-Efficient Linear-Log-Sum-Exp, Backward Pass (Algorithm 3)

The backward pass receives $\nabla\text{LSE} \in \mathbb{R}^N$ (the gradient of the loss with respect to each element of the LSE vector) and must compute:

∇E⊤=(S⋅∇LSE) C\nabla E^\top = (S \cdot \nabla\text{LSE}) \, C

∇C⊤=(S⋅∇LSE)⊤E\nabla C^\top = (S \cdot \nabla\text{LSE})^\top E

where $S = \text{softmax}(C^\top E) \in \mathbb{R}^{|V| \times N}$ is the full softmax matrix, $\cdot$ denotes row-wise element-wise multiplication (broadcasting $\nabla\text{LSE} \in \mathbb{R}^N$ across the $|V|$ vocabulary entries for each position), and $\hat{S} = S \cdot \nabla\text{LSE}$ is the effective weight matrix for the gradient computation.

The challenge. Computing $\nabla E$ and $\nabla C$ requires, in principle, the full $|V| \times N$ matrix $\hat{S}$, which is the same size as the logit matrix that the forward pass avoided materializing. The backward pass is two matrix multiplications ($\hat{S} C$ and $\hat{S}^\top E$) with an intermediate non-linear operation (the softmax) applied to the $C^\top E$ product.

Algorithm 3 structure. CCE solves this by recomputing the logit matrix in SRAM, applying the softmax, and immediately using the resulting $\hat{S}$ block for the gradient accumulation—all in a single pass, without ever writing $\hat{S}$ to global memory:

  1. Recomputation of logits: for each block pair $(n, v)$, perform the same tiled matrix multiplication as in the forward pass: $A_{nv} = C_v^\top E_n$, producing a $V_B \times N_B$ logit block in SRAM.

  2. Softmax computation: compute $S_{nv} = \exp(A_{nv} - \text{LSE}_n)$. Crucially, the $\text{LSE}_n$ vector was already computed and stored in global memory during the forward pass, so the softmax normalization constant is available without recomputation. The subtraction $A_{nv} - \text{LSE}_n$ is the standard numerically stable softmax operation.

  3. Gradient filtering check: evaluate whether the entire block $S_{nv}$ contains only negligible values. The condition is $\text{all}(S_{nv} < \varepsilon)$ where $\varepsilon = 2^{-12}$, the smallest bfloat16 value that is not truncated during floating-point addition (see the gradient filtering section below). If the condition is true, skip the gradient accumulation for this block—the contribution of these softmax values to $\nabla E$ and $\nabla C$ would be zero after bfloat16 truncation anyway.

  4. Gradient accumulation (if block is non-negligible): for each $D_B$-sized sub-block along the feature dimension, perform two outer-product updates:

    • $\nabla E_{n,d}^\top \;\;+\!\!=\;\; (S_{nv} \cdot \nabla\text{LSE}_n) \; C_{v,d}$
    • $\nabla C_{v,d}^\top \;\;+\!\!=\;\; (S_{nv} \cdot \nabla\text{LSE}_n)^\top \; E_{n,d}$

These updates are performed using locking thread-safe atomic additions into the global gradient buffers $\nabla E$ and $\nabla C$. The locking is necessary because multiple CUDA blocks processing different vocabulary ranges $v$ will all contribute to the same embedding gradient positions $\nabla E_n$.

What this achieves: the backward pass produces the full $D \times N$ gradient $\nabla E$ and the $D \times |V|$ gradient $\nabla C$ using only the input buffers ($E$, $C$, $\text{LSE}$, $\nabla\text{LSE}$), the output gradient buffers (which are pre-allocated by the optimizer regardless of the loss implementation), and SRAM. The memory footprint is $O(N)$ for the LSE vector beyond the mandatory gradient buffers—the same asymptotic complexity as the forward pass.

Why recomputation: the alternative would be to store the softmax matrix $S$ during the forward pass and reuse it during the backward pass. But storing $S$ requires $O(N \times |V|)$ memory—exactly the memory that CCE is designed to avoid. Recomputation is the standard space-time tradeoff used throughout deep learning (e.g., activation checkpointing, FlashAttention): spend extra computation to avoid storing large intermediate tensors. In CCE's case, the recomputation cost is partially offset by the gradient filtering, which skips the most expensive part (the outer-product gradient updates) for the vast majority of blocks.

Why the combined backward kernel (Algorithm 4): the paper notes that "in practice, we found it to be easier and more memory-efficient to merge the indexed matrix-multiplication backward implementation with the backward pass of the linear-log-sum-exp operator." The indexed matrix multiplication backward pass computes the gradient of the loss with respect to $\nabla E$ and $\nabla C$ for the correct-token logit term $(C^\top E)_x$. This gradient is:

∂(C⊤E)x∂E=Cxand∂(C⊤E)x∂C=Ex\frac{\partial (C^\top E)_x}{\partial E} = C_x \quad \text{and} \quad \frac{\partial (C^\top E)_x}{\partial C} = E_x

where $C_x$ and $E_x$ are the indexed classifier columns and embeddings respectively. The combined kernel (Algorithm 4) fuses this with the LSE backward pass by replacing the softmax matrix $S_{nv}$ in Algorithm 3 with the full cross-entropy gradient with respect to the logits:

Gnv=1[v=xn⊤]−SnvG_{nv} = \mathbb{1}[v = x_n^\top] - S_{nv}

where $\mathbb{1}[v = x_n^\top]$ is an indicator matrix that is 1 at the position of the correct token and 0 elsewhere. This single $G_{nv}$ matrix captures both the positive gradient from the correct-token term and the negative gradient from the log-sum-exp term, and the rest of the gradient accumulation proceeds identically. The fusion avoids a separate pass over the data and eliminates the need to store intermediate per-token gradients.

Gradient Filtering

Gradient filtering is the mechanism that makes CCE's backward pass fast enough to be competitive despite the added cost of recomputing the logits. The core observation is that the softmax matrix $S$ over a large vocabulary is extremely sparse, and values below the bfloat16 truncation threshold cannot affect the gradient accumulation regardless.

The threshold. The paper chooses $\varepsilon = 2^{-12}$ as the filtering threshold. This value is derived from the properties of bfloat16 arithmetic:

  • bfloat16 uses a 7-bit fraction (mantissa) plus 1 implicit bit.
  • When adding two bfloat16 numbers $a$ and $b$ where $|a| < |b|$, the mantissa of the smaller number $a$ is shifted right until its exponent matches $b$'s exponent. If the exponent difference exceeds 7 (the number of mantissa bits), the shift pushes all bits of $a$'s mantissa past the least significant bit position of the sum, and $a$ is effectively rounded to zero.
  • Since softmax values are in $[0, 1]$, the largest possible value in a gradient accumulation is 1, which has exponent $2^0$. A value smaller than $2^{-7}$ would, in principle, be truncated when added to 1. However, the paper includes "5 extra bits above the fractional size" to account for rounding rules and the blocking strategies used in summation, arriving at $\varepsilon = 2^{-12}$.

What gradient filtering does: before performing the expensive outer-product gradient updates for a block $(n, v)$, the kernel checks whether $\text{all}(S_{nv} < \varepsilon)$—i.e., whether every softmax value in this $V_B \times N_B$ block is below the truncation threshold. If so, the block is skipped entirely: no $\nabla E$ updates, no $\nabla C$ updates, no global memory atomic operations. The computation for that block reduces to the matrix multiplication $A_{nv} = C_v^\top E_n$ (which must be performed anyway to determine the softmax values) and the filtering check.

Empirical sparsity (Figure 3). The paper measures the sorted softmax probabilities for the average token prediction in a Gemma 2 (2B) model. The results show:

  • The probability of the most likely token is approximately $10^{-1}$ to $10^{-2}$.
  • Probabilities decay rapidly: by approximately the 50th most likely token (out of 256,000), the probability has fallen below $2^{-12}$, the gradient filtering threshold.
  • On a log-log plot, there is a roughly linear relationship between log-rank and log-probability for the top $10^5$ tokens.
  • In practice, less than 0.02% of softmax elements are non-zero (above the filtering threshold). For a vocabulary of 256,000, this means fewer than approximately 50 tokens per position have non-trivial gradients.

This extreme sparsity explains why gradient filtering achieves a 3.5× speedup over the unfiltered version (Table 1, row 1 vs. row 7): the unfiltered backward pass must perform $|V| / V_B$ outer-product updates per batch block $n$, regardless of whether the softmax values are significant. With filtering, only a handful of vocabulary blocks per position contain any significant softmax values—the vast majority are skipped entirely after the cheap filtering check.

Why this is lossless. The filtering is not an approximation; it is an exact computation to within bfloat16 numerical precision. Any softmax value below $2^{-12}$, when multiplied by a gradient element $\nabla\text{LSE}_n$ (which is at most on the order of 1), produces a product below $2^{-12}$. When this product is atomically added to a bfloat16 gradient accumulator that already contains values up to roughly 1, the addition has no effect—the small value is truncated to zero during the floating-point addition. Filtering these values out simply avoids performing arithmetic whose result would be discarded by the hardware anyway. This is why the paper can claim "no loss of precision" for the basic CCE variant during fine-tuning (Figure 4): the result is bitwise identical to a computation that includes all softmax values but accumulates in bfloat16.

Block-level vs. element-level filtering. The filtering is performed at the block level ($\text{if all}(S_{nv} < \varepsilon)$) rather than at the individual element level. This is a practical constraint of the Triton framework: "control flow must be specified at the block level and therefore our... gradient filtering [is] constrained to operate at the block level as well." A CUDA implementation could filter at finer granularity (e.g., individual warps or threads), potentially skipping partially-populated blocks and achieving further speedups. The paper suggests this as a direction for future optimization.

Vocabulary Sorting

Vocabulary sorting is a complementary optimization that increases the effectiveness of gradient filtering by clustering tokens with similar logit magnitudes together, thereby making blocks more likely to be either entirely empty (skipped) or densely populated (processed efficiently).

The problem. Without sorting, the vocabulary order is arbitrary—determined by the Byte Pair Encoding construction process, which groups tokens by merge frequency rather than by semantic or statistical similarity. This means that a block of $V_B$ consecutive vocabulary entries might contain a mix of common tokens (with high average logits and significant softmax values) and rare tokens (with low average logits and negligible softmax values). Such a block would be partially populated, and the entire block must be processed (since not $\text{all}(S_{nv} < \varepsilon)$) even though only a fraction of its entries contribute meaningfully to the gradient. This wastes computation on the sparse regions within the block.

The solution. The paper proposes to sort the vocabulary by the average logit of each token, computed during the forward pass. Specifically, during the forward LSE computation, each CUDA block that processes a vocabulary block $v$ can atomically accumulate the sum (and count) of logits observed for each token in that block. After the forward pass, these sums are divided by the counts to produce an average logit per token. The vocabulary is then reordered so that tokens with similar average logits are adjacent. This requires a temporary buffer of size $O(|V|)$—approximately 1 MB for the largest contemporary vocabularies, which is negligible compared to the tens of gigabytes the logit matrix would otherwise consume.

Effect on block population. After sorting, tokens with high average logits (which will tend to have larger softmax values and thus contribute to gradients) are clustered together, as are tokens with low average logits (which will tend to have negligible softmax values). This means a given block $V_B$ is more likely to be either entirely high-logit (densely populated, and thus processed efficiently with minimal wasted work) or entirely low-logit (entirely below the filtering threshold, and thus skipped entirely). The number of partially-populated blocks—where some entries contribute and others don't—is minimized.

Quantified impact. Table 1 (row 1 vs. row 6) shows that without vocabulary sorting, the backward pass takes 115 ms compared to 100 ms with sorting—a 15% (15 ms) increase. This is a meaningful but not dominant effect compared to the 3.5× improvement from gradient filtering, consistent with sorting being a refinement that increases the hit rate of the filtering mechanism rather than a primary source of speedup on its own.

Why average logit: the average logit is a natural proxy for "how often this token contributes to the gradient." Tokens with high logits are those the model frequently predicts as likely; their corresponding softmax values will often be above the filtering threshold. Tokens with very low logits are those the model rarely considers probable; their softmax values will almost always fall below the threshold. Sorting by this metric groups tokens by their expected contribution to the gradient computation, which directly optimizes the block-level sparsity pattern that gradient filtering exploits.

The Combined Backward Pass (Algorithm 4)

Algorithm 4 merges the indexed matrix multiplication backward pass with the LSE backward pass into a single kernel. The combined kernel replaces $S_{nv}$ from Algorithm 3 with $G_{nv}$:

Gnv=1[v=xn⊤]−SnvG_{nv} = \mathbb{1}[v = x_n^\top] - S_{nv}

where $\mathbb{1}[v = x_n^\top]$ is an indicator matrix: element $(j, i)$ is 1 if vocabulary entry $v_j$ (the $j$-th token in the current vocabulary block) matches the ground-truth token $x_{n_i}$ (the correct token for the $i$-th batch element in the current batch block), and 0 otherwise.

What $G_{nv}$ represents: this is the gradient of the cross-entropy loss with respect to the logit matrix $C^\top E$. For a single position with correct token $x$ and softmax distribution $s$, the gradient is $e_x - s$ where $e_x$ is the one-hot vector at position $x$. The first term comes from the derivative of the correct-token logit (which increases the logit of the correct answer), and the second term comes from the derivative of the log-sum-exp (which decreases all logits in proportion to their softmax probability). The combined kernel computes $G_{nv}$ directly from the recomputed logits and the known ground-truth indices, then uses $G_{nv} \cdot \nabla\text{CEL}_n$ (where $\nabla\text{CEL}$ is the gradient of the scalar loss with respect to the log-probability) in place of $S_{nv} \cdot \nabla\text{LSE}_n$ from Algorithm 3 for the outer-product updates.

Why fusion is beneficial: the indexed matrix multiplication gradient and the LSE gradient share the same access pattern—both need $E$, $C$, the recomputed logits, and the ground-truth indices—and both produce updates to $\nabla E$ and $\nabla C$. Computing them in separate kernels would require loading $E$ and $C$ from global memory twice and performing the logit recomputation twice. Fusing them into a single kernel amortizes the memory loads and the recomputation cost over both gradient contributions.

Gradient filtering with $G_{nv}$: the combined kernel applies the same block-level filtering check, but on $|G_{nv}|$ rather than $S_{nv}$. A block is skipped if $\text{all}(|G_{nv}| < \varepsilon)$. This is a slight relaxation compared to filtering on $S_{nv}$ alone, since $G_{nv}$ includes the one-hot indicator which can create isolated non-zero entries in blocks that would otherwise be all-zeros under $S_{nv}$. However, since each position has exactly one correct token, the indicator contributes at most $N_B$ non-zero entries across all vocabulary blocks for a given batch block—a negligible overhead that does not meaningfully affect the sparsity pattern.

Numerical Stability Variants

The paper identifies two sources of numerical error that affect pretraining but not fine-tuning:

Source 1: Gradient filtering prevents gradient flow to rare tokens during pretraining. During fine-tuning, the model starts from a pretrained checkpoint where all vocabulary entries already have reasonable classifier vectors. Even rare tokens have received gradient updates during pretraining and sit in a reasonable region of the embedding space. During pretraining from scratch, however, gradient filtering applied to $\nabla C$ means that tokens with very low softmax probabilities receive no gradient updates at all during many training steps—their gradient contribution is filtered out because their softmax values are below $\varepsilon$ on every example where they are not the ground truth. These tokens, which include rare vocabulary entries that appear infrequently in the training data, fail to learn reasonable representations, which harms overall model quality. This is not an issue in fine-tuning because the token representations are already learned.

Source 2: bfloat16 summation in global memory loses precision. CCE performs the gradient accumulation $\nabla E_{n,d}^\top \;\;+\!\!=\;\; \hat{S}_{nv} C_{v,d}$ using atomic additions in the bfloat16 data type (the final gradient type). For large reductions—particularly the $\nabla C$ accumulation, which sums over all $N$ batch positions, each contributing a $D_B \times V_B$ outer product—bfloat16's limited precision (7 mantissa bits) causes truncation errors that accumulate over many atomic additions. During fine-tuning with small gradient magnitudes, this error is negligible. During pretraining with larger gradient magnitudes and many more training steps, the accumulated error can measurably degrade convergence.

CCE-Kahan (Table 1, rows 8-10): this variant addresses Source 2 by replacing the bfloat16 atomic additions with Kahan summation (Kahan, 1965). Kahan summation maintains a running compensation term that captures the low-order bits lost during each floating-point addition, effectively doubling the precision of the accumulation without changing the data type. The cost is increased memory: the compensation term requires an additional buffer of the same size as the gradient, roughly doubling the memory for $\nabla E$ and $\nabla C$ during the backward pass (from ~1 GB to ~2 GB for the Gemma 2 (2B) configuration). The time cost is modest: the backward pass takes 114 ms with Kahan vs. 100 ms for basic CCE (Table 1, row 8 vs. row 1), a 14% increase. The forward pass is unaffected since Kahan summation only applies to the gradient accumulation in the backward pass.

CCE-Kahan-FullC (Table 1, row 9): this variant addresses both Source 1 and Source 2 by combining Kahan summation with disabling gradient filtering for the $\nabla C$ computation. The $\nabla E$ gradient filtering remains active (its accumulation is over the vocabulary dimension, which is typically large and benefits from filtering), but the $\nabla C$ computation processes all blocks regardless of softmax magnitude. This ensures that even rare tokens with very low softmax probabilities receive gradient updates during pretraining, allowing them to learn meaningful representations. The cost is a substantial slowdown: the backward pass takes 268 ms with CCE-Kahan-FullC vs. 100 ms for basic CCE (Table 1, row 9 vs. row 1), a 2.7× increase. However, the paper notes that this slowdown is "often offset by the larger batch sizes CCE-Kahan-FullC enables"—for Mistral NeMo, the larger batch size enabled by the memory savings reduced total training time by 2 hours (16%) compared to torch.compile despite the per-step slowdown.

CCE-Kahan-FullE (Table 1, row 10): for completeness, the paper also tests a variant that disables gradient filtering only for $\nabla E$ (keeping filtering for $\nabla C$). This variant performs similarly to CCE-Kahan-FullC (247 ms vs. 268 ms), confirming that the $\nabla C$ filtering is the primary source of pretraining instability. The paper focuses on CCE-Kahan-FullC as the recommended pretraining variant because it is the more conservative choice that guarantees no gradient starvation for any parameter.

Practical guidance. The paper's experiments (Figures 4 and 5) establish a clear two-tier recommendation:

  • For fine-tuning: Use basic CCE. Gradient filtering with $\varepsilon = 2^{-12}$ produces training curves indistinguishable from torch.compile across all four tested models (Gemma 2 2B, Phi 3.5 Mini, Qwen 2.5 7B, Mistral NeMo), with no loss of convergence quality and maximum memory savings.
  • For pretraining: Use CCE-Kahan-FullC. The Kahan summation prevents accumulation errors from bfloat16 reduction, and disabling gradient filtering for $\nabla C$ ensures that rare tokens receive gradient updates. This variant matches torch.compile's validation perplexity curves identically across all four tested models (Figure 5), while still providing substantial memory savings over the baseline—the temporary buffers for Kahan summation are typically smaller than the memory that would have been used for the logit matrix, so the net memory reduction remains significant.

Why the memory increase for Kahan variants is tolerable. The paper notes that "the increased memory usage of CCE-Kahan-FullC vs. CCE is due to temporary buffers used in the backward pass. The size of these buffers is typically less than the amount of free memory needed to rematerialize activations when using activation/gradient checkpointing. Thus CCE-Kahan-FullC often shares the same memory saving benefits as CCE." In other words, the Kahan compensation buffers are small enough that they fit within the memory headroom that activation checkpointing already requires, so the effective memory reduction for the full training loop remains substantial even with Kahan summation enabled.

Summary of Design Choices and Their Justifications

  • Decomposition into indexed term + LSE term rather than computing full log-softmax: avoids materializing the $|V| \times N$ logit matrix by separating the vocabulary-sparse operation (indexed dot product) from the vocabulary-dense operation (log-sum-exp reduction), enabling each to be optimized independently with different memory access patterns.

  • Blockwise SRAM-resident computation with online log-sum-exp reduction rather than chunked global-memory approaches: eliminates the memory-latency tradeoff inherent in chunking—chunking reduces memory by processing a subset of vocabulary entries at a time but incurs kernel launch overhead and repeated global memory loads. SRAM-resident blocks achieve both minimal memory and minimal latency by keeping all intermediate data on-chip.

  • Recomputation of logits in the backward pass rather than storing the softmax matrix: applies the standard space-time tradeoff to the loss layer. The recomputation cost is offset by gradient filtering, which skips the expensive gradient accumulation for the vast majority of blocks.

  • Gradient filtering at $\varepsilon = 2^{-12}$ rather than an approximate or learned threshold: the threshold is derived from bfloat16 hardware properties, not heuristics. Values below this threshold are provably truncated during accumulation regardless, so filtering them loses no information. This makes the optimization mathematically exact for bfloat16 gradients, which is why fine-tuning loss curves match exactly.

  • Block-level filtering rather than element-level: a practical concession to the Triton framework's control-flow constraints. The paper acknowledges this as a limitation and suggests that a CUDA implementation with finer-grained filtering could yield further speedups.

  • Vocabulary sorting by average logit rather than random or frequency-based ordering: directly optimizes the block-sparsity pattern that gradient filtering exploits. Tokens with similar logit magnitudes are clustered, making blocks more likely to be either entirely skippable or densely productive.

  • Kahan summation for pretraining rather than always-on or always-off: a targeted solution to the specific numerical errors that matter during pretraining (accumulation precision over many steps). Fine-tuning doesn't need it because gradient magnitudes are smaller and the training duration is shorter.

  • Disabling gradient filtering for $\nabla C$ during pretraining rather than a higher threshold: addresses the gradient starvation problem for rare tokens without fundamentally changing the filtering mechanism. The $\nabla C$ computation only covers $D \times |V|$ parameters, which is small relative to the total model size, so the computational cost of processing all blocks is acceptable.

  • Triton implementation rather than direct CUDA: enables rapid experimentation and easy integration with PyTorch autograd, but the paper notes that Triton's block-level control flow limits optimization granularity and that a CUDA implementation could potentially achieve further speedups through warp-level or thread-level filtering decisions.

4. Key Insights and Innovations

Innovation 1: The Logit Matrix Is a Memory Artifact, Not a Computational Necessity

The paper's most intellectually distinctive contribution is the recognition that the full |V| × N logit matrix—the intermediate tensor that standard cross-entropy implementations materialize, store, and operate on—is an artifact of how we've chosen to implement the loss, not a mathematical requirement of the computation itself. This is a conceptual reframing that changes how one thinks about the classification head: from an inevitable memory hog that must be endured, chunked, or parallelized across GPUs, to a problem of algorithm design where the intermediate representation can be eliminated entirely through kernel fusion and online reduction.

Prior to CCE, the dominant approaches to the vocabulary-memory problem all accepted the necessity of the logit matrix in some form. Chunked implementations (Torch Tune, Liger Kernels) reduce the peak size by processing vocabulary subsets sequentially but still write logit chunks to global memory—the O(N × |V|) memory dependence is reduced by a constant factor (the number of chunks) but not eliminated. Sequence and model parallelism distribute the matrix across GPUs, trading memory for communication but again treating the logit matrix as fundamentally necessary. Vocabulary reduction techniques (Grave et al., 2017; Yu et al., 2023) avoid the large-|V| regime entirely, solving the memory problem by changing the model architecture rather than the loss computation. In all cases, the logit matrix is treated as a given—something to be managed, not something to be questioned.

CCE's key reframing comes from the mathematical decomposition ℓ = (C^⊤ E)_x − log Σ_j exp(C_j^⊤ E). By recognizing that the indexed term needs only N specific logit entries and the log-sum-exp term can be computed incrementally over vocabulary blocks with a running reduction, CCE demonstrates that the full logit matrix is never needed in its entirety at any point in the computation. Every operation that standard implementations perform on the stored logit matrix (indexing, exponentiation, summation, softmax gradient computation) can instead be performed on the fly as each logit block is computed in SRAM, immediately consumed for the reduction or gradient accumulation, and then discarded. The logit matrix transitions from being a stored intermediate to being a transient quantity that exists only fleetingly in on-chip memory.

This reframing is not merely an implementation detail—it is a diagnostic insight with implications beyond cross-entropy. It suggests that many deep learning layers that currently materialize large intermediate tensors (contrastive loss matrices, full attention weight matrices before FlashAttention, pairwise distance computations) might similarly be amenable to SRAM-resident recomputation with online reduction, provided the downstream operations can be expressed in a form amenable to incremental accumulation. The paper draws the explicit parallel to FlashAttention, and together these works establish a design pattern: IO-aware algorithms for deep learning are not about doing less computation, but about reorganizing computation so that large intermediates never touch global memory. This is a fundamental contribution to the systems-for-ML design philosophy, not an incremental refinement of existing loss implementations.

The evidence for this insight's power is Table 1: CCE achieves a memory footprint for the loss computation (1 MB) that is 24,000× smaller than the baseline (24 GB) while maintaining comparable or better latency (145 ms vs. 143 ms for torch.compile on loss+gradient). This is not a tradeoff—it is a simultaneous improvement in both memory and speed, which is the hallmark of an algorithmic insight rather than an engineering optimization.

Innovation 2: Gradient Filtering as Exact Numerical Exploitation of Hardware Precision Limits

The second conceptual contribution is the recognition that the sparsity of the softmax distribution over a large vocabulary can be exploited in a mathematically exact way, not approximately, by aligning the filtering threshold with the hardware truncation boundary of the gradient accumulation data type. This is a fundamentally different approach to sparsity than the typical machine learning pattern of pruning, thresholding, or approximating.

The standard approach to exploiting sparsity in deep learning is to introduce approximation: set small weights to zero (pruning), use low-rank approximations, apply dropout, or employ top-k selection on softmax outputs. These are all lossy operations—they change the mathematical result, and their justification rests on the empirical observation that the approximation error does not harm model quality. The risk is always that the approximation interacts poorly with some aspect of training dynamics, requiring careful tuning of thresholds, schedules, or compensation mechanisms.

CCE's gradient filtering takes a fundamentally different approach: the threshold ε = 2⁻¹² is not a heuristic chosen by observing loss curves or validation perplexity; it is derived from the hardware specification of bfloat16 floating-point addition. The reasoning (detailed in Section 3.4 and Appendix E) is that when two bfloat16 numbers are added and their exponents differ by more than 7 (the number of mantissa bits), the smaller number is shifted to zero during the alignment step and contributes nothing to the sum. Since softmax values lie in [0, 1] and gradient accumulators can hold values up to roughly 1, any softmax value below 2⁻¹² is provably truncated when added to the accumulator—regardless of the specific values involved. Filtering these values out is not an approximation; it is an exact computation to within the precision of the data type.

This is a conceptual innovation in how to think about numerical optimization: rather than treating floating-point precision as a source of error to be managed (e.g., via mixed-precision training, loss scaling, or Kahan summation), CCE treats it as a source of information about which computations are unnecessary. The hardware's truncation behavior becomes a guarantee—a contract that says "values below this threshold cannot affect the result"—and the algorithm exploits that guarantee to skip work. This is a form of zero-cost sparsity that requires no tuning, no scheduling, and no verification beyond confirming that the data type and accumulation pattern match the assumptions.

The significance of this innovation extends beyond cross-entropy. Any deep learning operation that produces a distribution over a large output space and accumulates gradients in a low-precision format (as is standard in mixed-precision training) could potentially apply the same principle: compute the distribution, check against the precision-derived threshold, and skip gradient updates for elements whose contribution would be truncated. The paper's Figure 3 provides the diagnostic tool for determining when this approach will be effective—a log-log plot of sorted softmax probabilities that reveals how quickly the distribution drops below truncation—and this diagnostic pattern could be applied to other domains.

The evidence that this is truly lossless for fine-tuning is Figure 4: across four different models (Gemma 2 2B, Phi 3.5 Mini, Qwen 2.5 7B, Mistral NeMo), the training loss curves for CCE and torch.compile are "nearly indistinguishable." This is the empirical confirmation that the theoretical argument (filtered values would be truncated anyway) holds in practice.

Innovation 3: The Fine-Tuning vs. Pretraining Numerical Sensitivity Distinction

The paper's third contribution is a diagnostic finding rather than a method: the discovery that the same gradient filtering that is provably lossless for fine-tuning causes measurable degradation during pretraining, and the identification of the specific mechanisms responsible. This distinction is not obvious a priori—one might expect that if gradient filtering is mathematically exact to bfloat16 precision, it should be equally valid in any training regime. The paper shows that it isn't, and in doing so reveals something about the different numerical demands of fine-tuning versus pretraining that has practical implications for any technique that exploits sparsity or low-precision accumulation.

The degradation in pretraining comes from two sources (Section 5.3):

Gradient starvation for rare tokens. During fine-tuning, the model starts from a pretrained checkpoint where all vocabulary entries have reasonable classifier vectors—they've already received gradient updates during pretraining and sit in a functionally useful region of the embedding space. Gradient filtering during fine-tuning may skip updates for rare tokens on individual batches, but those tokens' representations are already learned and stable; an occasional skipped update does not cause them to drift or collapse. During pretraining from scratch, however, tokens with very low softmax probabilities never receive gradient updates because their softmax values are always below ε on examples where they are not the ground truth. These tokens—which include rare vocabulary entries that appear infrequently in the training corpus—fail to learn meaningful representations entirely. The model's classifier for these tokens remains close to its random initialization, which harms overall language modeling quality because even rare tokens need reasonable representations for the few times they do appear.

This is not a flaw in the threshold derivation—each individual skipped update would indeed be numerically truncated if performed. The problem is cumulative: over many training steps, the absence of any gradient flow to rare tokens prevents them from ever reaching a state where they would receive non-trivial gradients. The gradient filtering creates a "rich get richer" dynamic where tokens that are already well-represented continue to improve while rare tokens stagnate. This is a form of emergent bias from a per-step lossless optimization—a subtlety that would be easy to miss without the careful pretraining experiments in Figure 5.

Accumulation precision in large reductions. The ∇C gradient computation sums contributions from all N batch positions (potentially thousands) into a single bfloat16 accumulator. Even with Kahan summation disabled, the bfloat16 additions progressively lose precision as the accumulator grows. During fine-tuning with smaller gradient magnitudes and fewer total steps, this error is negligible. During pretraining from scratch with larger learning rates and many more steps, the accumulated precision loss measurably affects convergence. The paper's identification of this as a distinct source of error—separate from the gradient filtering mechanism—is important because it suggests that the two fixes (Kahan summation and FullC) address independent problems.

The practical outcome is a two-tier recommendation that is genuinely novel: basic CCE for fine-tuning, CCE-Kahan-FullC for pretraining. This is not a compromise or a weakness; it is a principled response to a real phenomenon. The paper demonstrates that getting this distinction right matters—Figure 5 shows that CCE-Kahan-FullC matches torch.compile's validation perplexity identically during pretraining across four models, while the basic CCE variant (not shown but implicitly tested during the authors' "initial experiments") does not. This finding establishes that memory-saving optimizations cannot be validated solely on fine-tuning benchmarks; pretraining from scratch stresses numerical precision in ways that fine-tuning does not, and optimization techniques must be evaluated in both regimes to be considered generally applicable.

Innovation 4: Vocabulary Sorting as a Mechanism for Amplifying Block-Level Sparsity

The fourth contribution is more incremental but still conceptually interesting: the recognition that the order of the vocabulary dimension is a free parameter that can be optimized to improve the effectiveness of block-level filtering, and that sorting by average logit (a statistic computable during the forward pass with negligible overhead) is an effective heuristic for this optimization. This transforms vocabulary ordering from an arbitrary byproduct of the tokenization process into a design choice that directly affects computational efficiency.

The default vocabulary ordering in LLMs is determined by the Byte Pair Encoding construction process: tokens are created by iteratively merging the most frequent byte pairs in the training corpus, and the resulting vocabulary order reflects merge frequency. This ordering has no relationship to the statistical properties that matter for gradient computation—whether a token's softmax probability will typically be above or below the filtering threshold. As a result, a block of consecutive vocabulary entries will contain a random mix of common and rare tokens, producing partially-populated blocks that gradient filtering cannot skip (because at least one entry has a significant softmax value) but where significant computation is wasted on the entries that don't.

Vocabulary sorting is a data layout optimization that increases the effective sparsity at the block granularity. By clustering tokens with similar average logits, the sorted vocabulary makes blocks more homogeneous: high-logit blocks are densely productive (every entry contributes to the gradient, so no computation is wasted), and low-logit blocks are entirely skippable (no entry exceeds ε, so the block is filtered). The number of partially-populated blocks—which must be processed but contain wasted work—is minimized.

What makes this idea intellectually interesting rather than merely a performance trick is that it reveals a degree of freedom in the classifier head that most practitioners treat as fixed. The vocabulary order is arbitrary; permuting it does not change the model's output or the loss computation in any mathematical sense, because the softmax is permutation-equivariant (reordering the vocabulary entries reorders the softmax probabilities correspondingly). But the computational efficiency of CCE's backward pass depends strongly on the permutation, because the block-level filtering check operates on contiguous blocks of the permuted vocabulary. This is a clean separation between the mathematical specification of the computation (which is permutation-invariant) and the implementation efficiency (which is permutation-sensitive), and the paper exploits this separation by choosing the permutation that maximizes efficiency.

The contribution is incremental because the speedup from vocabulary sorting is modest (15% on the backward pass, per Table 1 row 1 vs. row 6) compared to the 3.5× from gradient filtering, and because the sorting mechanism itself (average logit, computed via atomic additions during the forward pass) is straightforward. But the conceptual pattern—identifying a permutation degree of freedom in a computation and optimizing it for hardware efficiency—is broadly applicable. Many deep learning operations have permutation symmetries (channel order in convolutions, head order in multi-head attention, sequence order in set operations) where reordering could improve cache locality, reduce warp divergence, or increase the effectiveness of block-level filtering. The paper provides a concrete example of this pattern being applied successfully, which may inspire similar optimizations elsewhere.

The evidence is clear if modest: the backward pass time drops from 115 ms to 100 ms (Table 1), and the paper notes that this improvement comes from increasing the proportion of blocks that are either entirely skipped or entirely productive. The temporary buffer required for the sorting statistics (≈1 MB for the largest vocabularies) is negligible compared to the memory savings CCE provides, making this a pure win with no tradeoff.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The primary computational benchmarks use the Alpaca dataset (Taori et al., 2023) for inputs and labels. For fine-tuning convergence experiments, the full Alpaca dataset is used. For pretraining validation perplexity experiments, the paper uses a 5% subset of the Open WebText corpus (Gokaslan et al., 2019) for training, with a held-out 0.25% subset of Open WebText (non-overlapping with the training split) for validation. The specific models and configurations for each experiment vary and are specified in the relevant subsections.

  • Base model(s). The runtime and memory benchmarks (Table 1, Table A3) use classifier head configurations drawn from a range of contemporary frontier models: Gemma 2 at 2B, 9B, and 27B scales (Rivière et al., 2024), Phi 3.5 Mini (Abdin et al., 2024), Mistral NeMo (Mistral AI Team, 2024), and Qwen 2.5 at 7B and 32B scales (Qwen Team, 2024). The fine-tuning convergence experiments (Figure 4) use the instruction-tuned variants of Gemma 2 2B, Phi 3.5 Mini, Qwen 2.5 7B, and Mistral NeMo. The pretraining convergence experiments (Figure 5) use the same four model architectures. This diverse set was chosen to demonstrate that CCE's benefits generalize across model scales (2B to 32B parameters), vocabulary sizes (32,064 for Phi 3.5 Mini to 256,128 for Gemma 2), and hidden dimensions (2,304 to 5,120), capturing the range of |V|/D ratios encountered in practice.

  • Metrics. The paper employs two complementary sets of metrics. For computational performance, the metrics are peak GPU memory footprint (in MB, measured via PyTorch's memory allocator, rounded to the nearest MB) and wall-clock time (in milliseconds, averaged over 5 random seeds) for three stages: the forward loss computation, the backward gradient computation, and the combined loss+gradient computation. Memory measurements distinguish between the loss and gradient phases because "intermediate buffers can often... be reused between the loss and gradient computation, resulting in lower peak memory consumption than the sum of the parts" (Table 1 caption). For training quality, the metrics are training loss (for fine-tuning; Figure 4) and validation perplexity on held-out text (for pretraining; Figure 5), both tracked over gradient steps.

  • Baselines. The paper compares CCE against five baselines, each representing a different point in the memory-latency design space. Baseline: the default PyTorch cross-entropy implementation (F.cross_entropy), which materializes the full logit matrix into global memory and is the standard in frameworks such as Torch Tune and HuggingFace Transformers. torch.compile: Baseline wrapped with torch.compile (Ansel et al., 2024), which applies kernel fusion and other graph-level optimizations. Torch Tune (8 chunks): the chunked cross-entropy implementation from the Torch Tune library (Torch Tune Team, 2024), configured with 8 chunks, combined with torch.compile. Liger Kernels: the fused forward-backward cross-entropy kernel from the Liger Kernels library (Hsu et al., 2024), which computes loss and gradient simultaneously to avoid storing intermediate logits. Lower bound: the theoretical minimum memory—only the output gradient buffers ∇E and ∇C—representing the memory that must be allocated regardless of the cross-entropy implementation (since the optimizer needs these gradients). Additionally, six CCE variants are evaluated: CCE (the basic configuration with gradient filtering and vocabulary sorting), CCE (No Vocab Sorting), CCE (No Grad. Filter), CCE-Kahan (Kahan summation for gradient accumulation), CCE-Kahan-FullC (Kahan summation + gradient filtering disabled for ∇C), and CCE-Kahan-FullE (Kahan summation + gradient filtering disabled for ∇E). For the training convergence experiments (Figures 4 and 5), the baseline is torch.compile cross-entropy, chosen because it is the fastest memory-efficient baseline while still using standard PyTorch operations (unlike Liger Kernels, which requires kernel-level implementation of loss transformations).

  • Generation budget / compute accounting. Memory and runtime are measured on a single NVIDIA A100-SXM4 GPU with 80 GB of HBM, using PyTorch 2.4.1 and CUDA 12.4. All measurements are taken for a fixed batch of 8,192 tokens with a vocabulary size of 256,000 and hidden dimension 2,304 for the primary Gemma 2 (2B) benchmarks (Table 1), with additional configurations evaluated in Tables A3, A1 and Figures A1, A2. The embedding matrix E and classifier matrix C are loaded from actual Gemma 2 (2B) Instruct weights during the Alpaca fine-tuning process, ensuring that memory access patterns and sparsity characteristics reflect realistic training conditions rather than synthetic data. For the "removing ignored tokens" analysis (Table A1), the batch is pre-filtered to exclude padding, system prompt, and user input tokens before the loss computation. For the maximum batch size analysis (Figure 1, Table A4), the paper assumes a 16-GPU fully-sharded data-parallel setup with 75 GB usable memory per GPU (80 GB minus a 5 GB buffer for libraries) and a global batch of 65,536 tokens, computing the maximum per-GPU batch size as (total memory − weights+optimizer+gradients memory) / (activation memory + logit memory per token).

  • Cross-validation / statistical protocol. All runtime and memory measurements are averaged over 5 random seeds to account for GPU scheduling variability, though the paper notes that many memory figures are "multiples of 1,000 due to dimensions chosen and PyTorch's allocation strategy," indicating deterministic allocation patterns. The training convergence experiments (Figures 4 and 5) are also averaged over 5 random seeds to capture training stochasticity, with 95% confidence intervals shown via shaded regions or p95 confidence ranges. For the fine-tuning experiments, each model is fine-tuned for 700 gradient steps; for pretraining, each model is trained for 1,500 gradient steps on the 5% Open WebText subset. The paper does not employ cross-validation for the computational benchmarks (since these are deterministic given hardware and software versions) or for the training experiments (since the goal is to demonstrate that loss curves match, not to select hyperparameters).

Main Quantitative Results

Computational Performance: CCE Achieves Near-Lower-Bound Memory with Competitive Latency

The headline result of the computational benchmarks is that CCE reduces the memory footprint of the cross-entropy computation to within 1–3 MB of the theoretical lower bound while maintaining latency comparable to or better than the fastest existing implementation. Table 1 provides the primary evidence for the Gemma 2 (2B) configuration (batch size 8,192, vocabulary size 256,000, hidden dimension 2,304).

Memory footprint. For the combined loss+gradient computation, CCE uses 1,164 MB of GPU memory. This compares to 28,000 MB for Baseline (a 24× reduction), 16,000 MB for torch.compile (a 13.7× reduction), 9,631 MB for Torch Tune with 8 chunks (8.3× reduction), and 1,474 MB for Liger Kernels (1.3× reduction). The lower bound—the memory required for the gradient buffers ∇E and ∇C alone—is 1,161 MB, meaning CCE's incremental memory overhead beyond the mandatory gradient storage is only 3 MB (0.26% overhead). The paper explicitly states this is "within 3 MB of the lowest possible memory consumption" (Appendix C.2, for the Gemma 2 2B case).

Breaking this down by phase: for the forward loss computation alone, CCE uses 1 MB of memory versus 24,000 MB for Baseline (a factor of 24,000×) and 4,000 MB for torch.compile. For the backward gradient computation alone, CCE uses 1,163 MB versus 16,000 MB for Baseline and 12,000 MB for torch.compile, with a lower bound of 1,161 MB—again, only 2 MB of overhead.

The key comparison to Liger Kernels deserves attention. Liger Kernels achieves 1,474 MB for the combined computation—competitive with CCE's 1,164 MB and dramatically better than the other baselines. However, Liger Kernels' memory usage grows with O(N × D) rather than CCE's O(N + |V|), meaning that for very large vocabularies or long sequences, CCE's asymptotic advantage becomes decisive. The paper also notes that Liger Kernels "computes the loss and gradient simultaneously, not in separate forward/backward passes" (Table 1, footnote 2), which is why it shows equal memory for loss and gradient in the table—both are computed in a single fused operation. This fusion prevents per-phase memory measurement and, as noted in Section 2, requires any loss transformation to be implemented inside the kernel.

Latency. For the combined loss+gradient computation, CCE takes 145 ms, which is essentially identical to torch.compile's 143 ms (within 2 ms, or a 1.4% difference). Baseline takes 208 ms, Torch Tune takes 169 ms, and Liger Kernels takes 304 ms. These numbers yield several observations:

  • CCE is 6% faster than Baseline (145 ms vs. 208 ms) for the combined computation, despite avoiding materialization of the logit matrix. The paper attributes this to the fact that CCE "does not write all the logits to global memory" and "is able to save time in other parts of the computation" (Section 5.1).
  • Liger Kernels is dramatically slower than all other methods (304 ms vs. 143–208 ms), more than doubling the Baseline latency and taking over 2× CCE's time. This is the penalty for aggressive chunking to minimize memory: kernel launch overhead dominates.
  • Torch Tune with 8 chunks achieves 169 ms, a 17% slowdown versus torch.compile (143 ms), demonstrating the memory-latency tradeoff inherent in chunked approaches. CCE achieves both lower memory (1,164 MB vs. 9,631 MB) and lower latency (145 ms vs. 169 ms) than Torch Tune—a simultaneous improvement on both axes.

The forward pass latency reveals an interesting pattern: CCE computes the loss in 46 ms, which is 6% faster than torch.compile's 49 ms and 44% faster than Baseline's 82 ms. This is because CCE's fused linear-log-sum-exp kernel computes the matrix multiplication and log-sum-exp reduction in a single pass through SRAM, avoiding the separate kernel calls and global memory writes that Baseline and torch.compile require for the logit matrix, softmax, and log operations. The backward pass, by contrast, takes 100 ms for CCE versus 92 ms for torch.compile—an 8.7% slowdown. This slowdown reflects the cost of recomputing the logits during the backward pass, which CCE must do to avoid storing the softmax matrix from the forward pass. However, as the paper notes in Appendix C.1 (Table A2), this recomputation cost (45 ms) is more than offset by time saved on other backward pass operations (saving 62 ms total on the cross-entropy gradient, softcapping gradient, and ∇C computation), resulting in a net speedup for the forward+backward combination.

Scaling with batch size. Figures A1 and A2 show how latency scales with the number of tokens for all methods across seven model configurations. The paper notes that "CCE behaves very similarly to Baseline and torch.compile" and, critically, "because CCE does not utilize chunking, it does not reach a point where the overhead of dispatching all the kernels becomes the dominating factor." This means CCE's latency scales roughly linearly with batch size, without the superlinear degradation that chunked approaches can exhibit at large batch sizes due to kernel launch overhead scaling with the number of chunks. The paper also observes that while CCE-Kahan-FullC "is slower than the Liger Kernel and Torch Tune baselines with a large number of tokens, it becomes more performant than those baselines as the number of tokens reduces" (Appendix C.2)—an important practical consideration for training setups with smaller per-GPU batch sizes.

Scaling with model size and vocabulary-to-hidden ratio. Table A3 extends the benchmarks to six additional model configurations spanning vocabulary sizes from 32,064 (Phi 3.5 Mini) to 256,128 (Gemma 2 at 9B and 27B) and hidden dimensions from 3,072 to 5,120. Across all configurations, CCE consistently achieves near-lower-bound memory: the combined loss+gradient memory ranges from 236 MB (Phi 3.5 Mini) to 2,325 MB (Gemma 2 27B), always within 0–3 MB of the lower bound. The latency picture is more nuanced and depends on the ratio |V|/D:

  • For the largest |V|/D ratio considered (Gemma 2 2B: |V|/D ≈ 111), CCE is slightly faster than torch.compile for the combined computation (145 ms vs. 143 ms, essentially tied).
  • For intermediate ratios (Qwen 2.5 7B: |V|/D ≈ 42; Mistral NeMo: |V|/D ≈ 26), CCE is 10–26% slower than torch.compile for the combined computation (136 ms vs. 121 ms for Qwen; 180 ms vs. 143 ms for Mistral NeMo), but the absolute differences are small (15–37 ms) relative to total training step times (seconds).
  • For the smallest |V|/D ratio (Phi 3.5 Mini: |V|/D ≈ 10.4), CCE is approximately 50% slower than torch.compile (34 ms vs. 22 ms), but the absolute difference is only 12 ms, which the paper characterizes as "largely negligible" and "only increases training time by one to two percent."

The paper explains this scaling behavior (Appendix C.2): "As |V|/D continues to decrease, Liger Kernels is able to make better use of the GPU. All other methods use two matrix multiplications to compute the gradient. The amount of work that can be performed in parallel to compute ∇E and ∇C is B × D and |V| × D, respectively. The amount of parallel work for CCE is B × |V|, thus increasing D increases the amount of work but not the amount of parallelism." This identifies a potential direction for future optimization: incorporating ideas from split-k matrix multiplication to expose more parallelism to CCE for large D. Nevertheless, the paper's empirical conclusion is that even for the most unfavorable |V|/D ratio tested, the latency overhead is small enough to be inconsequential in practice—and it is the price paid for a memory reduction of 12.6× (3,006 MB for Baseline vs. 236 MB for CCE on Phi 3.5 Mini).

Gradient Filtering and Vocabulary Sorting: Empirical Validation of the Sparsity Assumption

The quantitative justification for gradient filtering's effectiveness appears in Figure 3 and the ablation rows of Table 1.

Softmax sparsity (Figure 3). The paper measures the sorted softmax probabilities for the average token prediction in a Gemma 2 (2B) model. The results show:

  • The probability of the most likely token is approximately 10⁻¹·⁵ (≈0.03), declining rapidly.
  • By the ~50th most likely token (out of 256,000), the probability has fallen below the bfloat16 cutoff of 2⁻¹².
  • On a log-log plot, there is a roughly linear relationship between log rank and log probability for the top 10⁵ tokens.
  • Across all tokens, "less than 0.02% of softmax elements are non-zero" in the sense of having values above the gradient filtering threshold.

This extreme sparsity—fewer than ~50 out of 256,000 vocabulary entries per position contribute non-negligible gradients—is the empirical foundation for the 3.5× speedup from gradient filtering.

Gradient filtering ablation (Table 1, row 7 vs. row 1). Without gradient filtering, CCE's backward pass takes 314 ms—a 3.1× increase over the filtered version's 100 ms. For the combined loss+gradient, the unfiltered version takes 357 ms versus 145 ms for the filtered version, a 2.5× slowdown. This confirms that the gradient filtering is the dominant source of CCE's speed, responsible for roughly two-thirds of the backward-pass performance. The paper notes that without filtering, the backward pass must perform |V|/V_B outer-product gradient updates for every batch block, even for vocabulary entries whose softmax contribution is numerically zero. With filtering, the vast majority of blocks are skipped after the cheap matrix multiplication and threshold check.

Vocabulary sorting ablation (Table 1, row 6 vs. row 1). Without vocabulary sorting, CCE's backward pass takes 115 ms versus 100 ms with sorting—a 15% increase. The combined loss+gradient takes 159 ms versus 145 ms, a 9.7% slowdown. This is a more modest effect than gradient filtering but still meaningful, consistent with sorting being a refinement that increases the hit rate of the filtering mechanism rather than a primary source of speedup. The paper explains that without sorting, the arbitrary vocabulary order from BPE produces blocks containing a mix of common and rare tokens, meaning many blocks are partially populated—some entries exceed the filtering threshold while others do not—and cannot be skipped, yet contain wasted computation on the sub-threshold entries.

Removing Ignored Tokens: A Complementary Optimization That Benefits All Methods

Table A1 applies a simple pre-filtering step to all methods: tokens that are ignored in the loss computation (padding, system prompts, user inputs) are removed before the logits+loss computation. The paper notes that "in all implementations we are aware of, the logits and loss for these ignored tokens is first computed and then set to zero. We notice that this is unnecessary."

The impact on CCE is substantial but not transformative for memory (CCE is already near the lower bound). CCE's combined loss+gradient time drops from 145 ms to 54 ms (a 2.7× speedup) and memory drops from 1,164 MB to 1,164 MB (unchanged, since CCE's memory is already minimal). However, the impact on other methods is dramatic:

  • Baseline drops from 208 ms / 28,000 MB to 75 ms / 12,826 MB—a 2.8× speedup and 2.2× memory reduction.
  • torch.compile drops from 143 ms / 16,000 MB to 53 ms / 7,337 MB.
  • Torch Tune drops from 169 ms / 9,631 MB to 77 ms / 6,157 MB.
  • Liger Kernels time is essentially unchanged (304 ms → 303 ms) because "due to heavy chunking... it is bound by kernel launch overhead, not computation."

This result demonstrates that removing ignored tokens is a simple, universally beneficial optimization that should be standard practice, and it narrows the latency gap between methods. After filtering, CCE (54 ms), torch.compile (53 ms), and Baseline (75 ms) are all within a narrow range, with CCE and torch.compile essentially tied. The key differentiator becomes memory: CCE uses 1,164 MB versus 7,337 MB for torch.compile and 12,826 MB for Baseline—factors of 6.3× and 11.0×, respectively.

Fine-Tuning Convergence: CCE Matches torch.compile Exactly

Figure 4 shows the training loss curves for fine-tuning four models (Gemma 2 2B, Phi 3.5 Mini, Qwen 2.5 7B, Mistral NeMo) on the Alpaca dataset using CCE (basic variant, with gradient filtering) versus torch.compile cross-entropy as the control. Across all four models and all 700 gradient steps, the loss curves for CCE and torch.compile are "nearly indistinguishable" (paper text). The curves overlay each other almost perfectly, with the 95% confidence intervals (shaded regions) showing substantial overlap throughout training.

Specific observations from Figure 4:

  • Gemma 2 2B (Figure 4a): Both methods start at a training loss of approximately 1.25–1.30, decline to approximately 0.80 by step 100, and then slowly converge to approximately 0.72–0.73 by step 700. The CCE and torch.compile curves are visually identical, with overlapping confidence intervals at all points.
  • Phi 3.5 Mini (Figure 4b): Similar pattern starting at ~1.25, declining to ~0.78 by step 100, and converging to ~0.70 by step 700. Again, curves are indistinguishable.
  • Qwen 2.5 7B (Figure 4c): Same convergence behavior, with final loss around 0.70–0.72.
  • Mistral NeMo (Figure 4d): Starting loss ~1.30, declining to ~0.82 by step 100, converging to ~0.74 by step 700. CCE and torch.compile overlay perfectly.

This result validates the paper's central claim that gradient filtering at ε = 2⁻¹² is mathematically exact for bfloat16 fine-tuning—the filtered values would be truncated during bfloat16 accumulation regardless, so skipping them changes nothing about the computed gradients or the resulting parameter updates. The empirical indistinguishability of the loss curves is the expected outcome given the theoretical argument, and it holds across four different model architectures, vocabulary sizes, and hidden dimensions.

An important nuance: these experiments use the basic CCE variant (with gradient filtering enabled for both ∇E and ∇C, and without Kahan summation). This means that for fine-tuning specifically, the most aggressive memory-saving configuration is sufficient—there is no need for the Kahan summation or FullC variants that increase memory and latency for pretraining. This is the empirical basis for the paper's two-tier recommendation.

Pretraining Convergence: CCE-Kahan-FullC Matches torch.compile; Basic CCE Does Not

Figure 5 shows the validation perplexity curves for pretraining the same four models from scratch on 5% of Open WebText. The control is torch.compile cross-entropy, and the CCE variant tested is CCE-Kahan-FullC (Kahan summation + gradient filtering disabled for ∇C). The paper reports that in "initial experiments using CCE for pretraining, we found that validation perplexity suffered," motivating the development of the Kahan-FullC variant, though these initial negative results are not explicitly plotted.

Across all four models and 1,500 gradient steps, CCE-Kahan-FullC produces validation perplexity curves that are identical to torch.compile within the 95% confidence intervals:

  • Gemma 2 2B (Figure 5a): Validation perplexity starts at approximately 180–200, declines rapidly to ~60 by step 500, and converges to ~45–50 by step 1,500. The two curves overlay with overlapping confidence intervals throughout.
  • Phi 3.5 Mini (Figure 5b): Starting perplexity ~180, declining to ~60 by step 500, converging to ~45 by step 1,500.
  • Qwen 2.5 7B (Figure 5c): Starting perplexity ~160–180, declining to ~55 by step 500, converging to ~40–42.
  • Mistral NeMo (Figure 5d): Starting perplexity ~160, declining to ~50 by step 500, converging to ~38–40.

The p95 confidence ranges (shaded regions) are shown for both methods and overlap completely in all panels, confirming that the observed differences are not statistically significant. This validates the paper's claim that CCE-Kahan-FullC is a numerically equivalent drop-in replacement for standard cross-entropy during pretraining, and that the two identified sources of pretraining error (gradient starvation for rare tokens and bfloat16 accumulation imprecision) are fully addressed by the combination of Kahan summation and disabling gradient filtering for ∇C.

The paper also makes an important practical note about this variant: "the increased computation time of CCE-Kahan-FullC vs. torch.compile is often offset by the larger batch sizes CCE-Kahan-FullC enables. In our experiments with Mistral NeMo, CCE-Kahan-FullC enabled doubling the batch size, thereby decreasing training time by 2 hours (16%) compared to torch.compile." This demonstrates that even though CCE-Kahan-FullC is slower per-step than torch.compile (268 ms vs. 143 ms for the loss+gradient in the Gemma 2 2B configuration; Table 1 row 9 vs. row 4), the memory savings enable larger batch sizes that reduce the total number of steps needed to process the training data, yielding a net reduction in wall-clock training time. This is the key practical argument for adopting CCE-Kahan-FullC in pretraining: the memory savings are not merely about fitting larger models on fixed hardware but about accelerating training throughput.

Ablation Studies and Robustness Checks

Gradient filtering threshold (implicit in the ε = 2⁻¹² design choice): The paper does not present a sweep over different filtering thresholds because the chosen threshold is derived from hardware properties rather than tuned empirically. The threshold ε = 2⁻¹² is justified theoretically (Appendix E) as the value below which a softmax element, when multiplied by a gradient of magnitude ≤1 and added to a bfloat16 accumulator, is truncated to zero during the floating-point alignment step. The empirical validation is indirect: the fine-tuning loss curves (Figure 4) match exactly, confirming that no gradient information is lost. There is no experiment testing whether a higher threshold (e.g., ε = 2⁻¹⁰) would maintain convergence while providing greater speedup, which could be informative for practitioners willing to trade some precision for additional throughput.

Kahan summation variants (Table 1, rows 8–10): The paper tests three numerical stability configurations. CCE-Kahan (Kahan summation for all gradient accumulations, gradient filtering enabled for both ∇E and ∇C) increases backward pass time from 100 ms to 114 ms (14% slowdown) and memory from 1,163 MB to 2,325 MB (roughly double, due to the compensation buffer). CCE-Kahan-FullC (Kahan summation + disable ∇C filtering) further increases backward time to 268 ms (2.7× over basic CCE) and memory to 2,326 MB. CCE-Kahan-FullE (Kahan summation + disable ∇E filtering) takes 247 ms and 2,326 MB. The comparison between FullC (268 ms) and FullE (247 ms) reveals that disabling ∇E filtering is slightly faster than disabling ∇C filtering, consistent with the idea that the ∇C computation involves larger reductions (summing over N) and thus benefits more from filtering. The paper focuses on CCE-Kahan-FullC as the recommended pretraining variant because "it is the more conservative choice that guarantees no gradient starvation for any parameter."

Removing ignored tokens (Table A1): This ablation, discussed in the Main Results above, demonstrates that a simple preprocessing step (filtering out padding and non-loss tokens before the loss computation) provides a 2.5–2.8× speedup across most methods without affecting the loss or gradient mathematically. For CCE specifically, the combined loss+gradient time drops from 145 ms to 54 ms. The paper notes that this optimization "represents a significant speed up for all methods but Liger Kernels," which is bound by kernel launch overhead rather than computation. This ablation also narrows the memory gap between chunked approaches and CCE—after filtering, Baseline uses 12,826 MB (vs. 28,000 MB unfiltered) and torch.compile uses 7,337 MB (vs. 16,000 MB)—but CCE remains far more memory-efficient at 1,164 MB.

Scaling with model configuration (Table A3 and Figures A1, A2): The paper benchmarks CCE across a comprehensive range of model configurations: Gemma 2 at 2B, 9B, and 27B scales (holding |V| = 256,000 constant while increasing D from 2,304 to 3,584 to 4,608); Phi 3.5 Mini (|V| = 32,064, D = 3,072); Mistral NeMo (|V| = 131,072, D = 5,120); and Qwen 2.5 at 7B and 32B (|V| = 152,064, D = 3,584 and 5,120). This sweeps |V|/D ratios from ~10 (Phi 3.5 Mini) to ~111 (Gemma 2 2B). The key robustness finding is that CCE's memory advantage is universal—for every configuration, CCE's memory is within 3 MB of the lower bound, representing a reduction of 5× (Phi 3.5 Mini: 3,006 MB → 236 MB) to 24× (Gemma 2 9B: 28,000 MB → 1,809 MB) versus Baseline. The latency penalty relative to torch.compile grows as |V|/D decreases (from essentially zero at |V|/D ≈ 111 to ~50% at |V|/D ≈ 10), but the absolute latency differences remain small (2–37 ms) compared to total training step times. Figures A1 and A2 further show that CCE's latency scales smoothly with batch size across all model configurations, without the superlinear degradation that affects chunked methods at large batch sizes.

Performance breakdown of the backward pass (Table A2): This detailed timing analysis decomposes the backward pass for both CCE and Baseline into five components: recomputation of the logits (CCE only), gradient of the loss with respect to logits, gradient filtering overhead, gradient with respect to softcapped logits, ∇E computation, and ∇C computation. The key finding is that CCE spends 45 ms (43.2% of its backward pass time before filtering savings) on recomputing the logits and applying the gradient filter—a cost that Baseline does not incur. However, CCE saves 30 ms on the cross-entropy gradient (4.7 ms vs. 35 ms), 12 ms on the softcapping gradient (4.7 ms vs. 17 ms), 5 ms on ∇E (31 ms vs. 37 ms), and 15 ms on ∇C (18 ms vs. 34 ms). Total savings: 62 ms, which more than offsets the 45 ms recomputation cost. The paper attributes these savings to CCE performing these operations on logits that are already in SRAM, avoiding the expensive global memory reads and writes that Baseline requires.

Critical Assessment

Do the Experiments Support the Claim That CCE's Memory Reduction Is "24,000×" (Loss) or Enables "1.5× to 10×" Batch Size Increases?

The 24,000× figure (24 GB → 1 MB for the loss computation) appears in the executive summary and is supported by Table 1 for the Gemma 2 (2B) configuration. However, this figure refers specifically to the forward loss computation only—the memory used during the forward pass to compute the scalar loss value. The more practically relevant figure is the combined loss+gradient memory, where CCE achieves a 24× reduction (28 GB → 1,164 MB), not 24,000×. The forward-only number is the most dramatic and headline-worthy, but it is the combined memory that determines training feasibility (since both forward and backward passes must fit in GPU memory during a training step). The paper is transparent about this distinction—Table 1 clearly shows separate columns for loss, gradient, and combined memory—but readers focusing only on the 24,000× headline may overestimate the practical benefit. The combined reduction is 24×, which is still transformative but an order of magnitude less dramatic.

The batch size increase claims (1.5× to 10×) from Figure 1 are calculated analytically rather than measured empirically. The paper computes the maximum batch size by dividing available GPU memory (75 GB per GPU on a 16-GPU setup, reserving 5 GB for libraries) by the per-token memory of activations plus logits, assuming a fixed global batch of 65,536 tokens. These numbers do not come from actually training the models with CCE and measuring the largest batch size that fits—they are extrapolations from the memory-per-token calculations. The actual achievable batch size in practice could differ due to memory fragmentation, framework overhead, or the need for additional scratch space that the analytical model does not account for. The paper acknowledges that the logit memory estimate "likely undercounts the amount of memory used for computing the probability distribution, as its common to also keep a copy of the logits in bfloat16 and, for models like Gemma 2 that use logit softcapping, an additional copy of the logits after softcapping may be needed" (Appendix D). This suggests the batch size increases in Figure 1 may be optimistic.

Additionally, the analysis assumes a specific parallelism configuration (16-GPU fully-sharded data-parallel with activation checkpointing). For different parallelism strategies (tensor parallelism, pipeline parallelism, sequence parallelism), the memory bottlenecks may shift, and CCE's impact on maximum batch size would differ. The paper discusses this in Section 6 ("CCE may enable better pipeline balancing or reducing the number of stages") but does not quantify the effect.

Do the Experiments Support the Claim That CCE Has "No Sacrifice in Speed or Convergence"?

Speed. The "no sacrifice in speed" claim is supported for the Gemma 2 (2B) configuration (Table 1), where CCE's combined loss+gradient time (145 ms) is essentially identical to torch.compile's (143 ms). However, this parity does not hold uniformly across all model configurations. For Mistral NeMo, CCE takes 180 ms versus torch.compile's 143 ms (a 26% slowdown). For Phi 3.5 Mini, CCE takes 34 ms versus 22 ms (a 55% slowdown). The paper argues that these absolute differences are small (37 ms and 12 ms, respectively) and represent only "one to two percent" of total training step time for even a 2B-parameter model. This argument is reasonable—a 12 ms increase on a training step that takes several seconds is a 0.3–0.6% throughput reduction—but it means the "no sacrifice" claim is conditional on the |V|/D ratio being large enough. For models with small vocabularies relative to their hidden dimension, there is a measurable latency cost, albeit a small one.

For the pretraining variant CCE-Kahan-FullC, the latency sacrifice is substantial: 268 ms versus 143 ms for torch.compile on the Gemma 2 (2B) configuration (Table 1, row 9 vs. row 4)—an 87% increase for the loss+gradient computation specifically. The paper argues that this is offset by the larger batch sizes CCE-Kahan-FullC enables, citing the Mistral NeMo experiment where total training time decreased by 16% despite the per-step slowdown. However, this claim rests on a single data point ("our experiments with Mistral NeMo") and is not systematically evaluated across models. It is possible that for some model configurations or training setups, the per-step slowdown is not fully compensated by batch size increases, resulting in a net throughput decrease. The paper does not provide the data to assess this tradeoff quantitatively for all tested models.

Convergence. The convergence claim is well-supported for fine-tuning (Figure 4): CCE and torch.compile produce identical loss curves across four models. The evidence for pretraining is also strong (Figure 5) but with an important caveat: the experiments use only 5% of Open WebText, training for only 1,500 steps. For pretraining from scratch on a full-scale corpus (trillions of tokens), the accumulated numerical differences between CCE-Kahan-FullC and torch.compile could potentially diverge in ways that are not visible in a 1,500-step run. The paper's theoretical argument—that Kahan summation addresses accumulation precision and FullC addresses gradient starvation—is sound, and the empirical match at 1,500 steps is encouraging, but it does not constitute proof that the methods would remain identical over a full pretraining run. A stronger validation would be to train a model from scratch to a meaningful downstream benchmark and compare final performance, though the computational cost of such an experiment is acknowledged to be prohibitive for an academic paper.

Do the Experiments Support the Claim That Gradient Filtering Is "Lossless"?

The theoretical argument that gradient filtering at ε = 2⁻¹² is lossless for bfloat16 accumulation is rigorous (Appendix E). The empirical support is the exact match of fine-tuning loss curves (Figure 4). However, there is a subtlety: the argument assumes that the gradient accumulator starts at a value up to ~1 in magnitude and that filtered contributions are added to it. In practice, the gradient accumulators in ∇E and ∇C are large tensors that are initialized to zero, not to 1. A softmax value of, say, 2⁻¹⁵, when multiplied by ∇LSE (typically ≤1) and added to a zero accumulator, would produce a non-zero result in bfloat16 (2⁻¹⁵ is above the minimum subnormal bfloat16 value). The truncation argument applies when adding to an accumulator that already contains larger values, which is true after the first few atomic additions. The paper's threshold derivation implicitly assumes the accumulator has grown large enough that values below 2⁻¹² are truncated, which is reasonable for the vast majority of the computation but may lose very small gradient contributions at the beginning of the accumulation process. This is a theoretical edge case that does not appear to affect training in practice (given Figure 4), but it means the "lossless" claim is better characterized as "lossless to within the precision that matters for training" rather than strictly bitwise-identical.

What Is Missing?

No comparison against a model-parallelism baseline for maximum batch size. Figure 1 shows that CCE increases maximum batch size dramatically compared to standard data-parallel training, but it does not compare against what could be achieved by applying vocabulary-parallel strategies (splitting the classifier head across GPUs) to the standard cross-entropy implementation. For practitioners with access to many GPUs, the relevant question is not "CCE vs. Baseline on the same GPUs" but "CCE vs. the best available alternative for training large-vocabulary models." A comparison showing CCE enabling larger batch sizes than vocabulary parallelism, or achieving the same batch size with fewer GPUs, would strengthen the practical case.

No end-to-end training throughput benchmarks. Table 1 and Table A3 measure the time for the loss+gradient computation in isolation, not the time for a complete training step including the forward pass through the backbone and the optimizer update. While the paper argues that the loss+gradient time is a small fraction of the total step time (and thus CCE's overhead is negligible), this is not empirically demonstrated. A plot showing training throughput (tokens/second) for full training steps with CCE versus torch.compile across model sizes would be more informative than the micro-benchmarks alone. The Mistral NeMo result (2-hour reduction, 16% faster) is the only end-to-end timing data point provided.

No sensitivity analysis for the gradient filtering threshold. The paper derives ε = 2⁻¹² from bfloat16 properties and validates it empirically through the fine-tuning loss curves, but it does not test whether a higher threshold would also work. Given that the softmax sparsity is extreme (Figure 3 shows probabilities dropping below 2⁻¹² by rank ~50 out of 256,000), it is plausible that a threshold of 2⁻¹⁰ or even 2⁻⁸ would still capture all practically significant gradient contributions while providing additional speedup. Exploring this tradeoff would be valuable for practitioners who might prefer a slight convergence degradation in exchange for faster training.

No evaluation on non-transformer architectures or non-language domains. All experiments use decoder-only transformer language models. The paper suggests in Section 6 that CCE "may prove beneficial for training very large models" and "could also be interesting to extend CCE to other classification problems where the number of classes is large, such as image classification and contrastive learning," but provides no empirical evidence in these domains. The gradient filtering mechanism depends on the softmax distribution being extremely sparse, which holds for large-vocabulary language modeling but may not hold for, e.g., ImageNet-21K classification where the effective number of classes per example is smaller and the softmax may be less sparse.

Limited exploration of CCE's interaction with other training optimizations. The paper tests CCE with mixed-precision training (bf16) and activation checkpointing, but does not systematically evaluate interactions with other common training techniques: gradient accumulation, gradient clipping, loss scaling, different optimizers (e.g., AdamW vs. SGD with momentum), or different parallelism strategies. Some of these interactions could be non-trivial—for instance, gradient accumulation across micro-batches might interact with the bfloat16 truncation argument that underlies gradient filtering, since the accumulator would grow larger with more micro-batches.

The full pretraining convergence experiment is relatively short. Training for 1,500 steps on 5% of Open WebText is a small fraction of what a full pretraining run would entail (typically hundreds of thousands to millions of steps). While the match between CCE-Kahan-FullC and torch.compile at 1,500 steps is encouraging, the possibility of subtle divergence over much longer training is not ruled out. This is a practical limitation imposed by computational cost, but it should be acknowledged as a limitation of the evidence rather than definitive proof of equivalence.

Overall, the experiments strongly support the paper's central claims about memory reduction and fine-tuning convergence, provide good evidence for pretraining convergence (with the caveat about training duration), and offer reasonable but not exhaustive support for the speed claims, which are configuration-dependent and less uniformly favorable than the memory claims. The most impactful open question is the end-to-end training throughput picture, which the Mistral NeMo example suggests is favorable but which is not systematically characterized.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Unaccounted for in All Reported Efficiency Gains

The assumption or constraint. The entire compute-optimal allocation framework depends on being able to estimate the difficulty of each prompt before deciding how to allocate the inference budget. The paper's method for doing so—generating 2048 samples per question, then averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted)—is enormously expensive. The authors acknowledge this upfront in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. The headline efficiency gains (4× over best-of-N, matching 256-generation performance with only 64 generations) are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, the total compute would be 2048 (for difficulty estimation) + N (for the actual strategy execution). For N = 16 or 64, the difficulty estimation cost dominates the budget by one to two orders of magnitude. The claimed 4× improvement is therefore best understood as an upper bound on achievable efficiency in a regime where difficulty can be estimated nearly for free—a regime that does not exist with the current method.

The practical implication is stark: if a practitioner deploys the system as described, with the recommended 2048-sample difficulty estimation step, the total per-prompt compute cost would be far higher than simply running best-of-N at the maximum N, defeating the purpose of the adaptive allocation. The paper acknowledges this as an exploration-exploitation tradeoff ("compute spent assessing difficulty versus compute spent solving the problem," Section 3.2) but makes no progress on resolving it.

What evidence exists in the paper. The 2048-sample protocol is described in Section 3.2. The paper states explicitly that this cost is not included in any of the budget calculations for Figures 4, 8, or 9. There is no experiment measuring total cost (estimation + execution) versus a baseline that avoids estimation. The predicted difficulty bins—which use the PRM's own scores rather than ground truth—are tested and shown to work nearly as well as oracle bins (Figures 4, 8), but this variant still requires the 2048 samples to produce the difficulty estimate; it only removes the need for ground-truth labels, not the sampling cost.

Mitigation status. The paper suggests but does not evaluate two directions: (1) "training models to directly predict difficulty of a question," flagged as future work in Section 8, and (2) the possibility of adaptive schemes that interleave difficulty assessment with problem-solving. Neither is implemented. The difficulty estimation gap is arguably the largest barrier between the paper's analytical results and practical deployment, and it remains entirely unaddressed.

Hard Problems Remain Essentially Unsolved—Test-Time Compute Cannot Compensate for Fundamental Capability Gaps

The assumption or constraint. The compute-optimal framework operates on the implicit premise that the base model's proposal distribution contains correct solutions at a non-trivial rate—that is, pass@1 > 0 for the given prompt. When this assumption fails, the entire framework collapses: no allocation of search, revisions, or their combination can extract a correct answer that does not exist in the model's output distribution.

The consequence. Across every method studied—search (Figure 3, right), revisions (Figure 7, right), and their compute-optimal combinations (Figures 4, 8)—the hardest difficulty bin (bin 5) shows near-zero improvement regardless of how much compute is allocated. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 shows ~2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling curve is essentially flat near 0–5% while the ~14× larger model achieves markedly higher accuracy (the stars in Figure 9 are well above the bin 5 scaling line).

The paper is explicit about this finding:

"on the hardest questions (bin 5), no method makes meaningful progress—the base model simply lacks the capability to produce correct solutions regardless of how the budget is allocated" (Section 5.3 discussion of Figure 3).

The practical consequence is that test-time compute is not a substitute for pretraining on genuinely hard problems. It amplifies existing capability but does not create it from nothing. For deployments where the problem distribution skews toward tasks beyond the base model's ability—out-of-distribution reasoning, novel problem types, or tasks requiring knowledge the model was not exposed to during pretraining—this approach offers no benefit. In the FLOPs-matched comparison (Section 7, bar charts in Figure 1), hard problems show a −52.9% relative disadvantage from test-time compute versus the larger model at R ≫ 1 for PRM search, and −37.2% for revisions, demonstrating that pretraining is strictly necessary for this difficulty tier.

What evidence exists in the paper. The per-difficulty-bin breakdowns in Figures 3 (right), 7 (right), and 9 provide extensive documentation of this failure mode. Bin 5 accuracy is consistently near-random across all methods and all budgets. The FLOPs-matched comparison (Figure 9) quantifies how much worse test-time compute performs relative to a larger pretrained model on hard problems.

Mitigation status. The paper does not attempt to solve this limitation—it acknowledges it as a fundamental boundary (Section 7 takeaway: "test-time compute amplifies existing capability but does not create it from nothing") and uses it to delineate when test-time compute should be preferred over pretraining. This is framed as a finding rather than a problem to be solved. The mitigation, implicitly, is to route hard problems to a larger model or to human review—but this requires knowing which problems are hard before attempting to solve them, which brings us back to the difficulty estimation problem above.

The ~14× Larger Model Baseline Is Weaker Than a Compute-Optimally Trained Model Would Be

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters. Crucially, this larger model is produced by scaling parameters only while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The paper acknowledges this departure from compute-optimal pretraining practice:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)

The consequence. A Chinchilla-optimal model (Hoffmann et al., 2022)—where both parameters and training data are scaled proportionally—would likely outperform a parameter-only-scaled model at the same total FLOPs budget. This means the reported advantages of test-time compute over pretraining (e.g., +27.8% on easy questions at R ≪ 1 for revisions, per the bar charts in Figure 1) may be overstated relative to what a properly compute-optimal larger model would achieve. The magnitude of this overstatement is unknown because the paper does not include a Chinchilla-optimal baseline.

Additionally, the larger model is evaluated with greedy decoding only—no majority voting, no best-of-N sampling, no search of any kind. This is a deliberately weak comparison point that the paper does not justify beyond noting that it represents "a canonical approach" (Section 7). A fairer comparison would give the larger model some test-time compute budget as well, since the question of interest is "should I spend my marginal FLOPs on making the model bigger or on making inference smarter," not "should I give my small model inference compute while giving my large model nothing."

What evidence exists in the paper. The FLOPs-matched comparison is presented in Figure 9 and the bar charts in Figure 1. The paper transparently states the parameter-only scaling choice in Section 7. However, there is no ablation comparing against a Chinchilla-optimal baseline, no discussion of how the results might change if the larger model received even a modest test-time compute budget (e.g., best-of-8), and no sensitivity analysis on the training-data-to-parameters ratio.

Mitigation status. The paper flags this as future work (Section 7) but provides no empirical characterization of how much the conclusions depend on this design choice. Given that the claimed substitution of test-time compute for pretraining is one of the paper's headline results and most practically significant claims, the weakness of the pretraining baseline is a consequential limitation. Practitioners evaluating the training-inference tradeoff for their own models should be aware that the ~14× figure is an upper bound on the substitution ratio, not a robust estimate of the general relationship.

Revisions and PRM Search Are Studied Independently—Their Combination Is Unexplored

The assumption or constraint. The paper studies two complementary mechanisms for improving test-time performance—PRM-guided search (Section 5) and iterative revisions (Section 6)—but evaluates them in isolation, never combining them into a single system. The paper acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions" (Section 8)

The consequence. The two mechanisms have complementary strengths that are well-documented in the paper's own experiments: revisions improve the proposal distribution (generating better candidates through sequential refinement), while PRM search improves candidate selection (finding the best among generated candidates through verifier-guided search). On medium-difficulty problems, the paper shows that both mechanisms help but in different ways—revisions help by making incremental improvements to nearly-correct answers, while beam search helps by exploring diverse solution paths. It is natural to expect that combining them—using the revision model as the proposal distribution within beam search, or using the PRM to decide which revision branches to pursue—could yield gains beyond either method alone.

The consequence is that the paper's reported performance numbers represent a lower bound on what a fully integrated system could achieve. The compute-optimal policy in Figures 4 and 8 selects between search strategies and revision strategies per difficulty bin, but never deploys both simultaneously on the same problem. This means the 4× efficiency gain over best-of-N may itself be improvable, and the ceiling for test-time compute scaling may be higher than the paper's results suggest.

The gap is particularly significant because the paper's own framework (Section 2) frames revisions (proposal modification) and search (verifier optimization) as two independent, complementary axes of test-time compute allocation. The fact that the paper never combines them leaves the central claim of complementarity empirically underdeveloped—the paper demonstrates that each axis works independently and that they have different difficulty-dependent optimal regimes, but does not demonstrate that they are additive when combined.

What evidence exists in the paper. All experiments in Sections 5 and 6 treat search and revisions separately. The difficulty-bin analyses show that search works best on medium problems (Figure 3, right) while revisions work best on easy problems (Figure 7, right), suggesting a natural division of labor that a combined system could exploit. But no experiment tests a combined system. The PRM is trained on base model outputs and does not transfer well to revision model outputs due to distribution shift (Appendix J, Figure 15a)—the paper acknowledges this and trains a separate ORM for revisions, but does not explore whether a PRM trained specifically on revision model trajectories could enable combined search+revision.

Mitigation status. The paper explicitly calls for future work on combining search with revisions (Section 8). The distribution shift issue with the PRM (Figure 15a) suggests that combining them may require retraining the verifier on revision model outputs, which is a non-trivial additional step. No partial results or pilot experiments are provided.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate with No Principled Solution

The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target answer (Section 6.1). During training, the model never sees examples where a correct answer appears in context and should be preserved rather than revised. This creates a mismatch at inference time: when the model produces a correct answer during a revision chain, the subsequent revision step may incorrectly "revise" it into a wrong answer.

The consequence. The paper reports that "approximately 38% of correct answers get converted back to incorrect ones" during sequential revision chains (Section 6.1). This means that simply taking the final revision in a chain—the most natural and intuitive way to use a revision model—frequently degrades correct answers. The paper mitigates this with post-hoc selection: using majority voting or verifier-based selection across the entire chain of revisions to pick the best answer from any point in the sequence, rather than always taking the last output.

This mitigation works—Figure 6 shows that sequential revision with within-chain selection outperforms parallel sampling—but it is fundamentally a patch for a training deficiency, not a solution. It means the revision model wastes a significant fraction of its sequential budget producing degraded answers that must then be filtered out by the selection mechanism. It also means the revision chain cannot be used in a streaming or low-latency setting where the most recent output is desired, because the most recent output is frequently wrong even when an earlier output was correct.

The 38% reversion rate also implies that longer revision chains have diminishing returns: each additional revision step has a ~38% chance of undoing a previously correct answer. The paper's Figure 6 (left) shows pass@1 at each step gradually improving over long chains, suggesting the net effect is still positive (more corrections than reversions), but the efficiency is clearly suboptimal—a significant fraction of the inference budget is spent on degradative revisions rather than productive ones.

What evidence exists in the paper. The 38% figure is reported in Section 6.1. The within-chain selection mechanism (majority voting or verifier-based) is described as the mitigation. The ReST^EM experiment in Appendix K (Figure 16) provides additional evidence of revision fragility: attempting to optimize the revision model with on-policy RL-style training caused performance to "substantially hurt," with fully sequential performance dropping to ~33.5% at 256 generations compared to ~38.5% at the optimal ratio. This suggests the revision training procedure is sensitive to data distribution in ways that are not fully understood.

Mitigation status. The paper's mitigation (within-chain selection) is effective but incomplete. It does not address the root cause: the model was never trained to recognize when no revision is needed. The authors do not propose a principled solution, such as training the model on trajectories that include correct-in-context examples with an explicit "no revision needed" target, or incorporating a confidence threshold that suppresses revision when the model's certainty is high. Section 8 does not flag this as a specific direction for future work, despite the quantitative significance of the 38% figure.

The Single Benchmark, Single Model Family Scope Leaves Generalization Unverified

The assumption or constraint. All experiments in this paper use a single benchmark (MATH, specifically the 500-question test split from Lightman et al., 2022) and a single model family (PaLM 2-S*, with the FLOPs-matched comparison using a larger model from the same family). The paper justifies this choice in Section 4:

"We believe this model is representative of the capabilities of many contemporary LLMs"

But this belief is not tested.

The consequence. Several aspects of the paper's findings could be model- or domain-specific in ways that affect their generality:

  • PRM quality and over-optimization behavior: The PRM is trained via Monte Carlo rollouts from PaLM 2-S* on MATH problems. A model with different calibration properties, different error patterns, or different output diversity might produce PRMs with different over-optimization thresholds. The finding that beam search degrades easy-problem performance at high budgets (Figure 3, right) might occur at different budget levels—or not at all—for other model families.

  • Revision model effectiveness: The revision model's ability to learn from incorrect in-context examples and produce targeted corrections (enabled by the edit-distance-based pairing in the training data) may depend on the base model's in-context learning capabilities, which vary substantially across model families. A model with weaker in-context learning might not benefit from the revision training procedure to the same degree.

  • Domain specificity: MATH consists of competition-level math problems requiring symbolic reasoning and multi-step deduction. The difficulty-dependent patterns—beam search hurting easy problems, revisions helping easy problems, sequential-to-parallel ratio optimality varying with difficulty—may not transfer to other reasoning domains (code generation, logical reasoning, scientific question-answering) or to tasks requiring factual recall rather than inference. The paper's theoretical framework (Section 2) is domain-agnostic, but the empirical findings are entirely MATH-specific.

  • Small test set for policy selection: The 500-question test set is split into five difficulty quintiles of ~100 questions each, then further split by two-fold cross-validation (Section 3.2). This means the compute-optimal policy is selected based on ~50 questions per fold per bin. The paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4, 8, 9), making it difficult to assess whether the observed policy choices are robust or an artifact of the small per-bin sample size.

What evidence exists in the paper. No experiments on any dataset other than MATH. No experiments with any model family other than PaLM 2. The paper does not report confidence intervals for the compute-optimal scaling curves in Figures 4, 8, and 9. The 500-question test set and the two-fold cross-validation protocol are described in Section 3.2.

Mitigation status. The paper does not address this limitation. Section 8 does not call for replication on other benchmarks or model families as future work (though it is a natural next step). The claim that PaLM 2-S* is "representative" is an assertion, not an empirical finding. Practitioners considering adopting the compute-optimal scaling approach for their own models and domains should treat the specific strategy recommendations (which search algorithm to use for which difficulty bin, what sequential-to-parallel ratio is optimal) as suggestive rather than prescriptive, and should plan to re-derive the optimal policy for their specific model-benchmark combination.

7. Implications and Future Directions

How This Work Changes the Landscape

CCE represents a systems-level reframing rather than a paradigm shift—it does not change what models are trained or what loss function is optimized, but it fundamentally changes the design space for how large-vocabulary models are trained by eliminating the O(N × |V|) memory bottleneck that has silently accumulated as vocabularies grew. The magnitude of the shift is best understood by what it enables: for Gemma 2 (2B), the loss layer drops from consuming 89% of training memory to consuming essentially zero incremental memory (3 MB beyond the mandatory gradient buffers). This is not an incremental 2× or 3× improvement—it is a qualitative change in the memory profile of LLM training, moving the classification head from being the dominant bottleneck to being a non-issue.

The reframing is methodological: CCE establishes that the logit matrix is an implementation artifact, not a mathematical requirement of the cross-entropy computation. Prior to CCE, every practical approach to the vocabulary-memory problem—chunking, vocabulary parallelism, vocabulary reduction—accepted the necessity of materializing at least a portion of the logit matrix into global memory. CCE demonstrates that the full matrix never needs to exist in global memory at any point; all operations can be performed on-the-fly in SRAM as logit blocks are computed, immediately consumed for reduction or gradient accumulation, and discarded. This is the same design philosophy that FlashAttention brought to self-attention, applied to the classification layer—and together, these two techniques mean that neither the attention mechanism nor the loss layer requires O(N^2) or O(N × |V|) global memory, removing the two largest intermediate tensors from the training pipeline.

This reframing resolves a latent tension in the LLM scaling community. On one side, vocabulary sizes have grown because larger vocabularies are genuinely beneficial—they compress sequences, improve comprehension, and are supported by scaling laws suggesting further expansion (Tao et al., 2024). On the other side, this growth has made the classifier head increasingly impractical to train, to the point where Gemma 2 (2B) could not fit a single 80K-token sequence on an 80 GB H100 GPU. Prior responses to this tension included accepting the memory cost (buying more GPUs), accepting architectural compromises (vocabulary reduction, which may hurt quality), or accepting latency penalties (chunking). CCE resolves the tension by making vocabulary size effectively free in memory terms—the loss layer's memory footprint becomes O(N) and decoupled from |V|. This means the community can continue to explore larger vocabularies without suffering a corresponding memory penalty, which in turn makes vocabulary scaling research (Tao et al., 2024) more practically actionable.

CCE also reconciles a contradiction in the systems-for-ML literature: the apparent tradeoff between memory and latency in cross-entropy implementations. Prior work consistently showed that reducing memory required increasing latency—Liger Kernels uses 95% less memory than Baseline but is 2.1× slower; Torch Tune's 8-chunk configuration uses 66% less memory but is 18% slower than torch.compile. CCE breaks this tradeoff by achieving simultaneously lower memory and lower latency than the chunked baselines, and latency essentially identical to the fastest existing implementation (torch.compile) for typical vocabulary-to-hidden-dimension ratios. Table 1 shows this clearly for Gemma 2 (2B): CCE uses 1,164 MB and 145 ms versus Torch Tune's 9,631 MB and 169 ms, and versus torch.compile's 16,000 MB and 143 ms. The mechanism is the SRAM-resident blockwise computation—by avoiding global memory writes for the logit matrix entirely, CCE saves both the memory for storing the matrix and the time for writing and reading it, yielding a genuine Pareto improvement over prior approaches.

The research directions that become more attractive as a result of CCE include:

  • Large-vocabulary pretraining at smaller scales. Before CCE, training models with 256K+ vocabularies was effectively restricted to organizations with large GPU clusters that could absorb the memory cost through aggressive parallelism. CCE makes large-vocabulary training feasible on smaller hardware configurations—a single GPU can now train with vocabularies that previously required model parallelism. This democratizes access to large-vocabulary LLM training and should accelerate research on vocabulary design, multilingual tokenization, and byte-level modeling.

  • IO-aware algorithm design for other "memory hog" layers. CCE and FlashAttention together establish a design pattern: identify layers where a large intermediate tensor is materialized solely to be reduced or indexed, then redesign the computation to perform the reduction incrementally in SRAM. This pattern is applicable to any layer with the structure reduce(f(X @ W)) where reduce is an associative operation (log-sum-exp, sum, max). Contrastive learning losses, large-scale retrieval scoring, and certain normalization layers are natural candidates.

  • Pipeline parallelism balancing. The paper specifically notes (Section 6) that "the classification head is currently an outlier, with a disproportionately high memory-to-computation ratio. CCE may enable better pipeline balancing or reducing the number of stages." By eliminating the classification head's memory bloat, CCE makes it possible to assign the head to the same pipeline stage resources as other layers, simplifying pipeline configuration and potentially reducing the number of stages needed—a practical systems benefit beyond raw memory savings.

The research directions that become less attractive include:

  • Chunking-based cross-entropy optimizations. Given that CCE achieves both lower memory and lower latency than chunked approaches (Torch Tune, Liger Kernels), and does so without the memory-latency tradeoff that chunking inherently imposes, further investment in chunking as a strategy for the cross-entropy layer seems difficult to justify. The chunking approach may still have value for other layers where the reduction operation is not amenable to online accumulation, but for cross-entropy specifically, CCE renders chunking obsolete.

  • Vocabulary reduction as a memory optimization. Techniques like hierarchical softmax (Grave et al., 2017) or byte-level tokenization (Yu et al., 2023) were partly motivated by the memory cost of large vocabularies. With CCE eliminating that memory cost, the case for vocabulary reduction rests entirely on its impact on model quality and sequence length efficiency, not on training feasibility. Researchers can now evaluate vocabulary design choices on their modeling merits alone, without the confounding factor of memory constraints.

Follow-Up Research This Work Enables

Training small-|V|/D models with CUDA-level gradient filtering granularity. The paper identifies that CCE's latency penalty relative to torch.compile grows as |V|/D decreases, reaching approximately 50% for Phi 3.5 Mini (|V|/D ≈ 10.4, 34 ms vs. 22 ms). The authors attribute part of this to Triton's block-level control flow constraint: "control flow must be specified at the block level and therefore our... gradient filtering [is] constrained to operate at the block level as well." A CUDA implementation with warp-level or thread-level filtering could skip partially-populated blocks at finer granularity, reducing wasted computation and potentially closing the latency gap for small vocabularies. A strong follow-up would implement CCE in CUDA with warp-level filtering granularity, benchmark across the same model configurations as Table A3 (particularly Phi 3.5 Mini, Gemma 2 9B, and Qwen 2.5 32B), and quantify the speedup relative to the Triton implementation and torch.compile. The key metric is whether CUDA-level filtering can bring CCE's latency to within 5–10% of torch.compile across all |V|/D ratios, which would make CCE strictly dominant (lower memory, no meaningful latency cost) rather than conditionally favorable.

End-to-end training throughput benchmarks with CCE across training scales. The paper provides micro-benchmarks for the loss+gradient computation in isolation but only one end-to-end training throughput result (the Mistral NeMo experiment where CCE-Kahan-FullC reduced total training time by 16% via larger batch sizes). A systematic evaluation is needed: train several models (Gemma 2 2B, Phi 3.5 Mini, Qwen 2.5 7B, Mistral NeMo) from scratch on a standard pretraining corpus (e.g., C4 or SlimPajama) for a fixed number of tokens (e.g., 10B) using both CCE-Kahan-FullC and torch.compile, measuring total wall-clock time, maximum per-GPU batch size, and final validation perplexity. The hypothesis from the paper is that CCE's memory savings enable larger batch sizes that more than compensate for any per-step slowdown, but this has only been demonstrated for one model. If the throughput benefit is confirmed across models, it would strengthen CCE's case for pretraining adoption; if it only holds for models with large |V|/D ratios, that would clarify the practical scope of CCE's advantage.

Dynamic vocabulary sorting via online logit statistics. The paper's vocabulary sorting uses average logits computed during the forward pass, which requires a temporary buffer and a separate sorting step. An online, streaming approach could be more efficient: maintain running exponential moving averages of token logits, update them during each forward pass, and periodically re-sort the vocabulary (or use the running estimates to guide block assignment without a full sort). This would eliminate the need for the atomic-addition-based statistics collection and the separate sorting pass, potentially reducing the overhead of vocabulary sorting from its current 15% backward pass latency increase. A strong experiment would compare: (1) static BPE order, (2) the paper's average-logit sort, and (3) an online EMA-based sort updated every N steps, measuring backward pass time and the fraction of blocks skipped by gradient filtering across a full training run.

Gradient filtering threshold sweep to identify precision-efficiency Pareto frontier. The paper derives ε = 2⁻¹² from bfloat16 truncation properties and validates it with exact loss curve matching for fine-tuning, but does not explore whether a higher threshold would also produce acceptable convergence. Given the extreme sparsity of the softmax (Figure 3 shows probabilities dropping below 2⁻¹² by rank ~50 out of 256,000), it is plausible that ε = 2⁻¹⁰ or even ε = 2⁻⁸ would provide additional speedup with minimal convergence degradation. A systematic sweep would fine-tune Gemma 2 (2B) on Alpaca with ε ranging from 2⁻¹⁴ to 2⁻⁶, measuring final training loss, backward pass time, and fraction of blocks filtered, to map out the precision-efficiency tradeoff. For practitioners willing to accept a 0.1–0.5% increase in final loss in exchange for 10–20% faster backward passes, a higher threshold could be an attractive option that the paper currently does not characterize.

Combining CCE with vocabulary-parallel training for extreme-scale models. For models with vocabularies in the 500K–1M range that may emerge in future work, even CCE's O(N) memory for the LSE vector may become a constraint, and vocabulary parallelism (splitting the classifier across GPUs) may still be necessary for the classifier parameters themselves (which grow as D × |V|). CCE should be compatible with vocabulary parallelism: each GPU could run CCE on its vocabulary shard, producing a partial LSE vector, and an all-reduce could merge the partial LSEs (log-sum-exp is not directly all-reducible, but the per-GPU max-and-sum approach from standard distributed softmax applies). A strong experiment would implement this combination, benchmark it against pure vocabulary parallelism with standard cross-entropy on a model with a 500K+ vocabulary, and measure whether CCE reduces the number of GPUs needed or enables larger per-GPU batch sizes within the parallel configuration. This would extend CCE's applicability to the regime where even the classifier parameters exceed single-GPU memory.

Applying CCE's gradient filtering to other large-classification losses. The paper suggests in Section 6 that CCE "could also be interesting to extend... to other classification problems where the number of classes is large, such as image classification and contrastive learning." The key question is whether the softmax sparsity that makes gradient filtering effective for language modeling also holds in other domains. A diagnostic experiment would measure the sorted softmax probability curve (analogous to Figure 3) for a trained ImageNet-21K classifier, a CLIP-style contrastive model with a large batch size, and a retrieval model with a large candidate set, and determine at what rank the probabilities drop below the bfloat16 cutoff. If the sparsity is similarly extreme (top ~50–100 classes out of tens of thousands), CCE's approach would transfer directly; if the softmax is more diffuse, the filtering threshold or the block-level strategy might need adjustment. This would establish the domain generality of CCE's core insight beyond language modeling.

Practical Applications and Downstream Use Cases

Single-GPU training of large-vocabulary LLMs. The most immediate practical application is enabling researchers and small organizations to train models with vocabularies of 128K–256K tokens on a single GPU—something that is effectively impossible with standard cross-entropy implementations. As documented in Figure 1, Gemma 2 (2B) with standard cross-entropy achieves a maximum batch size of only ~1.1M tokens on a 16-GPU setup, meaning ~70K tokens per GPU—barely enough for a single 80K-token sequence. With CCE, the per-GPU batch size ceiling rises to ~10.6M tokens, a 9.5× increase, meaning even a single GPU can handle reasonable batch sizes for large-vocabulary models. For a graduate student or independent researcher with a single 80 GB A100 or H100 GPU, CCE makes it possible to fine-tune (and potentially pretrain at small scale) models like Gemma 2 (2B) or Llama 3 (8B) without model parallelism, dramatically lowering the hardware barrier to large-vocabulary LLM research.

Cost reduction for large-scale pretraining via batch size increases. For organizations running large-scale pretraining, CCE's memory savings translate directly to cost reduction. The paper's Mistral NeMo experiment demonstrates the mechanism: CCE-Kahan-FullC enabled doubling the batch size per GPU, which reduced total training time by 2 hours (16%) compared to torch.compile. Extrapolating: if a 1,000-GPU training run for a 70B-parameter model with a 128K vocabulary can double its per-GPU batch size (reducing the number of gradient accumulation steps or enabling the same global batch size with fewer GPUs), the savings could reach tens of thousands of GPU-hours. The batch size increases in Figure 1—1.3× to 14.9× depending on the model—suggest this benefit applies broadly, with the largest gains for models where the logit matrix currently dominates memory (Gemma 2 2B at 14.9×, GPT-2 at 11.9×). For models where the logit matrix is a smaller fraction of total memory (Llama 2 13B at 1.3×), the benefit is more modest but still translates to 30% larger batch sizes, which can improve training efficiency or reduce the GPU count needed.

Fine-tuning large-vocabulary models on consumer hardware. For practitioners fine-tuning models like Llama 3 (8B, |V| = 128K) or Qwen 2.5 (7B, |V| = 152K) on consumer GPUs with 24 GB of VRAM (e.g., RTX 4090), the logit matrix for even modest batch sizes can exceed available memory. With standard cross-entropy, a batch of 2,048 tokens and a 128K vocabulary requires 1 GB just for the fp32 logit matrix—more after accounting for bf16 copies and intermediate buffers—which can push the total memory over the 24 GB limit. CCE's near-zero memory overhead for the loss layer removes this pressure entirely, allowing the full 24 GB to be used for model parameters, optimizer states, and activations. Combined with the paper's recommendation to use basic CCE for fine-tuning (no Kahan summation needed, maximum memory savings), this makes fine-tuning of large-vocabulary models on consumer hardware practical without the memory-induced batch size constraints that currently force aggressive gradient accumulation.

Pipeline-parallel training with balanced stages. In large-scale training setups using pipeline parallelism, each pipeline stage should ideally have similar memory and computation requirements to avoid bottlenecks. The classification head has traditionally been an outlier with disproportionately high memory demands relative to its computation, forcing pipeline planners to allocate extra resources to the final stage or accept underutilization elsewhere. CCE eliminates this imbalance by reducing the classification head's memory to near the lower bound, making its memory-to-computation ratio comparable to other transformer layers. For a 70B-parameter model with a 128K vocabulary trained with 8-stage pipeline parallelism, this could mean the difference between needing a dedicated high-memory stage for the head (which might require different GPU configurations than the other stages) versus assigning the head to the same stage type as the final transformer layers—simplifying cluster provisioning and potentially reducing the number of pipeline stages needed.