CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations

TL;DR

CHEHAB RL introduces a deep reinforcement learning framework to automate and optimize Fully Homomorphic Encryption code, significantly improving execution speed and noise reduction over existing methods.

Contribution

This work presents the first RL-based approach to automatically optimize FHE code, reducing reliance on heuristics and enabling efficient vectorization of scalar FHE programs.

Findings

Code generated is 5.3x faster in execution

Noise accumulation is reduced by 2.54x

Compilation time is 27.9x faster than Coyote

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