AERO: Softmax-Only LLMs for Efficient Private Inference

Authors: Nandan Kumar Jha, Brandon Reagen
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Private inference, LLMs, Architectural Optimization, Entropy Regularization

Abstract:
The pervasiveness of proprietary language models has raised privacy concerns for users' sensitive data, emphasizing the need for private inference (PI), where inference is performed directly on encrypted inputs. However, current PI methods face prohibitively higher communication and latency overheads,  primarily due to nonlinear operations. In this paper, we present a comprehensive analysis to understand the role of nonlinearities in transformer-based decoder-only language models.  We introduce AERO, a four-step architectural optimization framework that refines the existing LLM architecture for efficient PI by systematically removing nonlinearities such as LayerNorm and GELU and reducing FLOPs counts. For the {\em first time}, we propose a Softmax-only architecture with significantly fewer FLOPs tailored for efficient PI.  Furthermore, we devise a novel entropy regularization technique to improve the performance of Softmax-only models. AERO achieves up to 4.23$\times$ communication and 1.94$\times$ latency reduction. We validate the effectiveness of AERO by benchmarking it against the state-of-the-art.

Decision: reject
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730506945334

thread_type:
official_review

<thread_object>
creator: reviewer_rkod

<content>
summary:
This paper proposes the AERO framework, a new method to optimize large language models (LLMs) in resource-constrained contexts, such as private inference (PI), by systematically reducing non-linear operations until Softmax becomes the model’s sole source of non-linearity, thereby enhancing computational efficiency. The main contributions of this work are as follows:

1. In the Softmax-only model, the authors apply weight normalization and scaling techniques to the linear layers within the feed-forward network (FFN), effectively preventing training collapse. This shift from activation normalization to weight normalization avoids non-linear calculations during the inference phase, thereby improving inference efficiency.

2. In the Softmax-only architecture, the authors merge two linear layers in the FFN into a single linear layer, reducing the FFN’s FLOPs by 8x without compromising model performance. Additionally, they prune the deeper FFN layers, further enhancing computational efficiency.

3. To address issues of entropy overload and entropy collapse that are prone to occur during training in the Softmax-only architecture, the authors propose an entropy regularization method. By penalizing extreme entropy values during training, this method ensures a balanced entropy distribution across attention heads, maintaining model stability and performance.

strengths:
The paper demonstrates a range of strengths in terms of originality, quality, clarity, and significance, as detailed below:

The AERO framework presents a novel approach by systematically reducing non-linear operations until Softmax becomes the model’s sole non-linearity, thereby enhancing computational efficiency. Its originality lies in the combination of several innovative techniques, including the use of weight normalization and scaling instead of traditional activation normalization, merging FFN layers to reduce FLOPs by 8x, and applying entropy regularization to address issues of entropy overload and entropy collapse during training.

The paper employs a rigorous research methodology, conducting experiments across multiple models (such as GPT-2 and Pythia) and various context lengths. Detailed ablation studies demonstrate the effectiveness of each component within the AERO framework, such as weight normalization and entropy regularization. These experimental results enhance the technical reliability of the paper and support its main conclusions.

The paper is well-structured, with clear explanations of its motivations, methods, and experimental results. Complex concepts, such as entropy regularization and the transition from activation normalization to weight normalization, are thoroughly explained, aiding readers in understanding the rationale and benefits behind these technical choices.

By improving the efficiency of LLMs in PI scenarios, the AERO framework addresses a critical issue in deploying LLMs in resource-constrained environments. This approach holds significant value in the field of privacy-preserving machine learning, potentially facilitating broader applications of LLMs in privacy-sensitive or resource-limited settings. The demonstrated reductions in FLOPs, communication costs, and inference latency contribute meaningfully to the field.

weaknesses:
1. Lack of Detail in Entropy Regularization Implementation:  
   While entropy regularization is an important part of the AERO framework, the paper does not provide sufficient detail on how the regularization thresholds (for addressing entropy overload and entropy collapse) are selected. A clearer explanation or additional experiments testing different threshold values could strengthen the understanding of this technique and its practical application.

2. Limited Comparison with Existing Non-linear Reduction Techniques:  
   The paper primarily focuses on its unique approach but provides limited comparison with other existing methods that aim to reduce non-linear operations in LLMs. Incorporating experiments or discussions comparing AERO with similar frameworks could help highlight its relative strengths and weaknesses, making the contributions clearer.

