Neural Oblivious Decision Ensembles for Deep Learning on Tabular Data をまとめます。

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

モデル構造

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}}$ として、予測値を以下のように表すことができます。

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 で表すことで、対応する重みを参照して、予測値を得ることができます。

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 の内挿となるように設計された活性化関数で、スパースな出力を得ることができます。

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

\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 を用いて、以下のように表されます。

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

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

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

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

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

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

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)$ の直積のみを考えると、

\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$ を得られた重みで足し合わせて、予測値を得ることができます。

\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 することで最終的な出力とします。

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

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

全体のアーキテクチャ

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

$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 します。このステップは非常に重要で、学習の安定性と収束の早さが大きく変化します。

重みの初期化

重みの初期化には 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)}$ 式の以下の部分にかかります。

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

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

実験

他モデルとの比較

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

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

Ablation study

特徴量選択の関数

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

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

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

特徴量重要度

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

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

学習時間/推論時間

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

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

参考リンク