Encryption-Friendly LLM Architecture

We are happy to hear that the reviewer found our work valuable and convincing. We respond to the individual comments in the following.

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.

We ask the reviewer to refer to our common response.

The description in 2.1 does not seem to be consistent with Figure 1.

Thank you for raising this point. We revisited Section 2.1 and Figure 1 but could not fully understand the reviewer's concern about consistency. Could the reviewer kindly clarify which aspect in unclear or inconsistent? We would certainly want to address this issue.

Why does the statement “LLM weights are protected in the strict cryptographic sense (line 149)” hold?

Our paper states that the LLM weights are not protected in the strict cryptographic sense. This is because although the pre-trained LLM resides on the server side, their weights are not encrypted and model theft attacks are not mitigated by our approach.

Although the paper proposes a privacy-protecting LLM architecture, the security considerations of the model, especially against attacks and model theft, is insufficient.

We agree with the reviewer, and this is a point we acknowledge in our paper. We emphasize that our focus is on protecting the privacy of the user and protecting the server (model) is a separate and very active thread of research.

One core idea, Replacing Softmax with Gaussian kernel to solve difficulty of evaluation is very similar to Skyformer [2]. 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?

Thank you for this precise question. We would like to use this response as an opportunity to convey the rationale for using GK.

First and foremost, we clarify that we were aware of Skyformer, but omitted it in discussing the prior work as a simple honest mistake. (We did reference similar prior work such as SOFT [1].) We have incorporated the reference in the updated paper.

The difference between the Skyformer and our contribution is as follows. First, we view Gaussian kernel from a different perspective from Skyformer. On the one hand, the original self-attention is formulated as

where $A_{ij}=\exp\left(q_i^\top k_j/\sqrt{p}\right)$ ($q_i$ and $k_j$ are the $i$-th and $j$-th row of the query and the key, respectively) and $D$ is a diagonal matrix with diagonal element $\exp\left(QK^T/\sqrt{p}\right)\cdot \mathbf{1}$. The authors of Skyformer view this as a normalization of $A$, the attention scores, by $D$.

Skyformer modifies Gaussian kernel as

where $D_Q$ and $Q_K$ are diagonal matrices with $\left(D_Q\right)\_{ii}=\exp\left(\frac{\|q_i\|^2}{\sqrt{p}}\right)$ and $\left(D_K\right)_{ii}=\exp\left(\frac{\|k_i\|^2}{\sqrt{p}}\right)$ where $q_i$ and $k_i$ mean the $i$-th row of the query and key matrices, respectively. From this modification, the authors of Skyformer view Gaussian kernel as another normalization version of the attention scores through $D_Q$ and $D_k$. On the other hand, we interpret Gaussian kernel as an attention scoring mapping `without normalization'. Second, Skyformer applies Nystr"om method to approximate the kernel to accelerate the calculation. However, we do not use such an approximation. Third, the motivation of Skyformer is to handle long sequence lengths (not a concern for us) while our motivation is to avoid normalization operations, which are difficult to implement under HE.

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.

We ask the reviewer to refer to our common response.

Does the use of LoRA and Gaussian Kernel affect the interpretability of the model? How does the author balance efficiency and explainability?

The interpretability and explainability of deep learning models are very important to using these models more safely. However, this issue is not a focus of our work, especially since we consider the setup where the entire computation is meant to be hidden behind homomorphic encryption.

[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.

We are happy to hear that the reviewer found our work valuable, particularly for addressing fine-tuning in the non-interactive setup. We respond to the individual comments in the following.

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.

Thank you for this precise point. We chose the encoder-only model BERT due to the computational constraints. (Decoder-only models require a larger minimum scale to ``make them work'' compared to BERT.) Whether our findings with BERT translate to decoder-only models should be tested in future work, but we are quite optimistic since the architectural forms between the two types of model are quite similar with the only difference being the causal masks in the attention layers (and we can straightforwardly incorporate causal masks into attention layers with Gaussian kernels).

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.

Our model is based on the semi-honest security model. The client owns the personal data and sends these with HE encrypted ciphertext to the server. The server conducts fine-tuning and inference under HE. In our model, the client will know only the inference result, and the server only knows the encrypted client data. As a result, the entire security model relies on the semantic security of the underlying CKKS scheme.

We have clarified the adversary model in our revised paper.

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.

We believe that homomorphic encryption (HE) will play a role in the future use of some language model applications. Our work contributes to this progress by initiating considerations of architecture design specifically tailored for HE. While the landscape of large language model (LLM) research is evolving rapidly, we anticipate that the findings of this work will remain relevant amidst these changes, since we are targeting fundamental components of the transformer architecture.

We are happy to hear that the reviewer found our work well-motivated and the results valuable. We respond to the individual comments in the following.

It is worth mentioning https://arxiv.org/pdf/2410.00433, which appeared after the submission deadline.

Thank you for providing us with the reference to the P3EFT paper. This work focuses on energy consumption, so the focus is different from ours, but there is certainly some overlap in the techniques, so we have included the reference in the updated manuscript.

I am not particularly impressed by the novelty of this paper.

Regarding the novelty of our work, we ask the reviewer to refer to our common response.

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

Yes, we will provide the source code for the experiments. All of the code for the experiments under plaintext will be provided as PyTorch code. For the experiments under HE we will provide an implementation for LoRA-style CCMM (Appendix D) and polynomial approximations in Python linked with the HE transformer backbone written in C++. Recall that our HE code is not high-level PyTorch code, so it is taking us some time to organize the code into a GitHub repo.

Also for reproducibility (and in lack of code), I am interested in understanding better the polynomial approximations used.

Our code for the polynomial approximations will provide full clarity on the precise details on the approximations we 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?

We clarify that we pre-train (in plaintext) a new transformer model that is compatible with the Gaussian kernel. Because of the computational constraints of HE, the language model is a manageable size in plaintext, so this pre-training is not a significant computational cost, compared to the cost of inference and fine-tuning under HE.

The range penalizing loss in the following is applied to the last 100 steps of pre-training.

Here, the loss term multiplied by $\lambda$ is added to $L_{\text{pre}}$ in the last 100 steps of the pre-training. During these 100 steps, the input domains of the non-polynomial functions in $\mathcal{F}$ are narrowed to $\left[-L_f, L_f\right]$ so that it lies in the approximation domains. Other papers such as [1] use this penalizing loss so that the inputs of the non-polynomial functions do not deviate for inference. We found that using this new loss term prevents the input range in the fine-tuning from being too wide.

[1] Itamar Zimerman, Moran Baruch, Nir Drucker, Gilad Ezov, Omri Soceanu, and Lior Wolf. Converting transformers to polynomial form for secure inference over homomorphic encryption. International Conference on Machine Learning, 2024.