Encryption-Friendly LLM Architecture

Authors: Donghwan Rho, Taeseong Kim, Minje Park, Jung Woo Kim, Hyunsik Chae, Ernest K. Ryu, Jung Hee Cheon
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Homomorphic Encryption, PPML, Encrypted Fine-tuning, Transformer Architecture

Abstract:
Large language models (LLMs) offer personalized responses based on user interactions, but this use case raises serious privacy concerns. Homomorphic encryption (HE) is a cryptographic protocol supporting arithmetic computations in encrypted states and provides a potential solution for privacy-preserving machine learning (PPML). However, the computational intensity of transformers poses challenges for applying HE to LLMs. In this work, we propose a modified HE-friendly transformer architecture with an emphasis on inference following personalized (private) fine-tuning. Utilizing LoRA fine-tuning and Gaussian kernels, we achieve significant computational speedups---6.94$\times$ for fine-tuning and 2.3$\times$ for inference---while maintaining performance comparable to plaintext models. Our findings provide a viable proof of concept for offering privacy-preserving LLM services in areas where data protection is crucial. Our code is available on GitHub.

Decision: accept (poster)
Reviewer scores: 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1730252438072

thread_type:
official_review

<thread_object>
creator: reviewer_essp

<content>
summary:
This paper starts from the poor performance problem of homomorphic encryption (HE) in transformer architecture and focuses on optimizing the speed of HE in it. Specifically, the scheme tries to avoid CCMM computations to solve the poor performance brought by full fine-tuning. In addition, this paper aims to address the difficulty of evaluating under HE with Softmax, the core idea of which is to replace Softmax with a Gaussian kernel. Finally, this paper carried out many experiments, and the results show that their scheme is comparable to existing schemes in terms of modeling performance, while at the same time, the computational speed has been significantly improved. The overall narrative of the paper is clear, and has good logic.

strengths:
This paper is oriented to the problem of inefficiency of transformer architecture under HE, although the existing research has produced richer results. The main contribution of this paper is to enhance the speed of transformer architecture under HE, and the authors have carried out many experiments to verify the rationality and advantages of the scheme. I think the experiments in this paper are full, and the advantages of this paper are elaborated in terms of speed and model performance, which speed is emphasized in this paper. The experimental results are thorough and well-analyzed. Overall, this paper seems to incorporate some of the SOTA approaches in the current field and apply them to a widely researched topic.

weaknesses:
1. Insufficient innovation. First, the topic chosen for this paper is a more widely studied one. Second, the solutions in this paper seem to be a direct combination and application of existing advanced schemes, and it is not intuitively obvious in the paper that the authors have improved on existing methods.
2. The description in 2.1 does not seem to be consistent with Figure 1. Furthermore, why does the statement “LLM weights are protected in the strict cryptographic sense (line 149)” hold?
3. Although the paper proposes a privacy-protecting LLM architecture, the security considerations of the model, especially against attacks and model theft, is insufficient.

questions:
1.	One core idea of this paper to solve the difficulty of evaluating Softmax under HE is to replace it with Gaussian kernel, however, this method is very similar to that in the literature “Chen, Yifan, et al. Skyformer: Remodel self-attention with gaussian kernel and nystr\" om method. Advances in Neural Information Processing Systems 34 (2021): 2122-2135”. However, this paper does not describe the difference with this paper or even cite this literature. Is there a difference in the performance of the two? Furthermore, the polynomial approximation method of this paper seems to be very common way.
2.	In Section 3, is the author's approach in solving Bottleneck 1 just an application of some existing methods? If not, please clarify the differences and improvements.
3.	Does the use of LoRA and Gaussian Kernel affect the interpretability of the model? How does the author balance efficiency and explainability?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730497928501

thread_type:
official_review

<thread_object>
creator: reviewer_f3ds

<content>
summary:
As LLMs surge, privacy issue becomes important considering government-level regulations. Among methods in privacy-preserving machine learning, homomorphic encryption (HE) can provide cryptographic security by computing over encrypted data directly without extra communication like MPC. However, HE is not efficient to compute matmul or non-poly operations in the transformer-level scale. This work focus on fine-tuning stage with LoRA and softmax variant to create HE-friendly transformer architecture.

strengths:
1. Compared to MPC approaches, n on-interactive property helps HE to be feasible to compute over large-scale LLMs without including considerable communication overhead among computing parties.
2. This work has a great focus on the fine-tuning stage to make LLMs secure for users, which also concentrates on the key components like attention layers in the transformer, and it is also combined with SoTA techniques like LoRA to make the process more efficient.
3. Writing with bottleneck-improvement pattern for LoRA and GK looks good for readers to figure out key ideas.

