ATLAS: Automated Approximation of Transformers for Efficient Homomorphic Inference in One Hour Jianhang Xie
Sicheng Tan
City University of Hong Kong Hong Kong, China [email protected]
Shandong University Qingdao, China [email protected]
Vishnu Naresh Boddeti
Zhichao Lu
Michigan State University East Lansing, MI, USA [email protected]
City University of Hong Kong Hong Kong, China [email protected]
Fully homomorphic encryption (FHE) provides strong cryptographic guarantees for private inference, but deploying transformer models under FHE remains prohibitively expensive. A key bottleneck is that non-linear operations such as softmax, normalization, and activation must be replaced with polynomial approximations compatible with the CKKS scheme, and the multiplicative depth consumed by these approximations dominates inference cost. Recent frameworks have advanced approximation techniques, yet all rely on manually configured approximation hyperparameters (e.g., number of iterations, polynomial degree), applied uniformly across all layers. While convenient, this uniform-configuration approach is overly rigid: different layers can tolerate different levels of approximation error without degrading predictive accuracy, and uniform configurations cannot exploit this variability to reduce latency. Allowing each layer to adopt its own configuration, however, causes the search space to explode with model depth, reaching roughly 1084 configurations for BERT/ViT (12 layers) and 10225 for LLaMA3 (32 layers), rendering manual exploration practically impossible. We present ATLAS, an automated framework that configures perlayer approximation settings by formulating the problem as a multiobjective optimization over latency and predictive accuracy. The resulting problem is inherently difficult: 1) competing objectives over a large decision space (120 or 320 variables for BERT/ViT or LLaMA3); 2) expensive evaluation, as each configuration takes 70– 1,000 seconds even in cleartext; and 3) sparse optimization signals, as 35–50% of candidate configurations yield numerically invalid solutions. ATLAS addresses these challenges through a two-stage optimization strategy that progressively relaxes layer-wise constraints, combined with surrogate models to accelerate evaluation. Compared to iterative softmax (Cho et al., CCS 2024) and THOR (Moon et al., CCS 2025), ATLAS reduces multiplicative depth by ∼35% and ∼17% with negligible accuracy loss, translating to ∼25% and ∼20% reduction in end-to-end inference latency, respectively, while completing the entire configuration search in one hour. As a postprocessing step compatible with existing approximation advancements, ATLAS generalizes across encoder-only (BERT), decoder-only (LLaMA3-8B), and vision (ViT) transformers, and complements parallel work on packing and matrix multiplication to substantially lower the barrier to deploying cryptographically secure inference. Code: https://github.com/jianhayes/ATLAS
Transformers
Best Config Config 3 Automatically Configure Approximation
Accuracy
arXiv:2607.23478v1 [cs.CR] 26 Jul 2026
Abstract
Polynomial Transformers Search in One-Hour
Config 2 Config 1 Latency FHE Evaluation
ATLAS
Service Provider
Network
Encrypt
Decrypt
Public Secret Key Key User’s Private Data
Figure 1: ATLAS can automatically design polynomial approximation configurations for FHE compatible transformers. ATLAS solutions span the trade-off between accuracy and latency for ciphertext inference and can be deployed on the cloud server to satisfy a range of customer requirements.
Keywords Fully Homomorphic Encryption, CKKS Scheme, Transformers, Secure Inference, Multi-Objective Optimization
1
Introduction
The widespread deployment of machine learning has created a growing tension between model utility and data privacy. In sensitive applications such as medical diagnosis, financial assessment, and legal document analysis, a client wishes to obtain predictions from a server-hosted model without revealing the underlying input data. Fully homomorphic encryption offers a compelling cryptographic solution to this problem: as shown on the right side of Figure 1, under an honest-but-curious server assumption, FHE allows the server to evaluate the model entirely on encrypted data, providing strong privacy guarantees without requiring the server’s trust beyond adherence to the protocol. Transformer architectures have become the dominant paradigm for such tasks, spanning language understanding (BERT [22], LLaMA [48]) and vision (ViT [24]). Deploying transformers under FHE, however, remains extremely expensive. The Cheon-Kim-Kim-Song (CKKS) scheme [10, 11, 14], which supports approximate arithmetic on real-valued data, is the most practical choice for neural network
48.22s
FHE BERT 1107.52 s
567.71 s
90.58s
241.18 s
24.09s 89.51s
34.40 s
11.83 s
48.14 s 24.06 s
ATLAS FHE BERT 792.10 s
475.38 s
123.9 s 61.57 s
18.32 s
MatrixMul
Softmax
Normalization
Activation
Others
Boot in So
Boot in Norm
Boot in Act
29.17 s 11.57 s
Figure 2: Runtime breakdown of end-to-end FHE BERT [22] on an RTX 4090. ATLAS achieves a 28.5% speedup, and saves 29.01 s, 16.08 s, and 60.34 s in Softmax, normalization, and activation approximations, respectively, which further reduces bootstrapping overhead by ∼117 s.
inference, but it imposes a fundamental constraint: operations are limited to polynomial arithmetic, and the number of sequential multiplications (the multiplicative depth) determines both the ciphertext parameter size and the overall latency. As shown in Figure 2, a baseline FHE implementation of a transformer devotes approximately 46% of inference time to non-linear operations alone. This overhead arises because non-linear functions such as softmax, layer normalization, and GELU are incompatible with FHE and must be replaced with polynomial approximations, each consuming significant multiplicative depth. A substantial body of recent work has sought to reduce this cost. One line of work targets the linear components of transformers, particularly ciphertext-ciphertext matrix multiplication (CCMM) and plaintext-ciphertext matrix multiplication (PCMM) [28, 35, 38]. A second line of work addresses approximations of non-linear operations such as Softmax, LayerNorm, and GELU directly, either through polynomial approximations or through iterative computations [16, 41, 54]. A third line of work seeks to improve the computational efficiency of core cryptographic primitives, including bootstrapping [3, 8, 9, 12, 17] and polynomial multiplication [31]. Despite this progress, a critical design decision remains unaddressed across all existing frameworks: how to set the hyperparameters governing each polynomial approximation. Current practice treats this as a manual, per-function problem. A practitioner selects an approximation for each non-linear function type (e.g., a degree 𝑑 minimax polynomial for GELU) and applies it uniformly across all layers. This approach is intuitive but suboptimal. The true objective is not to approximate individual functions accurately in isolation, but to preserve the end-to-end task accuracy of the pretrained transformer model under the constraints of FHE. Optimizing for this global objective opens the possibility of heterogeneous, per-layer configurations in which different layers use different approximation configuration hyperparameters, trading local approximation fidelity for global depth reduction. The challenge is that this search space is vast. With multiple tunable hyperparameters per operator and dozens of layers per model, the space of valid configurations spans 1084 to 10225 possibilities depending on the architecture. We present ATLAS, a framework-agnostic post-processing method that automatically discovers depth-efficient approximation configurations for any FHE-compatible transformer. ATLAS formulates
the problem as multi-objective optimization over multiplicative depth and end-to-end accuracy, and solves it through a two-stage evolutionary search. ATLAS treats any existing FHE transformer framework as a black box, operating entirely post-training on the frozen pretrained model and requiring no retraining, fine-tuning, or changes to its packing strategy or encryption parameters. An overview of ATLAS is shown in Figure 1. The contributions of this paper are: (1) Problem formulation. We identify the selection of perlayer polynomial approximation hyperparameters as a critical and underexplored bottleneck in FHE transformer inference, and formulate it as a multi-objective combinatorial optimization problem over multiplicative depth and end-toend accuracy. (2) ATLAS. We present a framework-agnostic post-processing method that solves this problem through a two-stage evolutionary search, designed to handle conflicting objectives, expensive evaluations, and a high density of invalid configurations. ATLAS is complementary to existing frameworks, including THOR [38], NEXUS [54], and MOAI [55], and can be applied on top of any FHE-compatible transformer without modifying the underlying encryption or packing strategy. (3) Empirical evaluation. Evaluated on BERT-base, LLaMA3 8B, and ViT-base, ATLAS reduces multiplicative depth by 35 percent and inference latency by 25 percent relative to expert-designed baselines, with negligible accuracy loss. The search completes in under 30 minutes for BERT and ViT, and under one hour for LLaMA-3 8B.
2 Preliminaries 2.1 FHE Scheme and Bootstrapping CKKS Scheme. The Cheon-Kim-Kim-Song (CKKS) [14] is a leveled homomorphic encryption scheme. Compared with previous FHE schemes, e.g., BGV [6], BFV [5], and TFHE [15], which only support integer number encryption, the CKKS can encrypt real and complex number calculations. Under RNS-CKKS, the plaintexts and ciphertexts are elements in a residue cyclotomic polynomial ring R Qℓ = Z Qℓ [𝑋 ]/(𝑋 𝑁 +1). Îℓ The modulus is Qℓ = 𝑖=0 𝑞𝑖 , where 0 ≤ ℓ ≤ 𝐷, the ℓ is level, the 𝐷 is level budget. 𝑁 is the polynomial degree in RNS-CKKS, and a ciphertext has 𝑁 /2 slot counts for single instruction multiple data (SIMD) processing. We denote the above procedure from cleartext 𝑢 to SIMD ciphertext Enc(𝑢) as Enc(·), and the opposite direction decryption-decoding algorithm is denoted as Dec(·). Specifically, assuming that homomorphic addition is ⊞ and homomorphic multiplication is ⊠, the operations can be described in the following with 𝑢 and 𝑣 as: Dec(Enc(𝑢) ⊞Enc(𝑣)) =Dec(Enc(𝑢)) +Dec(Enc(𝑣)) ≈𝑢 +𝑣 (1) Dec(Enc(𝑢) ⊠Enc(𝑣)) =Dec(Enc(𝑢)) ×Dec(Enc(𝑣)) ≈𝑢 ×𝑣 Bootstrapping. Leveled homomorphic encryption schemes like CKKS only support a finite number of homomorphic multiplications, each of which consumes one level due to rescaling. So, when the level of a ciphertext becomes zero, the decryption would fail, so we need to perform a bootstrapping [10] operation to reset to a high level if it is too low to do a computation. The bootstrapping allows
us to evaluate circuit of arbitrary depth, as it homomorphically evaluates the decryption circuit and raises the modulus from Q0 to Q𝐷 by leveraging the isomorphism R𝑞0 R𝑞0 ×R𝑞1 × · · · ×R𝑞𝐷 [4]. A freshly encrypted ciphertext starts with 𝐷 levels, but bootstrapping consumes 𝐾 levels and reduces it to 𝐷 − 𝐾 levels. As bootstrapping requires a lot of key switching operations (KSO), it becomes the most time-consuming operation in RNS-CKKS.
2.2
Transformers
Transformer [49] models are typically encoder-only (BERT [22]), decoder-only (e.g., GPT [43] and LLaMA [48]), or vision encoder (ViT [24]). Despite this distinction, the structure of the Transformer layer is similar in both encoders and decoders: a stack of 𝐿 layers. A basic Transformer layer 𝑓 (𝒙) , includes Attention, Multilayer Perceptron (MLP), and LayerNorm, as shown in Equation (2). 𝑓 (𝒙) = LayerNorm2(MLP(LayerNorm1(Attention(𝒙)))),
(2)
where 𝒙∈R𝑛×𝑑 is the input embedding, 𝑛 is the number of tokens and 𝑑 is the hidden size. For attention, the weight matrices are 𝑾𝑄 , 𝑾𝐾 , 𝑾𝑉 , 𝑾𝑂 ∈ R𝑑 ×𝑑 , yielding 𝑸=𝒙𝑾𝑄 , 𝑲 =𝒙𝑾𝐾 , and 𝑽 =𝒙𝑾𝑉 . For standard multi-head attention (MHA) with ℎ heads, 𝑸, 𝑲 , and 𝑽 are partitioned into 𝑸ℎ , 𝑲ℎ , 𝑽ℎ ∈R𝑛×𝑑𝑘 , where 𝑑𝑘 =𝑑/ℎ is the head dimension. The attention computation for a single head is shown in Equation (3). 𝑸ℎ 𝑲ℎ𝑇 (3) Attention(𝑸ℎ , 𝑲ℎ , 𝑽ℎ ) = Softmax √ 𝑽ℎ 𝑑𝑘 Finally, the outputs of all heads are concatenated and projected via 𝑾𝑂 to produce the final attention result. For the MLP, let 𝑑 ff be the intermediate size. The MLP weights consist of two matrices 𝑾up ∈ R𝑑 ×𝑑ff , 𝑾down ∈ R𝑑ff ×𝑑 , with an activation 𝜎 (·) placed between them. Typically 𝜎 (𝑥)=GELU(𝑥). For the LLaMA MLP, 𝜎 (𝑥)=SiLU(𝑥) and it adopts a gated structure. For normalization, standard Transformers use LayerNorm [2]. Modern LLMs such as LLaMA employ RMSNorm [53] to reduce computational cost. The placement of normalization depends on the specific architecture.
2.3
FHE Transformers
Since homomorphic matrix multiplication is nearly lossless in precision, the precision loss arises from the approximation of non-linear components. Thus, constructing FHE inference requires a primary focus on these non-linearities. For Transformers, the non-linear components include Softmax in Attention, normalization, and activation in the MLP. Compared to CNNs [1, 32] where only ReLU requires approximation, building a Transformer inference system under FHE is considerably more challenging. Currently, the approximation of these components [16, 38, 44, 52, 54, 55] relies on expert-designed heuristics and manually configured hyperparameters. Softmax Approximation. The standard softmax function is deÍ fined as Softmax(𝒙)𝑖 =exp(𝑥𝑖 )/ 𝑗 exp(𝑥 𝑗 ), which contains the nonlinear exponential exp(𝑥) and inverse 1/𝑥. Some works replace the Attention with HE-friendly alternatives, such as Gaussian-kernel Attention [44], activation-based Attention [56], and BPMax [40], but these replacements all require model retraining.
Recent works [52, 54, 55] directly approximate exp(𝑥) and 1/𝑥 separately over a given interval. For example, exp(𝑥) is approxi𝑟 mated via limit (1 + 𝑥/2𝑟 ) 2 [54, 55] for 𝑥 ∈ [−2𝑟 , 0], or via a Chebyshev polynomial [52]; and 1/𝑥 via Goldschmidt iteration [54, 55] or a Chebyshev-initialized Goldschmidt refinement [52]. To improve numerical stability, the row maximum is subtracted: Softmax(𝒙)=Softmax(𝒙 −max𝑖 𝑥𝑖 ). In FHE, dynamically computing max𝑖 𝑥𝑖 is prohibitively expensive; thus, a statistical hard-coded constant maximum 𝑐 is used instead. Although this separate approximation incurs low multiplicative depth, the effective domains of the 1/𝑥 and exp(𝑥) approximations jointly constrain the softmax input range. From the NEXUS [47] codebase, we measure the usable softmax interval to be roughly [−2.19, 0]. As a result, the valid input range may be narrower than the range of 𝑥𝑖 − 𝑐, leading to approximation failure. A state-of-the-art alternative employs an iterative softmax [16] with normalize-and-square strategy. It assumes that inputs after subtracting the maximum are non-positive, i.e., 𝒙∈[−𝑀, 0] 𝑛 . The input is scaled by 1/2𝑘 and recovered via 𝑘 iterations. Consequently, the effective input range of the iterative softmax can be estimated as [−2𝑘 ln 𝑛, 0]. For its default setting 𝑘=5 and 𝑛=256, this gives a lower bound of roughly 𝑀≈177. So this mitigates the limited-domain problem. THOR [38] also proposes a similar square-and-normalize approach that reduces the required number of iterations. The precision of the above methods depends on the approximation degree of the inverse square root (or inverse) polynomial and the number of iterations 𝑘. In iterative softmax [16], these hyperparameters are manually configured based on expert heuristics. Normalization Approximation. LayerNorm requires the inverse √ square root (invsqrt): LayerNorm(𝑥𝑖 )=𝑤𝑖 ·(𝑥𝑖 −𝜇)/ 𝜎 2 + 𝜖+𝑏𝑖 , where 𝜇 and 𝜎 2 denote the mean and variance of input 𝒙, 𝑤𝑖 and 𝑏𝑖 are affine transform coefficients. RMSNorm is analogous by omitting the mean. Nearly all existing works [52, 54, 55] approximate invsqrt via Newton’s method [42], often combined with Goldschmidt iterations [27] for high-precision refinement. For example, the √ NEXUS [54] codebase [47] uses 𝜈=4 Newton steps to compute 1/ 𝑥, followed by 𝛾=2 Goldschmidt steps; MOAI [55] and ARION [52] adopt a similar variant. The effectiveness of LayerNorm and RMSNorm under FHE hinges on whether the input variance (or root mean square) lies within the valid approximation interval of the invsqrt. At the same time, the numbers of Newton and Goldschmidt iterations in the above methods remain hand-tuned fixed values. Activation Approximation. The GELU function is defined as GELU(𝑥) = 𝑥 ·Φ(𝑥), where Φ(𝑥) is the cumulative distribution function for Gaussian distribution, or by its tanh approximation. Early FHE Transformers [19, 54] and some interactive methods [23, 36, 39] use a piecewise GELU approximation, which employs a sign function approximated by minimax composition [33, 34] for segments selection. However, the piecewise GELU approximation is valid only for a small domain, e.g., 𝑥 ∈ [−8, 8] claimed in NEXUS [54], which is insufficient to cover the input range encountered during Transformer inference. To address this limitation, recent works [25, 38, 52, 55] explicitly define an approximate GELU with polynomials. For example, MOAI [55] uses a degree-23 polynomial; ARION [52]
adopts a degree-255 Chebyshev polynomial; and THOR [38] employs a composite polynomial with degree-31 and degree-27. For LLaMA’s SiLU, MOAI [55] does not specify a degree, deferring to Orion [25]. However, the polynomial degrees are largely determined by heuristics and manual tuning in all these works.
In the context of secure inference, the following threat model is generally assumed [1, 32, 54]: a customer uploads encrypted data to the Cloud; the Cloud service provider then processes the ciphertext using neural networks deployed on its servers and returns the encrypted output to the customer; finally, the customer decrypts the result locally using a secret key. Under this model, the Cloud provider cannot access the sensitive information contained in the customer’s input or output data, while the client remains unaware of the details of neural networks.
where 𝑀𝑢𝑙_𝐷𝑒𝑝𝑡ℎ(·) computes the total multiplicative depth and 𝜖 controls the tolerance on output deviation of a FHE transformer model 𝑓˜𝝀 from its cleartext counterpart 𝑓 . Unfortunately, choosing an appropriate 𝜖 in Equation (5) is also not straightforward, as different transformer models and tasks require different levels of precision (e.g., ViT can tolerate large deviation in outputs without dropping accuracy on ImageNet-1K classification compared to BERT on GLUE benchmark). An overconstrained 𝜖 leads to sub-optimal run-time latency, while an underconstrained 𝜖 leads to a noticeable drop in model performance. As opposed to adaptively/carefully tuning 𝜖, we propose to formulate the problem as a multi-objective optimization problem minimize𝝀 𝑀𝑢𝑙_𝐷𝑒𝑝𝑡ℎ( 𝑓˜𝝀 ), 𝑀𝐴𝐸 (𝑓 , 𝑓˜𝝀 ) , (6) 1 ∑︁ where 𝑀𝐴𝐸 (𝑓 , 𝑓˜𝝀 ) = 𝑓 (𝑥) − 𝑓˜𝝀 (𝑥) , |D| 𝑥 ∈ D
3
to simultaneously balance run-time latency and approximation precision, circumventing the need of choosing an 𝜖.
2.4
Threat Model
ATLAS: Adapting Transformers for FHE
ATLAS is an automated, framework-agnostic post-processing method that designs efficient polynomial approximation configurations for FHE transformers. Rather than applying a single set of handtuned hyperparameters uniformly across all layers, ATLAS assigns each layer its own approximation hyperparameters. We adopt a multi-objective evolutionary search algorithm to navigate the combinatorial space of layerwise approximation hyperparameters and discover configurations that are (i) significantly more efficient than hand-tuned hyperparameters, and (ii) span the accuracy-efficiency trade-off. Specifically, we adopt a two-stage evolutionary procedure that progressively relaxes layer-wise constraints.
3.1
Problem Formulation
Given a cleartext transformer 𝑓 , we would like to develop its polynomial approximation 𝑓˜𝝀 by choosing appropriate hyperparameters 𝝀 = (𝝀𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 , 𝝀𝑛𝑜𝑟𝑚 , 𝝀𝑎𝑐𝑡 ) to configure softmax 𝝀𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 , normalization 𝝀𝑛𝑜𝑟𝑚 , and activation function 𝝀𝑎𝑐𝑡 approximations. The goal is to minimize latency overhead without sacrificing security, while ensuring that 𝑓˜𝝀 exhibits similar predictive performance under FHE. This can be mathematically formulated as: minimize𝝀 𝐿𝑎𝑡 ( 𝑓˜𝝀 , D), (4) subject to 𝐴𝑐𝑐 (𝑓 , D) − 𝐴𝑐𝑐 ( 𝑓˜𝝀 , D) ≤ 𝜖𝑎𝑐𝑐 , where 𝐿𝑎𝑡 (·) measures runtime latency of the approximated model 𝑓˜𝝀 under FHE, 𝐴𝑐𝑐 (·) computes its accuracy over a representative set of data samples D, and 𝜖𝑎𝑐𝑐 signifies the error tolerance which is typically a small number approaching zero. Directly solving Equation (4) imposes two main challenges: (1) measuring latency reliably and consistently on hardware is not trivial; (2) 𝐴𝑐𝑐 (·) metric may not always be sensitive to deviations in model outputs due to approximation error. Accordingly, a more practical formulation of the problem is as follows. minimize𝝀 𝑀𝑢𝑙_𝐷𝑒𝑝𝑡ℎ( 𝑓˜𝝀 ), (5) subject to 𝑓 (𝑥) − 𝑓˜𝝀 (𝑥) ≤ 𝜖 ∀𝑥 ∈ D,
3.2
Search Space and Encoding
A decision vector 𝝀 specifies how softmax, normalization, and activation functions are approximated in each of the 𝐿 Transformer layer (𝐿 = 12 for BERT-Base and ViT-Base, 𝐿 = 32 for LLaMA38B). The decision variables are grouped by operators (i.e., softmax, normalization, and activation) and arranged in ascending layer order (i.e., from Layer 1 to 𝐿), as follows. 𝝀 = 𝑝 1(1) , . . . , 𝑝 5(1) , · · · , 𝑝 1(𝐿) , . . . , 𝑝 5(𝐿) , | {z } 𝝀𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 (1) (1) (1) (1) (𝐿) (𝐿) (𝐿) (𝐿) , 𝜈 𝑎𝑡𝑡𝑛 , 𝛾𝑎𝑡𝑡𝑛 , 𝜈𝑚𝑙𝑝 , 𝛾𝑚𝑙𝑝 , · · · , 𝜈 𝑎𝑡𝑡𝑛 , 𝛾𝑎𝑡𝑡𝑛 , 𝜈𝑚𝑙𝑝 , 𝛾𝑚𝑙𝑝 | {z } 𝝀𝑛𝑜𝑟𝑚 (1) (𝐿) 𝑝𝑎𝑐𝑡 , · · · , 𝑝𝑎𝑐𝑡 , | {z } 𝝀𝑎𝑐𝑡
(7)
resulting in a total of 5𝐿 + 4𝐿 + 𝐿 = 10𝐿 integer variables. Softmax 𝝀𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 Encoding. We adopt the iterative method proposed by Cho et al. [16] for softmax approximation, with one modification to use Chebyshev polynomials instead of power basis for coefficient estimation. Algorithm 1 outlines our procedure. Algorithm 1: Iter. Softmax(𝑥; 𝑘, 𝑝 1, · · · , 𝑝𝑘 ) in ATLAS Input: 𝑥 ∈ [−𝑀, 0] 𝑛 , 𝑘 ∈ Z>0 , (𝑝 1, · · · , 𝑝𝑘 ) Output: 𝑦 ≈ Softmax(𝑥) 𝑘 1 𝑦 ← ChebyPolyEvalExp(𝑥/2 , 15); 2 for 𝑗 ← 1 to 𝑘 do Í 3 𝑎 ← ChebyPolyEvalInvsqrt( 𝑛𝑖=1 𝑦𝑖2, 2𝑝 𝑗 − 1) ; // Chebyshev invsqrt with degree 2𝑝 𝑗 − 1 4 𝑦 ← (𝑎 · 𝑦) 2 ; 5 end 6 return 𝑦;
𝝀𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 =(𝑝 1(1) , · · · , 𝑝 5(1) , · · · , 𝑝 1(𝐿) , · · · , 𝑝 5(𝐿) ).
Activation 𝝀𝑎𝑐𝑡 Encoding. The non-linear activation function (e.g., GELU or SiLU) is approximated by a single Chebyshev polynomial [25, 52]. The degree directly determines the multiplicative depth consumed; we therefore include the degree 𝑝𝑎𝑐𝑡 in the search space for exploring the accuracy–depth trade-off. Each layer 𝑖 can select an integer 𝑝𝑎𝑐𝑡 ∈[1, 9], yielding a polynomial degree of 2𝑝𝑎𝑐𝑡 −1. The full activation configuration vector is then:
800 600 400 200
Algorithm 2: Newton–Gold InvSqrt(𝑥; 𝜈, 𝛾) in ATLAS Input: Input 𝑥; Newton steps 𝜈, Goldschmidt steps 𝛾 √ Output: 𝑦 ≈ 1/ 𝑥
0
77.5 ±0.1 s 117.9 ±0.0 s BERT-Base
ViT-Base
LLaMA3-8B
(a) Expensive evaluation
(8)
Normalization 𝝀𝑛𝑜𝑟𝑚 Encoding. Both the Attention and MLP modules in a Transformer layer use normalization (e.g., LayerNorm and RMSNorm). Algorithm 2 outlines the Newton–Goldschmidt approximation of the invsqrt used in normalization. The two loop counters 𝜈 and 𝛾 directly control the accuracy–depth trade-off: larger values improve the approximation quality at the cost of additional multiplicative depth. We therefore adopt the default setting from the NEXUS codebase [47] (𝜈 = 4, 𝛾 = 2) as a reference and treat both (𝜈, 𝛾) counters as searchable integer parameters.
1006.9 ±0.1 s
1000
% of sampled configs with finite MAE
(1) (2) (𝐿) 𝝀𝑎𝑐𝑡 = (𝑝𝑎𝑐𝑡 , 𝑝𝑎𝑐𝑡 , · · · , 𝑝𝑎𝑐𝑡 ).
Inference Time for 1,000 samples (s)
Given an input 𝒙 ∈ [−𝑀, 0] 𝑛 , the algorithm first divides the input by 2𝑘 and applies a 15-degree Chebyshev exponential approximation ChebyPolyEvalExp. It then applies 𝑘 iterations, each consisting of: (i) summing the squared intermediate values, (ii) evaluating a ChebyPolyEvalInvsqrt that approximates the invsqrt, and (iii) multiplying and squaring the result. The key hyperparameters are the total number of iterations 𝑘 and the degree of the Chebyshev polynomial 𝑝 in each iteration. We embed the choice of 𝑘 into the degree variables by allowing the 𝑝 𝑗 to be zero, indicating that the 𝑗-th and all its subsequent iterations are skipped. In summary, each layer 𝑖 ∈1,··· ,𝐿 is assigned five integer variables (𝑝 1(𝑖 ) , 𝑝 2(𝑖 ) , 𝑝 3(𝑖 ) , 𝑝 4(𝑖 ) , 𝑝 5(𝑖 ) ), where 𝑝 1(𝑖 ) ∈ [1, 7] and 𝑝 2(𝑖 ) , . . . , 𝑝 5(𝑖 ) ∈ [0, 7]. A non-zero value determines the polynomial degree for the invsqrt function approximation in an iteration, and a zero value omits that iteration and all the subsequent ones. The upper bound is set to 7 following the settings from [16], and the lower bound is chosen such that at least one iteration is executed. The full softmax configuration contributing 5𝐿 integer variables is then encoded as:
(10)
70 60 50
48.1 ±1.5%
45.9 ±1.6% 35.6 ±1.5%
40 30 BERT-base
ViT-base
Llama3-8B
(b) Fraction of finite MAE
Figure 3: Key challenges in searching for FHE approximation configurations. (a) Even a single cleartext evaluation of MAE is time-consuming, and FHE evaluation is orders of magnitude slower. (b) A large portion of the search space produces invalid (NaN/∞) MAE, resulting in a sparse feasible region.
3.3
Search Challenges
The search space defined in Section 3.2 is inherently combinatorial. For a single layer, the number of unique configurations is: |S𝑙𝑎𝑦𝑒𝑟 | = |S𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 | × |S𝑛𝑜𝑟𝑚 | × |S𝑎𝑐𝑡 | ! 5 ∑︁ 𝑘 = 7 × (4 · 2 · 4 · 2) × 9 (11) 𝑘=1
= 19 607 × 64 × 9 = 11 293 632 ≈ 1.13 × 107,
−4 · 𝑥 + 0.129 ; 1 𝑦 ← −1.29×10
// initial guess for 𝑖 ← 1 to 𝜈 do 3 𝑦 ← 𝑦 · 1.5 − 0.5 · 𝑥 · 𝑦 2 ; // Newton Iter 4 end 5 𝑎, 𝑏 ← 𝑥𝑦, 𝑦/2; 6 for 𝑖 ← 1 to 𝛾 do 7 𝑎, 𝑏 ← 𝑎(1.5 − 𝑎𝑏), 𝑏 (1.5 − 𝑎𝑏) ; // Goldschmidt Iter 8 end 9 return 𝑦 ← 2𝑏; 2
For each layer 𝑖 ∈1,··· ,𝐿 , four integer variables are assigned: at(𝑖 ) tention Newton iterations 𝜈 𝑎𝑡𝑡𝑛 ∈ [1, 4], attention Goldschmidt (𝑖 ) (𝑖 ) iterations 𝛾𝑎𝑡𝑡𝑛 ∈ [1, 2], MLP Newton iterations 𝜈𝑚𝑙𝑝 ∈ [1, 4], and (𝑖 ) MLP Goldschmidt iterations 𝛾𝑚𝑙𝑝 ∈ [1, 2]. The normalization variables of layer 𝑖 and the full configuration vector with 4𝐿 integer variables are encoded as: (1) (1) (1) (1) (𝐿) (𝐿) (𝐿) (𝐿) 𝝀𝑛𝑜𝑟𝑚 = (𝜈 𝑎𝑡𝑡𝑛 , 𝛾𝑎𝑡𝑡𝑛 , 𝜈𝑚𝑙𝑝 , 𝛾𝑚𝑙𝑝 , · · · , 𝜈 𝑎𝑡𝑡𝑛 , 𝛾𝑎𝑡𝑡𝑛 , 𝜈𝑚𝑙𝑝 , 𝛾𝑚𝑙𝑝 ).
(9)
where |S𝑠𝑜 𝑓 𝑡𝑚𝑎𝑥 | = 19 607, |S𝑛𝑜𝑟𝑚 | = 64, and |S𝑎𝑐𝑡 | = 9. When configurations are assigned per-layer, the total search volume grows to |S𝑙𝑎𝑦𝑒𝑟 |𝐿 , yielding 11 293 63212 ≈ 4.3 × 1084 for 𝐿=12 (BERT/ViT) and 11 293 63232 ≈ 4.9 × 10225 for 𝐿=32 (LLaMA3-8B). Beyond the sheer size of this space, the optimization must overcome several additional fundamental difficulties. • NP-Hard structure. Selecting per-layer approximation hyperparameters under a depth–accuracy budget is a constrained integer program that generalizes multi-dimensional knapsack and per-layer precision assignment, both NP-hard. No polynomial-time exact algorithm is therefore expected, motivating heuristic, population-based search rather than exhaustive or exact optimization. • Conflicting objectives. The two objectives—multiplicative depth and MAE—are inherently conflicting: higher-degree polynomials and more iterations reduce MAE but consume more depth, and vice versa. Finding a single solution that simultaneously minimizes both objectives is generally infeasible; instead, we seek a set of Pareto-optimal solutions that represent the efficient trade-off between the two.
Pareto Configs
Stage 1 Pareto
Random Init
Selection
Crossover
Mutation
Evaluation
Stage 2 Pareto
Survival
λ1
Softmax Norm Act
HV
Data subset surrogate Early-Layer proxy Distributed Eval λN
Depth
All layers share the same configuration
λ1′ λ2′
Acceleration Techniques
λ1′ HV
λ2
Softmax Norm Act Data subset surrogate Distributed Eval
HV ′
λ
′ N
MAE
λ2
MAE
λ1
MAE
Init
λN
Depth
Stage 1 – All layers share the same configuration
Each layer has its own configuration
λ1
′′
λ2′′
λ2′
HV ′
Rejected Acceleration Techniques
NSGA-II survival selection
λN′′
HV ′′ λ ′ N
Depth
Stage 2 – Each layer has its own configuration
Figure 4: ATLAS is a two-stage process. Stage 1 identifies a set of efficient layer-uniform approximations, and Stage 2 refines these into heterogeneous layer-wise configurations to further improve the depth–accuracy trade-off. • Expensive evaluation. Each candidate configuration must be instantiated as a model and evaluated to obtain its MAE and depth. Even in the cleartext polynomial setting used during search, a single function evaluation (FE) is costly for computing MAE. As shown in Figure 3(a), a single FE under cleartext already takes 77.5 s for BERT, 118 s for ViT, and 1007 s for LLaMA; under end-to-end FHE, the latency is even higher. This cost compounds in population-based methods, where every generation evaluates an entire population, making naive deployment of evolutionary search prohibitive without aggressive evaluation acceleration. • Sparse feasible region. A large fraction of the search space yields configurations that produce NaN or infinite MAE, rendering them useless. Figure 3(b) shows that the proportion of configurations with finite MAE is below 50% for all three models. This sparse signal makes the optimization landscape highly discontinuous, further complicating the search.
3.4
Two-Stage Search with NSGA-II
In ATLAS, to address the challenges outlined in Section 3.3, we develop a two-stage evolutionary optimization strategy based on the Non-dominated Sorting Genetic Algorithm II (NSGA-II) [20]. The core idea is to decompose the search: first explore the compact LayerProblem, where all 𝐿 layers share the same configuration; then refine them into per-layer settings by solving NetworkProblem, where each layer can have different configurations. As shown in Figure 4, the first stage is computationally efficient and produces a set of layer-uniform yet high-quality configurations that serve as a warm start for the second stage, thereby drastically accelerating convergence. 3.4.1 Preliminaries on NSGA-II. ATLAS adopts NSGA-II [20] to handle the two competing objectives—multiplicative depth and MAE— through Pareto dominance: 𝝀 1 dominates 𝝀 2 if it is no worse in both objectives and strictly better in at least one. The non-dominated set forms the Pareto front, whose quality we measure by the hypervolume (HV) [57] with respect to a fixed reference point; a larger HV in the (depth, MAE) space indicates a front that achieves lower depth, lower error, or a better trade-off between the two. NSGA-II evolves a population of solutions (i.e., configuration vectors) over generations through four operators: selection, crossover, and mutation produce new candidate solutions from the current population, and survival selection then determines which solutions advance to the next generation. The survival step is the core of
NSGA-II: given a combined pool of current and newly generated solutions, non-dominated sorting partitions the pool into fronts F1, F2, . . ., admitting solutions front by front; when a front cannot be fully admitted, crowding distance—the objective-space distance to neighbors within the same front—serves as a tiebreaker, favoring solutions in less densely populated regions. Together, the two criteria balance convergence toward the Pareto front with diversity along it, as illustrated in Figure 4. The main evolutionary operators are summarized below. • Selection. Promising solutions, referred to as parents, are chosen via binary tournament selection: two individuals are randomly sampled, and the one with better non-dominated rank is kept. This process repeats until enough parents are selected. • Crossover. A standard two-point crossover is used to exchange sub-components between two parents at each crossover point to create a new set of solutions, referred to as offspring. • Mutation. An integer step mutation operator then perturbs each variable in the offspring solution by +1 or -1 following a binomial distribution with a probability proportional to one over the number of variables, rounding to the nearest integer within the bounds. 3.4.2 Two-Stage Framework. However, directly applying NSGA-II to the 10𝐿-variable NetworkProblem suffers from the curse of dimensionality: the initial population is extremely sparse, and the search wastes many evaluations on infeasible or poor regions. We observed that a configuration where all layers are identical already yields a reasonable baseline; evaluating it on the LayerProblem is fast, and the HV grows quickly. In contrast, the per-layer search requires tens of thousands of FEs to reach a comparable HV. This motivates an efficient two-stage design: • LayerProblem — all 𝐿 layers share the same configuration by constraining: 𝒑 (1) = 𝒑 (2) = · · · = 𝒑 (𝐿) , where 𝒑 = (𝑝 1, 𝑝 2, 𝑝 3, 𝑝 4, 𝑝 5 ) (1) (2) (𝐿) (1) (2) (𝐿) 𝜈 𝑎𝑡𝑡𝑛 = · · · = 𝜈𝑚𝑙𝑝 = 𝜈 𝑎𝑡𝑡𝑛 = · · · = 𝜈 𝑎𝑡𝑡𝑛 , 𝜈𝑚𝑙𝑝 = 𝜈𝑚𝑙𝑝 (12) (1) (2) (𝐿) (1) (2) (𝐿) 𝛾𝑎𝑡𝑡𝑛 = 𝛾𝑎𝑡𝑡𝑛 = · · · = 𝛾𝑎𝑡𝑡𝑛 , 𝛾𝑚𝑙𝑝 = 𝛾𝑚𝑙𝑝 = · · · = 𝛾𝑚𝑙𝑝 (1) (2) (𝐿) 𝑝𝑎𝑐𝑡 = 𝑝𝑎𝑐𝑡 = · · · = 𝑝𝑎𝑐𝑡 . • NetworkProblem — each layer has its own configuration, i.e. the 10𝐿-variable bounded integer space represented by 𝝀 in (7).
Stage 1. NSGA-II searches the single-layer configuration that is applied identically to all 𝐿 layers. With a moderate budget (e.g., 1/10 of the total budget ), the algorithm produces a Pareto set of high-quality layer-uniform configurations 𝝀𝑙𝑎𝑦𝑒𝑟 . Stage 2. These 𝝀𝑙𝑎𝑦𝑒𝑟 , together with a small subset of randomly created configurations, form a seeded initial population of the NetworkProblem for NSGA-II. Because 𝝀𝑙𝑎𝑦𝑒𝑟 already resides in promising regions of the objective space, the search starts from a favorable basin and efficiently refines per-layer configurations. The Stage-2 budget (e.g., 9/10 of the total budget ) is set to match the total evaluation budget.
3.5
Accelerating the Search
Even with the two-stage strategy, each evaluation call remains expensive. To further reduce the cost of the search, we introduce two surrogate models that approximate the true objectives at a fraction of the cost and a distributed evaluation technique. Data Subset Surrogate. A single evaluation over the full dataset dominates the search time. We observe that the MAE computed on a small random subset of the data is highly correlated with the full-dataset MAE. For a given configuration, one can view the subset MAE as an estimate of the true full-dataset MAE. That is, following Equation (6), we are not evaluating all 𝑥 ∈ D, but take a sampled subset of D. To assess the reliability of this surrogate, we measure the Kendall 𝜏 rank correlation between the two quantities: 𝜏 close to 1 indicates that the ordering of configurations is largely preserved. Figure 5 (top right) reports 𝜏 and the resulting speed-up as a function of subset size. For BERT, a subset of 50 IMDB sentences achieves 𝜏 ≈ 0.85 with ∼470× speed-up, sufficient for reliable ranking. For ViT, even 10 ImageNet validation images yield 𝜏 ≈ 0.936 with ∼440× speed-up. Figure 5 (bottom left) and Figure 5 (bottom right) illustrate this correlation for the chosen sizes: each point represents one configuration whose evaluation is performed on a random subset, with the corresponding full-dataset MAE on the vertical axis. In practice, we set the proxy subset size to 10 for all three models (BERT, ViT, and LLaMA) throughout the search to achieve maximum speed-up.
Figure 5: MAE at an early layer as a surrogate for MAE at the last layer, and the correlation between the MAE on the full and subset of the dataset, are shown on BERT and ViT.
Early-Layer Proxy. In the LayerProblem, all layers share the same configuration; consequently, the MAE measured at intermediate Transformer layers is strongly correlated with the final-layer MAE. Figure 5 (top left) plots Kendall 𝜏 between early layer-𝑖 MAE and last layer MAE against the resulting speed-up. For BERT, layer 4 already gives 𝜏 = 0.927 with ∼3× speed-up; for ViT, layer 7 achieves 𝜏 = 0.914 with ∼1.8× speed-up. Figure 5 (middle left) and Figure 5 (middle right) show the corresponding relations between last layer MAE and layer-𝑖 MAE. Stopping the forward pass at an early layer avoids the computation of all subsequent Transformer layers, yielding an acceleration proportional to 𝐿/𝑖. This proxy is only valid when all layers are identical (Stage 1) and is disabled in Stage 2, where per-layer configurations vary independently. For the surrogate models of LLaMA3-8B, we use a data subset surrogate with 10 IMDB sentences, with 𝜏 = 0.905 and ∼454× speed-up, and a layer-9 MAE as the early-layer proxy, with 𝜏 = 0.924 and ∼8.5× speed-up. The detailed experiments are shown in Appendix E.
ATLAS automatically approximates pre-trained transformers for efficient homomorphic encryption through evolutionary search. Here, we design experiments to answer the following research questions (RQs):
Distributed Evaluation. Evaluation is further accelerated by distributing the population across multiple GPUs, exploiting the natural parallelism of evolutionary search: the fitness of every individual in a generation can be computed independently. Overall Speed-Up. Combining the above techniques reduces the effective cost of a single evaluation by a factor of several hundred to over one thousand in Stage 1. In Stage 2, we drop the earlylayer proxy as different layers may use different configurations, yet the search is already warm-started in a promising region of the objective space. Together, these techniques make the ATLAS tractable even for large models such as LLaMA3-8B.
4
Experimental Validation
• RQ1: Can ATLAS discover transformer approximation hyperparameters that improve the accuracy–depth (inference latency) trade-off achieved by existing hand-tuned secure transformer inference solutions? (Section 4.2) • RQ2: What insights does ATLAS reveal about the approximation hyperparameters that improve the accuracy–depth? (Section 4.3)
• RQ3: How efficient and generalizable is ATLAS at approximating pre-trained transformers that span the accuracy-efficiency trade-off? (Section 4.4)
4.1
Experimental Setup
Datasets. During search, both BERT and LLaMA use 10 randomly sampled IMDB [37] training sentences as the proxy subset for MAE evaluation. For the task evaluation, BERT is assessed on three GLUE benchmark tasks [50]: SST-2, RTE, and QNLI; LLaMA is evaluated on MMLU [29], GSM8K [18], and HumanEval [7]. Specifically, we evaluate MMLU on 64 sampled test examples from its abstract algebra subset. ViT uses 10 random ImageNet-1K [21] images as the search proxy, and the full ImageNet-1K validation set for final Top-1 accuracy measurement. Hardware, Searching and FHE Libraries. The multi-objective search is implemented on top of PyMOO 0.6.1.6, a Python framework for evolutionary optimization on cleartext polynomial model evaluation, with HuggingFace and Transformers 4.57.1 providing off-the-shelf model architectures and weights, and our tensor backend is the PyTorch 2.6.0 + CUDA 12.4. The FHE backend is Phantom-FHE [51], a GPU-accelerated C++ homomorphic encryption library that offers an interface similar to Microsoft SEAL [45] but operates on CUDA-enabled devices. Our FHE inference system additionally relies on LibTorch 2.6.0 + CUDA 12.4 for model weights loading. All search runs and FHE inference latency measurements were conducted on a server equipped with the NVIDIA RTX 4090 48 GB GPUs and an AMD EPYC 7343 16-core 32-thread processor with 256 GB RAM. We implemented the ATLAS FHE inference system in CUDA/C++ on top of NEXUS-EndtoEnd [47] and Phantom-FHE [51], under the RNS-CKKS scheme. The system adopts the ciphertext and plaintext row-packing strategies as well as the PCMM and CCMM from NEXUS-EndtoEnd, and extends the CCMM and PCMM operations to meet the requirements of ViT and LLaMA. Bootstrapping placement in our FHE inference system follows a strategy similar to that of THOR [38]: a level check is placed before every heavy homomorphic operator, and bootstrapping is triggered only when the remaining level of a ciphertext is insufficient for the next operation. Each bootstrapping acts on the 32,768 slots, and the reported #Boot is the number of such 32K-slot bootstrappings consumed by a full forward pass. Search Parameters. The two-stage NSGA-II shares a unified set of hyperparameters across all three architectures. In LayerProblem (Stage 1), the population size is 48 and the algorithm runs for 50 generations, yielding 48 · (50 + 1) = 2,448 FEs. In NetworkProblem (Stage 2), the population size is 96 and the algorithm runs for 225 generations, yielding 96 · (225 + 1) = 21,696 FEs. The total FE budget per run is therefore 24,144 across all models. Both crossover and mutation operators are consistent across stages: crossover is a two-point homogeneous operation with probability 0.9; mutation applies an integer step perturbation to each individual with probability 0.9, while the per-variable mutation rate is set to 1/𝑛 var . We use two objective metrics: MAE (defined in Equation (6)) and the total approximation multiplicative depth. The total approximation multiplicative depth is the sum of the multiplicative depths of all non-linear approximation components in Transformers; the
depth of homomorphic matrix multiplications is excluded because it is fixed for a given architecture and does not affect the search. RNS-CKKS Parameters. We follow the cryptographic parameterization as prior FHE inference works [1, 32, 54]. The cyclotomic polynomial degree is 𝑁 = 216 ; the ciphertext budget is 𝐷 = 28; bootstrapping consumes 𝐾 = 14 levels, leaving 𝐷 − 𝐾 = 14 multiplicative levels after a refresh. The base, special, and bootstrapping moduli each use 51 bits, while the default modulus uses 46 bits [13]. This configuration guarantees 128-bit security. Baselines. Architecture Baseline: We benchmark the proposed ATLAS on three representative Transformer architectures with opensource weights in HuggingFace: the encoder of BERT-base (𝐿=12), the vision encoder of ViT-Base (𝐿=12), and the decoder of LLaMA3-8B (𝐿 =32). For BERT and LLaMA, the number of tokens is set to 128; the original ViT-Base (Patch16-224) sequence length of 197 is padded to 256 to match the slots. Approximation Baseline: For comparison, we reproduce the non-linear approximation recipes of NEXUS [54] and THOR [38] within our FHE inference system. For the iterative softmax of Cho et al. [16], we pair it with NEXUS [54] LayerNorm setting (𝜈=4, 𝛾=2) and a Chebyshev GELU [25, 52] of degree 511 to obtain a high-precision complete end-to-end pipeline. All baselines are evaluated under the same cryptographic and hardware settings. Because the level of input ciphertext, homomorphic MM, packing strategies, and bootstrapping placements vary substantially across frameworks, our reproducing end-to-end latencies of NEXUS [54] and THOR [38] reported here reflect only the approximation configuration and are not directly comparable to those in the original papers; we deliberately adopt a simple row-packing homomorphic MM instead of the highly optimized alternatives (discussed in Section 5.2) to ensure a fair, isolated evaluation of the approximation recipes alone.
4.2
Main Results of ATLAS
For each neural network architecture, ATLAS obtains a set of MAE– depth Pareto configurations with two-stage search (as shown in Figure 4) on a small set of proxy data (e.g., IMDB). Then, ATLAS evaluates these configurations on each downstream task to build the accuracy–depth Pareto front (or accuracy–latency Pareto front, see Appendix F for details). 4.2.1 BERT. The hand-tuned high-precision baseline using iterative softmax [16] (denoted IS) operates at total approximation depth 1356 with 588 bootstrappings, yielding a latency of 1107.5 s on RTX 4090 and accuracies of 93.2% on SST-2, 69.0% on RTE, and 91.3% on QNLI. Figure 6 reports the accuracy–depth Pareto fronts discovered by ATLAS on the three GLUE tasks: SST-2, RTE and QNLI. The fronts are steep: a modest increase in total approximation multiplicative depth yields configurations that match or exceed the IS in accuracy, while reducing latency substantially. For example, on SST-2, S4 comes within 0.2 pp of the IS at 93.0% accuracy and 23.6% lower latency; on RTE, R4 surpasses IS with 70.0% accuracy and 25.6% lower latency. The best configuration differs by task—S4 for SST-2, Q4 for QNLI, R4 for RTE—because different tasks have different preferences for the approximation configuration. This taskspecific preference is exactly what a hand-crafted recipe cannot