Nieuw Aandachtsmechanisme Vermindert KV-Cache voor Efficiënte LLM's

Inleiding tot Efficiënte Inferentie in Grote Taalmodellen

De toenemende populariteit van grote taalmodellen (LLM's) en de opkomst van nieuwe inferentieparadigma's hebben de uitdaging van efficiënte grootschalige inferentie naar voren gebracht. Een aanzienlijk knelpunt is de Key-Value (KV) cache binnen traditionele aandachtmechanismen. Deze cache groeit lineair met de batchgrootte en de sequentiële lengte, wat leidt tot een 'geheugenvreter' die de schaalbaarheid en uitbreiding van LLM's belemmert.

Hoewel varianten zoals MQA, GQA en MLA zijn ontwikkeld om dit probleem aan te pakken, ondervinden ze vaak moeilijkheden om de prestaties te behouden onder strikte geheugenbeperkingen. Bovendien introduceren ze complexiteiten die technische uitdagingen en compatibiliteitsproblemen veroorzaken.

Multi-matrix Factorization Attention (MFA)

Een recent artikel van onderzoekers van Stepes, Tsinghua University en andere instellingen introduceert een nieuw aandachtsmechanisme: Multi-matrix Factorization Attention (MFA) en zijn variant MFA-Key-Reuse (MFA-KR). Deze mechanismen verlagen de kosten van inferentie van taalmodellen aanzienlijk en verbeteren tegelijkertijd de prestaties.

MFA en MFA-KR presteren niet alleen beter dan MLA, maar evenaren ook de prestaties van traditionele MHA, terwijl het KV-cachegebruik met tot wel 93,7% wordt verminderd. MFA is ontworpen met het oog op eenvoud, gemakkelijke reproductie, lage gevoeligheid voor hyperparameters en compatibiliteit met verschillende Pos-embedding methoden.

MFA Aanpak en Analyse

Het onderzoeksteam analyseerde het algemene ontwerp en de capaciteit van aandachtmechanismen en identificeerde twee cruciale dimensies met betrekking tot capaciteit. Deze analyse leidde tot de ontwikkeling van nieuwe analytische methoden en ontwerpprincipes. Ze introduceerden het concept van Generalized Multi-Head Attention (GMHA) als een overkoepelend kader voor het begrijpen van verschillende MHA-varianten.

De onderzoekers verkenden ook de berekening en opslag van sleutelwaarden vanuit een inferentieperspectief en onderzochten de modelcapaciteit vanuit een decompositieperspectief. Fully Parameterized Bilinear Attention (FPBA) werd vastgesteld als de theoretische bovengrens van prestaties. Ze ontdekten dat MHA en zijn varianten laag-rang decomposities zijn van FPBA.

Vergelijking met MQA en MLA

De analyse richtte zich op twee representatieve verbeteringsschema's: Multi-Query Attention (MQA) en Multi-Head Latent Attention (MLA).

• MQA gebruikt een agressievere parameter-sharing strategie, waarbij alle aandachtshoofden dezelfde set sleutelwaardeparameters delen. Dit vermindert het geheugengebruik, maar kan de expressiviteit van het model beïnvloeden.
• MLA introduceert een gedeelde latente ruimte voor parametercompressie, maar de werkelijke expressieve kracht wordt beperkt door de kleinste dimensie. Dit betekent dat het verhogen van tussenliggende dimensies de prestaties niet significant verbetert.

MFA Kerninnovaties

De ontwikkeling van MFA werd gedreven door het doel om een aandachtmechanisme te creëren dat het resourceverbruik minimaliseert en de theoretische prestatielimieten benadert. Het ontwerp van MFA omvat drie belangrijke innovaties:

1. Het significant verhogen van het aantal en de dimensie van aandachtshoofden om de modelcapaciteit te maximaliseren.
2. Het toepassen van een agressieve laag-rang decompositiestrategie om de parameterefficiëntie te behouden terwijl het aantal aandachtshoofden en hun dimensies worden uitgebreid.
3. Het gebruik van een enkel sleutelwaardesysteem om het geheugengebruik minimaal te houden, zelfs bij een verhoogde modelcomplexiteit.

Capaciteitsmeting en Vergelijking

Om MFA en andere aandachtmechanismen verder te analyseren, introduceerde het team twee belangrijke metrieken:

• Total Effective Rank (TER): Het product van het aantal aandachtshoofden en de Factorization rank per head (FRH).
• Shared Latent Subspace Dimension (SLSD): De dimensie van de verborgen ruimte die door alle aandachtshoofden wordt gedeeld.

MFA behaalt hogere SLSD en TER in vergelijking met MQA. In vergelijking met MLA behaalt MFA een kleinere KV-cachegrootte en een hogere TER met vergelijkbare parameterbudgetten, terwijl een vergelijkbare SLSD wordt behouden. In vergelijking met traditionele MHA heeft MFA een hogere TER, hoewel de SLSD kleiner is.

Experimentele Resultaten

Uitgebreide experimenten werden uitgevoerd om de prestaties van de nieuwe architectuur op grotere schaal te evalueren. Modellen variërend van 1 miljard tot 7 miljard parameters en trainingsdata van 10 miljard tot 1 biljoen werden getest. MFA vertoonde schaalvermogen vergelijkbaar met traditionele MHA en behield uitstekende prestaties, zelfs op grotere schaal.

Hoewel MFA-KR iets lagere prestaties liet zien, kwam de schaalbaarheid ervan overeen met die van MHA. De geheugenbesparende voordelen van MFA en MFA-KR bleven toenemen met de modelgrootte, waarbij MFA 87,5% geheugenbesparing bereikte en MFA-KR het geheugengebruik terugbracht tot 6,25% op de grootste schaal.

Ablatie Studies

Ablatiestudies valideerden de effectiviteit van MFA en MFA-KR. Hun prestatievoordelen werden ook bevestigd bij verschillende gangbare positionele coderingsmethoden.

Vooruitzichten

MFA biedt aanzienlijke verbeteringen met een eenvoudig ontwerp, waardoor het geheugenknelpunt in LLM-inferentie effectief wordt aangepakt zonder extra technische complexiteit toe te voegen. Het integreert naadloos in het bestaande Transformer-ecosysteem, waardoor de toepassing van LLM's in verschillende scenario's wordt versneld.