Split Learning

Lo split learning è un metodo di addestramento distribuito collaborativo che divide una rete neurale in corrispondenza di un layer di taglio scelto, tra il proprietario dei dati, cioè il client, e un server. Il client esegue sui propri dati privati i layer precedenti al taglio e invia solo le attivazioni risultanti, talvolta chiamate dati frammentati; il server esegue i layer rimanenti, calcola la funzione di perdita e invia il gradiente nel punto di taglio affinché il client possa aggiornare i propri layer. I dati di addestramento grezzi non lasciano mai il client e nessuna delle due parti possiede il modello completo. Otkrist Gupta e Ramesh Raskar del MIT Media Lab lo hanno introdotto nell'ottobre 2018 nell'articolo Distributed Learning of Deep Neural Network over Multiple Agents; il successivo lavoro del gruppo, Split Learning for Health, pubblicato nel dicembre 2018, ha proposto configurazioni che consentono agli ospedali di addestrare insieme modelli usando tipi diversi di dati dei pazienti, come scansioni per l'analisi di immagini mediche, senza condividere i dati grezzi.

Come funziona lo split learning#

Ogni passaggio di addestramento attraversa due volte il punto di taglio: le attivazioni vanno dal client al server durante il passaggio in avanti e il gradiente della funzione di perdita rispetto a tali attivazioni torna indietro durante la retropropagazione. Poiché per proseguire basta un solo gradiente, grazie alla regola della catena gli aggiornamenti del client sono identici a quelli ottenuti addestrando il modello non suddiviso. Una simulazione a processo singolo in PyTorch mostra lo scambio:

import torch
from torch import nn

torch.manual_seed(0)
client = nn.Sequential(nn.Linear(784, 256), nn.ReLU())  # layer precedenti al taglio, accanto ai dati privati
server = nn.Sequential(nn.Linear(256, 10))  # layer successivi al taglio
images, labels = torch.randn(32, 784), torch.randint(0, 10, (32,))

# Client: passaggio in avanti fino al layer di taglio, poi invia solo le attivazioni
activations = client(images)
received = activations.detach().requires_grad_()

# Server: completa il passaggio in avanti e propaga all'indietro fino al punto di taglio, poi invia il gradiente al client
nn.functional.cross_entropy(server(received), labels).backward()

# Client: prosegue la retropropagazione a partire dal gradiente ricevuto
activations.backward(received.grad)
split_grad = client[0].weight.grad.clone()

# Lo stesso passaggio sul modello non suddiviso fornisce al client gradienti identici
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))

Stampa Same client gradients as unsplit training: True. Poi ciascuna parte aggiorna il proprio ottimizzatore. Con più client, il protocollo originale li addestra a turno e ognuno riceve i pesi più recenti lato client dal client precedente.

La configurazione cambia in base a ciò che deve restare privato. Nella configurazione standard qui sopra, le etichette vengono inviate al server; una configurazione a U mantiene anche gli ultimi layer e le etichette sul client; una configurazione verticale è pensata per organizzazioni che possiedono caratteristiche diverse degli stessi record:

Split learning e federated learning#

Anche il federated learning mantiene i dati sulle macchine dei rispettivi proprietari, ma ogni client addestra una copia completa del modello e invia aggiornamenti dei pesi da mediare. Lo split learning invia invece attivazioni e gradienti, quindi i client hanno bisogno di molta meno potenza di calcolo e non possiedono mai il modello completo; la comunicazione cresce in base alla quantità di dati, anziché alle dimensioni del modello, e i client si addestrano uno dopo l'altro, non in parallelo. SplitFed learning, proposto da Chandra Thapa e colleghi nel 2020, combina i due approcci eseguendo in parallelo i layer lato client e mediandoli secondo il metodo federato.

Lo split learning verticale è usato anche in ambito finanziario, dove una banca e un partner con caratteristiche diverse degli stessi clienti addestrano un unico modello, come nel tutorial di SecretFlow sul marketing bancario, senza scambiare dati privati. Il parallelismo di pipeline divide una rete per profondità allo stesso modo, ma tra GPU di un unico cluster affidabile, per poter caricare in memoria un modello grande.

Rischi per la privacy#

Attivazioni e gradienti contengono comunque informazioni sui dati. Le attivazioni dei primi layer possono essere invertite per ricostruire immagini riconoscibili e, quando le etichette sono sul server, i gradienti restituiti dal server possono rivelarle al client. Al convegno ACM CCS 2021, Dario Pasquini, Giuseppe Ateniese e Massimo Bernaschi hanno dimostrato che un server malevolo può dirottare l'addestramento tramite i gradienti che restituisce, indirizzando i layer del client verso una configurazione che consente al server di ricostruire gli input privati. Le difese aggiungono rumore con la privacy differenziale, penalizzano la quantità di informazioni sull'input rivelate dalle attivazioni o rilevano gradienti insoliti, ciascuna con un certo costo in termini di accuratezza. Le distribuzioni affiancano allo split learning controlli di sicurezza dei dati, come l'autenticazione e il trasporto cifrato, poiché il protocollo disciplina solo ciò che attraversa il punto di taglio.

Scelta del layer di taglio#

Il taglio determina il compromesso tra calcolo lato client, traffico ed esposizione: un taglio iniziale richiede meno calcolo al client, ma invia attivazioni grandi che somigliano ancora all'input; un taglio più avanzato invia mappe delle caratteristiche più astratte, ma richiede hardware client più potente, aspetto importante per i dispositivi di IA edge. Le connessioni skip complicano la scelta, perché occorre inviare ogni tensore che attraversa il taglio. Se si divide Ultralytics YOLO26 dopo il suo backbone, è necessario inviare tre mappe delle caratteristiche, perché la configurazione del modello fornisce al neck le uscite dei layer 4, 6 e 10: per YOLO26n, in imgsz=640, contengono 1.126.400 valori, quasi quanto i 1.228.800 dell'immagine di input.