2022年2月2日

PySpark と Pandas UDF を使用した SHAP 計算のスケーリング

動機

機械学習 (機械学習)、特にディープラーニング (DL) モデルの意思決定への応用が普及するにつれて、ブラックボックスの中身を理解し、そのようなモデルの出力に基づいて主要なビジネス上の意思決定を正当化することがますます重要になっています。例えば、ML モデルが顧客のローン申請を却下したり、ピアツーピアレンディングで特定の顧客に信用リスクを割り当てたりした場合、なぜその決定が下されたのかをビジネスのステークホルダーに説明することは、モデルの採用を促進する強力なツールとなり得ます。多くの場合、解釈可能な機械学習は、単なるビジネス要件ではなく、なぜ特定の決定や選択肢が顧客に提示されたのかを理解するための規制要件でもあります。SHapley Additive exPlanations (SHAP) は、説明可能な AI を実現し、ビジネス課題の解決における機械学習モデルやニューラルネットワークの成果に対する信頼を確立するために活用できる重要なツールです。

SHAP は、ゲーム理論に基づくモデル説明のための最新のフレームワークです。このアプローチでは、データセット内の各データポイントについて、モデルの特徴量とモデルの出力との間の線形関係を見つけます。このフレームワークを使用すると、モデルの出力をグローバルまたはローカルに解釈できます。グローバルな解釈可能性は、各特徴量が結果にプラスまたはマイナスのどちらにどの程度貢献するかを理解するのに役立ちます。一方、ローカルな解釈可能性は、特定の観測値に対する各特徴量の効果を理解するのに役立ちます。

データサイエンスコミュニティで広く採用されている最も一般的なSHAPの実装は、シングルノードマシンで実行されます。つまり、利用可能なコアの数に関係なく、すべての計算を単一のコアで実行します。そのため、分散コンピューティング機能を利用できず、シングルコアの制約に縛られてしまいます。

この投稿では、特にローカルな解釈可能性のために、複数のマシンにまたがってSHAP値の計算を並列化する簡単な方法を紹介します。次に、データセットの行と列の数が増えるにつれて、このソリューションがどのようにスケールするかを説明します。最後に、SparkでSHAP計算を並列化する際に、何が有効で何を避けるべきかについての我々の知見をいくつか紹介します。

シングルノード SHAP

説明可能性を実現するため、SHAP はモデルを Explainer に変換し、それに Explainer を適用することで個々のモデル予測を説明します。SHAP 値の計算には、Python の一般的なものを含め、さまざまなプログラミング言語でいくつかの実装があります。この実装では、各観測値の説明を取得するために、モデルに適した Explainer を適用できます。次のコードスニペットは、TreeExplainer を Random Forest Classifier に適用する方法を示しています。

この手法は少量のデータボリュームではうまく機能しますが、数百万のレコードに対する機械学習モデルの出力を説明する場合、実装がシングルノードであるため、うまくスケールしません。例えば、下の図1は、シングルノードマシン（4コア、30.5 GBのメモリ）上でレコード数を増やしていった場合の、SHAP値の計算における実行時間の増加を示しています。100万行50列を超えるデータ形状ではマシンがメモリ不足に陥ったため、図ではそれらの値が欠落しています。ご覧のとおり、実行時間はレコード数にほぼ比例して増加しており、これは実世界のシナリオでは持続可能ではありません。例えば、機械学習モデルがある予測を行った理由を理解するために10時間も待つことは、多くのビジネスシーンにおいて非効率的であり、容認できるものではありません。

シングルノードでのSHAP計算の実行時間 — 図1: シングルノード SHAP 計算の実行時間

この問題を解決する方法の1つとして、近似計算を使用することが考えられます。shap_values メソッドで approximate 引数を True に設定できます。そうすることで、ツリーの下位の分割はより高い重みを持つことになり、SHAP 値が正確な計算と一致するという保証はありません。これにより計算は高速化されますが、モデル出力の説明が不正確になる可能性があります。さらに、approximate 引数は TreeExplainers でのみ利用可能です。

別のアプローチとして、Apache Spark™ などの分散処理フレームワークを利用して、複数のコアで Explainer の適用を並列化する方法があります。

PySparkによるSHAP計算のスケーリング

SHAP 計算を分散させるため、PySpark でこの Python 実装と Pandas UDF を使用しています。kddcup99 データセットを使用して、侵入または攻撃と呼ばれる不正な接続と、正常な接続を区別できる予測モデルであるネットワーク侵入検知器を構築しています。このデータセットは、侵入検知の目的には欠陥があることが知られています。しかし、この記事では、基盤となる機械学習モデルのセマンティクスではなく、SHAP 値の計算にのみ焦点を当てています。

