🤖 machine learning

CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations

This paper introduces CHEHAB RL, a deep reinforcement learning framework that automates the optimization of Fully Homomorphic Encryption code by learning to apply rewriting rules for vectorization, resulting in significantly faster execution, reduced noise growth, and much quicker compilation compared to state-of-the-art methods.

Original authors: Bilel Sefsaf, Abderraouf Dandani, Abdessamed Seddiki, Arab Mohammed, Eduardo Chielle, Michail Maniatakos, Riyadh Baghdadi

Published 2026-01-28

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Bilel Sefsaf, Abderraouf Dandani, Abdessamed Seddiki, Arab Mohammed, Eduardo Chielle, Michail Maniatakos, Riyadh Baghdadi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Locked Box" Problem

Imagine you have a very valuable secret, like your bank password or a medical record. You want to send it to a stranger (a cloud server) to do some math on it, but you are terrified they will peek inside and steal your secret.

Fully Homomorphic Encryption (FHE) is a magical solution to this problem. It's like putting your data in a locked, transparent box. You can give this box to the stranger, and they can shake it, hit it, and mix the contents around (perform calculations) without ever opening the lock or seeing what's inside. When they are done, they give the box back to you, and you unlock it to see the result.

The Catch: This magic is incredibly slow and messy. Doing math on a locked box takes thousands of times longer than doing it on regular data. It's like trying to bake a cake while wearing thick, clumsy oven mitts and a blindfold.

The Problem: Writing the Recipe is Hard

To make this "locked box" math faster, experts have to write very specific, low-level code. It's like trying to write a recipe for a cake while wearing those oven mitts.

Too Complex: You have to decide exactly how to pack the ingredients (data) into the box so they don't get mixed up.
Too Slow: If you pack them wrong, the math takes forever, or the "noise" (static) in the box gets so loud that the final result is garbage.
The Search Trap: Finding the perfect way to pack the data is like looking for a needle in a haystack. Existing tools try to find the needle by checking every single spot one by one (a "brute force" search). This takes a long time and often gets stuck in a "good enough" spot that isn't actually the best.

The Solution: CHEHAB RL (The "Smart Apprentice")

The authors of this paper built a new tool called CHEHAB RL. Instead of trying to check every single possibility, they taught a computer how to learn using a technique called Reinforcement Learning (RL).

Think of it like training a smart apprentice chef:

The Goal: The apprentice's job is to rearrange the ingredients in the locked box so the math happens as fast as possible with the least amount of "noise."
The Training: Instead of giving the apprentice a rulebook, they let the apprentice try millions of different ways to rearrange the ingredients.
- If a rearrangement makes the math faster, the apprentice gets a "treat" (a reward).
- If it makes things slower or noisier, they get a "frown" (a penalty).
The Dataset: To teach the apprentice, the authors didn't just use random junk. They used a Large Language Model (LLM)—a super-smart AI trained on real-world code—to generate thousands of realistic "practice recipes." This is like giving the apprentice practice on real cakes rather than random piles of flour and eggs.

How It Works: The "Rewriting" Game

The core of the system is a game of Lego.

You have a structure built from Lego bricks (your encrypted code).
You have a set of rules (like "swap these two bricks" or "combine these three into one").
The RL Agent (the apprentice) looks at the structure and decides: "If I apply Rule A here, and then Rule B there, will the whole thing become more stable and faster?"

The agent learns a policy (a strategy) to apply these rules in the perfect order. It doesn't just look at one step; it looks at the whole picture to find the best path.

The Results: Speed, Silence, and Efficiency

The paper compares their "Smart Apprentice" (CHEHAB RL) against the current best tool, called Coyote. Here is what happened:

Execution Speed (The Cake Baking): The code generated by CHEHAB RL ran 5.3 times faster than Coyote. It was like switching from a slow, manual oven to a high-speed convection oven.
Noise (The Static): The "noise" in the locked box grew 2.54 times less with CHEHAB RL. This means the final result is much clearer and less likely to be ruined.
Compilation Time (The Recipe Writing): Perhaps the most surprising result: It took 27.9 times less time for CHEHAB RL to write the optimized code than it took Coyote to search for it. The apprentice learned the strategy quickly, whereas Coyote was still frantically checking the haystack.

Why This Matters

The paper claims that this is the first time Reinforcement Learning has been used to solve this specific problem of optimizing FHE code.

