プライバシー保護計算フレームワークSecretFlowを使ってみた

SecretFlowとは

プライバシー保護 x データ処理 + 機械学習 x Python

近年セキュリティインシデント頻発の中に、世の中の規制も厳しくなっていて、ユーザー／顧客データをより安全に取り扱う必要があります。ただし、それを実現するのは簡単ではなくて、

1. プライバシー保護計算のための技術（例えば：準同型暗号）がわかりにくい
2. 処理速度が低いまたは強いハードウェアが必要
3. GitHubの中に一番有名なレポジトリはMicrosoftのSEALだけど、C++なので手軽に実装できない

などの理由があるけど、一方でSecretFlowは、

1. PythonなのでColab／ノートブック上でも使える
2. 暗号化だけではなくて、デバイス管理などフレームワーク全体提供される

参考リンク

GitHub

ドキュメント

日本語の説明資料がまだ少ない気がしていて、もしかしてこれしかない

下に紹介する手法に関連する連合学習の説明

やってみた

SecureBoost

SecureBoostは、準同型暗号を使って、学習時のデータを保護してXGBoost（機械学習のアルゴリズム）を使う手法です。

今回はチュートリアル通りにやってみたけど、その中にいくつか重要なポイントをピックアップして説明します。

ソースコード＆説明

Colab環境でやってみて、最初はライブラリのインストールです。大きなフレームワークなので少し時間がかかります。

!pip install secretflow

---

次に必要なライブラリをインポートして、動作を確認します。

import spu
from sklearn.metrics import roc_auc_score

import secretflow as sf
from secretflow.data import FedNdarray, PartitionWay
from secretflow.device.driver import reveal, wait
from secretflow.ml.boost.sgb_v import (
    Sgb,
    get_classic_XGB_params,
    get_classic_lightGBM_params,
)
from secretflow.ml.boost.sgb_v.model import load_model
import pprint

pp = pprint.PrettyPrinter(depth=4)

print('SecretFlowのバーション：{}'.format(sf.__version__))
# SecretFlowのバーション：1.10.0b1

---

そしてデバイスの初期化およびいくつかの環境設定です。

今回は垂直連合学習の手法を使う予定で、aliceとbob二つの端末を仮定しています。各自にデータを持たせて、データを保護するまま中央サーバに集約して学習する形です。

また heu_config の 'he_parameters' の 'schema' の値は、暗号化のアルゴリズムとして、概念検証のために、処理速度が早い ou (Okamoto-Uchiyama) を選択したけど、本番環境ではより安全なPaillier系がサイトの方で推奨されています。

アルゴリズムに関する詳しい説明は、こちらにあります。

alice_ip = '127.0.0.1'
bob_ip = '127.0.0.1'
ip_party_map = {bob_ip: 'bob', alice_ip: 'alice'}

_system_config = {'lineage_pinning_enabled': False}
sf.shutdown()
sf.init(
    ['alice', 'bob'],
    address='local',
    _system_config=_system_config,
    object_store_memory=5 * 1024 * 1024 * 1024,
)

# SPU
cluster_def = {
    'nodes': [
        {'party': 'alice', 'id': 'local:0', 'address': alice_ip + ':12945'},
        {'party': 'bob', 'id': 'local:1', 'address': bob_ip + ':12946'},
    ],
    'runtime_config': {
        'protocol': spu.spu_pb2.SEMI2K,
        'field': spu.spu_pb2.FM128,
    },
}

# HEU
heu_config = {
    'sk_keeper': {'party': 'alice'},
    'evaluators': [{'party': 'bob'}],
    'mode': 'PHEU',
    'he_parameters': {
        'schema': 'ou',
        'key_pair': {
            'generate': {
                'bit_size': 2048,
            },
        },
    },
    'encoding': {
        'cleartext_type': 'DT_I32',
        'encoder': "IntegerEncoder",
        'encoder_args': {"scale": 1},
    },
}

---

パラメータの設定が完了したら、デバイスをセットアップします。

alice = sf.PYU('alice')
bob = sf.PYU('bob')
heu = sf.HEU(heu_config, cluster_def['runtime_config']['field'])

---

データソースは sklearn のデフォルトデータソースを使います。

今回は垂直連合学習なので、x の30カラムを15カラムずつ alice と bob に保存させます。

from sklearn.datasets import load_breast_cancer

ds = load_breast_cancer()
x, y = ds['data'], ds['target']

v_data = FedNdarray(
    {
        alice: (alice(lambda: x[:, :15])()),
        bob: (bob(lambda: x[:, 15:])()),
    },
    partition_way=PartitionWay.VERTICAL,
)
label_data = FedNdarray(
    {alice: (alice(lambda: y)())},
    partition_way=PartitionWay.VERTICAL,
)

---

続いて XGB のセットアップと訓練です。

こちらでJAXの問題かもしれなくて、エラーが出たので、調べたところ、環境変数 ENABLE_PJRT_COMPATIBILITY を設定して解決できます。

import os
os.environ['ENABLE_PJRT_COMPATIBILITY'] = 'true'

params = get_classic_XGB_params()
params['num_boost_round'] = 3
params['max_depth'] = 3

sgb = Sgb(heu)
model = sgb.train(params, v_data, label_data)

---

訓練したあとに、1回推論して結果を検証してみます。

こちらで推論した yhat は保護されているので、reveal を使って明文化して、y と比較します。reveal はセキュリティのため、検証環境以外に使うのは極めて注意が必要です。

yhat = model.predict(v_data)
print(yhat)
# <secretflow.device.device.pyu.PYUObject object at 0x7bdfc40da9b0>

yhat_reveal = reveal(yhat)
print(yhat_reveal[:3])
# [[0.19071478], [0.19071478], [0.19071478]]

print(f"auc: {roc_auc_score(y, yhat_reveal)}")
# auc: 0.9952235611225622

こちらの結果は 0.9952 になっているので、正しく訓練と推論ができていることが確認できました。