Technische Samenvatting: Regulariseren of Localiseren: Wanneer de K/V-cache Geometrie tijdens de Training Voordeel Biedt bij Kwantisatie

Probleemstelling

Transformer-taalmodellen die getraind worden met cross-entropy (CE) ontwikkelen doorgaans een anisotrope verborgen-toestand geometrie, waarbij representaties concentreren langs enkele gedeelde richtingen. Hoewel Joint-Embedding Predictive Architectures (JEPA) doelstellingen hebben geïntroduceerd zoals SIGReg (een geschetste Epps–Pulley normaliteitstest) om embeddings richting isotropie ($N(0, I)$) te duwen, hebben eerdere pogingen om SIGReg op Large Language Models (LLMs) toe te passen nulresultaten opgeleverd. Deze mislukkingen traden op wanneer SIGReg werd toegepast als een hulp bij niet-standaard stromen of alleen op de toestanden van de laatste laag tijdens fine-tuning.

Er blijft een kritische kloof bestaan: Verandert SIGReg de geometrie van de standaard autoregressieve (AR) pretraining, en belangrijker nog, levert deze vormgeving voordeel op voor de Key/Value (K/V) cache tijdens post-hoc kwantisatie? Bestaand werk richt zich op post-hoc compressie of architecturale veranderingen, waardoor de vraag of distributionele regularisatie tijdens de training van de standaard K/V-cache geometrie de robuustheid tegen kwantisatie kan verbeteren, grotendeels onverkend blijft.

Methodologie

De auteurs hanteren een experimenteel ontwerp in drie fasen met behulp van 110M-parameter Llama-stijl modellen die zijn voorgetraind op 10B FineWeb tokens.

1. Fase 1: Dose-Response en Verborgen Toestand Karakterisering

• Doel: Bepalen of het toevoegen van SIGReg aan de CE-loss tijdens pretraining de anisotropie van de verborgen toestand vermindert.
• Opzet: De auteurs doorlopen een sweep van de regularisatie-gewicht $\lambda$ ($0, 0.001, 0.01, 0.1, 1.0$).
• Metriek: Pairwise-cosine anisotropie en perplexiteit.
• Bevinding: $\lambda=0.01$ wordt geïdentificeerd als het optimale werkpunt, wat de anisotropie significant vermindert met minimale perplexiteitskosten.

2. Fase 2: Bridge Probes (Adaptatie vs. Retrofit)

Om te bepalen of de verandering in geometrie doorwerkt naar de K/V-cache en hoe dit te bereiken, testen de auteurs drie specifieke probes uitgaande van een $\lambda=0.01$ checkpoint:

• A1 (Adaptatie): Het deblokkeren van de trunk en het voortzetten van de training waarbij SIGReg direct wordt toegepast op zowel de verborgen toestanden ($h$) als de K/V-projecties ($K_\ell, V_\ell$) over alle lagen.
• A2 (Frozen Prediction): Het trainen van per-layer predictoren op een bevroren trunk om te zien of de K/V-geometrie hersteld kan worden vanuit enkel de verborgen toestanden.
• A3 (Frozen Decoding): Het trainen van decoders om de K/V te reconstrueren vanuit de laatste verborgen toestand van de bevroren trunk.
• Controle: Een "matched continuation" waarbij de trunk wordt voortgezet voor hetzelfde token-budget maar zonder de K/V-regularisatieterm ($\lambda_{KV}=0$).

3. Fase 3: Kwantisatie Evaluatie

De auteurs evalueren de resulterende modellen onder post-hoc fake kwantisatie (kwantiseren $\to$ dekwantiseren) van de K/V-cache.

• Geteste schema's: Symmetrische group-free kwantisatie (per-token vs. per-channel), asymmetrische grids (met zero-points), en gemengde arrangementen (KIVI-stijl: per-channel K, per-token V).
• Metriek: $\Delta$NLL (Negative Log-Likelihood) ten opzichte van de eigen full-precision referentie van het model.

Belangrijkste Resultaten

1. Verborgen Toestand Geometrie

Bij $\lambda=0.01$ vermindert SIGReg de pairwise-cosine anisotropie van de verborgen toestand met 38% over drie gepaarde seeds. De perplexiteit stijgt met minder dan 0.35%, en de zero-shot capaciteiten (LAMBADA, HellaSwag, etc.) blijven stabiel zonder significante degradatie.

