AERO: LLMs Apenas com Softmax para Inferência Privada Eficiente

AERO: Softmax-Only LLMs for Efficient Private Inference

October 16, 2024

Autores: Nandan Kumar Jha, Brandon Reagen

cs.AI

Resumo

A pervasão de modelos de linguagem proprietários tem levantado preocupações com a privacidade dos dados sensíveis dos usuários, enfatizando a necessidade de inferência privada (IP), onde a inferência é realizada diretamente em entradas criptografadas. No entanto, os métodos de IP atuais enfrentam sobrecargas de comunicação e latência proibitivamente mais altas, principalmente devido a operações não lineares. Neste artigo, apresentamos uma análise abrangente para compreender o papel das não linearidades em modelos de linguagem baseados em decodificadores do tipo transformer. Introduzimos AERO, um framework de otimização arquitetural de quatro etapas que aprimora a arquitetura existente de LLM para IP eficiente, removendo sistematicamente não linearidades como LayerNorm e GELU e reduzindo contagens de FLOPs. Pela primeira vez, propomos uma arquitetura apenas com Softmax com significativamente menos FLOPs adaptada para IP eficiente. Além disso, desenvolvemos uma nova técnica de regularização de entropia para melhorar o desempenho de modelos apenas com Softmax. AERO alcança redução de comunicação de até 4,23 vezes e redução de latência de 1,94 vezes. Validamos a eficácia do AERO comparando-o com o estado-da-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.