今回の論文
今回取り上げるのは、Erik Wijmans、Brody Huval、Alexander Hertzberg、Vladlen Koltun、Philipp Krähenbühlによる「Cut Your Losses in Large-Vocabulary Language Models」です。2024年11月にarXivで公開され、ICLR 2025で発表されました。公開元はAppleで、分野はLLMの学習システム、GPUカーネル最適化、メモリ効率のよい損失計算です。論文URLは https://arxiv.org/abs/2411.09009 です。
この論文を選んだ理由は、LLMを学習するときに見落としやすい「最後の損失計算」が、語彙の拡大とともに大きなメモリ制約になる点を、数学を変えずに解決しているからです。モデルを量子化したり小さくしたりせず、logitを一時的に置かないという発想は、限られたGPUで継続学習やファインチューニングを行う開発にそのまま役立ちます。
どんな技術か
Cut Cross-Entropy（CCE）は、LLMの出力層で使う交差エントロピー損失を、巨大な全語彙分のlogit行列をGPUのグローバルメモリに保存せず計算する手法です。通常は、バッチ中の全トークン数を N、語彙数を V とすると、N × V 個のスコアを一度に作ります。語彙が10万規模で、長いシーケンスや大きなバッチを扱うと、この中間テンソルだけで数十GBに達します。
CCEは正解トークンのスコアだけを個別に求め、正規化に必要な全語彙の和は小さなブロックごとに走査して集約します。つまり、全スコアを「計算はするが保存しない」設計です。交差エントロピーそのものは変えないため、近似損失ではありません。FlashAttentionがAttention行列をmaterializeしない考え方を、出力層のsoftmaxへ持ち込んだ技術と捉えると理解しやすいです。
課題
LLMの規模拡大ではパラメータ数に注目が集まりがちですが、学習時には語彙数も重要です。大きな語彙は、テキストを短いトークン列にし、多言語やコードなどを効率よく扱える可能性があります。一方で、各位置について全語彙を候補にするsoftmaxの中間結果が急増します。
通常の実装では、隠れ状態 H と分類器重み W の行列積から全logit Z = HW^T を作り、log-softmaxとbackwardのために保持します。この Z のサイズは N × V です。GPUメモリに収まらなければ、バッチを小さくする、勾配累積を増やす、語彙を縮める、GPUを増やす、といった選択になり、試行速度やコストを損ねます。
この問題は、語彙が大きいほど、またシーケンスを詰め込みGPUを有効利用したいほど顕著です。LoRAのような省パラメータ学習でも、出力層まで勾配を通す設定ではlogitテンソルが残るため、アダプタだけを軽くしても学習全体が軽くなるとは限りません。
用語解説
- logit
- softmaxの前にモデルが各語彙トークンへ出す生のスコアです。CCEはこの全語彙分のスコアをまとめてGPUメモリに置かないことが核心です。
- 交差エントロピー損失
- 正解トークンに高い確率を与えられたかを測るLLM学習の標準的な損失です。CCEは損失の定義を置き換えるのではなく、同じ値と勾配を得る計算経路を変えます。
- log-sum-exp
- 全logitを数値的に安定に足し合わせ、softmaxの分母を求める演算です。語彙ブロックを順番に読んでも集約できるため、CCEのオンライン計算を支えます。
- materialize
- 計算途中の行列を実体のテンソルとしてメモリに確保することです。CCEは巨大な`N × V` logitをmaterializeしないことで、メモリ量の増え方を変えます。
- GPUのSRAMとグローバルメモリ
- SRAMはチップ上の高速・小容量メモリ、グローバルメモリは大容量だがアクセスが相対的に遅い領域です。CCEは小さなタイルをSRAMで処理し、巨大な中間値をグローバルメモリへ書き出しません。
技術の仕組み
損失を二つの必要な量へ分解する
各トークン位置 i の正解ラベルを y_i、logitを z_i とすると、交差エントロピーは概念的に -z_i[y_i] + logsumexp(z_i) と書けます。ここで最終的に必要なのは正解ラベルの一つのlogitと、全語彙に対するlog-sum-expだけです。z_i 全体を配列として残しておく必要はありません。
CCEはまず、正解ラベルの行だけを選んで内積を計算するindexed matrix multiplicationを行います。次に分類器重みを語彙方向のタイルに分け、各タイルのlogitをSRAM内で計算します。タイルごとの最大値と指数和をオンライン更新すれば、全タイルを処理した後に正確なlog-sum-expを得られます。
forwardでは作って捨て、backwardで再構成する
通常のcross-entropyは巨大なlogitやlog確率を保持し、backwardで再利用します。CCEはforwardでタイルの中間値を外部メモリに残さず、backward時に必要なタイルを再計算します。再計算の演算は増えても、巨大なテンソルの読み書きを省けるため、メモリ節約だけでなく実行時間も保ちやすくなります。
勾配計算ではsoftmax確率が必要です。CCEはタイルを再計算してsoftmax値を求め、隠れ状態と分類器重みへの勾配を蓄積します。全語彙を候補に含める点は通常のsoftmaxと同じなので、サンプリングsoftmaxのような候補近似ではありません。
小さすぎる勾配を省く工夫
論文はsoftmax分布の裾が非常に小さいことにも着目します。数値精度以下で勾配への寄与が無視できる要素は、分類器重みの勾配計算を省略できます。これは計算を減らすための実装上の最適化であり、論文では閾値を使った場合も収束曲線が通常実装とほぼ一致するかを検証しています。
ただし、この省略を有効にするか、数値補償をどうするかは精度形式や再現性要件に依存します。完全な数値的一致が必要な検証では、論文で扱うKahan和などの安定化と、勾配フィルタを切る設定を検討すべきです。
実験と結果
メモリと実行時間の比較
論文はGemma 2 2Bなどのデコーダ専用LLMを対象に、通常の実装、torch.compileを用いた実装、チャンク化ベースの実装とCCEを比較しました。A100環境のGemma 2 2Bベンチマークでは、lossのforwardに必要なメモリは通常実装の約24GBに対してCCEは約1MBでした。forwardとbackwardを合わせた分類器ヘッドのメモリも、約28GBから約1GBへ下がっています。
時間面でも、同じベンチマークでCCEのforwardは46.4ms、通常実装は81.9ms、forwardとbackwardの合計は134.8msと通常実装の208.8msより短い結果でした。ここでの重要点は、メモリ節約のためにCPUへ退避したり、単純に語彙を分割して遅くなったりする方式ではなく、オンチップ処理と融合カーネルで速度も維持していることです。
より大きいバッチと収束の確認
16台の80GB GPUを使うfully sharded data parallel設定とactivation checkpointingを前提にした比較では、CCEによりモデルごとの最大バッチサイズを1.5倍から10倍へ増やせると報告されています。特に小〜中規模モデルでも大語彙では出力層が支配的になり得ることが可視化されました。
また、Gemma 2 2B、Phi 3.5 Mini、Qwen 2.5 7B、Mistral NeMoのファインチューニングで、通常実装とほぼ同じ損失推移を示しました。事前学習の検証でも、数値安定性を強めた設定でvalidation perplexityが一致することを確認しています。したがって結果は「メモリだけ小さくした近似法」ではなく、標準交差エントロピーを実用的な精度で実装し直せていることを示します。
何に使える？
限られたGPUでの継続事前学習
多言語コーパス、コード、専門記号を重視して語彙を大きくしたい場合、CCEは出力層のメモリ制約を緩和します。同じGPUでもシーケンス数を増やせれば、通信やoptimizer更新の回数を相対的に減らし、学習効率を改善できる可能性があります。
長文のSFT・ドメイン適応
社内文書や長いコードを使うSFTでは、コンテキスト長とバッチサイズの両方がGPUメモリを押し上げます。Attentionやactivationの対策を施した後もlossがボトルネックなら、CCEを足す価値があります。特にモデル本体を変更せずに適用できるため、既存モデルの比較実験を崩しにくい利点があります。
語彙設計の実験
トークナイザの語彙を増やすと、シーケンスを短くできる一方で出力softmaxが重くなります。CCEはこのトレードオフの片側を小さくします。トークン化の粒度を変える研究や、多言語・コード混在モデルで、語彙数を理由に試せなかった設定を探索する用途に向きます。
開発や事業へのヒント
メモリ内訳を計測してから最適化する
学習がOOMになると、先にモデル量子化やGPU増設を検討しがちです。しかし、loss用logitが主因なら対策は出力層にあります。プロファイラで、重み、activation、KV cacheではなく、[batch_tokens, vocab_size] のテンソルがどれほど占めるかを見る習慣が重要です。
「保存しない計算」を設計原則にする
CCEの本質は、最終出力に不要な巨大な中間値を保存しないことです。類似の考え方は、RAGのreranker、対照学習の類似度行列、画像分類の大クラスsoftmaxにも応用を検討できます。ただし適用先では、オンライン集約とbackward再計算が数学的に同値かを必ず確認します。
導入は小さな検証から始める
公開実装はLlama、Phi 3、Mistral、Gemma 2系への統合を提供しています。対象モデル、GPU、PyTorch/Tritonの組み合わせが合う場合は、まず固定バッチでloss・勾配・wall-clock・最大バッチを比較するのが現実的です。メモリに余裕がある小語彙モデルでは、導入コストに見合わないこともあります。
限界
CCEが消せるのは主にN × Vのlogit中間テンソルであり、分類器重み、optimizer state、activation、隠れ状態の勾配まで無料になるわけではありません。語彙が小さい、またはバッチが小さい場面では削減幅が限られます。
実装はGPUカーネルへの依存が強く、公開実装の高速なCCEはAmpere以降のGPU、PyTorch 2.4以降、Triton 3.0以降を前提にしています。macOSや古いGPUではtorch.compile実装を使えますが、同じ性能特性を期待すべきではありません。また、Transformers統合では学習時にlogitsを返さないため、loss以外にlogitを直接使う独自の学習ループは設計を見直す必要があります。
勾配フィルタは数値的に無視できる成分を省く最適化です。論文の検証では収束を損ねませんでしたが、異なる精度、極端な温度設定、特殊な損失重み付けで常に同じとは限りません。厳密な再現実験や新しいドメインでは、フィルタの有無を含むablationを行うのが安全です。
よくある質問
Q. CCEはsoftmaxを近似する技術ですか？
A. いいえ。全語彙をブロックごとに走査して正確なlog-sum-expを求めるため、損失の定義は通常のfull softmaxと同じです。勾配フィルタを有効にした場合だけ、数値精度以下の寄与を実装上省きます。
Q. LoRAだけを学習する場合にも効果はありますか？
A. 効果があり得ます。LoRAで学習可能パラメータを減らしても、通常は出力層を通してlossを計算するため、大語彙・大バッチならlogitテンソルが残ります。ただしLM headを凍結しても損失計算自体は必要なので、実際のメモリ内訳を測って判断してください。
Q. CCEを使えば語彙数を無制限に増やせますか？
A. いいえ。N × Vの中間テンソルは避けられますが、全語彙の分類器重みと全語彙への計算は依然必要です。語彙を増やすと計算量や重み・optimizer stateは増えるため、トークナイザ品質とのトレードオフは残ります。
Q. 推論時の生成速度も速くなりますか？
A. 主対象は学習時のlossとbackwardです。通常の自己回帰推論では正解ラベルへのcross-entropyを計算しないため、生成のdecode高速化には直接効きません。学習基盤のメモリ対策として評価してください。
Q. 導入時に最初に確認すべきことは何ですか？
A. 対象モデルの統合対応、GPU世代、PyTorch/Tritonのバージョン、そしてloss以外でlogitsを利用していないかを確認してください。そのうえで同一入力におけるloss・勾配の差、最大バッチ、実行時間を通常実装と比較するのが確実です。
今日の学び
この論文は、大語彙LLMの学習で交差エントロピーが作る巨大なlogitテンソルを課題にしています。Cut Cross-Entropyは、正解logitとオンラインlog-sum-expへ計算を分解し、タイルをSRAMで処理して中間行列を保存しないことで解きました。開発上のヒントは、OOMをモデルサイズだけの問題と決めつけず、最終的に必要な値と一時的に保存している値を分けて観察することです。