Meeting in the Middle: A Co-Design Paradigm for FHE and AI Inference

1. El Problema: Tensión entre Utilidad y Privacidad

El despliegue de grandes modelos neuronales como servicios alojados en la nube genera un conflicto fundamental:

• Privacidad del usuario: Los usuarios deben enviar entradas sensibles (a menudo en texto plano) para obtener inferencias.
• Propiedad del proveedor: Los proveedores deben ejecutar pesos de modelos propietarios en entornos que podrían tener vulnerabilidades de software o canales laterales.

Limitaciones de las soluciones actuales:

• Entornos de Ejecución Confiables (TEEs): Son la solución predominante en producción, pero sufren de vulnerabilidades de canales laterales, requieren confiar en fabricantes de hardware y no ofrecen garantías criptográficas a los usuarios.
• Cifrado Homomórfico (FHE) Genérico: Aunque ofrece garantías criptográficas, es impráctico para la inferencia de redes neuronales. Evaluar un transformador modesto requiere millones de operaciones no lineales y estructurales que son catastróficamente costosas bajo parámetros FHE estándar.

2. Metodología: El Paradigma de "Co-diseño"

Los autores proponen que la solución no reside en adaptar un modelo genérico a un esquema de cifrado genérico, sino en un enfoque de "co-diseño" o "reunión en el medio". La premisa central es que la inferencia de IA es estructuralmente predecible:

• La topología del circuito es fija.
• Las formas de los tensores son conocidas en tiempo de compilación.
• Las distribuciones de valores están acotadas.
• No existe ramificación dependiente de datos.

Esta predictibilidad permite especializar mutuamente tanto el esquema de cifrado como la arquitectura del modelo. El enfoque se organiza en dos ejes complementarios:

A. Especialización de Esquemas FHE para Inferencia

En lugar de diseñar FHE para computación genérica, se adaptan las primitivas criptográficas a las garantías estructurales de la IA:

• Optimización Offline: Al no haber ramificación dependiente de datos, toda la topología del circuito puede precomputarse. Esto permite optimizaciones agresivas a nivel de circuito que son imposibles en FHE dinámico.
• Ajuste de Parámetros: Se propone diseñar conjuntamente el módulo de texto plano, la estrategia de escalado y la codificación. Esto permite cumplir con los requisitos de precisión manteniendo la profundidad del circuito y el crecimiento del ruido dentro de un conjunto de parámetros elegido, aprovechando el rango dinámico acotado de las activaciones y pesos (tras entrenamiento con cuantización).
• Empaquetado SIMD y Transformaciones Lineales: Dado que las formas de los tensores son fijas, se pueden diseñar estrategias de empaquetado (packing) que alineen los "slots" del cifrado con la estructura de multiplicación de matrices del modelo. Esto reduce drásticamente la necesidad de rotaciones criptográficas y cambio de claves, operaciones usualmente muy costosas.
• Programación Estática de Bootstrapping: Para un circuito de inferencia fijo, se puede estimar estáticamente dónde se acumula el ruido más rápido. Esto permite precomputar un programa de bootstrapping óptimo y estático que minimice globalmente el número total de operaciones de bootstrapping.

B. Restricción de la Inferencia para FHE

La comunidad de IA debe rediseñar sus arquitecturas para ser "amigables con FHE":

• Búsqueda de Arquitectura Neural (NAS) Criptográfica: Se propone modificar el paradigma NAS tradicional. En lugar de solo optimizar por precisión y latencia en texto plano, se deben incluir restricciones criptográficas específicas (profundidad multiplicativa, conteo de rotaciones) como objetivos primarios. Esto guía la búsqueda hacia arquitecturas con capas anchas y poco profundas, penalizando operaciones de alto grado polinómico.
• Replanteamiento de No Linealidades: Funciones como ReLU o GeLU son costosas de aproximar homomóricamente. Los autores proponen entrenar modelos sin estas no linealidades estándar, sustituyéndolas por activaciones polinómicas de bajo grado soportadas nativamente.
• Aproximaciones y Entrenamiento: Se demuestra que es posible aproximar operaciones no lineales (como softmax) con errores mínimos. Además, mediante Entrenamiento Consciente de la Cuantización (QAT) dirigido directamente a los parámetros de FHE, la red aprende a operar nativamente dentro de las restricciones del módulo de texto plano.
• Dispersión Estructurada: Si los modelos se entrenan imponiendo patrones de dispersión (sparsity) estructurados que se alineen con la estrategia de empaquetado SIMD, se pueden omitir físicamente operaciones enteras durante la evaluación homomórfica.

3. Resultados y Evidencia Empírica

El documento presenta una prueba de concepto que valida la viabilidad de este enfoque:

• Aproximación de Softmax: Los autores derivaron una aproximación experimental para softmax sobre un rango de logits acotado.
  • Precisión: Error observado inferior al 0.5% bajo una métrica específica.
  • Rendimiento: En una prueba de micro-rendimiento (timming) que solo evalúa la función softmax, la evaluación homomórfica de esta aproximación fue más rápida que una implementación escalar en texto plano escrita en Rust en el mismo CPU (sin SIMD/AVX ni GPU).
  • Nota: Esta comparación excluye la generación de claves, cifrado/descifrado y bootstrapping, centrándose únicamente en la eficiencia de la evaluación de la función no lineal.

4. Contribuciones Clave

1. Marco Unificado: Presenta un marco que integra la especialización del cifrado y la restricción del modelo, superando los enfoques aislados previos.
2. Optimización Estática: Demuestra cómo la naturaleza estática de la inferencia permite optimizaciones (como el bootstrapping y las rotaciones) que son imposibles en computación genérica.
3. Rediseño de NAS: Propone cambiar los objetivos de la búsqueda de arquitecturas neuronales para priorizar la eficiencia criptográfica (profundidad y rotaciones) junto con la precisión.
4. Viabilidad de Aproximaciones: Muestra que la tolerancia de los modelos de aprendizaje automático a la cuantización y perturbaciones permite el uso de aproximaciones polinómicas agresivas sin sacrificar significativamente la precisión.

5. Significado e Impacto

Este trabajo es significativo porque desafía la noción de que la inferencia privada con FHE es inherentemente ineficiente. Al proponer un co-diseño, los autores sugieren que se puede alcanzar un "óptimo en el medio" que supera la integración ingenua de modelos SOTA con esquemas FHE genéricos.

El impacto potencial es doble:

• Para la Criptografía: Permite diseñar esquemas FHE más eficientes al eliminar la necesidad de soportar computación genérica, enfocándose en cargas de trabajo estructuradas.
• Para la IA: Fomenta el desarrollo de nuevas arquitecturas de redes neuronales que son nativamente eficientes en entornos cifrados, abriendo la puerta a servicios de inferencia en la nube que sean verdaderamente privados, verificables y seguros, sin depender de la confianza en hardware o TEEs.

En resumen, el papel argumenta que la barrera actual no es tecnológica, sino de enfoque: la solución requiere que tanto los criptógrafos como los ingenieros de IA acepten restricciones mutuas para lograr un sistema viable.