Entropie-geleide aandacht voor private LLM's.

Entropy-Guided Attention for Private LLMs

January 7, 2025

Auteurs: Nandan Kumar Jha, Brandon Reagen

cs.AI

Samenvatting

De alomtegenwoordigheid van eigen taalmodellen heeft kritieke privacyzorgen doen rijzen, waardoor vooruitgang op het gebied van privé inferentie (PI) noodzakelijk is, waar berekeningen rechtstreeks op versleutelde gegevens worden uitgevoerd zonder de gevoelige informatie van gebruikers prijs te geven. Hoewel PI een veelbelovende oplossing biedt, wordt de praktische implementatie ervan belemmerd door aanzienlijke communicatie- en latentie-overhead, voornamelijk voortkomend uit niet-lineaire bewerkingen. Om dit aan te pakken, introduceren we een informatietheoretisch kader om de rol van niet-lineariteiten in alleen-decoder taalmodellen te karakteriseren, waarbij een principiële basis wordt gelegd voor het optimaliseren van transformer-architecturen die zijn afgestemd op de eisen van PI. Door gebruik te maken van Shannon's entropie als een kwantitatieve maat, onthullen we de eerder onverkende dubbele betekenis van niet-lineariteiten: naast het waarborgen van trainingsstabiliteit zijn ze cruciaal voor het behouden van diversiteit in aandachtsmechanismen. Specifiek vinden we dat het verwijderen ervan twee kritieke faalmodi activeert: 'entropiecollaps' in diepere lagen die de training destabiliseert, en 'entropische overbelasting' in eerdere lagen die leidt tot onderbenutting van de representatiecapaciteit van Multi-Head Attention (MHA). We stellen een entropie-geleid aandachtsmechanisme voor, gekoppeld aan een nieuwe entropieregularisatietechniek om entropische overbelasting te verminderen. Daarnaast verkennen we PI-vriendelijke alternatieven voor laagnormalisatie om entropiecollaps te voorkomen en de training van LLM's met verminderde niet-lineariteiten te stabiliseren. Onze studie overbrugt de kloof tussen informatietheorie en architectonisch ontwerp, waarbij entropiedynamica wordt vastgesteld als een principiële leidraad voor het ontwikkelen van efficiënte PI-architecturen. De code en implementatie zijn beschikbaar op 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}.