エクスペリメントで構築した 2 つのモデルは、列数の違いによるソリューションのスケーラビリティを示すために、10 および 50 の特徴量を持つデータセットでトレーニングされた、単純なランダムフォレスト分類器です。元のデータセットの列数は 50 未満であるため、目的のデータ量に達するようにこれらの列の一部を複製している点にご注意ください。実験したデータ量は 4MB から 1.85GB の範囲です。

コードを詳しく見ていく前に、Spark Dataframes と UDF の仕組みについて簡単に概説します。Spark Dataframesはクラスター全体に行単位で分散されます。行の各グループはパーティションと呼ばれ、各パーティションは（defaultで）1つのコアで操作できます。これが、Spark が根本的に並列処理を実現する仕組みです。PandasはSHAPに簡単に入力でき、パフォーマンスも高いため、Pandas UDFは自然な選択肢です。ベクトル化 UDF とも呼ばれる pandas UDF は、Apache Arrow を使用してデータ転送を最適化することにより、Python UDF よりも優れたパフォーマンスを実現します。

以下のコードスニペットは、PySpark で Pandas UDF を使用して Explainer の適用を並列化する方法を示しています。「calculate_shap」という pandas UDF を定義し、この関数を mapInPandas に渡します。このメソッドは、並列化されたメソッドを PySpark データフレームに適用するために使用されます。この UDF を使用して、SHAP パフォーマンステストをランします。

図2は、100万行10列のデータについて、シングルノードマシンと、それぞれサイズが2、4、8、16、32、64のクラスターでの実行時間を比較しています。すべてのクラスターの基盤となるマシンは同様です（4コア、30.5GBのメモリ）。興味深いことに、並列化されたコードはクラスター内の全ノードのすべてのコアを活用します。そのため、サイズ2のクラスターを使用するだけで、パフォーマンスがほぼ5倍向上します。

データ量の増加に応じたスケーリング

SHAPの実装方法により、追加の行よりも追加の特徴量の方がパフォーマンスに大きな影響を与えます。これで、SparkとPandas UDFを使用することでSHAP値をより高速に計算できることがわかりました。次に、追加の特徴量/列がある場合にSHAPがどのように動作するかを見ていきます。

直感的に、データサイズが大きくなるということは、SHAP アルゴリズムが処理しなければならない計算が増えることを意味します。図 3 は、16 ノードのクラスターで、行と列の数が異なる場合の SHAP 値の実行時間を示しています。行をスケーリングすると実行時間がほぼ正比例して増加し、つまり行数を 2 倍にすると実行時間もほぼ 2 倍になることがわかります。列数のスケーリングは実行時間と比例関係にあり、列を 1 つ追加すると実行時間は約 80% 増加します。

これらの観察結果（図2と図3）から、データが多ければ多いほど、実行時間を妥当な範囲に保つために計算を水平方向にスケール（ワーカーノードを追加）できるという結論に至りました。

行数と列数が異なる場合の 6 ノード並列 SHAP 計算の実行時間 — 図3: さまざまな行と列の数における16ノード並列SHAP計算実行時間

並列化を検討すべき時とは？

私たちが答えたいと考えた問いは次のとおりです。並列化する価値があるのはどのような場合か。計算量が増える可能性を承知の上で、SHAP 計算を並列化するために PySpark をいつから使い始めるべきか。クラスタサイズを2倍にすることがSHAP計算の実行時間の改善に与える影響を測定するエクスペリメントを行いました。このエクスペリメントの目的は、問題に対してより多くの水平リソース（つまり、ワーカーノードの追加）を投入することを正当化するデータサイズを明らかにすることです。

データ列数 10、行数を 10、100、1000、...、最大 1,000 万として、SHAP 計算を実行しました。各行数について、クラスターサイズを2、4、32、64として、SHAP計算の実行時間を4回測定しました。実行時間比とは、より大きいクラスターサイズ（4 と 64）での SHAP 値の計算実行時間を、ノード数が半分のクラスターサイズ（それぞれ 2 と 32）で同じ計算を実行した場合の時間で割った比率です。

図4はこのエクスペリメントの結果を示しています。主なポイントは次のとおりです。

