【論文メモ】Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data

[Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data](https://arxiv.org/abs/1909.06312) をまとめます。

![スクリーンショット 2021-04-02 5.20.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/1e12c113-61ff-de00-8c01-401ecc369858.png)

本論文は [Catboost](https://catboost.ai/) で有名な Yandex から出た論文で、Neural Oblivious Decision Ensembles (NODE) という、テーブルデータ用の Deep Learning モデルを提案しています。
実験結果で他の GBDT モデルを outperform しており、「テーブルデータの universal framework になること間違いなし」だそうです。
実装は[こちら](https://github.com/Qwicen/node)。

## モデル構造

- NODE レイヤー
- 全体のアーキテクチャ
- 学習・推論

の順に紹介します。

### NODE レイヤー

NODE レイヤーは決定木をベースにしているため、先に決定木を定式化しておきます。
深さ $d$ の決定木に関して、 それぞれの分割で使用される特徴量を $f \in \mathbb{R}^{d}$ 、分割の閾値を $b \in \mathbb{R}^{d}$ 、$2^{d}$ 個の末端ノードに紐付いた $d$ 階の重みテンソルを $R \in \mathbb{R}^{\underbrace{2 \times 2 \times 2 \times \cdots}_{d}}$ として、予測値を以下のように表すことができます。

```math
h(x)=R\left[\mathbb{1}\left(f_{1}(x)-b_{1}\right), \ldots, \mathbb{1}\left(f_{d}(x)-b_{d}\right)\right]
```

ここで、$\mathbb{1}\left(\cdot\right)$ はヘヴィサイド関数を表しています。
それぞれのデータが左右どちらに割り当てられたかを 0 or 1 で表すことで、対応する重みを参照して、予測値を得ることができます。

```math
h(x) = R \left[0, 1, 1, \ldots, 0, 1\right]
```

この出力を微分可能にするには、特徴量の選択 $f_i$ と、その分割 $\mathbb{1} \left(f_i(x) - b_i \right)$ を微分可能な関数に置き換える必要があります。
NODE ではこれらを置き換えるのに、 $\alpha$-entmax を使用します。

$\alpha$-entmax は、softmax と sparsemax の内挿となるように設計された活性化関数で、スパースな出力を得ることができます。

<img width="550" alt="スクリーンショット 2021-02-12 19.49.01.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/87ad9a56-dae2-07d9-bade-fc39e21c1722.png">

これを用いて、特徴量選択のための関数 $f$ は、あるひとつの特徴量を選択するようなハードな割り当てを行うのではなく、それぞれの特徴量の重み付き足し合わせとして表現されます。
学習可能な feature selection matrix $F \in \mathbb{R}^{d \times n}$ を用いて、以下のように表されます。

```math
\hat{f}_{i}(x)=\sum_{j=1}^{n} x_{j} \cdot \operatorname{entmax}_{\alpha}\left(F_{i j}\right) \tag{1}
```

同様に、分岐においても、各データを左右どちらかのノードにハードに割り当てるのではなく、左30%・右70%などのように、両方のノードに重み付きで属するように表されます。これは2クラスの $\alpha$-entmax を用いて、以下のように表されます。

```math
\sigma_{\alpha}\left(x\right)=\operatorname{entmax}_{\alpha}\left(\left[x, 0\right]\right)
```

以上を用いて、各分割における出力は

```math
c_i \left(x\right) = \sigma_{\alpha}\left(\frac{\hat{f}_{i} \left(x\right) - b_i}{\tau_i} \right)
```

となります。ここで、$\tau_i$ という学習可能なパラメータを用いて、特徴量ごとのスケールの違いを吸収させています。

![スクリーンショット 2021-04-02 5.20.31.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/bd6b0d5d-bdb4-05ce-fceb-bbda7a769e70.png)

上の図では ${F}_2$ や ${F}_3$ は複数描かれていますが、別の重みを学習させているわけではないことに注意してください。分割関数をノードごとに学習させるのではなく、深さごとに学習させており、左に流れたデータも右に流れたデータも、同じ関数を用いてさらに分割されることになります。

すべての分岐を通した結果、各データは $2^d$ 個の末端ノードすべてに、重み付きで属することになります。
これを式で書くと以下のようになります。

```math
C(x)=\left[\begin{array}{c}
c_{1}(x) \\
1-c_{1}(x)
\end{array}\right] \otimes\left[\begin{array}{c}
c_{2}(x) \\
1-c_{2}(x)
\end{array}\right] \otimes \cdots \otimes\left[\begin{array}{c}
c_{d}(x) \\
1-c_{d}(x)
\end{array}\right] \\
```

ここで、$C\left(x\right)$ はそのデータの各末端ノードへ属する強さを表す重みです。
これは、たとえば $c_1 \left( x \right)$ と $c_2 \left( x \right)$ の直積のみを考えると、

```math
\left[\begin{array}{c}
c_{1}(x) \\
1-c_{1}(x)
\end{array}\right] \otimes\left[\begin{array}{c}
c_{2}(x) \\
1-c_{2}(x)
\end{array}\right] = \begin{pmatrix}
c_1\left(x\right)c_2(x) & c_1\left(x\right) \left(1 - c_2\left(x\right) \right) \\
\left(1 - c_1\left(x\right)\right) c_2\left(x\right) & \left(1 - c_1\left(x\right)\right)\left(1 - c_2\left(x\right) \right) \\
\end{pmatrix}
```

と書けることから、各成分がそれぞれのノードへの所属率を表していることがわかります。

最後に、それぞれの末端ノードに紐付いた値 $R$ を得られた重みで足し合わせて、予測値を得ることができます。

```math
\hat{h}(x)=\sum_{i_{1}, \ldots i_{d} \in\{0,1\}^{d}} R_{i_{1}, \ldots, i_{d}} \cdot C_{i_{1}, \ldots, i_{d}}(x)
```

この $R$ を response tensor と呼びます。
NODE レイヤーでは並列に $m$ 個の木を学習させ、それらの予測値を concat することで最終的な出力とします。

```math
\left[\hat{h}_{1}(x), \ldots, \hat{h}_{m}(x)\right]
```

マルチクラス分類の場合は単純に、各末端ノードにクラス数分の重みを紐付けることで表現できます。

### 全体のアーキテクチャ

NODE は、$k$ 個の NODE レイヤーを直列に積んだ構成をとります。

![スクリーンショット 2021-04-02 5.20.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/b666dd24-c36b-dd2e-6921-2fef8c275110.png)

$i$ 番目のレイヤーの入力には、もとの入力に加えて、$i - 1$ 番目までのレイヤーの出力も concat して渡します。
最終的な予測値は、すべてのレイヤーの出力の単純平均で得ることができます。

ここで、後続レイヤーの入力として使用するための、各レイヤーの出力の次元は、クラス数 ${C}$ より大きくてもよく、末端ノードの重み $R$ を $(d+1)$ 次元のテンソル $R \in \mathbb{R}^{\underbrace{2 \times 2 \times 2 \times \cdots}_{d} \times l}$ とすることで、任意の次元数 $l$ の出力を持つことができます。
最終的な出力の際は、$\hat{h} \left(x\right) \in \mathbb{R}^l$ の最初から $C$ 次元分の値を使用します。

### 学習・推論

##### 前処理

前処理として、入力特徴量を [Quantile transform](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.QuantileTransformer.html) します。このステップは非常に重要で、学習の安定性と収束の早さが大きく変化します。

##### 重みの初期化

重みの初期化には data-aware initialization を使用します。
具体的には、閾値 $b$ は最初のバッチのランダムな値を選んで初期化します。
また、スケールパラメータ $\tau$ は、最初のバッチのすべてのサンプルが、 $\sigma_\alpha$ による変換の線形な部分に入るように初期化されます。
feature selection matrix ${F}$ は一様分布 $U(0, 1)$、response tensor は正規分布 $N(0, 1)$ に従って初期化されます。

##### Training

実験では $c = 5$ つの連続した checkpoint の平均を重みとして使用しました。また、hold-out validation データを使用した early stopping を行いました。

##### Inference

学習時間は、その大部分が ${(1)}$ 式の以下の部分にかかります。

```math
\operatorname{entmax}_{\alpha}\left(F_{i j}\right)
```

推論時にはこのパラメータは固定なので、事前に計算して保存しておくことで、推論時間を早めることができます。

## 実験

### 他モデルとの比較

実験は、別々のドメインの6種類のデータに対して行いました。
分類タスクについては cross-entropy loss、回帰タスクについては MSE loss を使用し、その値を評価指標としています。

![スクリーンショット 2021-04-02 14.51.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/e0774211-f0dd-5c83-af78-59d7cf8c01cc.png)

Table 1 は、 validation データに対して early stopping のみを行うデフォルトパラメータの実験、Table 2 はハイパーパラメータもチューニングする実験です。NODE のデフォルトパラメータとしては、レイヤー数1、深さ6のモデルを使用しました。
ほとんどのデータで NODE が他の GBDT モデルを outperform していることがわかります。

### Ablation study

##### 特徴量選択の関数

次の4つの関数について性能を比較しました。

- Softmax
- Gumbel-Softmax
- Sparsemax
- Entmax (${\alpha = 1.5}$)

![スクリーンショット 2021-04-02 15.01.22.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/0a5016ce-bb16-49a2-7fc7-b0a28d37d632.png)

entmax がすべての実験で最も良い性能を発揮していることがわかります。

##### 特徴量重要度

学習させたモデルの内部表現の特徴を調べるために、レイヤーごとに、「特徴量をシャッフルした際にどれだけスコアが下がるか (permutation importance)」と「最終的な予測値にどれだけ貢献しているか」の2つの指標について値を計測しました。

![スクリーンショット 2021-04-02 15.26.06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/f7177591-f686-b5af-dc40-d8ab72cb2a9d.png)

結果を見ると、入力に近いレイヤーでは permutation importance が高く、出力に近いレイヤーでは予測値への貢献度が高いことがわかります。
これらから、モデルの入力に近い部分ではデータの良い表現ベクトルを得て、それを出力に近い部分で利用する、という役割があることが伺えます。

##### 学習時間/推論時間

学習・推論時間についても比較を行いました。レイヤー数8、深さ6のモデルを使用しています。

![スクリーンショット 2021-04-02 15.38.54.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/271209/024e4650-1e0f-d109-a4f3-cfb333574e0f.png)

推論時間に関しては、他とあまり変わらない速度で実行できていることがわかります。

## 参考リンク

- NODE : https://arxiv.org/abs/1909.06312
- $\alpha$-entmax : https://www.aclweb.org/anthology/P19-1146
- data-aware initialization : https://arxiv.org/abs/1511.06422