2. K/V Cache Geometrie en de "Retrofit" Mislukking

• Geen Automatische Propagatie: De reductie in anisotropie in de verborgen toestanden verplant zich niet automatisch naar de K/V-cache. De frozen-trunk probes (A2 en A3) falen in het herstellen van de geometrie, wat aangeeft dat K/V-anisotropie niet louter een functie is van de uiteindelijke verborgen toestand.
• Directe Regularisatie is Vereist: Alleen A1 (voortgezette training met directe K/V-regularisatie) slaagt erin de cache-anisotropie te verminderen, waarbij een reductie van 94% in de gemiddelde cache-anisotropie wordt bereikt.
• Isolatie van Mechanisme: De matched continuation zonder de K/V-term ($\lambda_{KV}=0$) reproduceert dit effect niet, wat de K/V-term isoleert als de noodzakelijke oorzaak.

3. Kwantisatie Voordeel

De voordelen van de geregulariseerde geometrie zijn sterk afhankelijk van het kwantisatieschema:

• Coarse Scales (Group-Free): Onder ongetransformeerde symmetrische group-free kwantisatie is het A1-geadapteerde model de enige conditie die de voorkeur geeft aan per-channel scaling boven per-token scaling. Onder 3-bit symmetrische per-channel kwantisatie ondervindt de baseline 4.3–7.9× meer $\Delta$NLL schade dan het A1-model.
• Fine-Grained Scales (KIVI-Stijl): Onder een volledige gesimuleerde KIVI-stijl configuratie (gemengd arrangement, zero-points en grouped scales met $g=32$), verdwijnt het voordeel. Alle modellen, inclusief de baseline, bereiken bijna gelijke prestaties.
• Opslag Overhead: Zelfs wanneer rekening wordt gehouden met de metadata-overhead van grouped zero-points (waardoor 3-bit grouped opslag vergelijkbaar wordt met 4-bit group-free), verdwijnt het voordeel van de training-tijd onder de token-lokale grouping en mixed scaling configuratie.

Betekenis en Claims

Het artikel beweert de eerste te zijn die een distributionele regularisatie van de standaard K/V-cache geometrie tijdens de training evalueert tegen post-hoc cache kwantisatie.

• Geometrische Controle: De auteurs demonstreren dat de K/V-cache geometrie expliciet gecontroleerd kan worden tijdens de training via directe regularisatie, maar dat dit effect niet voortkomt uit enkel de verborgen-toestand regularisatie en niet kan worden "geretrofit" via frozen-trunk probes.
• Conditionele Bruikbaarheid: De interventie tijdens de training biedt een tastbaar voordeel alleen wanneer de kwantiteitsschalen grof zijn (over lange token-extents verspreid). Het voordeel is specifiek voor de interactie tussen de geregulariseerde geometrie en het kwantisatieschema; het generaliseert niet naar moderne, fijnmazige configuraties (zoals KIVI-stijl met zero-points en lokale grouping).
• Interpretatie van Metriek: Het artikel verheldert dat de pairwise-cosine anisotropie metriek de gemiddelde-richting uitlijning meet (gedomineerd door gedeelde gemiddelden) in plaats van covariantie whitening. De resultaten tonen een dissociatie aan: terwijl de gemiddelde uitlijning daalt, gedragen spectrale eigenschappen (zoals top-1 share) zich anders voor K en V lagen, wat suggereert dat de regularisator primair werkt op de gemiddelden/rogue richtingen.

Beperkingen

De auteurs beperken hun claims expliciet:

• Regime: Resultaten zijn beperkt tot 110M parameters en de FineWeb corpus; grotere modellen met verschillende cache-lay-outs (GQA/MQA) zijn niet getest.
• Simulatie: Kwantisatie wordt gesimuleerd (fake quantization) zonder metingen van wall-clock tijd of geheugen.
• Scope van Probes: Negatieve resultaten voor frozen probes zijn gebaseerd op specifieke capaciteiten en budgetten; een "middenweg" (bijv. LoRA-schaal adaptatie) blijft ongetest.
• Geen "Eerste Isotropie" Claim: De auteurs beweren niet de eersten te zijn die voor isotropie regulariseren (zij citeren CosReg, SimCTG, etc.), maar zijn de eersten die specifiek de K/V-cache geometrie targeten voor kwantisatie-robuustheid.