Search papers, labs, and topics across Lattice.
This paper introduces Homomorphic Encryption-Aware Training (HEAT), a novel fine-tuning method that allows for learnable iteration counts for nonlinearity approximations in fully homomorphic encryption (FHE) settings. By co-adapting these iteration counts with model weights during training, HEAT significantly reduces the number of required iterations and bootstraps, leading to improved efficiency in encrypted language model inference. The results show a 3.1x reduction in iterations, a 1.6x reduction in bootstraps, and a 1.4x decrease in end-to-end latency while enhancing decode agreement over existing methods.
Tailoring iteration counts for nonlinearity approximations can cut inference latency in encrypted language models by over 40%.
Fully homomorphic encryption (FHE) allows a server to run a language model directly on encrypted user prompts, but current approaches remain prohibitively slow. Ciphertexts natively support only addition, multiplication, and rotation, and multiplications may be composed only to a bounded depth before a costly bootstrapping operation is needed to continue. Every nonlinearity must therefore be approximated by an iterative method, and each iteration uses multiplications. A higher iteration count buys precision but exhausts the available depth faster and triggers more bootstraps, which dominate latency. Existing approaches fix the iteration counts uniformly across the model rather than tailoring them to each site's error tolerance. We introduce Homomorphic Encryption-Aware Training (HEAT), a fine-tuning method that makes the per-nonlinearity iteration counts learnable, enabling them and the model weights to co-adapt during training. HEAT optimizes iterations with respect to the task objective, allowing the model to adapt to approximation errors encountered during inference without architectural changes or retraining from scratch. On encrypted GPT-2 decoding, HEAT reduces iterations by $3.1\times$, bootstraps by $1.6\times$, and end-to-end latency by $1.4\times$, while improving decode agreement over the calibrated baseline.
