LightGBMのver. 4.0.0で登場した量子化ですが、今はLightGBMも4.3.0が出て量子化についての不具合修正もすすんで安定して動くようになってきました。

論文をベースに、「どういう計算をしているのか」という理論面の概要と、実際に動かして「どれだけ計算速度や予測精度やモデルのファイルサイズが変わるのか」を見ていきたいと思います。

個人的に特に気になっているのはモデルのファイルサイズの減少です。

量子化の概要

LightGBMのような最近のGBDTで使われている決定木では、葉の出力 $w$ は誤差関数の2次のテイラー近似をもとに、以下のように計算されます（このあたりはChen & Guestrin, 2016が比較的わかりやすいです）

ここで $g_i$ は誤差関数の勾配、 $h_i$ は誤差関数の二次の微分です。

この $g_i,h_i$ を32bitや64bitのfloatではなく4bitなどの低ビット幅の整数で保持しよう、というのが量子化です。

LightGBMの開発者たちによるGBDTの量子化についての研究（Shi et al., 2022）では、量子化したGBDTの予測精度に影響を与えるポイントとして

整数への丸め込みをどうするか（rounding strategy）
最終的な予測値を本来のビット幅で出すか（leaf-value refitting）

があると報告されています。

Rounding Strategy

シンプルに思いつく丸め込みは四捨五入のように最も近い整数へ丸め込む方法（round-to-nearest: RN）です。

しかしShi et al. (2022) では、RNよりも精度を上げる丸め込みの方法として、確率的な丸め込み（stochastic rounding）という方法を提案しています。

式だとややわかりにくいかもと思い、x = np.linspace(0.001, 0.999, 1000)の範囲でStochastic Roundingを行った結果をグラフにしてみました。

グレーの点は $x$ を10のbinに分割し、それぞれのbinにおける平均値と標準偏差をエラーバーにしたものです。

$x$ の値が大きくなるほど、切り上げられる確率が上がっていることがわかります。

Leaf-Value Refitting

ランキング学習など一部タスクにおいて、木の成長は量子化した勾配で行い、最終的な予測値を元の勾配で算出する方法が予測精度を高めることがわかっているようです。

実際、Shi et al. (2022) の実験結果においてもランキングや回帰でrefitを実行したほうが精度が高くなる傾向がみられました。

しかし、refitするとそのぶん計算量が増えるというデメリットがあります。

試してみる

環境

lightgbm==4.3.0
Google Colab（無料のやつ、os.cpu_count()は2だったので2スレッドのCPU）

実験内容

scikit-learnのmake_regression で0から99までの100個のseedで回帰タスクのデータを生成して、量子化の有無による処理時間や予測精度などの差を見てみます。

長すぎるので全部のコードは載せませんが、データは以下のように5万レコード生成し、ランダムに選んだ10,000レコードをtestにして残った40,000のうち32,000をtrainに、8,000をearly-stoppingのためのvalidationにしました。

from sklearn.datasets import make_regression
X, y = make_regression(n_samples=50_000, n_features=10, noise=0.5, random_state=seed)

ハイパーパラメータは以下のような形に設定しました。

params = {
    'device_type': 'cpu',
    'num_threads': os.cpu_count(),
    'objective': 'mse',
    'metric': 'rmse',
    'num_leaves': 31,
    'learning_rate': 0.1,
    'feature_fraction': 0.8,
    'verbose': 0,
    'seed': 0,
    'deterministic': True,
}
num_boost_round = 100_000
stopping_rounds = 100

以下の3条件での結果を比較します。

量子化なし：上記のハイパーパラメータのみ
量子化あり・Renewなし：上記に加え、use_quantized_grad=Trueを設定（量子化を有効化）
量子化あり・Renewあり：上記に加え、quant_train_renew_leaf=Trueを設定（Renewは論文における「Refitting」のこと）

測るのは以下4つになります

訓練時間：モデルの訓練に要した時間。wall-time
推論時間：テストデータの予測に要した時間。wall-time
RMSE：テストデータに対する予測値のRoot Mean Squared Error
モデルのサイズ：joblibでデフォルトの圧縮率で保存したモデルのファイルサイズ。単位はメガバイト。具体的には以下の方法で取得。

model_path = Path("model.joblib")
with open(model_path, "wb") as f:
   joblib.dump(model, f)
model_size_mb = model_path.stat().st_size / 1024**2

実験結果

100回の実験結果の平均・標準偏差の表と、分布の図は以下になります。

訓練時間（Training time）と推論時間（Inference time）とModel Sizeは「量子化あり・Renewなし」が大幅に小さくなっています。

一方でRMSEは「量子化あり・Renewあり」が平均としては最も小さくなっています（分布としては「量子化なし」とそこまで大差はないですが）。

まとめ

今回の実験では

訓練時間・推論時間・モデルのサイズは 「量子化あり・Renewなし」<「量子化なし」< 「量子化あり・Renewあり」
予測精度（RMSE）は 「量子化あり・Renewあり」<「量子化なし」< 「量子化あり・Renewなし」

といった関係が見られました。

精度面では「量子化あり・Renewあり」が最もよいものの、Renewすると訓練時間・推論時間・モデルサイズが増加するため、量子化の高速化のメリットは失われてしまいます。

損失関数やデータセットと量子化には相性のようなものがありそうですが、もしRenewが不要なくらい量子化での精度損失が小さい損失関数・データセットであれば量子化の高速化・軽量化のメリットを享受できそうですね。

盆暗の学習記録

データサイエンス，エンジニアリング，ビジネスについて日々学んだことの備忘録としていく予定です。初心者であり独学なので内容には誤りが含まれる可能性が大いにあります。

LightGBMの量子化を試してみる