CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations

Fully Homomorphic Encryption (FHE) enables computations directly on encrypted data, but its high computational cost remains a significant barrier. Writing efficient FHE code is a complex task requiring cryptographic expertise, and finding the optimal sequence of program transformations is often intractable. In this paper, we propose CHEHAB RL, a novel framework that leverages deep reinforcement learning (RL) to automate FHE code optimization. Instead of relying on predefined heuristics or combinatorial search, our method trains an RL agent to learn an effective policy for applying a sequence of rewriting rules to automatically vectorize scalar FHE code while reducing instruction latency and noise growth. The proposed approach supports the optimization of both structured and unstructured code. To train the agent, we synthesize a diverse dataset of computations using a large language model (LLM). We integrate our proposed approach into the CHEHAB FHE compiler and evaluate it on a suite of benchmarks, comparing its performance against Coyote, a state-of-the-art vectorizing FHE compiler. The results show that our approach generates code that is $5.3\times$ faster in execution, accumulates $2.54\times$ less noise, while the compilation process itself is $27.9\times$ faster than Coyote (geometric means).

💡 Research Summary

The paper introduces CHEHAB RL, a novel framework that uses deep reinforcement learning (RL) to automate the optimization of Fully Homomorphic Encryption (FHE) programs. Writing efficient FHE code is notoriously difficult because it requires deep cryptographic expertise and the search for optimal sequences of program transformations is combinatorial and often intractable. Existing FHE compilers such as HECO, CHET, Coyote, and Porcupine either focus only on structured (loop‑based) code or rely on hand‑tuned heuristics, integer linear programming (ILP), or exhaustive program synthesis to explore the design space. These approaches suffer from scalability issues and can become trapped in poor local optima.

CHEHAB RL reframes FHE code optimization as a sequential decision‑making problem modeled as a Markov Decision Process. The RL agent’s action space is hierarchical: first it selects a rewriting rule (e.g., commutativity, common‑factor extraction, vectorization) and then chooses a concrete location in the intermediate representation (IR) where the rule should be applied. Because the rule set is large and each rule may be applicable at many positions, the authors design a multi‑discrete hierarchical policy consisting of a policy network (rule selection) and a position network (location selection). This architecture outperforms a flat action space and enables efficient exploration of the massive transformation space.

Running actual FHE code to obtain a reward is prohibitively slow, so the authors devise an analytical reward function that estimates execution cost, circuit depth, and multiplicative depth—key factors influencing both latency and noise growth. The reward combines weighted terms for multiplications, rotations, and additions, allowing fast, differentiable feedback during training while still reflecting the true performance objectives of FHE.

A major obstacle is the lack of publicly available datasets of optimizable FHE programs. To address this, the authors build a synthesis pipeline powered by a large language model (LLM) that generates 15,855 diverse FHE expressions, ranging from simple arithmetic to neural‑network layers. Experiments show that training the RL agent on this LLM‑generated corpus yields significantly better policies than training on randomly generated programs.

The trained agent is integrated into the CHEHAB compiler and evaluated against Coyote on a benchmark suite of twelve real‑world FHE kernels (including neural‑network inference, matrix multiplication, and statistical functions). The evaluation measures three metrics: execution time, accumulated noise, and compilation time. CHEHAB RL achieves a geometric‑mean speed‑up of 5.3× in execution, reduces noise by 2.54×, and speeds up compilation by 27.9× compared to Coyote. Importantly, the approach works for both structured and unstructured code, demonstrating robust vectorization capabilities even on code without loops.

The contributions are: (1) formalizing FHE code optimization as a sequential decision problem and proposing an RL‑based solution; (2) introducing CHEHAB RL, the first RL system to automatically vectorize and optimize FHE code for both structured and unstructured programs; (3) creating a large LLM‑synthesized dataset for training; (4) empirically showing substantial performance gains over the state‑of‑the‑art compiler; and (5) releasing the framework to the community.

Limitations include specialization to the BFV scheme and specific parameter settings, which may affect generalization to other schemes like CKKS. The analytical reward, while efficient, is an approximation and may not capture all hardware‑level nuances. Future work is suggested in extending support to multiple FHE schemes, incorporating hardware‑feedback into the reward, applying meta‑learning for policy transfer, and scaling training with distributed RL infrastructure.