3. Absence of Optimization with Specific Privacy-Preserving Techniques in Real-world PI Applications:  
   Although the paper demonstrates AERO’s efficiency improvements in terms of FLOPs, communication cost, and inference latency, it lacks evaluations in real-world PI application scenarios, especially in conjunction with specific privacy-preserving techniques, such as multi-party computation or homomorphic encryption. Including benchmarks that evaluate AERO’s performance when combined with these privacy techniques could further showcase the method’s practicality and effectiveness in actual applications.

questions:
1.Could the authors provide more detail on how the entropy regularization thresholds were selected for addressing entropy overload and entropy collapse? Were these values determined empirically, or was there a specific criterion used? Further insights here would help clarify the robustness of the regularization approach.

2.The paper highlights the novelty of reducing non-linear operations to a Softmax-only model, but could the authors provide additional comparisons with other existing methods beyond LayerNorm-free design? (Are there any more SOTA that aim to achieve similar goals?)  

3.While weight normalization was shown to prevent training collapse, how does it affect the model's generalization ability in different contexts? 

4.Could the authors elaborate on how AERO performs in real-world private inference scenarios, especially when integrated with specific privacy-preserving techniques like homomorphic encryption or multi-party computation?

5. Realizing the limited computational resource and time, the reviewer is still curious about AERO's performance on large models other than GPT-2 and Pythia-70M.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730565905505

thread_type:
official_review

<thread_object>
creator: reviewer_aqto

<content>
summary:
This paper studies the efficient private inference of LLM and proposes an algorithm using only Softmax nonlinear functions. The problem of privacy inference studied is of practical value. The proposed scheme can greatly reduce the nonlinear operation in LLM and has practical application value.

strengths:
The proposed scheme can greatly reduce the nonlinear operation in LLM and has practical application value.

weaknesses:
The introduction does not fully introduce the motivation of the research, and does not elaborate on whether the proposed algorithm can solve the challenges mentioned in the introduction. The experimental results lack more detailed evaluation of test accuracy and other indicators.

questions:
1. The introduction of the paper emphasizes the importance and challenge of private inference. However, how the proposed AERO addresses these challenges is not fully described. Why reducing nonlinear operations can help LLMS perform private inferences needs to be explained in more detail.

2. How much gain and benefit will the proposed scheme bring to private inference? This requires further objective numerical evaluation.

3. The author describes in the limitation that this work mainly discusses the PPL performance of the model. What is the test performance on the actual NLP task? This needs to be explained in the experimental results section.

4. The content of the appendix is not in good shape, and there is no text introduction under many sub-headings, which makes it difficult for readers to obtain information from them.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688278650

thread_type:
official_review

<thread_object>
creator: reviewer_kstg

<content>
summary:
This paper presents an optimization framework to make large language models more efficient for private inference by minimizing non-linear operations such as LayerNorm and GELU. The proposed architecture, AERO, includes a Softmax-only model that reduces both communication and latency overheads. A novel entropy regularization technique is introduced to prevent training instabilities and entropic overload in the model.

strengths:
1. This paper highlights an important research area—private inference (PI)—and extends the study of PI for LLMs by providing valuable insights into the impact of LayerNorm, a topic that has not been extensively explored before.

2. The insights into entropic overload and the proposed solution, entropy regularization, are both novel contributions.

weaknesses:
1. Replacing GeLU with ReLU has been proposed in the work [1]. Despite this paper proposes insightful analysis, this method is not totally novel.
[1] Dake Chen, Yuke Zhang, Souvik Kundu, Chenghao Li, and Peter A Beerel. Rna-vit: Reduceddimension approximate normalized attention vision transformers for latency efficient private inference. In IEEE/ACM International Conference on Computer Aided Design (ICCAD), 2023.

2. Although AERO achieves notable reductions in communication overhead and latency, there is a trade-off in terms of higher perplexity compared to baseline models. As Iron [2] provides efficient private inference protocols for layernorm, I would suggest the authors to compare the PPL and PI savings with Iron.

