Research map / Architectures and efficiency / RNNs and normalization

Cut Your Losses in Large-Vocabulary Language Models

ICLR 2025 · Architectures and efficiency › RNNs and normalization · approach: Neural scaling laws · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Cut Cross-Entropy (CCE), a fused Triton kernel that computes LLM cross-entropy loss and its gradient without materializing the full logit matrix, plus gradient filtering that skips negligible softmax blocks.

Problem

The cross-entropy loss layer, whose memory scales with batch tokens times vocabulary size, dominates LLM training memory for models with large vocabularies.

Approach

Compute the correct-token logit directly and the log-sum-exp over the vocabulary on the fly in on-chip SRAM via a fused matmul-plus-reduction kernel, so logits are never written to global memory; the backward pass recomputes logits blockwise and skips blocks whose softmax entries fall below bfloat16 precision, with vocabulary sorting to increase block sparsity.

Key results

• For Gemma 2 (2B), CCE reduces loss-computation memory from 24000 MB to 1 MB and loss+gradient memory from 28000 MB to 1164 MB, with a loss+gradient time of 145 ms versus 208 ms for baseline.
• Gradient filtering gives about 3.5x speedup in the backward pass without loss of precision; without it the backward pass is 3.4x slower.
• Loss and validation perplexity curves for CCE match torch.compile when fine-tuning on Alpaca and pretraining on OpenWebText (CCE-Kahan-FullC).

Measured effects reported in the paper

Surprising findings

• The log-probabilities of the cross-entropy layer account for up to 89% of the memory footprint of Gemma 2 (2B) training, larger than the rest of the model's memory combined in some settings.
• Softmax probabilities vanish so quickly that fewer than 0.02% of elements are non-zero in evaluated frontier models, enabling skipping of most gradient blocks.

Limitations

• Triton's block-level control flow limits gradient filtering and the thread-safe log-add-exp to block granularity; a CUDA implementation may do better.
• Without Kahan summation, gradient filtering on the classifier gradient harmed pretraining perplexity; Kahan variants use more memory (e.g. 3558 MB vs 1809 MB for loss+gradient on Gemma 2 9B) and time.
• For small vocabulary-to-hidden-dimension ratios (e.g. Phi 3.5 Mini), CCE is about 50% slower than torch.compile on loss computation, though memory is much lower.

Experimental setup

Datasets: Alpaca OpenWebText

Metrics: peak memory (MB) wall-clock time (ms) training loss curves validation perplexity maximum batch size

Baselines: Baseline PyTorch cross-entropy torch.compile Liger Kernels Torch Tune chunked cross-entropy

Builds on: FlashAttention Liger Kernels torch.compile Kahan summation

Closest papers by meaning

• Constrained Entropic Unlearning: A Primal-Dual Framework for Large Language Models (NeurIPS 2025 · similarity 0.88)
• Mitigating Memorization in Language Models (ICLR 2025 · similarity 0.88)
• Cross-model Control: Improving Multiple Large Language Models in One-time Training (NeurIPS 2024 · similarity 0.87)
• Over-Tokenized Transformer: Vocabulary is Generally Worth Scaling (ICML 2025 · similarity 0.87)
• Exploiting Vocabulary Frequency Imbalance in Language Model Pre-training (NeurIPS 2025 · similarity 0.87)
• EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models (NeurIPS 2026 · similarity 0.87)