title: Cut Your Losses in Large-Vocabulary Language Models
summary: 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.
paper_authors._id: 6736d4060cbcd485282f342c
paper_authors.name: Erik Wijmans
paper_authors.status: extracted_confirmed
paper_authors.statusLastChangedAt: 2024-11-15T22:47:34.444Z
paper_authors._id: 664ce8c217586a9634370552
paper_authors.avatarUrl: /avatars/b8f96ed0dc321c833d654eab24a2443d.svg
paper_authors.fullname: Erik Wijmans
paper_authors.type: user
paper_authors.user: erikwijmans
paper_authors._id: 6736d4060cbcd485282f342d
paper_authors.name: Brody Huval
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-11-15T13:59:30.106Z
paper_authors._id: 6685b9b6eff6f8ec51725397
paper_authors.avatarUrl: /avatars/2425f72d0d6be1dadc6f7797a94687ba.svg
paper_authors.fullname: Brody Huval
paper_authors.type: user
paper_authors.user: brodyhuval
paper_authors._id: 6736d4060cbcd485282f342e
paper_authors.name: Alexander Hertzberg
paper_authors._id: 6736d4060cbcd485282f342f
paper_authors.name: Vladlen Koltun
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-11-15T13:59:45.191Z
paper_authors._id: 6707b1a9356f027ca986d6b1
paper_authors.avatarUrl: /avatars/d32eae81b27dc4f7090196d410e96876.svg
paper_authors.fullname: Vladlen Koltun
paper_authors.type: user
paper_authors.user: vkoltun
paper_authors._id: 6736d4060cbcd485282f3430
paper_authors.name: Philipp Krähenbühl
paper_title: Cut Your Losses in Large-Vocabulary Language Models
paper_summary: 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.
paper_ai_keywords: cross-entropy logit matrix token vocabulary cross-entropy in the loss computation Cut Cross-Entropy (CCE) log-sum-exp matrix multiplications flash memory memory footprint Gemma 2 (2B) model classifier head sparsity softmax gradient computation