Atenção Guiada por Entropia para LLMs Privados

Entropy-Guided Attention for Private LLMs

January 7, 2025

Autores: Nandan Kumar Jha, Brandon Reagen

cs.AI

Resumo

A prevalência de modelos de linguagem proprietários tem levantado preocupações críticas de privacidade, exigindo avanços em inferência privada (IP), onde cálculos são realizados diretamente em dados criptografados sem revelar informações sensíveis dos usuários. Embora a IP ofereça uma solução promissora, sua implementação prática é prejudicada por substanciais sobrecargas de comunicação e latência, principalmente decorrentes de operações não lineares. Para lidar com isso, introduzimos um framework teórico da informação para caracterizar o papel das não linearidades em modelos de linguagem de decodificação única, estabelecendo uma base fundamentada para otimizar arquiteturas de transformadores adaptadas às demandas da IP. Ao alavancar a entropia de Shannon como uma medida quantitativa, descobrimos o significado dual das não linearidades anteriormente não explorado: além de garantir estabilidade de treinamento, elas são cruciais para manter a diversidade de cabeças de atenção. Especificamente, descobrimos que a remoção delas desencadeia dois modos críticos de falha: o "colapso de entropia" em camadas mais profundas que desestabiliza o treinamento, e a "sobrecarga entrópica" em camadas mais iniciais que leva à subutilização da capacidade representacional da Atenção Multi-Cabeça (AMC). Propomos um mecanismo de atenção guiado por entropia combinado com uma técnica de regularização de entropia inovadora para mitigar a sobrecarga entrópica. Além disso, exploramos alternativas amigáveis à IP para normalização de camada a fim de prevenir o colapso de entropia e estabilizar o treinamento de LLMs com menos não linearidades. Nosso estudo preenche a lacuna entre teoria da informação e design arquitetônico, estabelecendo a dinâmica de entropia como um guia fundamentado para desenvolver arquiteturas eficientes de IP. O código e a implementação estão disponíveis em https://github.com/Nandan91/entropy-guided-attention-llm{entropy-guided-llm}.

English

The pervasiveness of proprietary language models has raised critical privacy concerns, necessitating advancements in private inference (PI), where computations are performed directly on encrypted data without revealing users' sensitive information. While PI offers a promising solution, its practical deployment is hindered by substantial communication and latency overheads, primarily stemming from nonlinear operations. To address this, we introduce an information-theoretic framework to characterize the role of nonlinearities in decoder-only language models, laying a principled foundation for optimizing transformer-architectures tailored to the demands of PI. By leveraging Shannon's entropy as a quantitative measure, we uncover the previously unexplored dual significance of nonlinearities: beyond ensuring training stability, they are crucial for maintaining attention head diversity. Specifically, we find that their removal triggers two critical failure modes: {\em entropy collapse} in deeper layers that destabilizes training, and {\em entropic overload} in earlier layers that leads to under-utilization of Multi-Head Attention's (MHA) representational capacity. We propose an entropy-guided attention mechanism paired with a novel entropy regularization technique to mitigate entropic overload. Additionally, we explore PI-friendly alternatives to layer normalization for preventing entropy collapse and stabilizing the training of LLMs with reduced-nonlinearities. Our study bridges the gap between information theory and architectural design, establishing entropy dynamics as a principled guide for developing efficient PI architectures. The code and implementation are available at https://github.com/Nandan91/entropy-guided-attention-llm{entropy-guided-llm}.