Resumo Técnico: CHEHAB RL: Aprendendo a Otimizar Computações de Criptografia Totalmente Homomórfica

1. Declaração do Problema

A Criptografia Totalmente Homomórfica (FHE) permite o processamento de dados criptografados sem a necessidade de descriptografia, preservando a privacidade. No entanto, a FHE permanece computacionalmente dispendiosa; por exemplo, uma multiplicação de texto cifrado pode ser $10^6$ vezes mais lenta que sua contraparte em texto simples. Essa lacuna de desempenho decorre da complexidade inerente das operações homomórficas, que envolvem multiplicações polinomiais de alto grau e um crescimento significativo de ruído.

Escrever código FHE eficiente é uma tarefa complexa e propensa a erros, que exige profundo conhecimento criptográfico. Embora compiladores tenham sido desenvolvidos para automatizar a geração de código FHE, as soluções existentes enfrentam limitações significativas:

• Código Estruturado vs. Não Estruturado: Compiladores como HECO e CHET focam em código estruturado (baseado em loops), enquanto outros como Coyote e Porcupine tentam lidar tanto com código estruturado quanto não estruturado (arbitrário, não baseado em loops).
• Complexidade do Espaço de Busca: Os compiladores de estado da arte atuais (ex: Coyote, Porcupine) enquadram a vetorização como um problema de busca combinatória. Eles dependem de heurísticas ajustadas manualmente, Programação Linear Inteira (ILP) ou síntese de programas para explorar um vasto espaço de design (ex: decidir quais subexpressões empacotar, layout de dados, inserção de rotação).
• Escalabilidade e Ótimos Locais: Neste espaço complexo, decisões locais têm efeitos não locais (ex: uma decisão de empacotamento que aumenta os custos de rotação futuros). Heurísticas simples frequentemente ficam presas em ótimos locais ruins, e a busca exaustiva torna-se computacionalmente intratável à medida que o tamanho do programa aumenta.

2. Metodologia

O artigo propõe o CHEHAB RL, um novo framework que utiliza Aprendizado por Reforço Profundo (Deep Reinforcement Learning - RL) para automatizar a otimização de código FHE. Em vez de busca combinatória, o problema é formulado como um processo de tomada de decisão sequencial onde um agente de RL aprende uma política para aplicar uma sequência de regras de reescrita.

2.1 Visão Geral do Framework

O CHEHAB RL é integrado ao compilador FHE CHEHAB. O fluxo de trabalho é o seguinte:

1. Entrada: Um programa escrito na Linguagem de Domínio Específico (DSL) do CHEHAB, que é reduzido a uma Representação Intermediária (IR) na forma de uma Árvore de Sintaxe Abstrata (AST) e um Grafo de Fluxo de Dados (DAG).
2. Otimização por RL: Um agente de RL percorre a IR, selecionando uma sequência de regras de reescrita de um Sistema de Reescrita de Termos (TRS) para vetorizar código escalar, simplificar expressões e minimizar ruído/latência.
3. Saída: A IR otimizada é usada para gerar código C++ visando o backend Microsoft SEAL (esquema BFV).

2.2 Componentes Técnicos Principais

A. Representação de Estado

Para representar a IR do programa para o agente de RL, os autores empregam a tokenização de Identificador e Invariante de Constante (ICI).

• Normalização: Variáveis são renomeadas (ex: $v_0, v_1$) e constantes (exceto 0 e 1) são abstraídas para tokens genéricos ($c_0, c_1$). Isso cria uma forma canônica invariante a nomes e valores de constantes específicos, colapsando programas semanticamente equivalentes na mesma representação.
• Codificação: A sequência tokenizada é processada por um codificador Transformer (4 camadas, 8 cabeças de atenção) para produzir um embedding de comprimento fixo de 256 dimensões. Esta arquitetura foi escolhida em vez de RNNs (LSTM/GRU) por sua capacidade superior de modelar dependências de longo alcance em programas FHE desenrolados.

B. Espaço de Ação

O espaço de ação é projetado para lidar com o grande número de regras e locais de aplicação.

• Política Hierárquica: Em vez de um espaço de ação plano (regra + localização), o agente usa uma política hierárquica de dois passos:
  1. Rede de Seleção de Regra: Um MLP seleciona qual regra de reescrita aplicar de um conjunto de 84 regras.
  2. Rede de Seleção de Localização: Um segundo MLP seleciona a localização específica (índice de correspondência) na IR onde a regra é aplicada.
• Terminação: Uma ação especial END permite que o agente encerre o episódio de otimização.

C. Modelo de Recompensa

O agente é treinado para minimizar uma função de custo composta que equilibra o tempo de execução e o crescimento do ruído.

• Função de Custo: $Cost(e) = w_{ops} \cdot C_{ops}(e) + w_{depth} \cdot D_{circuit}(e) + w_{mult} \cdot D_{mult}(e)$.
  • $C_{ops}$: Soma ponderada dos custos de operação (Soma Vetorial: 1, Multiplicação Vetorial: 100, Rotação: 50, Operações Escalares: 250).
  • $D_{circuit}$: Profundidade do circuito.
  • $D_{mult}$: Profundidade multiplicativa (crítica para o crescimento de ruído no BFV).
