PaperlayerView exact revision on arXiv

Source-linked AI summary

HEAT: Faster Fully Homomorphic Inference via Approximations-Weights Co-Adaptation

Alessandro Zirilli, Davide Marincione, Evgenios M. Kornaropoulos, Giuseppe Ateniese, Emanuele Rodolà

arXiv:2609.01730v1cs.CRcs.AI

TL;DR

FHE enables private language-model inference but is slowed by bounded multiplicative depth, iterative nonlinear approximations, and bootstrapping. HEAT fine-tunes per-site iteration counts jointly with model weights, and encrypted GPT-2 decoding reduces iterations, bootstraps, and latency while improving agreement with plaintext outputs.

  • Problem

    FHE inference is slow because bounded multiplicative depth makes iterative nonlinear approximations consume levels and trigger costly bootstrapping, while iteration counts are usually fixed across sites.

  • Method

    HEAT fine-tunes model weights jointly with learnable per-site iteration counts, using task loss and depth regularization to adapt the deployed circuit to approximation errors.

  • Results

    3.1× fewer iterations, 1.6× fewer bootstraps, and 1.4× lower end-to-end latency were achieved on encrypted GPT-2 while decode agreement improved over the calibrated baseline.

  • Takeaways & Limitations

    HEAT reduces encrypted GPT-2 inference cost while remaining more faithful to its plaintext reference, without changing the parent architecture.

  • Takeaways & Limitations

    Evidence covers one 124M backbone, one cryptographic parameterization, and one modality, so other tasks or modalities may shift where HEAT reduces latency.

Abstract

from arXiv · show

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.

1 Introduction

FHE protects outsourced language-model inference by keeping prompts and outputs encrypted, but limited ciphertext operations make inference extremely slow. HEAT fine-tunes iteration counts jointly with model weights to reduce this cost without changing the architecture.

  • FHE lets a server evaluate a language model on an encrypted prompt while only the client can recover the output.
  • ≈59% of token wall time and 71% of bootstrap time belong to the two iterative nonlinearities, where HEAT concentrates its savings.
  • Finite multiplicative depth forces nonlinearities into iterative approximations whose multiplications consume the budget and trigger costly bootstrapping.
  • HEAT learns per-site iteration counts jointly with model weights, preserving the parent architecture while reducing approximation depth.
  • 3.1× fewer total iterations and ≈54 s/tok to ≈38 s/tok decoding latency are reported for encrypted GPT-2.

2 Related Works

Encrypted inference systems are shaped by both nonlinear approximation and ciphertext packing. HEAT optimizes the deployed circuit’s multiplicative depth jointly with model weights while retaining existing packing and approximation choices.

  • Encrypted inference cost depends on nonlinear approximations and tensor packing, which determines rotations, multiplications, and ciphertext count.
  • HEAT builds on Cachemir’s autoregressive decoding layout, THOR’s nonlinear approximations, and an encrypted arg max.
  • Unlike architecture-redesign approaches, HEAT learns the deployed circuit’s multiplicative depth jointly with the model weights.

3 Method

HEAT makes approximation depth learnable at each nonlinear site and trains it against the task objective, allowing weights and solvers to co-adapt. A staged fine-tuning procedure then freezes the learned counts into a fixed encrypted circuit.

  • Method: Iterative approximation counts are tailored to each site according to their effect on global task loss rather than a universal fidelity target.
  • Training dynamics: The training dynamics show stable validation loss while the depth term rises, learned counts descend, and cooldown disables the term.
  • Differentiable iteration selection: A learnable distribution over iteration counts makes the approximation differentiable and lets its parameters co-adapt with model weights.
  • Iteration count penalty: The KL-based iteration penalty favors shallow approximations while the HEAT loss balances task accuracy against iteration reduction.
  • Fine-tuning schedule: HEAT uses warm-up, joint optimization, and cooldown phases before freezing each site to its learned mode for deployment.
  • Fine-tuning schedule: The deployed count is shared across inputs because data-dependent execution would leak information through latency.

4 Experiments

