Cut Cross Entropy from first principles - Omkaar Kamath

If you’ve ever tried training a 70B model and watched your GPU memory bar hit 99%, you’ve already met the enemy: Cross Entropy Loss. It’s silently eating half your VRAM, and Apple’s new Cut Cross Entropy (CCE) paper could help out. Cross Entropy Loss is commonly calculated by materializing all the logits, which, given the large vocab size and number of batch tokens, leads to exploding VRAM usage. The green pie segment shown in the image below is proportion of memory used by the Cross Entropy loss calculation... notice how with CCE green pie segments seemingly vanish. Following is an image from Apple's paper showing ratio of memory consumed by different parts of a training run with and without CCE. The CE Loss formula = -log(prob(Correct Token)) Probability of a token comes from softmax over the vocab dimension of the logits tensor. Putting all this together and a few math tricks, we get loss = -correct_token + LSE The most naive form of Cross Entropy loss calculation would be: That wou

Cut Cross Entropy from first principles - Omkaar Kamath Cut Cross Entropy from first principles If you’ve ever tried training a 70B model and watched your GPU memory bar hit 99%, you’ve already met the enemy: Cross Entropy Loss. It’s silently eating half your VRAM, and Apple’s new Cut Cross Entropy (CCE) paper could help out. Cross Entropy Loss is commonly calculated by materializing all the logits, which, given the large vocab size and number of batch tokens, leads to exploding VRAM usage. The green pie segment shown in the image below is proportion of memory used by the Cross Entropy loss ca

saved by

• Curtis Chong

related reading

• frontier model training methodologies | Alex Wa's Blogdjdumpling.github.io
• Transformer Math 101 | EleutherAI Blogblog.eleuther.ai
• microgptkarpathy.github.io
• The Little Book of Deep Learningfleuret.org
• Understanding Memorization via Loss Curvaturegoodfire.ai
• State of RL for reasoning LLMs | A. Weersaweers.de
• A Gentle Introduction to Cross-Entropy for Machine Learning - MachineLearningMastery.commachinelearningmastery.com
• Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
• Overleaf Examplearxiv.org
• Reiner Pope – The math behind how LLMs are trained and serveddwarkesh.com
• [2604.18002] Neural Garbage Collection: Learning to Forget while Learning to Reasonarxiv.org
• How is LLaMa.cpp possible?finbarr.ca