Resumen Técnico: TIGER – Inversión de Gradientes de Transformers mediante la Optimización de la Distancia del Subespacio de Embeddings

Planteamiento del Problema

El Aprendizaje Federado (FL) permite el entrenamiento colaborativo de modelos compartidos manteniendo los datos locales de los clientes. Sin embargo, los ataques de inversión de gradientes han demostrado que las actualizaciones del modelo (gradientes) pueden filtrar información suficiente para reconstruir entradas privadas de los clientes. Aunque estos ataques han tenido éxito con datos de imágenes, la reconstrucción de texto de Modelos de Lenguaje de Gran Escala (LLMs) basados en Transformers sigue siendo un desafío.

Los métodos de inversión de texto existentes enfrentan limitaciones significativas en entornos de FL realistas:

1. Enfoques basados en optimización: Los métodos que optimizan entradas ficticias para igualar gradientes completos son computacionalmente costosos e inestables para modelos modernos, particularmente con tamaños de lote grandes.
2. Enfoques basados en subespacios: Trabajos previos (por ejemplo, DAGER) explotan la estructura de bajo rango de los gradientes de las capas lineales de los Transformers para identificar un subespacio que contenga los embeddings de las capas reales. Sin embargo, estos métodos dependen de pruebas de membresía discretas (verificar si los tokens candidatos residen en el subespacio). Este enfoque es frágil ante el ruido numérico (por ejemplo, de cuantización o Privacidad Diferencial) y escala pobremente para modelos de codificador con atención no causal, donde la recuperación secuencial de tokens es imposible.

Consecuentemente, los ataques existentes a menudo fallan o subestiman los riesgos de privacidad en entornos de FL defendidos que involucran grandes lotes, gradientes cuantizados o ruido.

Metodología: TIGER

Los autores presentan TIGER (Transformer Data Inversion from Gradients via Embedding-subspace Reconstruction), un ataque de inversión de datos de Transformer continuo que transforma la señal del subespacio en un objetivo diferenciable. En lugar de buscar sobre tokens discretos o igualar gradientes completos, TIGER optimiza directamente los embeddings de tokens continuos para minimizar su distancia a los subespacios inducidos por el gradiente.

Mecanismo Central

TIGER aprovecha la observación de que, para una capa lineal en un Transformer, el rango de filas de las entradas de la capa coincide con el espacio de columnas del gradiente de peso correspondiente (bajo condiciones de rango leves).

• Señal de Subespacio: Dados los gradientes observados $\frac{\partial \mathcal{L}}{\partial W^l_Q}$, el atacante computa el espacio de columnas $S_l$.
• Objetivo Continuo: En lugar de probar candidatos discretos, TIGER optimiza una matriz de embeddings de entrada bruta continua $\bar{Z}_0$. El objetivo es asegurar que los estados ocultos $\bar{Z}_l$ inducidos por el paso hacia adelante (forward pass) se alineen con los subespacios observados $S_l$.
• Función de Pérdida: El objetivo principal es una pérdida de distancia de rango hacia adelante ($\ell_{forw}$), que minimiza la distancia al cuadrado entre los estados ocultos recuperados normalizados y su proyección sobre el subespacio inducido por el gradiente.

Adaptaciones Específicas de la Arquitectura

TIGER emplea estrategias distintas para modelos solo-decodificador (decoder-only) y solo-codificador (encoder-only):

1. Modelos Solo-Decodificador (Causales):

   • Explota las máscaras de atención causal para recuperar tokens secuencialmente (token por token).
   • Dado que la representación en la posición $i$ depende solo de los tokens precedentes, el ataque optimiza los embeddings uno por uno.
   • Deduplicación: Para evitar que múltiples secuencias en un lote colapsen al mismo primer token, se añade una pérdida de deduplicación ($\ell_{dedup}$). Esta penaliza el embedding del primer token actual si se alinea con el subespacio ya ocupado por los primeros tokens previamente recuperados en el lote.

2. Modelos Solo-Codificador (No Causales):

   • La recuperación secuencial causal es inviable porque todos los tokens en una secuencia son interdependientes.
   • TIGER realiza una optimización conjunta de los embeddings brutos de todo el lote.
   • Alineación Bidireccional: Para evitar soluciones degeneradas (por ejemplo, que todos los tokens recuperados colapsen hacia una única dirección), TIGER introduce una pérdida de distancia de rango hacia atrás ($\ell_{back}$). Esta pérdida impone que las direcciones de la base inducida por el gradiente también deben residir dentro del rango de los estados ocultos recuperados. Esto asegura que el lote recuperado cubra el subespacio observado completo, desalentando soluciones de tokens repetidos o secuencias duplicadas.