HEAT is evaluated on encrypted GPT-2 decoding against a calibrated baseline using end-to-end encrypted runs. It substantially reduces solver iterations, bootstraps, and latency while improving agreement with plaintext logits.

  • Setup: The experiments compare a GPT-2 (124M) HEAT checkpoint with a pretrained model using per-site iteration counts calibrated to plaintext relative error 10^-4.The pipeline is fine-tuned on OpenWebText and evaluated with CKKS at 128-bit security.
  • Latency: 228 solver iterations per forward versus 712 for the calibrated baseline yields a 3.1× reduction, with iterations retained where precision demands them.The reduction is non-uniform across approximation sites.
  • Latency: 326 bootstraps per token versus 513 for the baseline yields a 1.6× reduction.The reduction is concentrated in the two iterative operators.
  • Latency: 37.6 s/tok versus 54.3 s/tok for the calibrated baseline and approximately 68 s/tok for EncryptedLLM demonstrates lower end-to-end latency.The measurements include data-movement overheads and are taken during actual encrypted runs.
  • Decode agreement: HEAT improves pooled top-1 agreement to 82.9% versus 74.0% for the baseline, while final-position KL is 6.6× lower.Across the sequence, HEAT remains flatter than the baseline in both KL and top-1 agreement.

5 Conclusions

HEAT makes approximation iteration counts learnable and jointly optimizes them with model weights. On encrypted GPT-2, it reduces computation and latency while remaining more faithful to plaintext, but the evidence is limited in scope.

  • Conclusion: HEAT turns the precision-latency trade-off of encrypted inference into a training objective by jointly optimizing iteration counts and model weights.The method co-adapts approximations and weights during fine-tuning.
  • Conclusion: On encrypted GPT-2, HEAT cuts forward-pass iterations by 3.1× and end-to-end latency by 1.4× while improving faithfulness to the plaintext reference.These results summarize the reported encrypted decoding evaluation.
  • Limitations: The evidence covers one 124M backbone, one cryptographic parameterization, and one modality.Other modalities and tasks may place latency in different sites and change which reductions HEAT achieves.

A Extended Related Works

The related-work discussion distinguishes interactive and non-interactive privacy-preserving inference and positions HEAT as a training-side method for non-interactive FHE. Prior approaches reduce cost through protocols, architectures, or adaptive-computation ideas, whereas HEAT keeps the deployed schedule input-oblivious.

  • HEAT positioning: HEAT targets non-interactive FHE from the training side without replacing model components.Its learned iteration schedule is optimized before deployment.
  • Interactive private inference: Interactive protocols split computation between client and server, but transformer inference can require tens of gigabytes of data transfer.MPC evaluates nonlinearities exactly while paying through repeated interaction and communication.
  • Non-interactive FHE inference: Non-interactive FHE sends one encrypted prompt for server-side evaluation, after which only the client decrypts the answer.The setting includes prior encrypted language-model systems such as EncryptedLLM.
  • Reducing FHE cost before deployment: Architecture-level methods replace expensive operators or distill models into encryption-friendly forms, changing the computed model and requiring retraining.Examples include Gaussian-kernel attention, polynomial Softmax variants, Powerformer, and Peregrine.
  • Adaptive computation, made oblivious: HEAT uses adaptive computation only during training and freezes each site to one count before deployment, keeping execution constant across inputs.This avoids encrypted branching and latency leakage from input-dependent iteration counts.

B Background

FHE enables encrypted model evaluation but imposes a finite multiplicative-depth budget. Since nonlinear functions must be implemented through iterative approximations, additional iterations improve accuracy while increasing depth consumption and bootstrapping cost.

  • Background: The cryptographic background covers encrypted inference, the threat model, and the numerical methods assumed by the main method.The server is honest-but-curious and processes encrypted prompts without seeing plaintext tokens or intermediate values.
  • Fully homomorphic encryption: CKKS supports approximate encrypted arithmetic, with ciphertext operations consisting natively of additions, multiplications, and rotations.Encrypted addition and multiplication decrypt approximately to the corresponding plaintext operations.
  • Computational budget: A finite multiplicative-depth budget is exhausted by successive multiplications, after which bootstrapping refreshes the ciphertext and restores the budget.Bootstrapping costs two to three orders of magnitude more than other primitives.
  • Iterative approximations: Nonlinear functions such as division, inverse square root, and activations are evaluated with iterative methods built from the limited native operations.Goldschmidt and Newton-Raphson recurrences are examples of these approximations.
  • Iterative approximations: For quadratically convergent methods, accuracy improves doubly exponentially with iteration count while depth grows linearly.Each additional iteration adds a fixed number of multiplications and a proportional share of refreshes.