- 行数が少ない場合、クラスターサイズを倍にしても実行時間は改善されず、Spark のタスク管理によるオーバーヘッドが加わるため、かえって悪化する場合もあります (したがって、実行時間比 > 1 となります)。
- 行数を増やすにつれて、クラスタサイズを2倍にすることがより効果的になります。1000万行のデータの場合、クラスターサイズを2倍にすると実行時間がほぼ半分になります。
- どの行数においても、クラスターサイズを32から64に倍増させるよりも、2から4に倍増させる方が効果的です（青線とオレンジ線の間のギャップに注目してください）。クラスターサイズが大きくなるにつれて、ノードを追加するオーバーヘッドも大きくなります。これは、パーティションごとのデータサイズが小さすぎるパーティションサイズになっているためです。少量のデータを処理するために個別のタスクを作成すると、より最適なデータ/パーティションサイズを使用する場合よりもオーバーヘッドが増加します。

さまざまなデータ量においてクラスターサイズを倍増させることが実行時間に及ぼす影響 — 図4：さまざまなデータ量において、クラスターサイズを2倍にすることが実行時間に与える影響

注意点

再パーティション化

上記で述べたように、Sparkはパーティションという概念を通じて並列処理を実装します。データは行のチャンクに分割され、各パーティションはデフォルトで単一のコアによって処理されます。Apache Sparkによってデータが最初に読み込まれるとき、クラスタで実行したい計算に最適なパーティションが必ずしも作成されるとは限りません。特に、SHAP値を計算する場合、データセットを再パーティション化することで、パフォーマンスを向上させられる可能性があります。

十分に小さいパーティションを作成することと、その作成のオーバーヘッドが計算を並列化するメリットを上回るほど小さくしすぎないことの間で、バランスを取ることが重要です。
パフォーマンステストでは、次のコードを使用して、クラスター内のすべてのコアを利用することにしました。

さらに大きなデータ量の場合は、パーティションの数をコア数の 2 倍または 3 倍に設定することをお勧めします。重要なのは、エクスペリメントして、自分のデータに最適なパーティショニング戦略を見つけ出すことです。

display()の使用

Databricks ノートブックで作業している場合、実行時間のベンチマークを行う際にはdisplay()関数の使用を避けた方がよいでしょう。display() を使用しても、完全な変換にかかる時間が必ずしもわかるわけではありません。これにはクエリーに挿入される暗黙的な行数制限があり、また、測定対象の操作（例: ファイルへの書き込み）によっては、結果をドライバーに収集する際に追加のオーバーヘッドが発生するためです。我々の実行時間は、Sparkのwriteメソッドを「noop」フォーマットで使用して測定しました。

まとめ

このブログ記事では、PySparkとPandas UDFでSHAP計算を並列化することによって高速化するソリューションを紹介しました。次に、データ量の増加、異なるマシンタイプ、構成の変更に対して、そのソリューションのパフォーマンスを評価しました。主なポイントは以下の通りです:

- - シングルノードの SHAP 計算は、行と列の数に比例して線形に増加します。
  - PySparkでSHAP計算を並列化すると、クラスター全体のすべてのCPUで計算が実行されるため、パフォーマンスが向上します。
  - ビッグデータボリュームが大きい場合、クラスターサイズを大きくする方が効果的です。データが小さい場合、この方法は効果的ではありません。

今後の作業

垂直スケーリング - このブログ記事の目的は、大規模なデータセットで水平スケーリングすることが、SHAP値の計算パフォーマンスをどのように向上させることができるかを示すことでした。クラスター内の各ノードが4コア、30.5 GBを持つという前提で始めました。将来的には、垂直スケーリングと水平スケーリングのパフォーマンスをテストすることも興味深いでしょう。例えば、4ノード（各4コア、30.5GB）のクラスターと2ノード（各8コア、61GB）のクラスターのパフォーマンスを比較するなどです。

シリアライズ/デシリアライズ - 前述のように、Python UDF ではなく Pandas UDF を使用する主な理由の 1 つは、Pandas UDF が Apache Arrow を使用して、JVM と Python プロセス間のデータのシリアライズ/デシリアライズを改善することです。Spark のデータパーティションを Arrow のレコードバッチに変換する際に、いくつかの最適化の可能性があります。Arrow のバッチサイズでエクスペリメントすることで、さらなるパフォーマンス向上が見込めます。

分散 SHAP 実装との比較 - このソリューションの結果を、Shparkley などの SHAP の分散実装と比較することは、興味深いでしょう。このような比較研究を行うにあたっては、まず第一に、両ソリューションの出力が比較可能であることを確認することが重要です。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事