• Estrutura de Recompensa:
  • Recompensa de Passo ($R_{step}$): Recompensa imediata baseada na porcentagem de redução de custo após cada ação.
  • Recompensa Terminal ($R_{final}$): Uma recompensa global baseada na redução total de custo do estado inicial para o final. Isso evita que o agente seja localmente ganancioso e incentiva sequências que resultem em circuitos globalmente ótimos.

D. Síntese de Dados de Treinamento

Reconhecendo a falta de conjuntos de dados de otimização FHE existentes, os autores sintetizaram um conjunto de dados de treinamento de 15.855 expressões únicas usando um Modelo de Linguagem Grande (LLM, Gemini 2.5 Flash).

• Guia do LLM: O LLM foi instruído com regras de sintaxe, exemplos de regras de reescrita e motivos de kernels do mundo real (ex: adição de matrizes, diferença de quadrados) para gerar padrões de código diversos e otimizáveis.
• Validação: As expressões geradas foram filtradas por validade sintática, unicidade (via tokenização ICI) e exclusão de programas de benchmark para garantir a generalização.

E. Algoritmo de Treinamento

O agente é treinado usando o algoritmo de Otimização de Política Próxima (PPO) em uma única CPU multi-core durante 2 milhões de timesteps.

3. Principais Contribuições

1. Formulação: A primeira formulação da otimização de código FHE como um problema de tomada de decisão sequencial solucionável via Aprendizado por Reforço.
2. Framework CHEHAB RL: Um novo sistema que vetoriza automaticamente códigos FHE estruturados e não estruturados, minimizando simultaneamente o crescimento de ruído e a latência do circuito. Ele escala melhor do que as abordagens de busca combinatória existentes.
3. Dataset Sintetizado por LLM: A criação e validação de um dataset de treinamento grande e diverso gerado por um LLM, demonstrando que treinar em tais dados produz agentes significativamente melhores do que treinar em programas aleatórios.
4. Ganhos de Desempenho: Evidência empírica mostrando que o CHEHAB RL gera códigos que são 5,3× mais rápidos em execução, acumulam 2,54× menos ruído e possuem um tempo de compilação 27,9× mais rápido que o compilador Coyote de última geração (médias geométricas).
5. Código Aberto: Lançamento do framework CHEHAB RL para a comunidade.

4. Resultados e Avaliação

Os autores avaliaram o CHEHAB RL contra o Coyote (um compilador FHE de vetorização de última geração) em uma suíte de benchmarks incluindo filtros de processamento de imagem, kernels de álgebra linear e polinômios irregulares gerados aleatoriamente.

• Tempo de Execução: O CHEHAB RL gerou códigos que foram, em média, 5,3× mais rápidos que o Coyote. Em casos específicos (ex: Regressão Linear 32), o aumento de velocidade atingiu 114×. Isso foi atribuído a menos operações e layouts de dados mais simples encontrados pelo agente de RL.
• Orçamento de Ruído: Os circuitos gerados consumiram 2,54× menos orçamento de ruído em média. Em benchmarks onde o Coyote esgotou o orçamento de ruído e falhou na execução, o CHEHAB RL produziu circuitos válidos e executáveis com sucesso.
• Tempo de Compilação: O processo de compilação baseado em RL foi 27,9× mais rápido que a abordagem de busca do Coyote. Embora o agente de RL incorra em overhead para benchmarks pequenos (devido às etapas de inferência de rede neural), ele demonstrou escalabilidade superior para circuitos maiores, onde a busca combinatória do Coyote torna-se proibitivamente lenta.
• Estudos de Ablação:
  • LLM vs. Dados Aleatórios: Agentes treinados em dados gerados por LLM produziram circuitos 13× mais rápidos em benchmarks específicos comparados aos treinados em dados aleatórios, destacando a importância de padrões de código realistas.
  • Design de Recompensa: Combinar recompensas de passo e terminais superou as recompensas apenas de passo, evitando ótimos locais.
  • Arquitetura: O codificador Transformer superou as bases de GRU em aprender representações de programas, e o espaço de ação hierárquico superou os espaços de ação planos em aprendizado de velocidade e recompensa.

5. Significância e Alegações

O artigo afirma que o CHEHAB RL representa uma mudança de paradigma na compilação FHE. Ao mover-se de uma busca combinatória baseada em heurísticas para a otimização de política baseada em aprendizado, o framework aborda os problemas de escalabilidade e de ótimos locais inerentes às abordagens anteriores.

Os autores enfatizam que este é o primeiro uso de RL para resolver a otimização FHE, navegando com sucesso pelos compromissos entre vetorização, acumulação de ruído criptográfico e custos de operação. O trabalho demonstra que agentes de RL podem aprender estratégias globais eficazes para transformação de código que superam heurísticas ajustadas manualmente e solvers de ILP, particularmente para código não estruturado. A integração de LLMs para síntese de datasets destaca ainda mais um caminho para gerar dados de treinamento de alta qualidade em domínios onde tais dados são escassos.

O artigo conclui que o CHEHAB RL permite a adoção prática da FHE ao reduzir significativamente tanto o tempo de execução das computações criptografadas quanto o tempo necessário para compilá-las, tudo isso preservando as garantias de segurança do esquema criptográfico subjacente.