It handles messy code: Unlike older tools that only work on neat, loop-based code, this system can handle messy, unstructured code (like a tangled ball of yarn).
It scales: As the problems get bigger, the "Smart Apprentice" gets better at handling them, while the old "brute force" search methods get slower and slower.

Summary Analogy

Imagine you are trying to organize a massive library of books that are all locked in glass cases.

Old Way (Coyote): You hire a team of librarians who try every possible arrangement of books to see which one is fastest to find. They get tired, take a long time, and often pick a "good" arrangement that isn't the best.
New Way (CHEHAB RL): You hire a genius librarian who watches thousands of videos of other libraries being organized. They learn the patterns of what works. When you give them a new library, they instantly know the best way to organize it, doing it in a fraction of the time and making the books much easier to find.

The paper concludes that by using this "learning" approach, we can make the "locked box" technology much more practical for real-world use, allowing us to do complex math on private data without waiting days for the result.

Technical Summary: CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations

1. Problem Statement

Fully Homomorphic Encryption (FHE) enables computation on encrypted data without decryption, preserving privacy. However, FHE remains computationally expensive; for instance, a ciphertext multiplication can be $10^6$ times slower than its plaintext counterpart. This performance gap stems from the inherent complexity of homomorphic operations, which involve high-degree polynomial multiplications and significant noise growth.

Writing efficient FHE code is a complex, error-prone task requiring deep cryptographic expertise. While compilers have been developed to automate FHE code generation, existing solutions face significant limitations:

Structured vs. Unstructured Code: Compilers like HECO and CHET focus on structured (loop-based) code, while others like Coyote and Porcupine attempt to handle both structured and unstructured (arbitrary, non-loop-based) code.
Search Space Complexity: Current state-of-the-art compilers (e.g., Coyote, Porcupine) frame vectorization as a combinatorial search problem. They rely on hand-tuned heuristics, Integer Linear Programming (ILP), or program synthesis to explore a vast design space (e.g., deciding which subexpressions to pack, data layout, rotation insertion).
Scalability and Local Optima: In this complex space, local decisions often have non-local effects (e.g., a packing decision increasing future rotation costs). Simple heuristics often get trapped in poor local optima, and exhaustive search becomes computationally intractable as program size increases.

2. Methodology

The paper proposes CHEHAB RL, a novel framework that leverages Deep Reinforcement Learning (RL) to automate FHE code optimization. Instead of combinatorial search, the problem is formulated as a sequential decision-making process where an RL agent learns a policy to apply a sequence of rewriting rules.

2.1 Framework Overview

CHEHAB RL is integrated into the CHEHAB FHE compiler. The workflow is as follows:

Input: A program written in the CHEHAB Domain-Specific Language (DSL), which is lowered to an Intermediate Representation (IR) in the form of an Abstract Syntax Tree (AST) and Dataflow Graph (DAG).
RL Optimization: An RL agent traverses the IR, selecting a sequence of rewriting rules from a Term Rewriting System (TRS) to vectorize scalar code, simplify expressions, and minimize noise/latency.
Output: The optimized IR is used to generate C++ code targeting the Microsoft SEAL backend (BFV scheme).

2.2 Key Technical Components

A. State Representation

To represent the program IR for the RL agent, the authors employ Identifier and Constant Invariant (ICI) tokenization.

Normalization: Variables are renamed (e.g., $v_0, v_1$ ) and constants (except 0 and 1) are abstracted to generic tokens ( $c_0, c_1$ ). This creates a canonical form invariant to naming and specific constant values, collapsing semantically equivalent programs into the same representation.
Encoding: The tokenized sequence is processed by a Transformer encoder (4 layers, 8 attention heads) to produce a fixed-length 256-dimensional embedding. This architecture was chosen over RNNs (LSTM/GRU) for its superior ability to model long-range dependencies in unrolled FHE programs.

B. Action Space

The action space is designed to handle the large number of possible rules and application locations.

Hierarchical Policy: Instead of a flat action space (rule + location), the agent uses a two-step hierarchical policy:
1. Rule Selection Network: An MLP selects which rewriting rule to apply from a set of 84 rules.
2. Location Selection Network: A second MLP selects the specific location (match index) in the IR where the rule is applied.
Termination: A special END action allows the agent to terminate the optimization episode.

C. Reward Model

The agent is trained to minimize a composite cost function that balances execution time and noise growth.

