Resumen Técnico: CHEHAB RL: Aprendizaje para la Optimización de Computaciones de Cifrado Totalmente Homomórfico

1. Declaración del Problema

El Cifrado Totalmente Homomórfico (FHE, por sus siglas en inglés) permite la computación sobre datos cifrados sin necesidad de descifrarlos, preservando la privacidad. Sin embargo, el FHE sigue siendo computacionalmente costoso; por ejemplo, una multiplicación de texto cifrado puede ser $10^6$ veces más lenta que su contraparte en texto plano. Esta brecha de rendimiento surge de la complejidad inherente de las operaciones homomórficas, que involucran multiplicaciones polinómicas de alto grado y un crecimiento significativo del ruido.

Escribir código FHE eficiente es una tarea compleja y propensa a errores que requiere una profunda experiencia criptográfica. Aunque se han desarrollado compiladores para automatizar la generación de código FHE, las soluciones existentes enfrentan limitaciones significativas:

• Código Estructurado vs. No Estructurado: Compiladores como HECO y CHET se centran en código estructurado (basado en bucles), mientras que otros como Coyote y Porcupine intentan manejar tanto el código estructurado como el no estructurado (arbitrario, no basado en bucles).
• Complejidad del Espacio de Búsqueda: Los compiladores actuales de vanguardia (por ejemplo, Coyote, Porcupine) enmarcan la vectorización como un problema de búsqueda combinatoria. Dependen de heurísticas ajustadas a mano, Programación Lineal Entera (ILP) o síntesis de programas para explorar un vasto espacio de diseño (por ejemplo, decidir qué subexpresiones empaquetar, disposición de datos, inserción de rotaciones).
• Escalabilidad y Óptimos Locales: En este espacio complejo, las decisiones locales suelen tener efectos no locales (por ejemplo, una decisión de empaquetado que aumenta los costos de rotación futuros). Las heurísticas simples a menudo quedan atrapadas en óptimos locales pobres, y la búsqueda exhaustiva se vuelve computacionalmente intratable a medida que aumenta el tamaño del programa.

2. Metodología

El artículo propone CHEHAB RL, un nuevo marco que aprovecha el Aprendizaje por Refuerzo Profundo (Deep Reinforcement Learning - RL) para automatizar la optimización de código FHE. En lugar de una búsqueda combinatoria, el problema se formula como un proceso de toma de decisiones secuenciales donde un agente de RL aprende una política para aplicar una secuencia de reglas de reescritura.

2.1 Descripción General del Marco

CHEHAB RL está integrado en el compilador FHE CHEHAB. El flujo de trabajo es el siguiente:

1. Entrada: Un programa escrito en el Lenguaje de Dominio Específico (DSL) de CHEHAB, que se reduce a una Representación Intermedia (IR) en forma de Árbol de Sintaxis Abstracta (AST) y Grafo de Flujo de Datos (DAG).
2. Optimización por RL: Un agente de RL recorre la IR, seleccionando una secuencia de reglas de reescritura de un Sistema de Reescritura de Términos (TRS) para vectorizar el código escalar, simplificar expresiones y minimizar el ruido/latencia.
3. Salida: La IR optimizada se utiliza para generar código C++ dirigido al backend de Microsoft SEAL (esquema BFV).

2.2 Componentes Técnicos Clave

A. Representación del Estado

Para representar la IR del programa para el agente de RL, los autores emplean la tokenización de Identificador e Invariante de Constante (ICI).

• Normalización: Las variables se renombran (por ejemplo, $v_0, v_1$) y las constantes (excepto 0 y 1) se abstraen a tokens genéricos ($c_0, c_1$). Esto crea una forma canónica invariante al nombre y a los valores constantes específicos, colapsando programas semánticamente equivalentes en la misma representación.
• Codificación: La secuencia tokenizada es procesada por un codificador Transformer (4 capas, 8 cabezales de atención) para producir un embedding de longitud fija de 256 dimensiones. Esta arquitectura fue elegida sobre las RNN (LSTM/GRU) por su capacidad superior para modelar dependencias de largo alcance en programas FHE desplegados.

B. Espacio de Acción

El espacio de acción está diseñado para manejar el gran número de reglas posibles y ubicaciones de aplicación.

• Política Jerárquica: En lugar de un espacio de acción plano (regla + ubicación), el agente utiliza una política jerárquica de dos pasos:
  1. Red de Selección de Reglas: Un MLP selecciona qué regla de reescritura aplicar de un conjunto de 84 reglas.
  2. Red de Selección de Ubicación: Un segundo MLP selecciona la ubicación específica (índice de coincidencia) en la IR donde se aplica la regla.
• Terminación: Una acción especial END permite al agente terminar el episodio de optimización.

C. Modelo de Recompensa

El agente es entrenado para minimizar una función de costo compuesta que equilibra el tiempo de ejecución y el crecimiento del ruido.

• Función de Costo: $Cost(e) = w_{ops} \cdot C_{ops}(e) + w_{depth} \cdot D_{circuit}(e) + w_{mult} \cdot D_{mult}(e)$.
  • $C_{ops}$: Suma ponderada de los costos de operación (Suma de Vectores: 1, Multiplicación de Vectores: 100, Rotación: 50, Operaciones Escalares: 250).
  • $D_{circuit}$: Profundidad del circuito.
  • $D_{mult}$: Profundidad multiplicativa (crítica para el crecimiento del ruido en BFV).
