EmbBERT: Attention Under 2 MB Memory

Probleemstelling

Transformer-architecturen, gebaseerd op het attention-mechanisme (zoals BERT), hebben de natuurlijke taalverwerking (NLP) revolutionair verbeterd. Echter, hun enorme geheugen- en rekenvereisten maken ze onbruikbaar voor ultra-beperkte apparaten zoals wearables en Internet-of-Things (IoT)-eenheden. Deze apparaten hebben vaak slechts enkele megabytes aan geheugen (bijvoorbeeld < 2 MB) en beperkte verwerkingskracht. Bestaande "kleine" taalmodellen, zoals MobileBERT of zelfs sterk gereduceerde varianten van BERT, zijn nog steeds te groot voor deze omgevingen. Er is een dringende behoefte aan een taalmodel dat specifiek is ontworpen om binnen een strikt geheugenbudget van 2 MB te opereren zonder in te leveren op prestaties.

Methodologie: De EmbBERT Architectuur

De auteurs introduceren EmbBERT, een Tiny Language Model (TLM) dat volledig opnieuw is ontworpen voor extreme efficiëntie. In plaats van bestaande modellen simpelweg te verkleinen, heeft het team de architecturale componenten van transformers herzien om het geheugengebruik te minimaliseren.

De kerncomponenten van EmbBERT zijn:

1. Nano Embedder Block: Een compacte invoerlaag die tokens omzet in dichte representaties. Deze gebruikt een "learned positional embedding" en projecteert tokens via een dense layer naar een gereduceerde ruimte ($rd$) voordat ze worden geprojecteerd naar de volledige dimensie ($d$). Dit vermindert de parametergrootte aanzienlijk ten opzichte van standaard embeddings.
2. Efficient Encoder Blocks: De kern van het model bestaat uit een reeks van $N$ blokken die twee parallelle paden combineren:
   • Efficient Attention: Een lichtgewicht variant van self-attention die expliciete Key- en Value-projectiematrices elimineert. In plaats van $3hd^2$ parameters voor projecties, gebruikt deze slechts $hd^2$. Dit reduceert de parameters met ongeveer 66% en verlaagt de rekenlast. Het model gebruikt slechts één attention-head ($h=1$).
   • Convolutional Skip Connection: Een lichtgewicht lokale token-mixer die gebruikmaakt van dieptewijze 1D-convoluties (depthwise 1D Conv) in plaats van zware feed-forward lagen.
3. Gewogen Aggregatie: In plaats van de output van de attention- en convolutiepaden simpelweg op te tellen, gebruikt EmbBERT een geleerde gewogen som (met een verschilterm: $\lambda_{EA} \Delta_{EA} - \lambda_{CS} \Delta_{CS}$) om de informatie te combineren.

Geheugenberekening:
Het totale geheugen ($M_{tot}$) wordt berekend als de som van de gewichten en activeringen, waarbij rekening wordt gehouden met de precisie (standaard 32-bit float, maar compatibel met kwantisatie). De formule is:
$$M_{tot} = (W_{emb} + N \cdot W_{enc}) \cdot P_w + \max(A_{emb}, A_{enc}) \cdot P_a$$
Door de architectuur te optimaliseren, past het volledige model (inclusief activeringen) binnen 2 MB.

Belangrijkste Bijdragen

1. EmbBERT Model: Een nieuw TLM dat specifiek is ontworpen voor microcontrollers, met een totaal geheugengebruik van slechts 2 MB.
2. Analyse van Geheugen en Rekenkracht: Een gedetailleerde wiskundige analyse van de geheugenvoetafdruk en complexiteit van elke component, wat een hulpmiddel biedt voor het afwegen van trade-offs in TinyML.
3. Kwantisatie en Schaalbaarheid: Het model is robuust tegen 8-bit post-training kwantisatie, wat het geheugengebruik verder verlaagt naar 781 kB. Daarnaast toont de architectuur schaalbaarheid van sub-megabyte tot tientallen megabytes.
4. TinyNLP Benchmark: De auteurs hebben een nieuwe, gespecialiseerde benchmark (TinyNLP) ontwikkeld met datasets die specifiek relevant zijn voor embedded NLP-taken (zoals sentimentanalyse en intentieherkenning), om eerlijke vergelijkingen mogelijk te maken.

Resultaten

De prestaties van EmbBERT zijn getest op de TinyNLP benchmark en de GLUE suite, en vergeleken met andere modellen binnen het 2 MB-budget (zoals BERT(2MB), MAMBA(2MB), NanoBERT) en grotere referentiemodellen (BERT-Tiny, 20 MB).

• TinyNLP Benchmark: EmbBERT behaalde een gemiddelde nauwkeurigheid van 87,19%. Dit is beter dan BERT-Tiny (20 MB), die 86,99% haalde, en aanzienlijk beter dan andere 2 MB-modellen (bijv. NanoBERT: 86,33%, BERT(2MB): 83,93%).
• GLUE Benchmark: EmbBERT behaalde een score van 63,50, opnieuw beter dan BERT-Tiny (63,16%) en alle andere 2 MB-modellen.
• Kwantisatie: Na 8-bit kwantisatie daalde de geheugenvoetafdruk naar 781 kB met een verwaarloosbare prestatiedaling (slechts -0,7% op GLUE), en zelfs een lichte verbetering op TinyNLP.
• Schaalbaarheid: Experimenten met schaalvariërende configuraties (van 0,5 MB tot 40 MB) tonen aan dat de EmbBERT-architectuur schaalbaar is en zelfs bij grotere maten (bijv. 40 MB) consistent beter presteert dan BERT-Tiny.

Betekenis en Conclusie

EmbBERT bewijst dat het mogelijk is om hoogwaardige NLP-taken uit te voeren op apparaten met extreem beperkte middelen (zoals wearables en IoT), waarvoor eerder geen taalmodellen beschikbaar waren. Het model overtuigt door:

• Efficiëntie: Het bereikt prestaties die vergelijkbaar zijn met modellen die 10x meer geheugen vereisen.
• Architecturale Innovatie: Door het combineren van compacte embeddings, geoptimaliseerde attention-mechanismen en convolutieve skip-connections, wordt een betere balans gevonden tussen parameteraantal en activeringsgeheugen dan bij "downscaled" bestaande modellen.
• Reproduceerbaarheid: Alle code, scripts en checkpoints zijn openbaar beschikbaar gemaakt.

Dit werk legt de basis voor de volgende generatie "TinyML" toepassingen, waarbij intelligente taalverwerking lokaal en in real-time kan plaatsvinden op apparaten met een geheugenbudget van slechts enkele megabytes, wat privacy verbetert en latentie verlaagt.