XGBoost 2.0でアップデートされたランキング学習を試す

昨年9月にリリースされたXGBoost 2.0で、ランキング学習（Learning to rank）がアップデートされたので、今回はこのランキング学習を実践してみたいと思います。

デモンストレーション

環境はGoogle Corabです。

インストール

import xgboost as xgb
print(xgb.__version__) 

執筆時点でGoogle Corabでは最新のリリースバージョン(2.0.3)がインストールされていますので、特にインストールする必要はありません。

ライブラリの読み込み

import pandas as pd
import numpy as np

ランキング学習のためのシミュレーションデータの生成

　from sklearn.datasets import make_classification

seed = 1
X, y = make_classification(random_state=seed)

make_classification()で分類問題のシミュレーションデータを生成します。

rng = np.random.default_rng(seed)
n_query_groups = 3
qid = rng.integers(0, 3, size=X.shape[0])

sorted_idx = np.argsort(qid)
X = X[sorted_idx, :]
y = y[sorted_idx]
qid = qid[sorted_idx]

各データポイントがどのクエリグループに属するかを示すクエリID(qid)の割り当てを行います。

学習

ranker = xgb.XGBRanker(tree_method="hist", lambdarank_num_pair_per_sample=8, objective="rank:ndcg", lambdarank_pair_method="topk")

XGBRankerは、XGBoostのランキング学習のためのモデルになります。

ranker.fit(X, y, qid=qid)

ランキング学習を行います。

df = pd.DataFrame(X, columns=[str(i) for i in range(X.shape[1])])
df["qid"] = qid
ranker.fit(df, y)

from sklearn.model_selection import StratifiedGroupKFold, cross_val_score

kfold = StratifiedGroupKFold(shuffle=False)
cross_val_score(ranker, df, y, cv=kfold, groups=df.qid)

StratifiedGroupKFoldで交差検証を行います。
ここでは、各グループ（クエリ）内でのクラスの比率が保持されるようにしています。

予測

scores = ranker.predict(X)
sorted_idx = np.argsort(scores)[::-1]
scores = scores[sorted_idx]

predict()で予測を行います。スコアに基づいて予測された関連性が最も高いものから最も低いものへと順序付けます。