• Estructura de Recompensa:
  • Recompensa de Paso ($R_{step}$): Recompensa inmediata basada en el porcentaje de reducción de costo después de cada acción.
  • Recompensa Terminal ($R_{final}$): Una recompensa global basada en la reducción total de costo desde el estado inicial al final. Esto evita que el agente sea localmente codicioso y fomenta secuencias que produzcan circuitos globalmente óptimos.

D. Síntesis de Datos de Entrenamiento

Reconociendo la falta de conjuntos de datos de optimización FHE existentes, los autores sintetizaron un conjunto de datos de entrenamiento de 15,855 expresiones únicas utilizando un Modelo de Lenguaje Grande (LLM, Gemini 2.5 Flash).

• Guía del LLM: El LLM fue instruido con reglas sintácticas, ejemplos de reglas de reescritura y motivos de kernels del mundo real (por ejemplo, suma de matrices, diferencia al cuadrado) para generar patrones de código diversos y optimizables.
• Validación: Las expresiones generadas fueron filtradas por validez sintáctica, unicidad (vía tokenización ICI) y exclusión de programas de referencia para asegurar la generalización.

E. Algoritmo de Entrenamiento

El agente es entrenado utilizando el algoritmo de Optimización de Política Próxima (PPO) en una sola CPU multinúcleo durante 2 millones de pasos de tiempo.

3. Contribuciones Clave

1. Formulación: La primera formulación de la optimización de código FHE como un problema de toma de decisiones secuenciales resoluble mediante Aprendizaje por Refuerzo.
2. Marco CHEHAB RL: Un nuevo sistema que vectoriza automáticamente el código FHE tanto estructurado como no estructurado, minimizando simultáneamente el crecimiento del ruido y la latencia del circuito. Escala mejor que los enfoques de búsqueda combinatoria existentes.
3. Conjunto de Datos Sintetizado por LLM: La creación y validación de un conjunto de datos de entrenamiento grande y diverso generado por un LLM, demostrando que entrenar con tales datos produce agentes significativamente mejores que entrenar con programas aleatorios.
4. Ganancias de Rendimiento: Evidencia empírica que muestra que CHEHAB RL genera código que es 5.3× más rápido en ejecución, acumula 2.54× menos ruido y tiene un tiempo de compilación 27.9× más rápido que el compilador Coyote (medias geométricas).
5. Código Abierto: Lanzamiento del marco CHEHAB RL a la comunidad.

4. Resultados y Evaluación

Los autores evaluaron CHEHAB RL frente a Coyote (un compilador de vectorización FHE de vanguardia) utilizando una suite de benchmarks que incluye filtros de procesamiento de imágenes, kernels de álgebra lineal y polinomios irregulares generados aleatoriamente.

• Tiempo de Ejecución: El código generado por CHEHAB RL fue, en promedio, 5.3× más rápido que el de Coyote. En casos específicos (por ejemplo, Regresión Lineal 32), la aceleración alcanzó las 114×. Esto se atribuyó a menos operaciones y disposiciones de datos más simples encontradas por el agente de RL.
• Presupuesto de Ruido: Los circuitos generados consumieron 2.54× menos presupuesto de ruido en promedio. En benchmarks donde Coyote agotó el presupuesto de ruido y falló en la ejecución, CHEHAB RL produjo exitosamente circuitos válidos y ejecutables.
• Tiempo de Compilación: El proceso de compilación basado en RL fue 27.9× más rápido que el enfoque de búsqueda de Coyote. Aunque el agente de RL incurre en una sobrecarga para benchmarks pequeños (deb debido a los pasos de inferencia de la red neuronal), demostró una escalabilidad superior para circuitos más grandes donde la búsqueda combinatoria de Coyote se vuelve prohibitivamente lenta.
• Estudios de Ablación:
  • LLM vs. Datos Aleatorios: Los agentes entrenados con datos generados por LLM produjeron circuitos 13× más rápidos en benchmarks específicos comparado con aquellos entrenados con datos aleatorios, resaltando la importancia de los patrones de código realistas.
  • Diseño de Recompensa: La combinación de recompensas de paso y terminal superó a las recompensas de solo paso, evitando los óptimos locales.
  • Arquitectura: El codificador Transformer superó a las bases de GRU en el aprendizaje de representaciones de programas, y el espacio de acción jerárquico superó a los espacios de acción planos en el aprendizaje de velocidad y recompensa.

5. Significado y Reivindicaciones

El artículo afirma que CHEHAB RL representa un cambio de paradigma en la compilación de FHE. Al pasar de una búsqueda combinatoria impulsada por heurísticas a la optimización de políticas basada en el aprendizaje, el marco aborda los problemas de escalabilidad y de óptimos locales inherentes a los enfoques anteriores.

Los autores enfatizan que este es el primer uso de RL para resolver la optimización de FHE, navegando con éxito los compromisos entre la vectorización, la acumulación de ruido criptográfico y los costos de operación. El trabajo demuestra que los agentes de RL pueden aprender estrategias globales efectivas para la transformación de código que superan a las heurísticas ajustadas a mano y a los solvers de ILP, particularmente para el código no estructurado. La integración de LLMs para la síntesis de conjuntos de datos resalta aún más una vía para generar datos de entrenamiento de alta calidad en dominios donde tales datos son escasos.

El artículo concluye que CHEHAB RL permite la adopción práctica de FHE al reducir significativamente tanto el tiempo de ejecución de las computaciones cifradas como el tiempo requerido para compilarlas, todo ello preservando las garantías de seguridad del esquema criptográfico subyacente.