Split Learning

スプリットラーニングは、選択したカット層でニューラルネットワークをデータ所有者であるクライアントとサーバーの間で分割する、協調型の分散学習手法です。クライアントは非公開データでカット層より前の層を実行し、その結果得られた活性化（smashed dataと呼ばれることもあります）だけを送信します。サーバーは残りの層を実行して損失を計算し、カット層の勾配を返します。クライアントはその勾配を使って自身の層を更新します。生の学習データがクライアントの外へ出ることはなく、どちらの側もモデル全体を保持しません。MIT Media LabのOtkrist GuptaとRamesh Raskarは、2018年10月に複数エージェント間でのディープニューラルネットワークの分散学習でこの手法を発表しました。同グループが2018年12月に発表した続編「Split Learning for Health」では、医用画像解析用のスキャン画像など、異なる種類の患者データを持つ病院が、生の記録を共有せずにモデルを共同学習するための構成を提案しました。

スプリットラーニングの仕組み#

学習ステップごとに、カット層を2回またぎます。順伝播では活性化がクライアントからサーバーへ送られ、逆伝播では活性化に対する損失の勾配が戻されます。連鎖律によれば、処理を続けるために必要な勾配は1つだけなので、クライアントの更新はモデルを分割せずに学習した場合と同一です。PyTorchを使った単一プロセスのシミュレーションで、データのやり取りを確認できます。

import torch
from torch import nn

torch.manual_seed(0)
client = nn.Sequential(nn.Linear(784, 256), nn.ReLU())  # カット層より前の層（非公開データの隣）
server = nn.Sequential(nn.Linear(256, 10))  # カット層より後の層
images, labels = torch.randn(32, 784), torch.randint(0, 10, (32,))

# クライアント：カット層まで順伝播し、活性化だけを送信します
activations = client(images)
received = activations.detach().requires_grad_()

# サーバー：順伝播を完了してカット層まで逆伝播し、その勾配を返します
nn.functional.cross_entropy(server(received), labels).backward()

# クライアント：返された勾配から逆伝播を続けます
activations.backward(received.grad)
split_grad = client[0].weight.grad.clone()

# モデルを分割しない場合も、同じステップでクライアントに同一の勾配が得られます
client.zero_grad()
nn.functional.cross_entropy(server(client(images)), labels).backward()
print("Same client gradients as unsplit training:", torch.equal(split_grad, client[0].weight.grad))

Same client gradients as unsplit training: Trueが出力されます。その後、各側がそれぞれのオプティマイザーを実行します。複数のクライアントがいる場合、元のプロトコルではクライアントを順番に学習させ、各クライアントは前のクライアントが更新した最新のクライアント側の重みを引き継ぎます。

非公開にする必要がある情報に応じて構成は変わります。上記の標準構成ではラベルをサーバーに送信します。U字型構成では最後の層とラベルもクライアント側に保持します。垂直型構成は、同じ記録について異なる特徴を持つ組織向けです。

スプリットラーニングと連合学習#

連合学習もデータを所有者のマシン上に保持しますが、各クライアントがモデル全体のコピーを学習させ、重みの更新を送信して平均化します。スプリットラーニングでは代わりに活性化と勾配を送信するため、クライアントに必要な計算量ははるかに少なく、モデル全体を保持する必要もありません。一方、通信量はモデルサイズではなくデータ量に応じて増え、クライアントは並列ではなく順番に学習します。Chandra Thapaらが2020年に提案したSplitFed学習は、クライアント側の層を並列実行して連合学習方式で平均化することで、両方の手法を組み合わせています。

垂直スプリットラーニングは金融にも利用されます。銀行と、同じ顧客について異なる特徴を持つパートナーが、SecretFlowの銀行マーケティングチュートリアルのように、非公開データを交換せずに1つのモデルを学習させます。パイプライン並列処理も同様にネットワークを深さ方向に分割しますが、大規模モデルをメモリに収めるため、信頼された単一クラスター内のGPU間で分割します。

プライバシーリスク#

活性化と勾配にもデータに関する情報が含まれます。初期層の活性化から認識可能な画像を復元できる場合があります。また、サーバーがラベルを保持している場合、返す勾配からクライアントにラベルが漏れる可能性があります。Dario Pasquini、Giuseppe Ateniese、Massimo BernaschiはACM CCS 2021で、悪意のあるサーバーが返す勾配を通じて学習を乗っ取り、サーバーが非公開入力を再構成できるようクライアントの層を誘導できることを示しました。防御策としては、差分プライバシーによるノイズの追加、入力について活性化が明らかにする情報量へのペナルティ設定、異常な勾配の検出などがありますが、いずれも精度とのトレードオフが生じます。プロトコルが規定するのはカット層を越えるデータだけなので、デプロイでは認証や暗号化通信などのデータセキュリティ対策とスプリットラーニングを併用します。

カット層の選択#

カット層の位置によって、クライアント側の計算量、通信量、情報漏えいリスクのバランスが決まります。早い位置で分割すればクライアント側の計算量は最小限になりますが、入力に似た大きな活性化を送信することになります。遅い位置で分割すると、より抽象化された特徴マップを送信できますが、クライアントのハードウェアにより大きな負荷がかかります。これはエッジAIデバイスで重要です。スキップ接続があると、カット層をまたぐすべてのテンソルを送信する必要があるため、選択が複雑になります。Ultralytics YOLO26をバックボーンの後で分割する場合、層4、6、10の出力がネックに入力されるため、3つの特徴マップを送信します。モデル構成においてimgsz=640を使用するYOLO26nでは、それらの値は合計1,126,400個であり、入力画像の1,228,800個にほぼ匹敵します。