Resumen Técnico: CGF-Softmax

Enunciado del Problema

El Cifrado Homomórfico (HE) permite el aprendizaje automático que preserva la privacidad al permitir el cálculo sobre datos cifrados sin necesidad de descifrarlos. Sin embargo, la implementación de operaciones no polinómicas, particularmente la función softmax central en las arquitecturas Transformer, sigue siendo un cuello de botella significativo en entornos de HE estrictamente no interactivos.

La formulación estándar del softmax, $\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}$, presenta tres desafíos principales en HE:

1. Inestabilidad Numérica: La función exponencial causa un crecimiento rápido, lo que lleva a desbordamientos. La inferencia en texto plano estándar mitiga esto restando el valor de entrada máximo ($x_{\max}$), pero la operación de máximo requiere lógica de comparación compleja que es incompatible con HE y hace necesarias aproximaciones polinómicas costosas.
2. Operaciones de División: El paso de normalización requiere división, que los esquemas de HE no soportan de forma nativa. Trabajos anteriores dependen de aproximaciones polinómicas iterativas (por ejemplo, el algoritmo de Goldschmidt o polinomios de Chebyshev) para funciones recíprocas, aumentando drásticamente la profundidad multiplicativa.
3. Costo Computacional: Las aproximaciones polinómicas de alto grado requeridas tanto para las funciones exponencial como de división consumen el presupuesto de niveles limitado de los esquemas de HE escalonados (como CKKS), a menudo haciendo necesarias costosas operaciones de reinicio (bootstrapping) para refrescar los textos cifrados.

Los enfoques existentes se dividen en dos categorías: aproximación del softmax (alta precisión pero alta profundidad debido a la división/máximo) y sustitución del softmax (baja profundidad pero degradación significativa de la precisión, especialmente en modelos a gran escala y benchmarks complejos).

Metodología: CGF-Softmax

Los autores proponen CGF-Softmax, una reformulación de la función softmax que elimina tanto la división homomórfica como las operaciones explícitas de resta del máximo. La innovación central radica en interpretar el denominador del softmax a través de la Función Generadora de Cumulantes (CGF).

Reformulación Teórica

Dado un vector de entrada $x$, los autores interpretan el denominador $\sum \exp(x_i)$ como una media muestral escalada de términos exponenciales. Al sustituir la media muestral por la media del conjunto (esperanza), el denominador se reemplaza por $n M_X(1)$, donde $M_X(t)$ es la Función Generadora de Momentos (MGF).

Utilizando la CGF, definida como $K_X(t) = \ln M_X(t)$, el softmax reformulado se expresa como:
$$\text{softmax}_{\text{CGF}}(x)_i = \exp(x_i - K_X(1) - \ln n)$$

Esta formulación ofrece tres ventajas críticas:

1. Eliminación de la Operación Máximo: En lugar de restar el valor máximo, la entrada se desplaza por $K_X(1) + \ln n$. Dado que $K_X(1)$ puede aproximarse utilizando cumulantes de bajo orden (media $\mu$ y varianza $\sigma^2$), este desplazamiento es adaptativo a los datos y computable mediante operaciones básicas de HE (suma y multiplicación) sin comparaciones.
2. Eliminación de la División: El término de normalización se absorbe en el exponente. El resultado es una única función exponencial, eliminando la necesidad de división homomórfica por completo.
3. Escalado del Dominio: El método soporta naturalmente el escalado del dominio. Al escalar el exponente por $1/2^k$, la función exponencial se evalúa sobre un intervalo más pequeño, y la escala original se recupera elevando al cuadrado $k$ veces. Esto evita el proceso iterativo de "normalizar y elevar al cuadrado" requerido por el softmax estándar, el cual implica divisiones repetidas.

Instanciación Práctica

Para la evaluación práctica de HE, los autores adoptan una aproximación de segundo orden de cumulantes:
$$K_X(1) \approx \mu + \frac{\sigma^2}{2}$$
Se omiten los términos de orden superior ($\kappa_j/j!$ para $j \ge 3$). Los autores argumentan que el denominador factorial suprime intrínsecamente estos términos, y el ajuste fino empírico reduce aún más la magnitud de los cumulantes de orden superior, haciendo que la aproximación sea lo suficientemente precisa para la inferencia.

