Split Learning

분할 학습은 선택한 컷 레이어에서 신경망을 데이터 소유자인 클라이언트와 서버 사이에 나누는 협업형 분산 학습 방법입니다. 클라이언트는 비공개 데이터에서 컷 이전의 레이어를 실행하고, 그 결과 활성화 값만 서버에 전송합니다. 이 값은 스매시드 데이터라고도 합니다. 서버는 나머지 레이어를 실행하고 손실을 계산한 다음 컷 지점의 그래디언트를 돌려보내 클라이언트가 레이어를 업데이트할 수 있도록 합니다. 원본 학습 데이터는 클라이언트 외부로 나가지 않으며, 어느 쪽도 모델 전체를 보유하지 않습니다. MIT Media Lab의 Otkrist Gupta와 Ramesh Raskar는 2018년 10월 여러 에이전트에 걸친 심층 신경망의 분산 학습 논문에서 이 방법을 소개했습니다. 연구팀은 2018년 12월 후속 논문인 Split Learning for Health에서 의료 영상 분석에 사용하는 스캔 이미지처럼 서로 다른 종류의 환자 데이터를 보유한 병원들이 원본 기록을 공유하지 않고 함께 모델을 학습할 수 있는 구성을 제안했습니다.

분할 학습 작동 방식#

각 학습 단계에서 컷 지점을 두 번 통과합니다. 순전파 중에는 활성화 값이 클라이언트에서 서버로 전달되고, 역전파 중에는 활성화 값에 대한 손실의 그래디언트가 클라이언트로 돌아옵니다. 연쇄 법칙으로 역전파를 계속하는 데 이 그래디언트 하나만 필요하므로, 클라이언트의 업데이트 결과는 분할하지 않은 모델을 학습했을 때와 같습니다. 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 은행 마케팅 튜토리얼과 같이 하나의 모델을 학습하면서 비공개 데이터를 교환하지 않습니다. 파이프라인 병렬 처리도 같은 방식으로 네트워크를 깊이 방향으로 나누지만, 대형 모델을 메모리에 올리기 위해 신뢰할 수 있는 하나의 클러스터에 있는 GPU 간에 분할합니다.

개인정보 보호 위험#

활성화 값과 그래디언트에도 데이터 관련 정보가 남아 있습니다. 초기 레이어의 활성화 값을 역산하면 식별 가능한 이미지를 복원할 수 있으며, 서버가 레이블을 보유한 경우 서버가 반환하는 그래디언트를 통해 클라이언트가 레이블을 알아낼 수 있습니다. Dario Pasquini, Giuseppe Ateniese, Massimo Bernaschi는 ACM CCS 2021에서 악의적인 서버가 반환하는 그래디언트를 통해 학습을 탈취하고, 서버가 비공개 입력을 재구성할 수 있는 형태로 클라이언트의 레이어를 유도할 수 있음을 보였습니다. 방어 방식은 차등 프라이버시를 통해 노이즈를 추가하거나, 활성화 값이 입력에 대해 드러내는 정보량을 제한하거나, 비정상적인 그래디언트를 탐지합니다. 각 방식은 정확도를 어느 정도 희생합니다. 프로토콜 자체는 컷 지점을 통과하는 데이터만 제어하므로, 배포 환경에서는 분할 학습을 인증 및 암호화된 전송 같은 데이터 보안 제어와 함께 사용합니다.

컷 레이어 선택#

컷 지점은 클라이언트 연산량, 트래픽, 데이터 노출 간의 균형을 결정합니다. 앞쪽에서 컷하면 클라이언트 연산량은 최소화되지만 입력을 닮은 큰 활성화 값이 전송됩니다. 뒤쪽에서 컷하면 더 추상적인 특징 맵을 전송하지만 클라이언트 하드웨어에 더 많은 연산을 요구하므로 엣지 AI 장치에서 중요합니다. 스킵 연결이 있으면 컷을 통과하는 모든 텐서를 전송해야 하므로 선택이 복잡해집니다. Ultralytics YOLO26을 백본 뒤에서 분할하면 세 개의 특징 맵을 전송해야 합니다. 모델 구성에서 레이어 4, 6, 10의 출력을 넥에 입력하기 때문입니다. imgsz=640의 YOLO26n은 1,126,400개의 값을 담고 있어 입력 이미지의 1,228,800개 값과 거의 같습니다.