pQuant: Towards Effective Low-Bit Language Models via Decoupled Linear Quantization-Aware Training

1. Het Probleem: Parameter Democratization

Hoewel kwantisatie (het verlagen van de precisie van modelgewichten) essentieel is voor het efficiënt inzetten van Large Language Models (LLMs) op randapparatuur (edge devices), lopen bestaande methoden voor extreem lage bitbreedtes (sub-2-bit) tegen fundamentele beperkingen aan.

• Achteruitgang in Nauwkeurigheid: Bestaande Quantization-Aware Training from scratch (QAT-Scratch) methoden, zoals BitNet, presteren vaak onder de 80% van de FP16-baselines, wat hun praktische toepasbaarheid beperkt.
• Schaalprobleem: Naarmate het model groter wordt, verbetert de prestatie van 1-bit modellen niet lineair, in tegenstelling tot FP16-modellen.
• De Kernoorzaak: Parameter Democratization: De auteurs identificeren een nieuw fenomeen: bij extreme kwantisatie (1-bit) wordt de gevoeligheid van alle parameters homogeen. In plaats dat een kleine subset van "gevoelige" parameters (die cruciaal zijn voor de output) wordt onderscheiden, worden alle parameters behandeld alsof ze even belangrijk zijn. Dit verlies aan structuur beperkt de expressiviteit van het model.

2. Methodologie: pQuant

Om dit probleem op te lossen, stellen de auteurs pQuant voor. Dit is een QAT-Scratch methode die de lineariteit van de laag ontkoppelt om parameter-specialisatie te forceren.

A. Ontkoppelde Lineaire Laag (Decoupled Linear Layer)

In plaats van alle gewichten naar 1-bit te kwantiseren, splitst pQuant de lineaire lagen (vooral in de Feed-Forward Networks) in twee gespecialiseerde takken:

1. Dominante 1-bit Tak: De meeste parameters (ongeveer 95-96%) worden gereduceerd tot 1-bit (INT1) voor maximale rekenefficiëntie.
2. Compacte Hoog-precisie Tak: Een klein subset van de parameters (ongeveer 4-5%) wordt bewaard in 8-bit (INT8). Deze tak is specifiek bedoeld om de meest "gevoelige" en informatieve parameters te behouden.

B. Feature Scaling en Dynamische Toewijzing

Een naïeve verdeling van gewichten werkt niet omdat de gevoelige parameters niet op vaste posities zitten in een willekeurig geïnitieerd model. pQuant lost dit op door:

• Feature Scaling: De uitgangen van de 1-bit en 8-bit takken worden vermenigvuldigd met leerbare scalars ($\alpha$ en $\beta$).
• Initiële Voorkeur: Tijdens training wordt $\alpha$ (voor de 8-bit tak) initieel veel groter ingesteld dan $\beta$ (bijv. $\alpha=2.0, \beta=0.2$). Dit dwingt het model om de belangrijkste representaties via de hoog-precisie tak te laten lopen en zorgt voor een sterkere gradiëntfeedback voor deze tak.
• Resultaat: Het model leert dynamisch welke parameters gevoelig zijn en wijst deze automatisch toe aan de 8-bit tak.

C. Schaalbaarheid via Experts (MoE)

Om de capaciteit verder te schalen zonder de inferentie-kosten te verhogen, wordt de 8-bit tak uitgebreid naar meerdere parallelle takken (experts). Een lichte router (Top-1 gating) selecteert per token de meest geschikte expert. Dit werkt als een Mixture of Experts (MoE) maar met een cruciaal verschil: de "experts" verschillen in precisie (8-bit vs 1-bit), niet alleen in gewichten.

3. Belangrijkste Bijdragen

• Identificatie van Parameter Democratization: Het paper introduceert en valideert het concept dat het homogeniseren van parametergevoeligheid de hoofdblokkade is voor 1-bit LLMs.
• Decoupled Architectuur: Een nieuwe architectuur die 1-bit efficiëntie combineert met een kleine, dynamisch toegewezen 8-bit buffer voor kritieke parameters.
• State-of-the-Art Prestaties: pQuant bereikt prestaties die dicht in de buurt komen van FP16-modellen, zelfs bij een gemiddelde bitbreedte van slechts ~1.35 bit.
• Efficiënte Schaalbaarheid: De methode toont aan dat het verhogen van de capaciteit via hoog-precisie experts leidt tot een schaalwettig gedrag dat veel dichter bij FP16-modellen ligt dan bij traditionele 1-bit modellen.

4. Resultaten

De experimenten werden uitgevoerd op modellen van 300M tot 2.6B parameters, getraind vanaf nul (QAT-Scratch).

• Prestaties:
  • pQuant (1.35-bit) op een 1.3B model behaalt een perplexiteit van 17.2 op WikiText-2, vergeleken met 21.8 voor BitNet (1-bit) en 16.9 voor BitNet1.58 (2-bit).
  • Het model presteert beter dan bestaande Post-Training Quantization (PTQ) methoden (zoals PTQ1.61) die op veel grotere modellen zijn getraind.
  • Op downstream taken (zoals ARC, BoolQ, etc.) overtreft pQuant 1-bit baselines met een significant marge en benadert de prestaties van 2-bit modellen, terwijl het minder bits gebruikt.
• Schaalbaarheid:
  • Bij het verhogen van het aantal 8-bit experts (N=8) sluit pQuant de prestatiekloof met FP16 LLaMA-2 aanzienlijk in, terwijl 1-bit BitNet hierin faalt.
• Efficiëntie:
  • Inferentie: pQuant is 38% sneller dan BitNet1.58 en 82% sneller dan FP16 LLaMA-2 op CPU (Apple M2).
  • Geheugen: Het model vereist 31% minder geheugen dan BitNet1.58 en 92% minder dan FP16-modellen, omdat de 1-bit gewichten in UINT8 formaat worden opgeslagen (8 parameters per byte).

5. Betekenis en Impact

pQuant is een doorbraak voor de implementatie van LLMs op randapparatuur (edge devices).

• Hardware Vriendelijkheid: Door het gebruik van 1-bit bitwise operaties voor het grootste deel van de berekening, wordt de rekenlast drastisch verlaagd. De toevoeging van een kleine 8-bit component is hardware-vriendelijk en vereist geen complexe floating-point hardware.
• Praktische Toepasbaarheid: Het bewijst dat het mogelijk is om modellen te trainen die bijna even goed presteren als FP16-modellen, maar met een fractie van het geheugen en de rekenkracht. Dit opent de deur voor krachtige AI op smartphones en IoT-apparaten zonder afhankelijkheid van de cloud.
• Nieuwe Richting: Het paper verschuift de focus van "alle parameters gelijk behandelen" naar "intelligente, dynamische toewijzing van precisie", wat een nieuwe richting aangeeft voor onderzoek in extreem lage bitbreedte kwantisatie.

Kortom, pQuant lost het expressiviteitsprobleem van 1-bit modellen op door een slimme hybride architectuur te introduceren, waardoor extreem efficiënte en toch hoogpresterende taalmodellen mogelijk worden.