Contribuciones Clave

1. Formulación CGF-Softmax: Una reformulación novedosa que elimina la necesidad de división homomórfica y resta de máximos, reduciendo sustancialmente la profundidad multiplicativa.
2. Análisis Teórico: El artículo establece propiedades clave del softmax-CGF, incluyendo:
   • Positividad: Salidas estrictamente positivas.
   • Preservación del Orden y de la Razón Pareada: Preserva el orden relativo de los logits y las razones exactas de salida pareadas ($\exp(x_i - x_j)$) del softmax estándar.
   • Invarianza al Desplazamiento: Al igual que el softmax estándar, es invariante a desplazamientos constantes.
   • Límite de Error Asintótico: Utilizando el teorema de Berry–Esseen, los autores demuestran que el error relativo converge a cero en probabilidad a medida que aumenta la dimensión de entrada $n$.
3. Validación Experimental: Evaluación extensa en Vision Transformers (ViT/DeiT) y Modelos de Lenguaje Grande (LLaMA-3.2-1B) que demuestra alta precisión con costos computacionales reducidos.

Resultados Experimentales

Los autores evaluaron CGF-softmax en ImageNet-1k (Vision Transformers) y múltiples benchmarks de NLP (Clinc150, Banking77, SST-2) utilizando el modelo LLaMA-3.2-1B.

• Profundidad Multiplicativa: CGF-softmax logra una profundidad multiplicativa de 7–10 para modelos a gran escala (por ejemplo, LLaMA-3.2-1B, ViT-Base). En contraste, los métodos de aproximación anteriores (por ejemplo, Cho et al., EncryptedLLM) requieren profundidades de 14–52 para modelos o tareas más pequeñas.
• Tiempo de Ejecución: En una matriz de entrada de 256×256, CGF-softmax redujo la latencia total de 107.65s (Cho et al.) a 3.25s, principalmente al eliminar el reinicio (bootstrapping) y reducir las multiplicaciones de texto cifrado (CMult).
• Precisión:
  • CGF-softmax mantiene la precisión de inferencia dentro de <1% de la línea base exacta de softmax en texto plano en todos los benchmarks.
  • Supera significativamente a los métodos de sustitución de softmax (por ejemplo, BPMax), que sufren caídas sustanciales de precisión en tareas complejas.
  • Logra una precisión cercana a la de texto plano mientras requiere 4–6× menos profundidad que las líneas base de aproximación de alta precisión.
• Comportamiento del Ruido: Al eliminar la división, CGF-softmax reduce el ruido acumulado. En profundidad 7, logra un ruido total de $3.05 \times 10^{-10}$, que es menor que las formulaciones estándar de softmax incluso a mayores profundidades.

Significado y Afirmaciones

El artículo afirma que CGF-softmax ofrece una solución escalable y eficiente para la inferencia segura en modelos basados en Transformer. Al conciliar alta precisión con una profundidad multiplicativa sustancialmente reducida, aborda el principal cuello de botella en el aprendizaje automático que preserva la privacidad (PPML) basado en HE.

Los autores enfatizan que, aunque los métodos de sustitución de softmax pueden ser suficientes para modelos pequeños en tareas simples, CGF-softmax es necesario para arquitecturas a gran escala y benchmarks complejos donde las sustituciones existentes no logran mantener la precisión. El método permite una inferencia totalmente externalizada y no interactiva en modelos grandes (como LLaMA-3.2-1B) con costos computacionales que eran previamente inalcanzables para inferencia HE de alta precisión.

El artículo nota modestamente que para modelos relativamente pequeños (por ejemplo, BERT) en tareas simples, los métodos de sustitución de baja profundidad existentes podrían seguir siendo suficientes, pero CGF-softmax ofrece un compromiso superior para escenarios más exigentes.