A training tweak halves cryptographic cost: HEAT jointly learns per-site solver iterations with model weights so encrypted GPT‑2 executes 3.1× fewer iterative steps (1.6× fewer bootstraps) and reduces end‑to‑end token latency from ≈54s to ≈38s while improving fidelity to its plaintext counterpart.

Main Finding

HEAT (Homomorphic Encryption-Aware Training) turns the multiplicative-depth / approximation trade-off in fully homomorphic encryption (FHE) inference into a trainable objective: it makes the per-site iteration counts of iterative approximations learnable and co-trains them with model weights. On encrypted GPT-2 (124M) this reduces solver iterations 3.1×, bootstraps per token 1.6×, and end-to-end decoding latency 1.4× (from ≈54.3 s/tok to ≈37.6 s/tok), while improving agreement with the plaintext model (top-1 agreement pooled over tokens rises from 74.0% to 82.9%). The method requires no architectural changes to the model.

Key Points

• Problem: FHE supports only additions/rotations/multiplications and multiplicative depth is limited; common nonlinearities (LayerNorm, Softmax, activations) are implemented by iterative polynomial/solver approximations whose iteration counts trade accuracy for depth (and thus bootstrapping cost). Prior work used a uniform (worst-case) iteration count across all sites, causing overprovisioning and many unnecessary bootstraps.
• HEAT idea: make the iteration count at every approximation site a learnable parameter and co-adapt it with weights so the model absorbs approximation error where possible and keeps precision where needed.
  • Differentiable iteration selection: attach a learnable discrete distribution p = (p0,...,pn) over iteration counts and use the expected approximation e_f(x) = Σ_i p_i F_i(x) (Fi are successive solver states). This makes the iteration policy differentiable.
  • Regularization: add a KL divergence from a negative-binomial prior (mode at low i) to encourage shallow counts and an activation-range squeeze term (Lrange) to make approximations easier.
  • Three-phase fine-tuning: (1) warm-up with Lrange only, (2) main HEAT optimization with Liter (the depth term) and prior ramp, (3) consolidation where each site is frozen to the mode of its learned distribution and weights are fine-tuned on the exact deployed circuit.
• Privacy consideration: learned per-site iteration counts are fixed at deployment for all inputs (no data-dependent halting) to avoid timing side-channel leakage.
• Empirical gains (GPT-2 124M, OpenWebText subset):
  • Iterations per forward: 228 (HEAT) vs 712 (calibrated baseline) → 3.1× reduction.
  • Bootstraps per token: 326 vs 513 → 1.6× reduction.
  • End-to-end latency: 37.6 s/tok vs 54.3 s/tok (baseline); also faster than EncryptedLLM ≈68 s/tok.
  • Fidelity to plaintext: pooled top-1 agreement 82.9% (HEAT) vs 74.0% (calibrated baseline); HEAT shows more stable KL divergence over decoding positions.
• HEAT checkpoints and code publicly released (GitHub + HuggingFace).

Data & Methods

• Model & data:
  • Base model: GPT-2 (124M).
  • Fine-tuning data: OpenWebText; total HEAT pipeline consumed ≈1.5×10^8 tokens (≈2% of pretraining corpus).
  • Compute: fine-tuning completed in under 7 hours on a single A100-64GB.
• FHE backend & cryptography:
  • Encrypted inference uses RNS-CKKS (approximate HE) at 128-bit security.
  • Implemented/evaluated on a FIDESlib-based GPU backend; measurements include data movement overheads.
• Metrics & evaluation:
  • Latency measured as seconds per generated token; bootstraps counted per token.
  • Fidelity measured as KL between encrypted logits and plaintext logits and top-k overlap/agreement across teacher-forced decoding (128 chains × 128 tokens).
  • Validation: mean over 32 held-out batches every 250 updates during fine-tuning.
• Implementation notes:
  • Iterative solvers for nonlinearities produce Fi states; expected output computed as a weighted sum using the learned distribution. Checkpointing/recomputation used to contain memory overhead.
  • Negative-binomial prior and KL penalty (Liter) steer per-site distributions toward low iteration counts while task loss retains precision where needed.
  • After training, each site’s iteration count is set to the learned mode for deterministic, constant-time deployment.

Implications for AI Economics

• Reduced per-query compute and wall-clock latency (1.4× end-to-end, 1.6× fewer bootstraps) lowers provider resource usage for FHE inference:
  • Bootstrapping is the dominant cost in FHE; reducing bootstraps has outsized operational savings (energy, GPU-hours, queuing).
  • Even modest latency cuts can materially reduce cost-per-query and improve throughput on existing hardware, improving the commercial viability of non-interactive private LLM inference.
• Low retraining barrier and integration cost:
  • HEAT is a fine-tuning method (no architecture redesign, no distillation), so providers can adopt it on pretrained models with modest extra training expense (the paper reports ≲7 hours on a single A100 for GPT-2-small).
  • This reduces time-to-market and engineering effort compared with proposals that require bespoke encryption-friendly model architectures or full retraining.
• Market and product implications:
  • Makes privacy-preserving inference (via FHE) more feasible for high-value, latency-tolerant use cases (e.g., sensitive enterprise queries, regulated data domains) that previously avoided FHE because of cost and latency.
  • Could shift pricing/packaging: providers may be able to offer lower-cost private-inference tiers or higher margins on private-inference products due to reduced computation.
  • Providers may prefer HEAT over heavier investments in re-architecting models purely for FHE, at least as an intermediate strategy.
• Limits and caveats for economic impact:
  • Absolute latency after HEAT remains large for interactive use (≈37.6 s/token on the evaluated stack). For many consumer-facing applications, that remains impractically high; HEAT primarily improves economics for privacy-sensitive, throughput-insensitive workloads or as part of a hybrid solution.
  • Results reported for one model size (GPT-2 124M), one cryptographic parameterization, and one modality (text). Gains and cost-effectiveness may differ for larger models, different packing strategies, other modalities, or alternate FHE parameter choices.
  • There is an upfront fine-tuning cost (compute, engineering), but the paper’s reported training cost was modest relative to pretraining; economic trade-offs will depend on provider scale and expected query volume.
• Strategic considerations for providers:
  • HEAT lowers the marginal cost of FHE inference per token, but full commercial adoption will likely require combining HEAT with other system and packing optimizations, hardware acceleration for bootstrapping, or selective use of cheaper privacy techniques (e.g., MPC, hybrid split computation) depending on use-case latency requirements.
  • Regulators or sectors demanding strong cryptographic privacy may value HEAT-enabled services more, increasing willingness-to-pay and offsetting remaining latency premiums.

Overall, HEAT is a practical, low-friction method to reduce expensive FHE operations by co-adapting model weights and per-site approximation depth—improving the cost-efficiency and fidelity of encrypted LLM inference and making some privacy-preserving deployments more economically attractive, while not yet making FHE inference broadly suitable for low-latency, consumer-scale interactive use.