Atención guiada por entropía para LLMs privados

Entropy-Guided Attention for Private LLMs

January 7, 2025

Autores: Nandan Kumar Jha, Brandon Reagen

cs.AI

Resumen

La omnipresencia de modelos de lenguaje propietarios ha suscitado preocupaciones críticas sobre la privacidad, lo que ha hecho necesario avanzar en la inferencia privada (PI), donde los cálculos se realizan directamente en datos encriptados sin revelar la información sensible de los usuarios. Aunque la PI ofrece una solución prometedora, su implementación práctica se ve obstaculizada por importantes sobrecargas de comunicación y latencia, que provienen principalmente de operaciones no lineales. Para abordar esto, introducimos un marco teórico de información para caracterizar el papel de las no linealidades en los modelos de lenguaje de solo decodificador, sentando así las bases para optimizar arquitecturas de transformadores adaptadas a las exigencias de la PI. Al aprovechar la entropía de Shannon como medida cuantitativa, descubrimos el doble significado de las no linealidades, previamente no explorado: más allá de garantizar la estabilidad del entrenamiento, son cruciales para mantener la diversidad de cabezas de atención. Específicamente, encontramos que su eliminación desencadena dos modos críticos de falla: el "colapso de entropía" en capas más profundas que desestabiliza el entrenamiento, y la "sobrecarga entrópica" en capas anteriores que conduce a una subutilización de la capacidad representativa de la Atención Multi-Cabeza (MHA). Proponemos un mecanismo de atención guiado por entropía junto con una novedosa técnica de regularización de entropía para mitigar la sobrecarga entrópica. Además, exploramos alternativas amigables con la PI para la normalización de capas con el fin de prevenir el colapso de entropía y estabilizar el entrenamiento de LLMs con menos no linealidades. Nuestro estudio cierra la brecha entre la teoría de la información y el diseño arquitectónico, estableciendo la dinámica de la entropía como guía fundamentada para el desarrollo de arquitecturas de PI eficientes. El código y la implementación están disponibles en 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}.