Split Learning

El aprendizaje dividido es un método de entrenamiento distribuido colaborativo que divide una red neuronal en una capa de corte elegida, entre el propietario de los datos, que actúa como cliente, y un servidor. El cliente ejecuta las capas anteriores al corte con sus datos privados y envía únicamente las activaciones resultantes, a veces llamadas datos fragmentados; el servidor ejecuta las capas restantes, calcula la pérdida y devuelve el gradiente en el punto de corte para que el cliente pueda actualizar sus capas. Los datos de entrenamiento sin procesar nunca salen del cliente, y ninguna de las partes posee el modelo completo. Otkrist Gupta y Ramesh Raskar, del MIT Media Lab, presentaron este método en octubre de 2018 en Aprendizaje distribuido de redes neuronales profundas entre varios agentes; en diciembre de 2018, el grupo publicó el trabajo complementario Split Learning for Health, que propuso configuraciones para que hospitales con distintos tipos de datos de pacientes, como escaneos para el análisis de imágenes médicas, entrenasen modelos juntos sin compartir registros sin procesar.

Cómo funciona el aprendizaje dividido#

Cada paso de entrenamiento cruza dos veces el punto de corte: las activaciones van del cliente al servidor durante la pasada hacia delante, y el gradiente de la pérdida con respecto a esas activaciones vuelve durante la retropropagación. Como la regla de la cadena solo necesita ese gradiente para continuar, las actualizaciones del cliente son idénticas a las del entrenamiento del modelo sin dividir. Una simulación en un único proceso con PyTorch muestra el intercambio:

import torch
from torch import nn

torch.manual_seed(0)
client = nn.Sequential(nn.Linear(784, 256), nn.ReLU())  # capas anteriores al corte, junto a los datos privados
server = nn.Sequential(nn.Linear(256, 10))  # capas posteriores al corte
images, labels = torch.randn(32, 784), torch.randint(0, 10, (32,))

# Cliente: pasada hacia delante hasta la capa de corte y envío exclusivo de las activaciones
activations = client(images)
received = activations.detach().requires_grad_()

# Servidor: finaliza la pasada hacia delante y propaga hacia atrás hasta el corte; después, devuelve ese gradiente
nn.functional.cross_entropy(server(received), labels).backward()

# Cliente: continúa la retropropagación a partir del gradiente recibido
activations.backward(received.grad)
split_grad = client[0].weight.grad.clone()

# El mismo paso con el modelo sin dividir da gradientes idénticos al cliente
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))

Imprime Same client gradients as unsplit training: True. A continuación, cada parte actualiza su propio optimizador. Con varios clientes, el protocolo original los entrena por turnos y cada uno recibe los últimos pesos del cliente anterior.

La configuración cambia según lo que se deba mantener en privado. En la configuración básica anterior se envían las etiquetas al servidor; una configuración en U también mantiene las últimas capas y las etiquetas en el cliente, y una configuración vertical está pensada para organizaciones que tienen distintas características de los mismos registros:

Aprendizaje dividido frente a aprendizaje federado#

El aprendizaje federado también mantiene los datos en los equipos de sus propietarios, pero cada cliente entrena una copia completa del modelo y envía actualizaciones de pesos para promediarlas. El aprendizaje dividido envía activaciones y gradientes, por lo que los clientes necesitan mucha menos capacidad de cálculo y nunca tienen el modelo completo; a cambio, las comunicaciones aumentan en función de la cantidad de datos, no del tamaño del modelo, y los clientes entrenan uno tras otro en vez de en paralelo. SplitFed, propuesto por Chandra Thapa y sus colegas en 2020, combina ambos enfoques ejecutando las capas del cliente en paralelo y promediándolas al estilo federado.

El aprendizaje dividido vertical también se utiliza en finanzas, donde un banco y un colaborador que poseen distintas características de los mismos clientes entrenan un único modelo, como en el tutorial de marketing bancario de SecretFlow, sin intercambiar datos privados. El paralelismo de canalización también divide una red por profundidad, pero lo hace entre GPU de un único clúster de confianza para que un modelo grande quepa en la memoria.

Riesgos para la privacidad#

Las activaciones y los gradientes siguen conteniendo información sobre los datos. Las activaciones de las primeras capas pueden invertirse para reconstruir imágenes reconocibles y, si el servidor tiene las etiquetas, los gradientes que devuelve pueden revelárselas al cliente. En ACM CCS 2021, Dario Pasquini, Giuseppe Ateniese y Massimo Bernaschi demostraron que un servidor malicioso puede secuestrar el entrenamiento mediante los gradientes que devuelve y dirigir las capas del cliente de forma que el servidor pueda reconstruir las entradas privadas. Las defensas añaden ruido mediante privacidad diferencial, penalizan la cantidad de información de entrada que revelan las activaciones o detectan gradientes inusuales; todas tienen cierto coste en precisión. Las implementaciones combinan el aprendizaje dividido con controles de seguridad de datos, como la autenticación y el transporte cifrado, ya que el propio protocolo solo regula qué cruza el punto de corte.

Cómo elegir la capa de corte#

El punto de corte establece el equilibrio entre el cómputo del cliente, el tráfico y la exposición: un corte temprano exige menos cómputo al cliente, pero envía activaciones grandes que todavía se parecen a la entrada; uno más tardío envía mapas de características más abstractos, pero requiere más capacidad del hardware cliente, algo importante para los dispositivos de IA en el borde. Las conexiones de salto complican la elección, porque hay que enviar todos los tensores que atraviesen el punto de corte. Si se divide Ultralytics YOLO26 después del backbone, hay que enviar tres mapas de características, ya que su configuración del modelo introduce en el cuello las salidas de las capas 4, 6 y 10: para YOLO26n en imgsz=640, contienen 1.126.400 valores, casi tantos como los 1.228.800 de la imagen de entrada.