Attention guidée par l'entropie pour les LLM privés

Entropy-Guided Attention for Private LLMs

January 7, 2025

Auteurs: Nandan Kumar Jha, Brandon Reagen

cs.AI

Résumé

La prévalence des modèles de langage propriétaires a soulevé des préoccupations critiques en matière de confidentialité, nécessitant des avancées en matière d'inférence privée (PI), où les calculs sont effectués directement sur des données chiffrées sans révéler les informations sensibles des utilisateurs. Bien que la PI offre une solution prometteuse, son déploiement pratique est entravé par des surdébits de communication et de latence substantiels, principalement dus aux opérations non linéaires. Pour y remédier, nous introduisons un cadre informationnel pour caractériser le rôle des non-linéarités dans les modèles de langage à décodeur uniquement, posant ainsi des bases fondées pour l'optimisation des architectures de transformateur adaptées aux exigences de la PI. En exploitant l'entropie de Shannon comme mesure quantitative, nous révélons la double importance des non-linéarités jusqu'alors inexplorée : au-delà d'assurer la stabilité de l'apprentissage, elles sont cruciales pour maintenir la diversité des têtes d'attention. Plus précisément, nous constatons que leur suppression déclenche deux modes d'échec critiques : l'« effondrement entropique » dans les couches plus profondes qui déstabilise l'apprentissage, et la « surcharge entropique » dans les couches plus précoces qui conduit à une sous-utilisation de la capacité de représentation de l'Attention Multi-Têtes (MHA). Nous proposons un mécanisme d'attention guidé par l'entropie associé à une nouvelle technique de régularisation de l'entropie pour atténuer la surcharge entropique. De plus, nous explorons des alternatives favorables à la PI à la normalisation des couches pour prévenir l'effondrement entropique et stabiliser l'apprentissage des LLM avec des non-linéarités réduites. Notre étude comble le fossé entre la théorie de l'information et la conception architecturale, établissant la dynamique de l'entropie comme guide fondé pour le développement d'architectures PI efficaces. Le code et l'implémentation sont disponibles sur 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}.