weaknesses:
1. When you mentioned SoTA LLMs, you should notice that decoder-based models have been proved very powerful in generative tasks. After iteration of the recent few years, BERT series is not as useful and prevalent as decoder models. Hence, the significance to protect BERT-based model is less essential in the current LLMs.
2. Although this work introduces how HE and CKKS work in the secure way, this work does not specify adversary model, such ability of adversaries, type of adversaries (e.g., semi-honest or malicious) and kind of attacks (e.g., member inference attack) this architecture counters with.
3. In the conclusion, this work is too vague on the future work. For example, how cryptographic community develops helps the improvement of this work (e.g., any change on HE). Also, how LLMs itself evolve may change security issue based on this work.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730565304148

thread_type:
official_review

<thread_object>
creator: reviewer_yq2m

<content>
summary:
This paper presents optimizations to LLM architectures to make them more friendly towards Homomorphic Encryption evaluation. Concretely, first they propose to use LoRA for fine-tuning, where the data used for fine-tuning is given in encrypted form. This is useful to avoid re-training a lot of parameters, since LoRA focuses on updating only a few weights. This minimizes the amount of homomorphic operations w.r.t. plain fine-tuning. They also observe that LoRA is useful for reducing the dimension of homomorphic matrix multiplications. Secondly, they replace Softmax-based with Gaussian Kernel attention. They show that this is a simpler function to evaluate in HE, leading to significant savings with little impact in accuracy.

