XGBoostのマルチクラス分類用目的関数の実装を読み解く

# TL;DR
stableなsoftmax関数を使ったCross Entropy Lossの微分値をCUDAを直叩きして計算してます！

# XGBoostとは？
言わずと知れたGradient Boostingの効率的な並列処理アルゴリズムおよび実装です。
[XGBoost: A Scalable Tree Boosting System](https://www.kdd.org/kdd2016/papers/files/rfp0697-chenAemb.pdf)
https://github.com/dmlc/xgboost

# マルチクラス分類問題どうなってるんだろう？
上の論文は話をシンプルにするために１変数の回帰問題を主に扱っています（データセットのラベル$y_i$が$y_i\in \mathbb{R}$）。
しかし、実際の問題ではマルチクラスの分類問題を取り扱いたいことが多々あります。そんな時、どんな計算が中でなされているんでしょうか？上記の論文だけではわかりませんでした。
Gradient Boostingの論文読む人には当たり前のようにわかっている前提なのかもしれませんが、マルチクラス分類問題において目的関数$l(y_i, \hat{y}^{(t-1)}_i)$がなんなのか？そしてその１階および２階の偏微分$g_i, h_i$がなんなのか？上記の論文ではほとんど説明がありません。

いいだろう、そしたらソース読んで理解してやる！というのが本項の趣旨です。

# 問題のソースコード
いきなりですが[核心のコードはここです](https://github.com/dmlc/xgboost/blob/master/src/objective/multiclass_obj.cu#L75)。

ここでやっている処理は、前段のPrediction $\hat{y}^{(t-1)}$ が与えられた時、次段のboostingを行うために目的関数に基づいて$g_i, h_i$を計算する、というものです。
論文では、式番号もなくさらっと

```math
g_i=\partial_{\hat{y}^{(t-1)}} l(y_i, \hat{y}^{(t-1)})\\
h_i=\partial_{\hat{y}^{(t-1)}}^2 l(y_i, \hat{y}^{(t-1)})\\
```

と書いてありますね。
該当するC++のコードを引用します。

```C++
    common::Transform<>::Init(
        [=] XGBOOST_DEVICE(size_t idx,
                           common::Span<GradientPair> gpair,
                           common::Span<bst_float const> labels,
                           common::Span<bst_float const> preds,
                           common::Span<bst_float const> weights,
                           common::Span<int> _label_correct) {
          common::Span<bst_float const> point = preds.subspan(idx * nclass, nclass);

          // Part of Softmax function
          bst_float wmax = std::numeric_limits<bst_float>::min();
          for (auto const i : point) { wmax = fmaxf(i, wmax); }
          double wsum = 0.0f;
          for (auto const i : point) { wsum += expf(i - wmax); }
          auto label = labels[idx];
          if (label < 0 || label >= nclass) {
            _label_correct[0] = 0;
            label = 0;
          }
          bst_float wt = is_null_weight ? 1.0f : weights[idx];
          for (int k = 0; k < nclass; ++k) {
            // Computation duplicated to avoid creating a cache.
            bst_float p = expf(point[k] - wmax) / static_cast<float>(wsum);
            const float eps = 1e-16f;
            const bst_float h = fmax(2.0f * p * (1.0f - p) * wt, eps);
            p = label == k ? p - 1.0f : p;
            gpair[idx * nclass + k] = GradientPair(p * wt, h);
          }
        }, common::Range{0, ndata}, device, false)
        .Eval(out_gpair, &info.labels_, &preds, &info.weights_, &label_correct_);
```
## 引用部分の構文

少しづつ読み解きます。
まず、引用部分全体の構文は `common::Transform<>::Init().Eval()` で、ファクトリ関数Initの第１引数の関数オブジェクトを第２引数の `common::Range` で並列で実行する、という処理をしています。Initの第３引数deviceはGPUのidで、そのGPUで実行されることになります。
Eval関数の5つの引数はそのままラムダ式の5つの引数として渡されます。
ちょっとまどろっこしいことをやっているのは、 `common::Transform` がCUDAのkernel関数呼び出しなどをwrapしているためです。

また、この関数オブジェクトは `XGBOOST_DEVICE` マクロによって `__host__ __device__`修飾されます。これによって、GPU上で実行されるdevice関数になります。
CUDAの関数呼び出しは `ホスト側関数（CPU）->kernel関数（GPU）->device関数（GPU）`という流れになりますが、ここでは  `common::Transform::Eval()->Evalが呼びだすkernel関数->関数オブジェクト` とCUDA固有部分がwrapされ、普通のC++コードで書けるように工夫されてます。
// `__host__` が付いているのはこの関数がhost上、つまりCPUによって実行することもできるようにしています。

以下のEval関数から追っていくと流れがわかると思います。
https://github.com/dmlc/xgboost/blob/master/src/common/transform.h#L72

ここまでまとめると、CUDAのdevice関数に5つの配列（ここでは `common::Span<>`型）を渡してあげて、`{0, ndata}` 各インプットデータ毎にidxを与えて並列実行させている・・・ということになります。

## device関数(lambda式)の引数
- `idx`: 各スレッドに渡されるインデックス。そのスレッドが受け持つデータを示す。0以上ndata未満の値がそれぞれのスレッドに渡される。
- `gpair`:  1階微分と２階微分をメンバにもつstructの配列。戻り値に相当する。
- `labels`: 各データの正解ラベル。論文で言う所の $y_i$。データ数N次元ベクトル。
- `preds`: トレーニングデータ数xクラス数(N x K)の行列。前段の予測値、論文で言う所の $\hat{y}^{(t-1)}$。
- `weight`: このweightは[こちらの記事で紹介されている](https://qiita.com/makotu1208/items/e58db3a376d70a763944)DMatrixのweight。
- `_label_correct`: 不正なラベルが貼られてた時にわかるようにflag

ここで、注目すべきなのは、入出力の次元です。入力$\hat{y}^{(t-1)}$はトレーニングデータの数(N) x クラス数(K)の次元を持つ行列になっています。つまり、前段のBoostingでは、トレーニングデータ毎に各クラスの予測値を計算している、ということになります。
また、device関数の出力である `gpair` も同様にN x K次元の行列になっており、こちらもトレーニングデータ毎に各クラスに対応する勾配・へシアンが計算されていることがわかります。

```C++
          common::Span<bst_float const> point = preds.subspan(idx * nclass, nclass);
```
この `point` 変数はidxが0以上ndata未満のrangeなことに注意すれば、あるトレーニングデータ$i$の各クラスの予測値（K次元ベクトル）であることがわかります。

## 関数の正体見たりsoftmax
以下のコメントにもありますが、ここで計算してるのはsoftmax関数ですね。

```C++
          // Part of Softmax function
          bst_float wmax = std::numeric_limits<bst_float>::min();
          for (auto const i : point) { wmax = fmaxf(i, wmax); }
          double wsum = 0.0f;
          for (auto const i : point) { wsum += expf(i - wmax); }
          auto label = labels[idx];

          for (int k = 0; k < nclass; ++k) {
            // Computation duplicated to avoid creating a cache.
            bst_float p = expf(point[k] - wmax) / static_cast<float>(wsum);
```

ただし、ここで予測値の中で最大の値であるwmaxを計算して `expf()` の中で引き算していることに注意してください。
これは、softmax関数のexpの中身全てに、ある定数`-wmax`を足しても結果が変わらないことを利用して、計算上発散を防ぐstableな実装です。
http://saitcelebi.com/tut/output/part2.html#numerical_stability_of_softmax_function

数式にすると、

```math
F(\hat{y}^{(t-1)}_{i,k})=\frac{e^{\hat{y}^{(t-1)}_{i,k}}}{\sum_k e^{\hat{y}^{(t-1)}_{i,k}}}=
\frac{e^{-w_{\rm max}}(e^{\hat{y}^{(t-1)}_{i,k}})}{e^{-w_{\rm max}}(\sum_k e^{\hat{y}^{(t-1)}_{i,k}})}
=\frac{\exp({\hat{y}^{(t-1)}_{i,k}-w_{\rm max}})}{\sum_k \exp({\hat{y}^{(t-1)}_{i,k}}-w_{\rm max})}
```

の一番右の計算をしているというわけです。

## Cross Entropy Lossの1階・2階微分
以上はsoftmax関数の計算でしたが、softmax関数は言わばシグモイド関数の多変数バージョン。それそのものは目的関数ではありません。
実際に最適化の対象になるのはクロスエントロピーロス関数、数式で表すと以下のような関数です。これは多変数の[LogitBoost](https://web.stanford.edu/~hastie/Papers/AdditiveLogisticRegression/alr.pdf)と同様の目的関数ですね。
Cross Entropy Lossは前述のsoftmax関数を$p_{i,k}=F(\hat{y}^{(t-1)}_{i,k})$とすると、

```math
L(y_i, \hat{y}^{(t-1)})=-\sum_k y_{i,k} \log(p_{i,k})
```

となります。
softmax関数、クロスエントロピーロス関数の微分については、以下のサイトを参考にするとわかりやすいです。
https://deepnotes.io/softmax-crossentropy

勾配を求めるところまではディープラーニングで頻出ですね。
softmax関数の微分は、長くなってきたので$o_{i,m}=\hat{y}^{(t-1)}_{i,m}$と変数名を変えて、

```math
\frac{\partial p_{i,l}}{\partial o_{i,m}}=p_{i,l}(\delta_{l,m}-p_{i,m})
```

と美しい形になります。 ここで、$l,m$共にクラスのインデックスです。また、$\delta_{l,m}$ はクロネッカーのデルタで$l,m$が同じクラスの時１になる関数です。
この結果を使うと、勾配$g_i$は

```math
\begin{align}
g_{i,l} &= \frac{\partial L(y_i, \hat{y}^{(t-1)})}{\partial o_{i,l}} \\
&= -\sum_m y_{i,m} \frac{\partial \log(p_{i,m})}{\partial o_{i,l}} \\
&= -\sum_m y_{i,m} \frac{\partial \log(p_{i,m})}{\partial p_{i,m}} \frac{\partial p_{i,m}}{\partial o_{i,l}} \\
&= -\sum_m y_{i,m} \frac{1}{p_{i,m}} p_{i,m}(\delta_{m,l}-p_{i,l}) \\
&= -\sum_m y_{i,m} (\delta_{m,l}-p_{i,l}) \\ 
&= p_{i,l} - \delta_{l,k^*}
\end{align}
```

最後は$k^*$が正解ラベルを表しています。$l$が正解ラベルと一致しているときに、クロネッカーデルタが１になる、という意味です。

$y_{i,m}$が$m=k^*$の時だけ１になる変数であるとするとこの変化が自然に導出されます。

次に2階微分です。

```math
\begin{align}
h_{i,l} &= \frac{\partial^2 L(y_i, \hat{y}^{(t-1)})}{\partial^2 o_{i,l}} \\
&= -\sum_m y_{i,m} \frac{\partial}{\partial o_{i,l}} \frac{\partial \log(p_{i,m})}{\partial o_{i,l}} \\
&= -\sum_m y_{i,m} \frac{\partial}{\partial o_{i,l}}(\delta_{m,l}-p_{i,l}) \\ 
&= -y_{i,m}\frac{\partial}{\partial o_{i,m}}(1-p_{i,m})+\sum_{l\neq m}y_{i,m} \frac{\partial p_{i,l}}{\partial o_{i,l}}\\
&= p_{i,l}(1-p_{i,l})
\end{align}
```

いやはや、綺麗になるもんです。最後は場合分けからの全パターン足し合わせた時、結果的に$m$が正解ラベルと一致しているときのみ値が残るので、正解ラベルによらずにある値を取ることがわかります。

該当するコードは、

```C++
            const bst_float h = fmax(2.0f * p * (1.0f - p) * wt, eps); // p = F(y_i,k)
            p = label == k ? p - 1.0f : p; // ここでpはsoftmax関数の出力からgradient（g_i,k）にアップデートされてます
            gpair[idx * nclass + k] = GradientPair(p * wt, h); // 戻り値行列に保存
```

あれ？だいたいあっているけど、2.0fが掛かっているのはなんでだろうか？

# まとめ
XGBoostのコード面白い。