Resumen Técnico: PowerSoftmax: Hacia Inferencia Segura de LLM sobre Datos Encriptados

Declaración del Problema

Los Modelos de Lenguaje Grande (LLM) modernos que preservan la privacidad a menudo dependen de la Encriptación Homomórfica (HE), específicamente el esquema CKKS, que permite realizar cálculos sobre datos encriptados sin desencriptarlos. Sin embargo, CKKS solo soporta operaciones polinómicas. Las arquitecturas Transformer estándar dependen en gran medida de componentes no polinómicos, siendo el más notable la función Softmax en la auto-atención, la cual involucra exponenciales, divisiones y operaciones logarítmicas (por ejemplo, para la estabilidad numérica mediante el truco log-sum-exp).

Los enfoques existentes para cerrar esta brecha enfrentan limitaciones significativas:

1. Aproximación Post-Entrenamiento (PTA): Aproximar directamente funciones no polinómicas en modelos pre-entrenados utilizando polinomios de alto grado. Esto a menudo resulta en una inferencia ineficiente debido a una profundidad multiplicativa alta y ruido acumulado.
2. Entrenamiento Consciente de la Aproximación (AAT): Reemplazar primitivas no polinómicas con alternativas más fáciles de aproximar (por ejemplo, atención punto a punto o ReLU escalada) antes del entrenamiento. Aunque efectivas para modelos más pequeños, estos métodos han luchado con la estabilidad y la escalabilidad, fallando en soportar LLMs con miles de millones de parámetros o arquitecturas profundas (por ejemplo, 32 capas).

El desafío central es diseñar una arquitectura Transformer que sea inherentemente compatible con las restricciones de HE (polinomios de bajo grado, rangos numéricos estables) mientras mantiene el poder expresivo, las capacidades de razonamiento y el aprendizaje en contexto (ICL) de los Transformers estándar basados en Softmax a la escala de miles de millones de parámetros.

Metodología

Los autores proponen PowerSoftmax, una variante novedosa y compatible con HE de la auto-atención diseñada para replicar los principios fundamentales de Softmax mientras permanece susceptible a la aproximación polinómica. La metodología involucra tres etapas distintas:

1. Diseño de Atención Compatible con HE

Los autores destilan las propiedades esenciales de Softmax: normalización a $[0, 1]$, amplificación de las diferencias de puntuación y preservación del orden. Reemplazan la función exponencial $\exp(x)$ en Softmax con una función de potencia $x^p$ (donde $p$ es un entero par positivo).
El mecanismo de atención propuesto se define como:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
Esta formulación evita las exponenciales y los problemas de división por cero inherentes a las aproximaciones polinómicas directas de Softmax.

2. Estabilidad del Entrenamiento y Robustez Numérica

Para abordar la inestabilidad del entrenamiento causada por la función de potencia (donde las entradas $>1$ explotan y las entradas $<1$ desaparecen), los autores introducen dos modificaciones clave:

• Variante Estable: Inspirados en el truco log-sum-exp, normalizan las entradas mediante una constante $c = \|x\|_\infty + \delta$ antes de aplicar la función de potencia. Esto asegura que las entradas permanezcan dentro de un rango manejable $(0, 1)$, previniendo desbordamientos/subdesbordamientos y preservando la precisión de punto flotante.
• División Acotada por $\epsilon$: Para manejar la operación de división, que es no polinómica y sensible cerca de cero, modifican el denominador a $\epsilon + \sum x_i^p$. Esto crea una función continua $1/\epsilon^2$-Lipschitz, haciendo que la aproximación polinómica posterior de la división sea significativamente más tratable y reduciendo la profundidad multiplicativa requerida.

3. Aproximación Agnóstica a la Longitud

La aproximación estándar de división (por ejemplo, algoritmo de Goldschmidt) lucha con la suma no acotada de puntuaciones de atención, que escala linealmente con la longitud de secuencia $L$. Los autores proponen un PowerSoftmax Agnóstico a la Longitud:
$$\text{Agnóstico a la Longitud}(x)_j = \frac{x_j^p}{L \cdot \text{Media}(x^p)}$$
Dado que $L$ (longitud de secuencia) es pública, el término $1/L$ puede calcularse directamente sin aproximación. El modelo aproxima la división de la media en lugar de la suma, asegurando que el dominio de aproximación permanezca acotado independientemente de la longitud de la secuencia.

4. Pipeline de Construcción (Algoritmo 1)

El pipeline completo para crear un LLM polinómico involucra:

