CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations

TL;DR

This work reframes FHE code optimization as sequential decision-making and introduces CHEHAB RL, an RL-guided term rewriting system that optimizes both structured and unstructured FHE code. It combines a Transformer-based state representation with a hierarchical action space and a dual-component reward (immediate and terminal) to efficiently learn rewrite policies that minimize operation counts, circuit depth, multiplicative depth, and noise. Training relies on a large, LLM-generated corpus of CHEHAB IR expressions to capture realistic motifs, paired with PPO for policy optimization; the approach is integrated into the CHEHAB compiler and evaluated against Coyote, achieving 5.3x faster execution, 2.54x less noise, and 27.9x faster compilation on a suite of benchmarks. The results demonstrate improved scalability and code quality, as well as insights from ablations on data generation, tokenization, and policy design, with a public release of CHEHAB RL. Overall, CHEHAB RL advances practical FHE deployment by delivering high-quality, scalable vectorization and noise-aware optimizations integrated into a full compilation pipeline.

Abstract

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).

Figures (13)

• Figure 1: Example of an unstructured scalar code.
• Figure 2: Vectorization approaches.
• Figure 3: CHEHAB RL overview (our contribution: RL TRS).
• Figure 4: Architecture of the Policy Network
• Figure 5: Semi-log plot of execution time of the benchmarks, comparing code generated using our compiler and Coyote.