Reduced Matrix Multiplication: Input-Adaptive Matrix-Product Reduction for LLM Inference

Technische Samenvatting: Reduced Matrix Multiplication (RMM) voor LLM-inferentie

Probleemstelling
Transformer-gebaseerde taalmodellen behalen sterke prestaties, maar brengen aanzienlijke inferentiekosten met zich mee door herhaalde hoogdimensionale matrixvermenigvuldigingen in attention- en feed-forward (MLP)-lagen. Hoewel eerder werk de efficiëntie van inferentie heeft aangepakt via gestructureerde pruning, low-rank benadering, tokencompressie of KV-cachebeheer, modificeren deze methoden doorgaans vaste modelstructuren of verkorten ze inputs/caches. Ze pakken niet direct aan of de gecontracteerde berekening binnen elke matrixproduct kan adaptief worden verminderd voor een specifieke input zonder de modelgewichten te wijzigen. Bestaande methoden voor activatie-schaarsheid (bijv. TEAL, CATS) slaan elementen met een lage magnitude over, maar richten zich op het schaar maken van verborgen toestanden (hidden states) in plaats van het reduceren van de gedeelde contractieassen van algemene matrixproducten.

Methodologie: Reduced Matrix Multiplication (RMM)
De auteurs stellen Reduced Matrix Multiplication (RMM) voor, een training-vrije, input-adaptieve inferentiemethode. RMM werkt op de verenigde vorm van Transformer-berekeningen, $Y = AB$, waarbij $A$ de activatiematrix vertegenwoordigt en $B$ de gewichtsmatrix (of intermediaire representatie) vertegenwoordigt.

1. Kernmechanisme: In plaats van het volledige product $AB$ te berekenen, selecteert RMM dynamisch een subset van indices $I$ langs de gedeelde contractiedimensie $d$ op basis van de magnitude van de huidige input-activaties. Het berekent het gereduceerde product $A_{:,I} B_{I,:}$.
2. Selectiestrategie: De methode maakt gebruik van activatie-bewuste dimensie-selectie. Voor elke feature-dimensie $j$ wordt een belangrijkheidsscore $s_j = \|A_{:,j}\|_2$ (de kolomnorm van de activatiematrix) berekend. De top-$k$ indices met de hoogste scores worden behouden, waarbij $k = \lceil \rho d \rceil$ en $\rho$ een door de gebruiker beheerde retentieratio is.
3. Theoretische basis: De auteurs bewijzen dat TopK-selectie op basis van kolomnorm minimax optimaal is voor het minimaliseren van de worst-case benaderingsfout over alle mogelijke gewichtsmatrices $B$ wanneer alleen $A$ bekend is op het moment van selectie. Dit biedt een principiële rechtvaardiging voor het gebruik van activatie-magnitudes om de reductie te sturen.
4. Toepassingsbereik: RMM wordt toegepast op:
   • Attention-lagen: Het reduceren van de feature-dimensie in $QK^\top$ (query-key interactie) en optioneel de token-dimensie in $PV$ (attention-value interactie).
   • MLP en Lineaire Projecties: Het reduceren van de verborgen dimensie in feed-forward projecties ($XW$).
5. Implementatie: De methode vereist het berekenen van feature-scores en het uitvoeren van top-$k$ selectie, wat lichtgewicht vectoroperaties zijn vergeleken met de dense matrixvermenigvuldigingen die ze vervangen. Er zijn aangepaste Triton-kernels geïmplementeerd om deze besparingen in de praktijk te realiseren.

Belangrijkste Bijdragen

