AERO: LLMs solo con Softmax para Inferencia Privada Eficiente

AERO: Softmax-Only LLMs for Efficient Private Inference

October 16, 2024

Autores: Nandan Kumar Jha, Brandon Reagen

cs.AI

Resumen

La omnipresencia de los modelos de lenguaje propietarios ha suscitado preocupaciones sobre la privacidad de los datos sensibles de los usuarios, enfatizando la necesidad de inferencia privada (PI), donde la inferencia se realiza directamente en entradas encriptadas. Sin embargo, los métodos actuales de PI enfrentan sobrecargas de comunicación y latencia prohibitivamente altas, principalmente debido a operaciones no lineales. En este documento, presentamos un análisis exhaustivo para comprender el papel de las no linealidades en los modelos de lenguaje basados en decodificadores de tipo transformador. Introducimos AERO, un marco de optimización arquitectónica de cuatro pasos que perfecciona la arquitectura existente de LLM para una PI eficiente al eliminar sistemáticamente no linealidades como LayerNorm y GELU y reducir el recuento de FLOPs. Por primera vez, proponemos una arquitectura solo con Softmax con muchos menos FLOPs diseñada para una PI eficiente. Además, ideamos una novedosa técnica de regularización de entropía para mejorar el rendimiento de los modelos solo con Softmax. AERO logra reducciones de hasta 4.23 veces en comunicación y 1.94 veces en latencia. Validamos la efectividad de AERO comparándola con el estado del arte.

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.