Privacy-Aware Split Inference with Speculative Decoding for Large Language Models over Wide-Area Networks

1. El Problema

La adopción de Modelos de Lenguaje Grandes (LLM) en entornos empresariales (salud, defensa, legal, finanzas) genera una tensión fundamental entre la capacidad del modelo y el cumplimiento normativo de privacidad.

• Dilema actual: Las organizaciones deben elegir entre usar modelos potentes alojados en la nube (arriesgando la exposición de datos sensibles) o ejecutar modelos más pequeños localmente (sacrificando capacidad).
• Limitaciones de las soluciones existentes: Enfoques como el cifrado homomórfico, la computación multipartita segura (MPC) o los entornos de ejecución de confianza (TEE) imponen sobrecargas computacionales prohibitivas (100-1000x más lentos).
• Desafío de la Inferencia Dividida (Split Inference): Aunque dividir el modelo entre un dispositivo local y la nube reduce la exposición de datos, la inferencia autoregresiva en LLMs requiere un viaje de ida y vuelta (round-trip) a la red por cada token generado. En redes de área amplia (WAN) con latencias típicas de 80-100 ms, esto limita el rendimiento a ~8-11 tokens/segundo, lo cual es insuficiente para aplicaciones interactivas.

2. Metodología

El sistema propone una arquitectura de inferencia dividida asimétrica combinada con decodificación especulativa adaptada a la red.

A. Arquitectura del Sistema

• División de Capas (Split Point): El modelo Transformer se divide de manera que las capas de embedding (entrada) y unembedding (salida/LM head) permanecen en el dispositivo local de confianza (GPU local). Las capas intermedias (cálculo intensivo) se ejecutan en una GPU en la nube no confiable.
  • Ventaja de privacidad: Los tokens crudos y sus embeddings nunca salen del dispositivo local. Solo se transmiten activaciones intermedias (vectores de alta dimensión en el espacio de representación aprendido), que son computacionalmente difíciles de invertir sin las matrices de embedding.
• Transporte: Se utiliza un protocolo WebSocket binario para transmitir las activaciones (vectores float16 de ~8 KB por token), evitando la sobrecarga de JSON y codificación Base64.
• Gestión de Caché: La caché de claves-valores (KV-Cache) se mantiene en el servidor de la nube para evitar transferir el estado completo de la secuencia a través de la red.

B. Decodificación Especulativa (Lookahead Decoding)

Para mitigar la latencia de la red, el sistema adapta la técnica de Lookahead Decoding:

• Mecanismo: En lugar de esperar un token por viaje de ida y vuelta, el sistema utiliza iteraciones de Jacobi para generar candidatos de n-gramas (secuencias de tokens) localmente.
• Amortización: Estos candidatos se verifican en paralelo en la nube. Si coinciden, se confirman múltiples tokens en un solo viaje de ida y vuelta.
• Máscara de Atención: Se implementan máscaras de atención explícitas para manejar la verificación de múltiples tokens especulativos sin errores de causalidad.

3. Contribuciones Clave

1. Sistema Desplegable: Un sistema práctico que logra 8-9 tokens/segundo en un enlace WAN de ~80 ms, con proyecciones de 15-19 tokens/segundo a 20 ms.
2. Primera Integración Lookahead-Split: Demuestra que la decodificación paralela estilo Jacobi puede amortizar la latencia de red, no solo la computación.
3. Ingeniería de Despliegue: Documenta hallazgos críticos sobre la arquitectura de redes de proveedores de nube (ej. el enrutamiento proxy de VAST.ai vs. SSH directo de RunPod) que afectan más la latencia que la distancia geográfica.
4. Validación de Escala: Pruebas exitosas en Mistral NeMo 12B (40 capas) con solo 4.9 GB de VRAM local, mostrando que el rendimiento se mantiene comparable al modelo de 7B.
5. Verificación de Calidad: Prueba formal de que la decodificación Lookahead produce una salida idéntica a la decodificación secuencial bajo greedy argmax, sin degradación de calidad.
6. Modelo de Descomposición RTT: Un modelo matemático que separa el tiempo de pared en latencia de red (RTT) y sobrecarga fija, permitiendo proyecciones precisas de rendimiento con <6.2% de error.
7. Evaluación de Privacidad Empírica: Experimentos de ataque de inversión que cuantifican la recuperabilidad de tokens según la profundidad de la división.

4. Resultados Experimentales

Rendimiento (Throughput)

• Mistral 7B (WAN ~80ms):
  • Decodificación secuencial: ~10.9 tok/s.
  • Con Lookahead: 13–17 tok/s (en enlaces de baja latencia) o 8.7–9.3 tok/s (en enlaces estándar).
• Mistral NeMo 12B:
  • Logra un rendimiento similar al de 7B (7.8–8.7 tok/s) a pesar de tener 1.7x más parámetros, debido a que la latencia de red domina sobre el tiempo de cómputo de la nube.
• Proyecciones: A 20 ms de RTT (co-locación regional), el sistema alcanzaría 15–19 tok/s, acercándose a la velocidad de inferencia nativa.

Privacidad y Ataques de Inversión

• Se evaluó la capacidad de un atacante para recuperar tokens originales a partir de las activaciones intermedias.
• División a 2 capas (predeterminado): Un atacante puede recuperar ~59% de los tokens.
• División a 8 capas: La precisión del ataque cae a ~35% con un impacto mínimo en el rendimiento (cada capa local adicional añade solo ~3 ms de latencia).
• Conclusión: La privacidad es un compromiso ajustable; aumentar las capas locales mejora significativamente la seguridad sin penalizar la velocidad.

Análisis de Escala

• La tasa de aceptación de tokens especulativos es consistente entre modelos (1.21–1.25 tokens/paso), independientemente del tamaño del modelo, dependiendo más de las estadísticas del lenguaje que de la capacidad del modelo.
• El texto estructurado (código) obtiene las mejores tasas de aceptación (hasta 1.57 tokens/paso), mientras que el texto creativo es menos predecible.

5. Significado e Impacto

Este trabajo demuestra que la inferencia dividida es viable para uso interactivo en LLMs, resolviendo el cuello de botella de la latencia de red mediante técnicas de especulación.

• Privacidad Arquitectónica: Ofrece una alternativa práctica a las APIs de nube tradicionales para sectores regulados. Aunque no ofrece garantías criptográficas (como TEE o MPC), garantiza que los tokens crudos nunca abandonan el dispositivo, lo cual es una mejora cualitativa significativa sobre el envío de texto plano.
• Eficiencia de Recursos: Permite ejecutar modelos grandes (hasta 12B y proyectado hasta 70B) en hardware de consumo (RTX 3090) con requisitos de memoria local bajos (4.9 GB para 12B).
• Guía de Implementación: Proporciona lecciones críticas sobre la ingeniería de redes (túneles SSH, elección de proveedores) que a menudo se pasan por alto en la literatura académica pero son determinantes para el rendimiento en el mundo real.

En resumen, el sistema presenta un camino práctico para que las organizaciones utilicen LLMs potentes en la nube sin sacrificar la confidencialidad de sus datos, logrando velocidades de generación que permiten aplicaciones en tiempo real.