AERO : LLM à Softmax uniquement pour une inférence privée efficace

AERO: Softmax-Only LLMs for Efficient Private Inference

October 16, 2024

Auteurs: Nandan Kumar Jha, Brandon Reagen

cs.AI

Résumé

La généralisation des modèles de langage propriétaires a soulevé des préoccupations en matière de confidentialité des données sensibles des utilisateurs, soulignant le besoin d'inférence privée (PI), où l'inférence est effectuée directement sur des entrées chiffrées. Cependant, les méthodes actuelles de PI sont confrontées à des surcoûts de communication et de latence prohibitifs, principalement en raison des opérations non linéaires. Dans cet article, nous présentons une analyse approfondie pour comprendre le rôle des non-linéarités dans les modèles de langage basés sur des décodeurs de transformateurs uniquement. Nous introduisons AERO, un cadre d'optimisation architecturale en quatre étapes qui affine l'architecture existante des modèles de langage à décodeur pour une PI efficace en supprimant systématiquement des non-linéarités telles que LayerNorm et GELU et en réduisant les calculs en FLOPs. Pour la première fois, nous proposons une architecture basée uniquement sur Softmax avec beaucoup moins de FLOPs adaptée à une PI efficace. De plus, nous concevons une nouvelle technique de régularisation de l'entropie pour améliorer les performances des modèles basés uniquement sur Softmax. AERO permet d'atteindre une réduction de la communication jusqu'à 4,23 fois et de la latence jusqu'à 1,94 fois. Nous validons l'efficacité d'AERO en le comparant aux technologies de pointe.

English

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 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.23times communication and 1.94times latency reduction. We validate the effectiveness of AERO by benchmarking it against the state-of-the-art.