Riepilogo Tecnico: CHEHAB RL: Apprendere l'Ottimizzazione di Computazioni in Crittografia Totalmente Omomorfica

1. Definizione del Problema

La Crittografia Totalmente Omomorfica (FHE) consente la computazione su dati criptati senza decrittazione, preservandone la privacy. Tuttavia, la FHE rimane computazionalmente onerosa; ad esempio, una moltiplicazione di ciphertext può essere $10^6$ volte più lenta della sua controparte in chiaro. Questo divario di prestazioni deriva dalla complessità intrinseca delle operazioni omomorfiche, che comportano moltiplicazioni polinomiali di alto grado e una crescita significativa del rumore.

Scrivere codice FHE efficiente è un compito complesso e soggetto a errori che richiede una profonda competenza crittografica. Sebbene siano stati sviluppati compilatori per automatizzare la generazione di codice FHE, le soluzioni esistenti affrontano limitazioni significative:

• Codice Strutturato vs. Non Strutturato: Compilatori come HECO e CHET si concentrano sul codice strutturato (basato su loop), mentre altri come Coyote e Porcupine tentano di gestire sia il codice strutturato che quello non strutturato (arbitrario, non basato su loop).
• Complessità dello Spazio di Ricerca: Gli attuali compilatori allo stato dell'arte (ad es., Coyote, Porcupine) inquadrano la vettorizzazione come un problema di ricerca combinatoria. Essi si affidano a euristiche calibrate manualmente, Programmazione Lineare Intera (ILP) o sintesi di programmi per esplorare un vastissimo spazio di progettazione (ad es., decidere quali sottoespressioni impacchettare, il layout dei dati, l'inserimento di rotazioni).
• Scalabilità e Ottimi Locali: In questo spazio complesso, le decisioni locali hanno spesso effetti non locali (ad es., una decisione di impacchettamento che aumenta i costi di rotazione futuri). Le semplici euristiche spesso rimangono intrappolate in scarsi ottimi locali, e la ricerca esaustiva diventa computazionalmente intrattabile all'aumentare della dimensione del programma.

2. Metodologia

Il paper propone CHEHAB RL, un nuovo framework che sfrutta il Deep Reinforcement Learning (RL) per automatizzare l'ottimizzazione del codice FHE. Invece della ricerca combinatoria, il problema è formulato come un processo decisionale sequenziale in cui un agente RL apprende una policy per applicare una sequenza di regole di riscrittura.

2.1 Panoramica del Framework

CHEHAB RL è integrato nel compilatore FHE CHEHAB. Il flusso di lavoro è il seguente:

1. Input: Un programma scritto nel linguaggio specifico del dominio (DSL) di CHEHAB, che viene abbassato a una Rappresentazione Intermedia (IR) sotto forma di Albero di Sintassi Astratta (AST) e Grafo di Flusso di Dati (DAG).
2. Ottimizzazione RL: Un agente RL attraversa l'IR, selezionando una sequenza di regole di riscrittura da un Sistema di Riscrittura di Termini (TRS) per vettorizzare il codice scalare, semplificare le espressioni e minimizzare rumore/latenza.
3. Output: L'IR ottimizzata viene utilizzata per generare codice C++ che punta al backend Microsoft SEAL (schema BFV).

2.2 Componenti Tecnici Chiave

A. Rappresentazione dello Stato

Per rappresentare l'IR del programma per l'agente RL, gli autori impiegano la tokenizzazione Identifier and Constant Invariant (ICI).

• Normalizzazione: Le variabili vengono rinominate (ad es., $v_0, v_1$) e le costanti (eccetto 0 e 1) vengono astratte in token generici ($c_0, c_1$). Ciò crea una forma canonica invariante rispetto ai nomi e ai valori specifici delle costanti, collassando programmi semanticamente equivalenti nella stessa rappresentazione.
• Codifica: La sequenza tokenizzata è elaborata da un encoder Transformer (4 strati, 8 teste di attenzione) per produrre un embedding a lunghezza fissa di 256 dimensioni. Questa architettura è stata scelta rispetto alle RNN (LSTM/GRU) per la sua superiore capacità di modellare dipendenze a lungo raggio in programmi FHE srotolati.

B. Spazio delle Azioni

Lo spazio delle azioni è progettato per gestire l'elevato numero di regole e posizioni di applicazione.

• Policy Gerarchica: Invece di uno spazio di azioni piatto (regola + posizione), l'agente utilizza una policy gerarchica a due fasi:
  1. Rete di Selezione della Regola: Un MLP seleziona quale regola di riscrittura applicare da un insieme di 84 regole.
  2. Rete di Selezione della Posizione: Un secondo MLP seleziona la posizione specifica (indice di match) nell'IR in cui la regola viene applicata.
• Terminazione: Un'azione speciale END permette all'agente di terminare l'episodio di ottimizzazione.

C. Modello di Ricompensa

L'agente è addestrato per minimizzare una funzione di costo composita che bilancia il tempo di esecuzione e la crescita del rumore.

• Funzione di Costo: $Cost(e) = w_{ops} \cdot C_{ops}(e) + w_{depth} \cdot D_{circuit}(e) + w_{mult} \cdot D_{mult}(e)$.
  • $C_{ops}$: Somma pesata dei costi delle operazioni (Addizione Vettoriale: 1, Moltiplicazione Vettoriale: 100, Rotazione: 50, Operazioni Scalari: 250).
  • $D_{circuit}$: Profondità del circuito.
  • $D_{mult}$: Profondità moltiplicativa (critica per la crescita del rumore in BFV).