1. Modificación Arquitectónica: Reemplazar Softmax con la variante estable de PowerSoftmax y entrenar desde cero (o mediante pre-entrenamiento continuo) con objetivos de minimización de rango.
2. Minimización de Rango: Aplicar una función de pérdida especializada para restringir los rangos de entrada de las capas no polinómicas (PowerSoftmax, GELU, LayerNorm) para asegurar que las aproximaciones polinómicas de bajo grado sean precisas.
3. Reemplazo Polinómico: Post-entrenamiento, reemplazar las operaciones no polinómicas restantes (división, raíz cuadrada inversa) con aproximaciones polinómicas de bajo grado (por ejemplo, método de Goldschmidt) para producir un modelo completamente polinómico.

Contribuciones Clave

1. Mecanismo PowerSoftmax: Una nueva variante de auto-atención que reemplaza la escalado exponencial con escalado polinómico, específicamente adaptada para entornos HE mientras mantiene los principios de diseño centrales de la atención.
2. LLMs Polinómicos Escalables: El desarrollo de los primeros LLMs polinómicos que superan los mil millones de parámetros (específicamente 1.4B parámetros con 32 capas). Estos modelos demuestran capacidades de razonamiento y aprendizaje en contexto (ICL) comparables a los transformers estándar del mismo tamaño.
3. Estrategias de Entrenamiento e Inferencia Estables: Introducción de variantes de entrenamiento numéricamente estables y estrategias de inferencia agnósticas a la longitud que superan los problemas de inestabilidad y escalabilidad de los enfoques anteriores de Transformer polinómico.
4. Evaluación Exhaustiva: Un desglose detallado de la latencia para cálculos encriptados y estudios de ablación que validan las decisiones de diseño, incluido el impacto del parámetro de potencia $p$ y la división acotada por $\epsilon$.

Resultados Experimentales

Los autores evaluaron su método en arquitecturas causales (GPT/Pythia) y bidireccionales (RoBERTa):

• Rendimiento: El modelo polinómico de 1.4B parámetros logró un rendimiento comparable a los modelos Pythia originales basados en Softmax en benchmarks de cero disparos y pocos disparos (por ejemplo, Lambada, PIQA, WinoGrande, ARC). Notablemente, el modelo demostró un rendimiento competitivo en benchmarks de razonamiento como ARC-Challenge, una capacidad previamente no observada en LLMs polinómicos.
• Modelos Bidireccionales: Una variante polinómica de RoBERTa-Base mostró solo una degradación de ~1% en comparación con el RoBERTa original basado en Softmax en tareas GLUE (SST-2, QNLI, MNLI), superando enfoques polinómicos anteriores como Nexus.
• Inferencia HE: Experimentos en el modelo Pythia 70M utilizando el framework HElayers confirmaron que la inferencia encriptada produce resultados casi idénticos a la inferencia en texto plano (pérdida MSE máxima de 0.005).
• Eficiencia: El método propuesto reduce significativamente la profundidad multiplicativa en comparación con alternativas de última generación (Nexus, Cho et al.). Para una secuencia de 128 tokens, PowerSoftmax redujo la latencia en un factor de ~9.7 en comparación con los enfoques Taylor/Goldschmidt y ~20 en comparación con los enfoques de estrategia cuadrada.

Significado y Afirmaciones

El artículo afirma un avance en el campo del aprendizaje automático que preserva la privacidad al demostrar que los LLMs completamente polinómicos pueden escalar al régimen de miles de millones de parámetros sin sacrificar las capacidades de razonamiento o aprendizaje en contexto.

• Escalabilidad: A diferencia de trabajos anteriores limitados a modelos pequeños o tareas de clasificación simples, este enfoque entrena exitosamente transformers profundos (32 capas) con 1.4 mil millones de parámetros.
• Seguridad No Interactiva: Al utilizar encriptación homomórfica totalmente (FHE) con un modelo completamente polinómico, la solución permite inferencia no interactiva, evitando la sobrecarga de comunicación y las vulnerabilidades potenciales asociadas con la computación multipartita (MPC) o protocolos interactivos.
• Practicidad: El análisis detallado de latencia y el despliegue exitoso de un modelo de 1.4B parámetros sugieren que la inferencia segura y privada de LLM se está volviendo prácticamente factible, allanando el camino para aplicaciones en sectores sensibles a los datos como la salud y las finanzas.

Los autores reconocen que, aunque los modelos funcionan bien, una evaluación completa de las capacidades generativas auto-regresivas en la decodificación secuencial sobre datos encriptados sigue siendo una dirección para trabajo futuro.