Inicialización y Optimización

• Inicialización: TIGER utiliza un prior Gaussiano dependiente de la posición ajustado a un corpus extenso (por ejemplo, WikiText) para inicializar los embeddings, proporcionando un punto de partida sólido para la optimización no convexa.
• Optimización: El ataque ejecuta múltiples optimizaciones independientes ($N_{init}$) con diferentes inicializaciones, utilizando la optimización Adam para minimizar la pérdida de distancia de rango.
• Discretización: Una vez que la optimización continua converge, los embeddings recuperados se proyectan al token de vocabulario más cercano mediante similitud de coseno.

Contribuciones Clave

1. Ataque TIGER: Un ataque de inversión de gradientes continuo, robusto y escalable para LLMs de tipo Transformer que recupera texto privado en entornos desafiantes (grandes lotes, gradientes cuantizados, ruido de DP).
2. Objetivo de Subespacio Diferenciable: Una formulación novedosa que convierte la estructura de bajo rango de los gradientes en un objetivo de subespacio de embeddings diferenciable. Esto evita el filtrado discreto frágil y se extiende a modelos solo-codificadores mediante una pérdida de alineación bidireccional.
3. Evaluación Exhaustiva: Experimentos extensos en modelos modernos (GEMMA-3-4B-IT y EMBEDDINGGEMMA-300M) demostrando que TIGER supera sustancialmente a ataques previos en entornos defendidos y de solo-codificador, manteniéndose competitivo en entornos de decodificador no defendidos.
4. Código Abierto: Una implementación de código abierto para facilitar la investigación futura sobre riesgos y defensas de privacidad.

Resultados Experimentales

Los autores evaluaron TIGER en lotes de WikiText con variaciones en tamaños de lote, longitudes de secuencia y niveles de defensa (ruido de gradiente y cuantización), midiendo el desempeño mediante ROUGE-1 y ROUGE-L.

• Robustez al Ruido (Decodificador): En entornos no defendidos, TIGER es competitivo o supera ligeramente a DAGER. Sin embargo, bajo ruido de gradiente de tipo Privacidad Diferencial ($\sigma = 10^{-3}$), DAGER colapsa a una calidad de reconstrucción cercana a cero, mientras que TIGER mantiene una recuperación significativa (ROUGE-1 > 70%).
• Desempeño del Codificador: En modelos solo-codificador, donde la recuperación causal es imposible, TIGER logra un ROUGE-1 de ~60% en lotes de 16 secuencias, superando significativamente a ataques de optimización previos (por ejemplo, LAMP).
• Cuantización: TIGER sigue siendo efectivo bajo precisión BF16, mientras que DAGER falla completamente debido a la falta de fiabilidad de sus pruebas discretas basadas en umbrales.
• Escalabilidad: TIGER escala a secuencias más largas y tamaños de lote más grandes de manera más efectiva que los métodos anteriores, particularmente en entornos de codificador.

Significancia y Reivindicaciones

El artículo afirma que TIGER revela vulnerabilidades de privacidad significativas en los LLMs de tipo Transformer que anteriormente se habían subestimado.

• Resiliencia de los Gradientes: El estudio demuestra que los gradientes de los transformers retienen información textual sustancial incluso después de perturbaciones comunes como el ruido aditivo (DP) o la reducción de la precisión numérica (cuantización).
• Limitaciones de las Defensas Actuales: Los resultados sugieren que no debe asumirse que el ruido moderado de gradiente o el cómputo de menor precisión eliminen los riesgos de reconstrucción.
• Amenaza Más Amplia: Al permitir reconstrucciones exitosas en entornos donde ataques previos fallan (grandes lotes, codificadores, entornos defendidos), TIGER expande el panorama de amenazas prácticas de la inversión de gradientes, resaltando la necesidad de evaluar y aplicar cuidadosamente las defensas de privacidad en despliegues de aprendizaje federado.

Los autores señalan limitaciones, específicamente que el ataque depende de la estructura de bajo rango de los gradientes, la cual disminuye cuando el número total de tokens excede la dimensión oculta. También reconocen que, aunque TIGER recupera contenido sustancial, no siempre produce una reconstrucción exacta ordenada, lo que sugiere potencial para enfoques híbridos que combinen la recuperación de subespacio continua con el refinamiento discreto.