# Cut Your Losses in Large-Vocabulary Language Models

arXiv: https://arxiv.org/abs/2411.09009 · HTML: https://mlsanity.com/paper/2411.09009/

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 39 references that are in the map.
- First posted on arXiv 13 Nov 2024.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers](https://mlsanity.com/paper/2305.07185/) (2023)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies](https://mlsanity.com/paper/2407.13623/) (2024)
- [Efficient softmax approximation for GPUs](https://mlsanity.com/paper/1609.04309/) (2016)
- [Speed Is All You Need: On-Device Acceleration of Large Diffusion Models via GPU-Aware Optimizations](https://mlsanity.com/paper/2304.11267/) (2023)

## Built on by (most cited first, 30 of 40)

- [Reinforcement Learning for Long-Horizon Interactive LLM Agents](https://mlsanity.com/paper/2502.01600/) (2025)
- [Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration](https://mlsanity.com/paper/2512.22382/) (2025)
- [FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling](https://mlsanity.com/paper/2502.14856/) (2025)
- [OVD: On-policy Verbal Distillation](https://mlsanity.com/paper/2601.21968/) (2026)
- [Arcee Trinity Large Technical Report](https://mlsanity.com/paper/2602.17004/) (2026)
- [Entropy-Preserving Reinforcement Learning](https://mlsanity.com/paper/2603.11682/) (2026)
- [TransAct V2: Lifelong User Action Sequence Modeling on Pinterest Recommendation](https://mlsanity.com/paper/2506.02267/) (2025)
- [Scaling Categorical Flow Maps](https://mlsanity.com/paper/2605.07820/) (2026)
- [A Course Correction in Steerability Evaluation: Revealing Miscalibration and Side Effects in LLMs](https://mlsanity.com/paper/2505.23816/) (2025)
- [Calibrated Language Models and How to Find Them with Label Smoothing](https://mlsanity.com/paper/2508.00264/) (2025)
- [Can LLMs Understand Unvoiced Speech? Exploring EMG-to-Text Conversion with LLMs](https://mlsanity.com/paper/2506.00304/) (2025)
- [Spherical Leech Quantization for Visual Tokenization and Generation](https://mlsanity.com/paper/2512.14697/) (2025)
- [The Design Space of Tri-Modal Masked Diffusion Models](https://mlsanity.com/paper/2602.21472/) (2026)
- [CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs](https://mlsanity.com/paper/2605.19269/) (2026)
- [DynaSpec: Context-aware Dynamic Speculative Sampling for Large-Vocabulary Language Models](https://mlsanity.com/paper/2510.13847/) (2025)
- [DiSTAR: Diffusion over a Scalable Token Autoregressive Representation for Speech Generation](https://mlsanity.com/paper/2510.12210/) (2025)
- [Efficient and Programmable Exploration of Synthesizable Chemical Space](https://mlsanity.com/paper/2512.00384/) (2025)
- [FlashSampling: Fast and Memory-Efficient Exact Sampling](https://mlsanity.com/paper/2603.15854/) (2026)
- [FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training](https://mlsanity.com/paper/2606.22932/) (2026)
- [Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss](https://mlsanity.com/paper/2608.03796/) (2026)
- [Distilled Continuous Diffusion Language Models Can Write Code in Few Steps---or One](https://mlsanity.com/paper/2609.04531/) (2026)
- [ABC: Achieving Better Control of Multimodal Embeddings using VLMs](https://mlsanity.com/paper/2503.00329/) (2025)
- [Memory-Efficient Backpropagation for Fine-Tuning LLMs on Resource-Constrained Mobile Devices](https://mlsanity.com/paper/2510.03425/) (2025)
- [Vocab Diet: Reshaping the Vocabulary of LLMs via Vector Arithmetic](https://mlsanity.com/paper/2510.17001/) (2025)
- [NanoQuant: Efficient Sub-1-Bit Quantization of Large Language Models](https://mlsanity.com/paper/2602.06694/) (2026)
- [Adjoint sharding for very long context training of state space models](https://mlsanity.com/paper/2501.00692/) (2025)
- [torchtune: PyTorch native post-training library](https://mlsanity.com/paper/2605.21442/) (2026)
- [Large Byte Model: Teaching Language Models About Compiled Code](https://mlsanity.com/paper/2606.02834/) (2026)
- [Spokes: Optimizing for Diverse Pretraining Data Selection](https://mlsanity.com/paper/2606.15216/) (2026)
- [Batch Size or Negatives? A Selection Rule for Memory-Constrained Recommender Training](https://mlsanity.com/paper/2608.11061/) (2026)

## Abstract (arXiv)

As language models grow ever larger, so do their vocabularies. This has shifted the memory footprint of LLMs during training disproportionately to one single layer: the cross-entropy in the loss computation. Cross-entropy builds up a logit matrix with entries for each pair of input tokens and vocabulary items and, for small models, consumes an order of magnitude more memory than the rest of the LLM combined. We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss without materializing the logits for all tokens into global memory. Rather, CCE only computes the logit for the correct token and evaluates the log-sum-exp over all logits on the fly. We implement a custom kernel that performs the matrix multiplications and the log-sum-exp reduction over the vocabulary in flash memory, making global memory consumption for the cross-entropy computation negligible. This has a dramatic effect. Taking the Gemma 2 (2B) model as an example, CCE reduces the memory footprint of the loss computation from 24 GB to 1 MB, and the total training-time memory consumption of the classifier head from 28 GB to 1 GB. To improve the throughput of CCE, we leverage the inherent sparsity of softmax and propose to skip elements of the gradient computation that have a negligible (i.e., below numerical precision) contribution to the gradient. Experiments demonstrate that the dramatic reduction in memory consumption is accomplished without sacrificing training speed or convergence.