[2] Meng Hao, Hongwei Li, Hanxiao Chen, Pengzhi Xing, Guowen Xu, and Tianwei Zhang. "Iron: Private inference on transformers." Advances in neural information processing systems 35 (2022): 15718-15731.

3. The acronym FFN is first introduced in line 70, but its full name, Feed-Forward Network, does not appear until line 108. This slight inconsistency in placement could lead to some confusion for readers.

questions:
The proposed AERO framework and experiments primarily focus on models with fewer than 1B parameters, limiting the insights into its performance and scalability on larger LLMs commonly used in industry. This might restrict the applicability of the findings to more demanding real-world scenarios. What are the authors' perspective on how AERO might perform with larger models and whether they anticipate similar benefits or new challenges at a larger scale?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732066849104

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (1/2)

comment:
We sincerely thank all reviewers for their comprehensive, thoughtful, and insightful feedback, as well as their encouraging evaluations. We deeply appreciate the time and effort invested in carefully reading and evaluating our work. We are particularly pleased that the reviewers recognized the novelty of our contributions and acknowledged the potential impact of our insights in advancing LLM private inference.

In response to the reviewers' thoughtful suggestions, we have revised our manuscript to include several new additions in the Appendix (**highlighted in blue** in the revised version): 

**C.3: Cryptographic Protocols for Linear and Nonlinear Operations:** Detailed analysis of cryptographic protocols for linear (FLOPs) and nonlinear operations (GELU, ReLU, LayerNorm, Softmax) used for LLM private inference. 

**D.7: Training Dynamics in Softmax-only Models with Fewer FFNs:** Analyzes how pruning deeper FFNs in Softmax-only architecture affects training stability, demonstrated as entropy collapses in deeper layers (Figure 16 and Figure 17). 

**D.8: Mitigating Over-Regularization with an Appropriate Threshold Margin:** Describes the strategies for selecting threshold margins to avoid over-regularization in our entropy-based regularization implementations (Figure 18 and Figure 19). 

**I: AERO Beyond Private Inference: Broader Impact and Promises**  Discusses the broader implications of the principled approach to architectural simplification in AERO, as well as the entropy regularization techniques. 

Additionally, as suggested by **Reviewer aqto**, we have added descriptions to contextualize and understand the experimental results presented in the Appendix. We have also updated the Related Work section (Appendix Section F) to include the prior work on the  nonlinearities's role in LLM

We would be glad to receive further feedback on the new results and welcome any perspectives on the discussion presented in Appendix I.

We sincerely hope that our revisions address all major concerns and questions raised by the reviewers. We remain eager to discuss and address any outstanding issues and further improve the work based on the reviewers' guidance.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The reviewers had reasonable arguments in favour of the paper and against it: on one hand they appreciated the approach of reducing non-linear operations and the structure of the work, on the other they had concerns with (a) higher perplexity and (b) limited comparison with prior work. All reviewers responded to the author feedback with very short responses indicating that they would not change their scores. No reviewers engaged in further discussion of the paper. After going through the paper and the discussion I have decided to vote to reject based on the above issues. Specifically, for (a) a reviewer pointed out that AERO had higher perplexities than baselines, and proposed an additional baseline. The authors argued that the baseline was out of scope as it has been shown to be much slower than Bumblebee. I agree with the authors on this. However, this sidesteps the concern about increased perplexity. Other tasks in the appendix show increased perplexity as well, and while this is to be expected it highlights a key downside to the approach: there is no way to tune the perplexity/communication/latency trade-off. If, after Step 4 of Figure 8, one ends up with a model that has a perplexity that is too high, they have no alternative. Whereas, methods that are aimed at approximations of non-linearities often have ways to trade-off the accuracy of the approximation with the communication/latency of the resulting network. This poses a serious threat to the applicability of the method. For (b), reviewers wanted to see comparisons with other methods for PI. While the above baseline was convincingly argued to be out of scope by the authors, the authors also argued that polynomial approximation methods have different design goals and face limitations including data-specific accuracy dependencies and narrow input ranges and so they are also out of scope. I disagree on this point: this only means that these methods cannot be applied on certain tasks, but this can be indicated in a results section. There is nothing stopping the authors from comparing against these methods. Given all of the above, I believe this work should be rejected at this time. Once these things and other issues mentioned in the reviews are addressed in an updated version, the work will be much improved.