strengths:
The problem statement is well motivated. Several works are currently exploring evaluation of LLMs under FHE, and the potential applications are also quite compelling. This is an extremely difficult task in terms of achieving viable efficiency, and any method that advances the state-of-the-art in this direction is welcome. The results achieved here show that several optimizations in other domains, that is, LoRA and the use of Gaussian Kernels, turn out to be useful for the evaluation of LLMs in FHE. I am not aware of this observation being made and explored in this depth in other papers (it is worth mentioning https://arxiv.org/pdf/2410.00433, which appeared after the submission deadline).

weaknesses:
I am not particularly impressed by the novelty of this paper. It uses existing FHE tools with existing ML optimizations. This may not be a weakness on its own given the positive results of combining these techniques, but I still think the improvement factors may not be big enough for these techniques to become "enablers" of private LLM applications in practice. Put differently, I am not convinced that the gains here are a significant enough to overcome the blockers that prevent LLMs + FHE from becoming more widely spread.

questions:
For reproducibility, can the authors comment on the source code? Whether they intend to make it public for validation?

Also for reproducibility (and in lack of code), I am interested in understanding better the polynomial approximations used. In Section E the authors talk about penalizing the model in a "pre-training stage". I am not sure I understand how this would work, especially in the context of secure inference (no fine-tuning). What do the authors mean exactly? The model owner retrains the model using this new loss function? Most prior works take a pre-trained model, changing only its non-linearities by polynomial while keeping its weights. What are the authors proposing to do here? Does this require changing the model's weights? Re-training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732507002218

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response

comment:
We sincerely thank all reviewers for their insightful comments. We are happy to hear that all reviewers found our work to be well-motivated and the writing to be clear.

&nbsp;



At the same time, however, there were some concerns about the simplicity of our approach, particularly the observation that our solutions are direct applications of existing schemes. Indeed, our contribution is not in the novelty of such mechanisms themselves. Rather, our contribution is identifying a novel application of these existing mechanisms.

&nbsp;

- ### **Use of LoRA**

The use of LoRA is not an obvious choice. In the standard plaintext setting, LoRA is used primarily to reduce the memory footprint, but the memory footprint is not a concern of our work under HE. Our contribution is finding a new purpose for LoRA: reduction in computational cost by avoiding large ciphertext-ciphertext matrix multiplications (CCMMs).

&nbsp;

- ### **Use of Gaussian kernel (GK)**
The use of Gaussian kernel (GK) is also not an obvious choice. In the standard plaintext setting, prior works such as SOFT [1] and Skyformer [2], use GK with a low-rank approximation to handle long sequence lengths, but sequence length is not a concern for us. Our contribution is finding a new purpose for GK: an attention layer that is inherently simpler to polynomially approximate compared to the softmax. (Prior work such as HETAL[3] devotes a significant amount of effort to adequately approximate the softmax function, but we circumvent this difficulty by using the GK.)

&nbsp;

Further, we quickly reiterate two points highlighting our contributions.

- This work represents a somewhat significant engineering effort. There is no existing library for implementing ML models using the CKKS scheme, so we had to develop our implementation code from scratch. This includes implementing backpropagation with LoRA, along with their corresponding CKKS-specific computations, in a multi-GPU setup with data parallelism.


- Our work is the first to successfully fine-tune an entire transformer model (as opposed to fine-tuning only the head) under HE. Considering the significant computational constraints imposed by HE, it was not previously evident that fine-tuning under such conditions would be practically feasible. Our work provides the first positive proof of concept in this direction.

&nbsp;  

In our view, our work represents meaningful steps of progress made through simple but non-obvious proposals. We kindly ask the AC and the reviewers to take this perspective into consideration.

&nbsp;

[1] Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. SOFT: Softmax-free transformer with linear complexity. Neural Information
Processing Systems, 2021.

[2] Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. Skyformer: Remodel self-attention with gaussian
kernel and nystr\"{o}m method. Neural Information Processing Systems, 2021.

[3] Seewoo Lee, Garam Lee, Jung Woo Kim, Junbum Shin, and Mun-Kyu Lee. HETAL: Efficient
privacy-preserving transfer learning with homomorphic encryption. International Conference on
Machine Learning, 2023b.
</content>
</thread_object>

</thread>


Metareview:
metareview:
All reviewers except one (EsSP) argued for accepting the paper. For this reviewer their main conerns were on #1 a lack of novelty, #2 incorrect descriptions, #3 insufficient model security considerations. Specifically for #1 the reviewer argued that (a) the topic is widely studied, (b) the solutions are a direct combination of existing methods, and (c) it is not clear that the authors have improved on existing methods. I disregard (a), as topic popularity alone should not count against a paper so long as (b) and (c) are dealt with. For (b) the authors argue that even though LoRA and the Gaussian kernel have been applied to LLM architectures, they have been used for very different purposes: LoRA to reduce memory, and the Gaussian kernel to handle long sequence lengths. They are using each in different, unintended ways: LoRA to avoid large ciphertext-ciphertext matrix multiplications, and the Gaussian kernel to find a better polynomial approximation, compared to prior work which directly tried to approximate the softmax function. The reviewer had not response to this. I am convinced by the authors’ arguments on both points. It is clear that LoRA would be useful for HE-friendly fine-tuning. And the observations that one need only approximate part of the Gaussian kernel and that this is well-approximated by an easy-to-compute polynomial approximation is brilliant. For (c), the reviewer did not elaborate what baselines should be compared against. In the paper the authors say “Only a few attempts, such as those by Lee et al. (2022b) and HETAL (Lee et al., 2023b), have explored fine-tuning, focusing exclusively on the classification head while leaving other key components like attention layers and feed-forward networks untouched. P3EFT (Li et al., 2024) also addresses fine-tuning of foundation models, but it concentrates more on energy consumption during fine-tuning.” The authors compare against HETAL and P3EFT was released the day of the ICLR deadline so a comparison against this, while I’d like to see it in the future, is out of scope for this submission. Concern #1 is resolved. For #2 the reviewer argued that Section 2.1 was not consistent with Figure 1. The authors revised both of these things but asked for further clarification to understand the inconsistency. The reviewer did not respond. I found no inconsistencies, so concern #2 is resolved. For #3 the reviewer argues that model security is not addressed. The authors agree and stated that their focus is on user security and that model security is separate. I agree and would further argue that one could apply recent developments in model security (e.g., differential privacy for LLMs) to achieve this. Including a new method for model security would create a very confusing conference submission for which most details would have to go into the appendix. This resolves concern #3. Overall the paper demonstrates an impressive new capability: fine-tuning an LLM (beyond just the heads) with HE, which outperforms recent work. The authors have agreed to release their code which will have a big impact on the community.  Given these things, I vote to accept. Authors: you’ve already made improvements to respond to reviewer changes, if you could double check their comments for any recommendation you may have missed on accident that would be great! The only additional request I would like to make is for you to run an experiment where you vary the approximation levels of LoRA and the Gaussian kernel to show the trade-off between computation speed and model performance. I would really recommend highlighting this aspect of your approach as it allows for much more varied use-cases beyond methods without such tunable knobs. After incorporating these changes the paper will make a nice contribution to the conference!