CatBoost論文のprediction shiftについて完全に理解する

CatBoostの論文における、prediction shiftについて調べる機会があったのでまとめてみました。
(CatBoost: unbiased boosting with categorical features https://arxiv.org/pdf/1706.09516.pdf
以下は、論文の4.1 Prediction shift、AppdendicesのA Proof of Theorem 1を元に、私の理解で構成しなおしたものです。

概要

勾配ブースティングの計算において、各ブースティングで同じ学習データを使うことにより、偏りが発生する。

シンプルな以下の前提で考える。

データセットの前提

データセットの前提は以下のとおり:

  • データの個数は$n$個、特徴量は $(s, t)$ の2つ
  • $s, t$はそれぞれベルヌーイ分布($p=0.5$)に従う二値変数
  • 目的変数 $y = c_1 s+ c_2 t$ であり、誤差は無い

つまり、このデータは以下のようになっている。

  • 4種類のデータ $(s, t) = (0, 0), (1, 0), (0, 1), (1, 1)$ から構成される。
  • 目的変数yの値はそれぞれ、$0, c_1, c_2, c_1 + c_2$である。$f^*(s, t)$とも表すことにする
  • 分布に従ってこれらのデータを生成するとき、現れたデータの個数を $\xi_{00}, \xi_{10}, \xi_{01}, \xi_{11}$ で表すことにする

学習の前提

学習の前提は以下のとおり:

  • 勾配ブースティング木で回帰タスクを解く、目的関数は二乗誤差。
  • 木の数は2、木の深さは1とする。
    $c_1$の方が$c_2$より大きく、1本目の木では$s$による分岐が行われ、2本目の木では$t$による分岐が行われると仮定する。

主張

1本目、2本目の木の学習に使うデータセットをそれぞれ$D_1$, $D_2$としたとき、

  • $D_1$と$D_2$が独立に生成されたサンプルであれば、学習されたモデルによるデータ$(s, t)$の予測値の期待値は、
    $f^*(s, t)$ となる。
  • $D_1$と$D_2$が同じであった場合、学習されたモデルによるデータ$(s, t)$の予測値の期待値は、
    $f^*(s, t) - \frac{1}{n-1} c_2 (t - \frac{1}{2} )$ となる。

(Threorem1の主張をまとめなおしたもの。なお、この記事では、論文中に出てくる微小な項である $O(1/2^{n})$ や、データが極めて偏った場合を除外する意味を持つAを省略する)

つまり、1本目と2本目の木の作成で同じデータセットを使うことにより、予測値の期待値に${1}/{(n-1)}$ に比例する偏りが発生する。 この偏りは、1本目と2本目の木で独立したデータを使って学習した場合には発生しない。

なお、prediction shiftの影響を受けないよう、CatBoostではデータ数が少ない場合にはordered boostingというアルゴリズムを使用するようになっている。

証明とポイント

このThreorem1の証明はそこそこの長さがあり、Appendicesの A Proof of Threorem 1に記述されている。

ポイント1

まず、この前提でデータセットが生成されていたとして、モデルを学習させたときにできる木のウェイトを求める。

1本目の木では、特徴量$s$を元に分岐が行われる。

  • 特徴量 $s=1$ であるデータの目的変数の平均は、$c_1 + c_2\xi_{11} / (\xi_{10} +\xi_{11})$
  • 特徴量 $s=0$ であるデータの目的変数の平均は、$c_2\xi_{01} / (\xi_{00} +\xi_{01})$

であり、目的変数の平均がその葉のウェイトとなるため、上記をまとめて、
データ $(s, t)$ に適用されるウェイト $h^{1}(s, t) = c_1 + c_2 \xi_{s1}/(\xi_{s0}+\xi_{s1})$ と表される。(論文11Pの式(7) )

2本目の木では、特徴量$t$を元に分岐が行われる。目的変数から1本目の木の予測値を差し引いた残差の平均がその葉のウェイトになる。
データ $(s, t)$ に適用されるウェイト $h^{2}(s,t)$ は、$t=0, t=1$それぞれにおいて、目的変数から1本目の木の予測値を差し引いた平均を計算し展開すると、

$t=0$のとき、 $$-c_2\left(\frac{\xi_{00}\xi_{01}}{(\xi_{00}+\xi_{01})(\xi_{00}+\xi_{10})} + \frac{\xi_{10}\xi_{11}}{(\xi_{10}+\xi_{11})(\xi_{00}+\xi_{10})}\right)$$

$t=1$のとき、 $$ c_2\left ( \frac{\xi_{00}\xi_{01}}{(\xi_{00}+\xi_{01})(\xi_{01}+\xi_{11})} + \frac{\xi_{10}\xi_{11}}{(\xi_{10}+\xi_{11})(\xi_{01}+\xi_{11})} \right ) $$

となる。(論文14Pの前半)

ポイント2

$D_1$, $D_2$に同じデータセットを使った場合の$h^{1}(s,t)+h^{2}(s,t)$ の期待値に偏りがあることを示す。

  • $h^1(s,t) + h^2(s,t)$ の期待値は、
    $f^*(s,t) - \frac{1}{n-1} c_2 \left(t-\frac{1}{2}\right) = c_1 s + c_2 t - \frac{1}{n-1} c_2 \left(t-\frac{1}{2}\right) $ である。
    これは以下の $h^1(s,t), h^2(s,t)$ の期待値の和として導かれる。
  • $h^1(s,t)$ の期待値は、$c_1 s + \frac{c_2}{2}$ である。
    (Proposition1参照。特にこの時点で偏りはない)
  • $h^2(s,t)$ の期待値は、$\ (2t - 1) \frac{c_2}{2} \left(1 - \frac{1}{n-1}\right)$ である。
    つまり、
    t=0のとき、 $\ - \frac{c_2}{2} \left(1 - \frac{1}{n-1}\right)$
    t=1のとき、 $\ \frac{c_2}{2} \left(1 - \frac{1}{n-1}\right)$
    である。(Proposition3のProof参照。 この$h^2(s,t)$の期待値に偏りがある。)

ここで、$h^1(s,t)$ の期待値の計算は特に難しくないが、$h^2(s,t)$ の期待値の計算はやや難しい。
対称性より、 $$\frac{\xi_{00}\xi_{01}}{(\xi_{00}+\xi_{01})(\xi_{00}+\xi_{10})}$$ のみ考えれば良く、他の項にも同様に適用できる。

この項の期待値を計算すると、$\frac{1}{4} \left(1 - \frac{1}{n-1}\right)$ となり、もう一つの項も同じ値のため2を乗じ、さらに係数$-c_2$もしくは$c_2$を乗じたものが$h^2(s,t)$となる。
この項の期待値が$\frac{1}{4}$ であれば偏りはないのだが、そうでなく $\frac{1}{4} \left(1 - \frac{1}{n-1}\right)$ となっていることがポイントで、この部分が偏りの発生源といえるだろう。
この計算はLemma3に記述されているとおり煩雑だが、分子である $\xi_{00}, \xi_{01}$ と分母が独立でなく、分子の値が大きくなるときほど分母が大きくなりやすい関係性にあるため、「フラットな」期待値よりも少し小さくなる、と感覚的に理解できる。