Technische Samenvatting: Density Field State Space Models (DF-SSM)

1. Probleemstelling

Large Language Models (LLMs), met name die met 1,3 miljard parameters in FP16-precisie, vereisen ongeveer 2,7 GB aan geheugen, wat de capaciteit van de meeste mobiele, embedded en edge-apparaten overschrijdt. Hoewel kwantisatie de modelgrootte kan verkleinen, resulteert extreme kwantisatie (1–4 bits) doorgaans in significante kwaliteitsdegradatie of vereist het training vanaf nul op enorme datasets (bijv. 100B+ tokens). Bovendien blijft de gewichtcompressie van State Space Models (SSMs) zoals Mamba-2, die architecturale voordelen bieden voor on-device implementatie (vaste grootte van de hidden states waardoor de lineair groeiende KV-cache wordt geëlimineerd), minder onderzocht vergeleken met Transformer-kwantisatie.

2. Methodologie: DF-SSM Pipeline

Het artikel introduceert Density Field State Space Models (DF-SSM), een framework ontworpen om SSMs te comprimeren tot een 1-bit scaffold met int8 low-rank correctie. De aanpak bestaat uit drie primaire stadia:

A. Density Field Weight (DFW) Training

De auteurs maken gebruik van quantization-aware distillation van een vooraf getrainde FP16 teacher (Mamba-2 1.3B).

• 17-Level Quantization: Gewichten worden geparametriseerd als continue latente waarden die via een lineaire clamp worden gemapt naar 17 discrete niveaus tussen -1 en +1.
• Quantization Annealing: Een annealing-parameter $\alpha$ interpoleert tussen continue en gekwantiseerde gewichten, waarbij deze lineair toeneemt van 0 naar 1 tijdens de training. Dit voorkomt een "catastrofale kwaliteitscurve" door het model toe te staan het taalmodelleringsdoel te leren met continue gewichten voordat het zich aanpast aan discrete beperkingen.
• Gradient Handling: Voor gradiënten wordt straight-through estimation gebruikt. De auteurs wijzen sigmoid mapping af ten gunste van lineaire clamping om gradiëntattenuatie bij extreme waarden te voorkomen.

B. Frozen Scaffold met LoRA Correctie

Na de training worden de gekwantiseerde scaffold-gewichten vastgezet op hun deterministisch afgeronde waarden.

• Low-Rank Adaptation: Een kleine, trainbare int8 LoRA-module ($W_A \in \mathbb{R}^{d_{out} \times r}, W_B \in \mathbb{R}^{r \times d_{in}}$ met rang $r=16$) wordt toegevoegd om residuele fouten te corrigeren.
• Architectuur: LoRA wordt toegepast op alle 48 lagen (zowel in_proj als out_proj).
• Mechanisme: De LoRA-correctie corrigeert niet louter per-gewicht afrondingsfouten (correlatie met de afrondingsresiduele is ~0), maar corrigeert input-gewogen outputfouten.
• Kwantisatie: Zowel de LoRA-matrices als de embedding/LM head-matrices worden gekwantiseerd naar int8 met per-channel/row schaling, wat effectief verliesloze perplexiteit bereikt (PPL 49.1 vs. 49.2 voor FP16).

C. Geoptimaliseerde Inference Pipeline

De auteurs ontwikkelen een aangepaste inference engine voor zowel GPU als CPU:

• GPU (CUDA):
  • Scaffold Matmul: 1-bit gepakte gewichten worden on-the-fly uitgebreid naar int8 en verwerkt via cuBLAS INT8 tensor cores (gebruikmakend van torch._int_mm), waarbij gebruik wordt gemaakt van 8x minder geheugenoverdracht en hogere INT8 throughput.
  • Stateful Ops: Custom CUDA-kernels verwerken de stateful shift register (convolutie) en SSM-recurrentie (warp-level reductie).
  • CUDA Graphs: De volledige 48-laagse forward pass wordt vastgelegd als een CUDA graph om kernel launch overhead te elimineren, wat cruciaal is voor batch size 1 prestaties.
• CPU (AVX-512):
  • Een bit-packed kernel breidt gewichten uit naar int8 voor VNNI instructies (_mm512_dpbusd_epi32), wat een hoge throughput bereikt met 8x minder geladen data vergeleken met standaard int8 gewichten.

D. Density Field State Representatie

Om de SSM hidden state te comprimeren, behoudt de auteur een $K \times K$ binair veld per state element ($K^2 \approx 4 \times d_{state}$). De cyclus omvat normalisatie, sigma-delta injectie (error feedback), binarisatie en popcount readout om continue waarden te herstellen, toegepast op chunk-grenzen.

3. Belangrijkste Resultaten

Model Efficiëntie en Grootte