C Implementation Details

The implementation fixes CKKS cryptographic parameters, benchmarks primitive costs, and makes learned halting distributions memory-efficient and differentiable. Training constrains approximation inputs during optimization and exports each site's modal iteration count.

  • Cryptographic settings: CKKS runs at 128-bit security with N=2^16, 215 slots, 28 RNS limbs, and a 53-bit uniform rescaling-prime tower.The scaling factor is Δ=2^53, with a 60-bit first modulus.
  • Primitive costs and precision: Bootstrap dominates other CKKS primitives by two to three orders of magnitude, motivating latency reductions through avoided refreshes.Primitive timings use 100 isolated repetitions on one A100-64GB.
  • Memory-efficient depth gradients: The backward pass recomputes solver states from saved inputs, keeping activation memory constant in the iteration count while preserving exact truncated-iteration gradients.Forward-mode tangents avoid a separate reverse sweep over intermediate states.
  • Halting distributions: Sequential halting logits define a valid per-site distribution, with floor states removed and the deepest state absorbing remaining probability mass.Gradients reach the logits through the halting distribution, expected approximation, and exported mode.
  • Initialization and schedule: Training uses calibrated-domain clamps with quadratic push-back gradients, while deployment exports each site's modal count and omits the clamps.The learned distribution may settle above the calibrated count when task demands require it.

D The Approximated Operators

The model approximates nonlinear operators with CKKS-compatible polynomial and iterative circuits, while HEAT learns halting depths for the iterative solver families. Softmax refinement retains structured depth variation across layers, whereas GELU and arg max use fixed schedules.

  • Iterative primitives: Goldschmidt reciprocal and Newton inverse square root converge quadratically, and each learned halting site wraps one of these iterative loops.Goldschmidt squares the reciprocal error at each iteration; Newton likewise converges quadratically.
  • Softmax: Softmax uses a midpoint shift, two scaling factors, and repeated square-and-normalize passes with Goldschmidt reciprocals of updated sums.The refinement loop is the Softmax halting family whose learned depth varies across layers.
  • LayerNorm: LayerNorm inverse square root uses a degree-(3,1) Remez rational initializer, with division implemented through a Goldschmidt reciprocal.The rational approximation is fitted over the calibrated variance domain.
  • GELU: GELU is a fixed-depth composition of two low-degree polynomials, so it has no halting distribution and contributes only a range guard to HEAT.The deployment evaluation uses the Paterson–Stockmeyer algorithm.
  • Encrypted arg max: Encrypted arg max uses a fixed five-round sharpening schedule shared across checkpoints and contributes approximately 9.3 s/token to every token.Its schedule is frozen because arg max is not a learned halting site.

E Additional Results

HEAT learns nonuniform solver depths across GPT-2's approximation sites, reducing total iteration counts while preserving deeper computation where approximation demands it. Softmax refinement is the main source of cross-layer variation.

  • Per-site iteration counts: LayerNorm families collapse to uniform shallow counts, while Softmax refinement retains 8 iterations at block 0 and 4–7 in later blocks.Newton increases to 2 iterations at blocks 0–1 despite calibrated counts of 1 there.
  • Per-site iteration counts: Softmax refinement separates early in the depth descent, with each layer settling at the depth its approximation demands rather than following a global schedule.Figure 4 traces this structured allocation across layers.
  • Halting laws: Expected learned depths are 3.44 ± 0.05, 1.79 ± 0.02, 2.72 ± 0.05, and 6.29 ± 0.90 for the four families in order.Softmax refinement's across-site spread is approximately 18× larger than any other family's.

F Societal Impacts

FHE-based inference may reduce exposure of sensitive user data, but combining strong cryptographic privacy with capable models raises deployment and societal concerns.

  • More capable models with stronger cryptographic privacy may bring both societal benefits and unintended consequences.
  • FHE can reduce exposure of sensitive user data through privacy-preserving inference.
  • Practical deployments may provide different privacy and security properties than FHE protocols demonstrated in controlled settings.Additional components, assumptions, and operational choices can affect those properties.
  • Naively adding FHE to an existing inference pipeline does not necessarily yield an end-to-end private system.
Loading 2609.01730v1…