• Struttura della Ricompensa:
  • Ricompensa per Passo ($R_{step}$): Ricompensa immediata basata sulla percentuale di riduzione del costo dopo ogni azione.
  • Ricompensa Terminale ($R_{final}$): Una ricompensa globale basata sulla riduzione del costo totale dal primo all'ultimo stato. Questo impedisce all'agente di essere localmente avido e incoraggia sequenze che producono circuiti globalmente ottimali.

D. Sintesi dei Dati di Addestramento

Riconoscendo la mancanza di dataset esistenti per l'ottimizzazione FHE, gli autori hanno sintetizzato un dataset di addestramento di 15.855 espressioni uniche utilizzando un Large Language Model (LLM, Gemini 2.5 Flash).

• Guida dell'LLM: L'LLM è stato istruito con regole di sintassi, esempi di regole di riscrittura e motivi di kernel reali (ad es., addizione di matrici, differenza del quadrato) per generare pattern di codice diversificati e ottimizzabili.
• Validazione: Le espressioni generate sono state filtrate per validità sintattica, unicità (tramite tokenizzazione ICI) ed esclusione di programmi benchmark per garantire la generalizzazione.

E. Algoritmo di Addestramento

L'agente è addestrato utilizzando l'algoritmo Proximal Policy Optimization (PPO) su una singola CPU multi-core per 2 milioni di timesteps.

3. Contributi Chiave

1. Formulazione: La prima formulazione dell'ottimizzazione del codice FHE come problema decisionale sequenziale risolvibile tramite Reinforcement Learning.
2. Framework CHEHAB RL: Un nuovo sistema che vettorizza automaticamente sia il codice FHE strutturato che quello non strutturato, minimizzando contemporaneamente la crescita del rumore e la latenza del circuito. Scala meglio degli approcci di ricerca combinatoria esistenti.
3. Dataset Sintetizzato da LLM: La creazione e la validazione di un dataset di addestramento ampio e diversificato generato da un LLM, dimostrando che l'addestramento su tali dati produce agenti significativamente migliori rispetto all'addestramento su programmi casuali.
4. Guadagni di Prestazioni: Evidenza empirica che mostra che CHEHAB RL genera codice che è 5,3 volte più veloce nell'esecuzione, accumula 2,54 volte meno rumore e ha un tempo di compilazione 27,9 volte più veloce rispetto al compilatore state-of-the-art Coyote (medie geometriche).
5. Open Source: Rilascio del framework CHEHAB RL alla comunità.

4. Risultati e Valutazione

Gli autori hanno valutato CHEHAB RL rispetto a Coyote (un compilatore FHE vettorizzante allo stato dell'arte) su una suite di benchmark che include filtri di elaborazione delle immagini, kernel di algebra lineare e polinomi irregolari generati casualmente.

• Tempo di Esecuzione: CHEHAB RL ha generato codice che era, in media, 5,3 volte più veloce di Coyote. In casi specifici (ad es., Linear Regression 32), l'accelerazione ha raggiunto le 114 volte. Ciò è stato attribuito a meno operazioni e a layout dei dati più semplici trovati dall'agente RL.
• Budget di Rumore: I circuiti generati hanno consumato in media 2,54 volte meno budget di rumore. Nei benchmark in cui Coyote ha esaurito il budget di rumore e ha fallito l'esecuzione, CHEHAB RL ha prodotto con successo circuiti validi ed eseguibili.
• Tempo di Compilazione: Il processo di compilazione basato su RL è stato 27,9 volte più veloce dell'approccio di ricerca di Coyote. Sebbene l'agente RL comporti un overhead per i benchmark piccoli (a causa dei passaggi di inferenza della rete neurale), ha dimostrato una scalabilità superiore per circuiti più grandi dove la ricerca combinatoria di Coyote diventava proibitivamente lenta.
• Studi di Ablazione:
  • LLM vs. Dati Casuali: Gli agenti addestrati su dati generati da LLM hanno prodotto circuiti 13 volte più veloci su specifici benchmark rispetto a quelli addestrati su dati casuali, evidenziando l'importanza di pattern di codice realistici.
  • Design della Ricompensa: La combinazione di ricompense per passo e terminali ha superato le ricompense basate solo sul passo, prevenendo gli ottimi locali.
  • Architettura: L'encoder Transformer ha superato i baseline GRU nel apprendere le rappresentazioni dei programmi, e lo spazio delle azioni gerarchico ha superato gli spazi delle azioni piatti nell'apprendimento della velocità e della ricompensa.

5. Significato e Rivendicazioni

Il paper sostiene che CHEHAB RL rappresenti un cambio di paradigma nella compilazione FHE. Passando da una ricerca combinatoria guidata da euristiche a un'ottimizzazione di policy basata sull'apprendimento, il framework affronta i problemi di scalabilità e degli ottimi locali inerenti ai precedenti approcci.

Gli autori sottolineano che questo è il primo uso di RL per risolvere l'ottimizzazione FHE, navigando con successo tra i compromessi tra vettorizzazione, accumulo di rumore crittografico e costi operativi. Il lavoro dimostra che gli agenti RL possono apprendere strategie globali efficaci per la trasformazione del codice che superano le euristiche calibrate manualmente e i solver ILP, particolarmente per il codice non strutturato. L'integrazione di LLM per la sintesi dei dataset evidenzia ulteriormente un percorso per la generazione di dati di addestramento di alta qualità in domini in cui tali dati sono scarsi.

Il paper conclude che CHEHAB RL abilita l'adozione pratica della FHE riducendo significativamente sia il tempo di esecuzione delle computazioni criptate che il tempo necessario per compilarle, il tutto preservando le garanzie di sicurezza dello schema crittografico sottostante.