• Compressie: Het resulterende model is 278 MB, een 9,7× reductie ten opzichte van de 2,7 GB FP16 teacher.
• Samenstelling: 56% Scaffold (1-bit), 37% Embedding/LM Head (int8), 4% LoRA (int8), 3% Overig (FP16/FP32).
• Trainingskosten: Distillatie vereist slechts 32M tokens en 6 uur op een enkele A100 GPU, vergeleken met 150B tokens voor het vanaf nul trainen van BitMamba-2.

Inference Prestaties

• GPU Versnelling: Bereikt 21,4× snellere inference op GPU (batch=1) ten opzichte van de geoptimaliseerde mamba-ssm referentie-implementatie. Bij batch=512 nadert de throughput de HBM-bandbreedte limiet van de A100 (~5.081 tok/s).
• CPU Versnelling: Bereikt een 2,9× versnelling op CPU (Intel Xeon, AVX-512) met 567 MB RAM verbruik.

Downstream Taak Prestaties

DF-SSM behoudt prestaties binnen 2–4 procentpunten van BitMamba-2 (een 1,58-bit model getraind vanaf nul op 150B tokens) en behoudt 78–95% van de prestaties van de FP16 teacher:

• BoolQ: 60,8% (Teacher: 64,2%)
• PIQA: 67,1% (Teacher: 73,2%)
• WinoGrande: 54,7% (Teacher: 60,9%)
• ARC-easy: 50,2% (Teacher: 64,1%)
• HellaSwag: 41,4% (Teacher: 59,9%) – Genoteerd als de outlier, waarschijnlijk door onvoldoende distillatie-tokens voor multi-step redeneren.

4. Interpretatie Bevindingen

Het artikel presenteert de eerste systematische analyse van de organisatie van feitelijke kennis in SSMs, wat drie verschillende verwerkingsfasen onthult:

1. Intent Classificatie (Layers 0–3): Werkt in een abstracte ruimte zonder vocabulaire alignment. Classificatie is syntactisch (gedreven door template structuur, bijv. "The capital of...") in plaats van semantisch. Categorieën met identieke templates clusteren onmiddellijk; gemengde-template categorieën doen dat niet.
2. Knowledge Retrieval (Layers 25–35): Factual associaties lokaliseren zich in een 5-laags venster (piek bij L32–L36). Causal intervention experimenten tonen aan dat het wisselen van hidden states in dit venster de feitelijke voorspellingen omkeert. Kennis bezet een laag-dimensionale subspace (10–16 dimensies) van de 2048-dimensionale state.
3. Output Formatting (Layers 36–47): De categorie structuur lost op naarmate het model convergeert naar gedeelde output representaties, ter voorbereiding op het emitteren van tokens.

Belangrijke Observatie: Het model vertoont goed georganiseerde interne kennisstructuren (schone clustering, causale lokalisatie) ondanks zwakke feitelijke recall (correcte antwoorden bereiken zelden de top-1 rang). Dit suggereert dat representatieve structuur vooraf kan gaan aan feitelijke kracht.

5. Betekenis en Claims

Het artikel claimt de volgende bijdragen:

• Efficiëntie: Een volledige pipeline om een 1,3B SSM te comprimeren tot een inzetbaar binair model in 6 uur met minimale distillatie data (32M tokens).
• Prestaties: Een 21,4× inference versnelling op GPU en 2,9× op CPU terwijl de nabijheid van de teacher-prestaties behouden blijft, primair gedreven door geheugen-bandbreedte optimalisatie en INT8 tensor core benutting.
• Architectuur: Een nieuwe mixed-precision architectuur (1-bit scaffold + int8 LoRA) die de kwaliteitscurves van uniforme extreme kwantisatie vermijdt.
• Interpretatie: De ontdekking van een drie-fasen verwerkingsmodel in SSMs en de bevinding dat vroege-laag classificatie syntactisch en template-gedreven is, opererend in een niet-vocabulaire-aligned ruimte.
• Implicatie: De resultaten suggereren dat distillatie efficiënt de organisatorische structuur van kennisrepresentaties kan overdragen, zelfs wanneer de feitelijke inhoud slechts gedeeltelijk behouden blijft, wat impliceert dat extra training tokens bestaande structuren kunnen invullen in plaats van ze te reorganiseren.

Erkende Beperkingen:

• De PPL degradatie (49.2 vs. 14.3) geeft aan dat er ruimte is voor verbetering via langere training of hogere LoRA ranks.
• De interpretatie bevindingen zijn specifiek voor één gedestilleerd model (1.3B Mamba-2); generalisatie naar andere schalen of architecturen getraind vanaf nul is onverifieerd.
• De "1-bit" label geldt voor de scaffold (56% van de gewichten); embeddings en LoRA blijven int8, consistent met andere binaire SSM methoden.