定義：
カットクロスエントロピー[CCE]は,Daniel M. Ziegler らによって 2023年 に提案された.
Ziegler, Daniel M., et al. "Cut cross-entropy: A compact cross-entropy loss for large vocabulary language models." arXiv preprint arXiv:2306.11606 (2023).
LLM の語彙サイズが大きくなると,損失層で使用されるメモリも大きくなり,多くの最近の言語モデルでは,この 1 つの層で使用されるメモリが学習時のメモリフットプリントの大部分を占めるようになってしまうという問題を抱える.
そこで,CCE は,損失とその勾配の計算が正解ラベルの対数確率のみに依存するようにすることで,この問題を解決することを目指す.
$N$ をサンプル数,$x_i$ を入力データ,$y_i$ を正解ラベル,$p(y_i | x_i)$ をモデルが $x_i$ を入力としたときに $y_i$ を出力する確率とする.このとき,従来のクロスエントロピー損失関数は以下のように定義される.\[L = -\frac{1}{N} \sum_{i=1}^N \log p(y_i | x_i) \]
CCEでは,この損失関数を以下のように分解する.\[L = -\frac{1}{N} \sum_{i=1}^N \left( \mathbf{I}(y_i)^\top \log \text{softmax}(\mathbf{C}^\top \mathbf{f}(x_i)) \right) \]
ここで,$\mathbf{I}(y_i)$ は,$y_i$ 番目の要素のみが1で,それ以外の要素は0であるような one-hot ベクトル.$\mathbf{C}$ は 分類器の重み行列,$\mathbf{f}(x_i)$ は モデルの出力ベクトルである.
この式は,以下の2つの操作に分解できる.
CCEでは,これらの操作を効率的に計算するために,TritonなどのGPUプログラミングフレームワークを用いてカスタムCUDAカーネルを実装している.
CCEの利点は,メモリ消費量を大幅に削減できることである.従来のクロスエントロピー計算においては,全語彙に対する確率を計算する必要があるため,メモリ消費量が大きくなっていた.この点,CCEでは,正解ラベルに対応する対数確率のみを計算すればよいため,メモリ消費量を大幅に削減できる.これにより,より大きなバッチサイズでLLMを訓練することが可能になり、訓練速度の向上に資する.
また,CCEは,勾配のフィルタリングと語彙のソートを組み合わせることによって,更なる性能向上を実現した.
Mathematics is the language with which God has written the universe.