Cost Function: $Cost(e) = w_{ops} \cdot C_{ops}(e) + w_{depth} \cdot D_{circuit}(e) + w_{mult} \cdot D_{mult}(e)$ $C os t (e) = w_{o p s} \cdot C_{o p s} (e) + w_{d e pt h} \cdot D_{c i r c u i t} (e) + w_{m u l t} \cdot D_{m u l t} (e)$ .
- $C_{ops}$ : Weighted sum of operation costs (Vector Add: 1, Vector Mul: 100, Rotation: 50, Scalar Ops: 250).
- $D_{circuit}$ : Circuit depth.
- $D_{mult}$ : Multiplicative depth (critical for noise growth in BFV).
Reward Structure:
- Step Reward ( $R_{step}$ ): Immediate reward based on the percentage reduction in cost after each action.
- Terminal Reward ( $R_{final}$ ): A global reward based on the total cost reduction from the initial to the final state. This prevents the agent from being locally greedy and encourages sequences that yield globally optimal circuits.

D. Training Data Synthesis

Recognizing the lack of existing FHE optimization datasets, the authors synthesized a training dataset of 15,855 unique expressions using a Large Language Model (LLM, Gemini 2.5 Flash).

LLM Guidance: The LLM was prompted with syntax rules, rewrite rule examples, and real-world kernel motifs (e.g., matrix addition, squared difference) to generate diverse, optimizable code patterns.
Validation: Generated expressions were filtered for syntactic validity, uniqueness (via ICI tokenization), and exclusion of benchmark programs to ensure generalization.

E. Training Algorithm

The agent is trained using the Proximal Policy Optimization (PPO) algorithm on a single multi-core CPU for 2 million timesteps.

3. Key Contributions

Formulation: The first formulation of FHE code optimization as a sequential decision-making problem solvable via Reinforcement Learning.
CHEHAB RL Framework: A novel system that automatically vectorizes both structured and unstructured FHE code while minimizing noise growth and circuit latency. It scales better than existing combinatorial search approaches.
LLM-Synthesized Dataset: The creation and validation of a large, diverse training dataset generated by an LLM, demonstrating that training on such data yields significantly better agents than training on random programs.
Performance Gains: Empirical evidence showing that CHEHAB RL generates code that is 5.3× faster in execution, accumulates 2.54× less noise, and has a compilation time 27.9× faster than the state-of-the-art Coyote compiler (geometric means).
Open Source: Release of the CHEHAB RL framework to the community.

4. Results and Evaluation

The authors evaluated CHEHAB RL against Coyote (a state-of-the-art vectorizing FHE compiler) on a suite of benchmarks including image processing filters, linear algebra kernels, and randomly generated irregular polynomials.

Execution Time: CHEHAB RL generated code that was, on average, 5.3× faster than Coyote. In specific cases (e.g., Linear Regression 32), the speedup reached 114×. This was attributed to fewer operations and simpler data layouts found by the RL agent.
Noise Budget: The generated circuits consumed 2.54× less noise budget on average. In benchmarks where Coyote exhausted the noise budget and failed to execute, CHEHAB RL successfully produced valid, runnable circuits.
Compilation Time: The RL-based compilation process was 27.9× faster than Coyote's search-based approach. While the RL agent incurred overhead for small benchmarks (due to neural network inference steps), it demonstrated superior scalability for larger circuits where Coyote's combinatorial search became prohibitively slow.
Ablation Studies:
- LLM vs. Random Data: Agents trained on LLM-generated data produced circuits 13× faster on specific benchmarks compared to those trained on random data, highlighting the importance of realistic code patterns.
- Reward Design: Combining step and terminal rewards outperformed step-only rewards, preventing local optima.
- Architecture: The Transformer encoder outperformed GRU baselines in learning program representations, and the hierarchical action space outperformed flat action spaces in learning speed and reward.

5. Significance and Claims

The paper claims that CHEHAB RL represents a paradigm shift in FHE compilation. By moving from heuristic-driven combinatorial search to learning-based policy optimization, the framework addresses the scalability and local-optima issues inherent in previous approaches.

The authors emphasize that this is the first use of RL to solve FHE optimization, successfully navigating the trade-offs between vectorization, cryptographic noise accumulation, and operation costs. The work demonstrates that RL agents can learn effective global strategies for code transformation that outperform hand-tuned heuristics and ILP solvers, particularly for unstructured code. The integration of LLMs for dataset synthesis further highlights a pathway for generating high-quality training data in domains where such data is scarce.

The paper concludes that CHEHAB RL enables the practical adoption of FHE by significantly reducing both the runtime of encrypted computations and the time required to compile them, all while preserving the security guarantees of the underlying cryptographic scheme.