• Input-Adaptieve Reductie: Introduceert een nieuw perspectief op inferentie-optimalisatie dat de gecontracteerde berekeningen van matrixproducten zelf vermindert, wat verschilt van het prunen van gewichten, het comprimeren van tokens of het beheren van caches.
• Training-Vrij en Controleerbaar: De methode vereist geen hertraining en biedt een vloeiende, voorspelbare nauwkeurigheid-efficiëntie afweging die wordt gecontroleerd door een eenvoudige retentieratio ($\rho$).
• Mechanistische Inzichten: Door middel van ablatie-studies onthullen de auteurs een structurele asymmetrie in Transformers:
  • Attention-zijde berekeningen (Q, K, V projecties, $QK^\top$, $PV$) zijn aanzienlijk meer reduceerbaar en robuust tegen agressieve pruning.
  • MLP-componenten zijn significant gevoeliger voor reductie, waarbij de "Up" projectie het meest gevoelig is en de "Down" projectie minder gevoelig is. Het prunen van de gehele MLP-block leidt tot ernstige prestatie-instorting.
• Generalisatie: Demonstreert dat het principe verder reikt dan alleen tekstmodellen naar multimodale vision-language inferentie (bijv. Qwen2.5-VL).

Experimentele Resultaten
De auteurs evalueerden RMM over modellen variërend van 1B tot 70B parameters (inclusief de LLaMA 3.1/3.2 en Qwen families) en diverse taken (QA, redeneren, samenvatten, lange context, en vision-language).

• Prestaties versus Statische/Random Baselines: Onder vaste retentieratio's (bijv. $\rho=0.5$), presteert RMM consistent beter dan statische pruning-methoden (SparseGPT, Wanda, SliceGPT) en random pruning. Statische methoden lijden onder instabiele degradatie, terwijl RMM coherente generatie en redeneervermogens behoudt.
• Schaalbaarheidstrends: Grotere modellen tolereren over het algemeen agressievere reductie. Zo blijft LLaMA 3.1 70B dicht bij de volledige modelprestaties bij $\rho=0.8$ op de meeste benchmarks, terwijl kleinere modellen sneller degraderen, vooral bij uitdagende taken zoals GSM8K en HumanEval.
• Robuustheid: RMM blijft stabiel in autoregressieve generatie en long-context settings (tot 30K tokens), waarbij een vloeiende degradatie wordt getoond in plaats van een abrupte fout.
• Runtime Efficiëntie: Benchmarks op NVIDIA A100 GPU's met aangepaste kernels laten zien dat computationele besparingen vertalen naar praktische wall-clock versnellingen, vooral bij langere sequentielengtes (bijv. 1.40x versnelling bij 4096 tokens). Opvallend genoeg liep de dense implementatie bij 4096 tokens vast door gebrek aan geheugen, terwijl RMM de inferentie voltooide.
• Compatibiliteit: De methode is compatibel met INT8 gewichtsquantisatie en breidt zich uit naar diverse vision-language model backbones.

Betekenis en Claims
Het artikel positioneert RRM als een schaalbare richting voor input-adaptieve inferentie-optimalisatie. De primaire betekenis ligt in het aantonen dat redundantie in Transformer-inferentie niet uniform verdeeld is; het is in plaats daarvan geconcentreerd in specifieke componenten (met name de attention-zijde operaties) en varieert per input. Door gebruik te maken van deze structurele asymmetrie en input-afhankelijke activatie-informatie, bereikt RMM controleerbare efficiëntiewinsten zonder de noodzaak van modelhertraining of structurele modificatie. De auteurs concluderen dat matrix-product-niveau adaptieve reductie een veelbelovende weg is voor efficiënte Transformer-inferentie, wat suggereert dat toekomstige methoden rekening moeten houden met component-specifieke redundantiepatronen.

Beperkingen
Het werk richt zich op een training-vrije formulering. De auteurs erkennen dat zij geen extensies buiten deze setting hebben verkend (bijv. het integreren van adaptieve reductie in pretraining doelstellingen) noch een gedetailleerde studie naar redundantiepatronen hebben uitgevoerd over alle componenten van complexe vision-language modellen buiten de initiële evaluatie. Daarnaast, hoewel aangepaste kernels versnellingen laten zien, blijft volledige integratie met alle inferentie-frameworks en gekwantiseerde backends toekomstig werk.