An Exploration of Mamba for Speech Self-Supervised Models

Probleemstelling

Hoewel Mamba (een Selective State Space Model) uitstekende prestaties heeft geleverd in taalmodellen door zijn lineaire complexiteit en efficiëntie bij lange sequenties, blijft de toepassing ervan in spraak-zelftoezichtende leermodellen (SSL) onderbelicht. Bestaande studies beperken zich vaak tot geïsoleerde taken en tonen aan dat Mamba in spraaktoepassingen vaak achterblijft bij Transformer-varianten, tenzij er hybride ontwerpen worden gebruikt die de lineaire schaalbaarheid tenietdoen.

De kernproblemen die dit paper adresseert zijn:

1. Rekenkosten: Transformers hebben een kwadratische rekencomplexiteit ($O(N^2)$) ten opzichte van de sequentielengte, wat leidt tot hoge kosten en geheugenproblemen (OOM - Out of Memory) bij lange audio-opnames.
2. Lange context: Het ontbreken van een unificatie van Mamba als een fundamenteel spraakmodel dat zowel voor lange context als real-time (streaming) toepassingen kan worden gebruikt.
3. Representatiekwaliteit: Het is onduidelijk of Mamba-architecturen even goede of betere fonetische en spreker-gerelateerde representaties leren dan Transformers.

Methodologie

De auteurs introduceren Mamba-based HuBERT, een model dat de Transformer-blokken in de standaard HuBERT-architectuur vervangt door Mamba-blokken, terwijl de feature-extractiemodules (CNN) behouden blijven.

• Architectuur:
  • Causale setting: Vergelijking tussen Mamba en Causal Transformer (met een onder-triagonale masker).
  • Bidirectionele setting: Vergelijking tussen standaard Transformer en twee varianten van Mamba: ExtBiMamba (twee volledig onafhankelijke Mamba-lagen) en InnBiMamba (gedeelde projecties, alleen SSM-modules gedupliceerd).
• Training:
  • Het model wordt getraind volgens het HuBERT-pipeline op het volledige 960-uurs LibriSpeech-dataset.
  • Twee iteraties: eerst met MFCC-doelen, daarna met pseudo-labels gegenereerd door de zesde laag van de eerste iteratie.
  • Training vond plaats op één NVIDIA V100 GPU (beperkte resources), wat resulteerde in een kleinere batchgrootte dan in de originele HuBERT-papers, maar voor eerlijke vergelijking werden ook de Transformer-baselines onder dezelfde beperkingen getraind.
• Evalueringsstrategie:
  • Fine-tuning: Voor Automatic Speech Recognition (ASR) in zowel lange context (document-level) als streaming (causale) scenario's.
  • Representatie-analyse: Gebruik van Phone Purity (k-means clustering) en Canonical Correlation Analysis (CCA) om fonetische en spreker-kenmerken te analyseren.
  • Downstream taken: Evaluatie op de SUPERB-benchmark (fonemenherkenning, sprekeridentificatie, emotieherkenning, intentieclassificatie).

Belangrijkste Bijdragen

1. Systematische exploratie: Dit is de eerste uitgebreide studie die Mamba-based HuBERT onderzoekt als zowel een fundamenteel spraakmodel als een feature-extractor.
2. Efficiëntiebewijs: Demonstratie dat Mamba-architecturen lineaire schaalbaarheid behouden, waardoor ze superieur zijn aan Transformers voor zeer lange audio-sequenties.
3. Superieure causale prestaties: Het bewijs dat Mamba-based modellen in causale settings (streaming) betere prestaties leveren dan Transformers, zelfs met minder parameters.
4. Representatiekwaliteit: De bevinding dat Mamba-modellen fonetische en spreker-gerelateerde features scherper onderscheiden dan Transformers, wat resulteert in hogere kwaliteit gekwantiseerde representaties.

Resultaten

1. Rekenefficiëntie en Lange Context

• MACs en RTF: De Multiply-Accumulate operations (MACs) per seconde voor Mamba blijven nagenoeg constant naarmate de sequentielengte toeneemt (van 5 sec tot 5 min). Bij Transformers stijgen deze exponentieel.
• Geheugen: Causale Transformers lopen tegen geheugenlimieten (OOM) aan bij sequenties langer dan 80 seconden. Mamba kan dit probleem niet hebben.
• Lange Context ASR: Bij het fine-tunen voor document-level ASR (TEDLIUM3) bereikte ExtBiMamba een Woordfoutpercentage (WER) van 11.08%. De Transformer kon deze taak niet uitvoeren vanwege geheugenproblemen. Op utterance-niveau presteerde Mamba iets slechter (13.37% vs 11.86%), maar de lange context gaf een significant voordeel.

2. Streaming ASR (Causale Setting)

• Een Mamba-based HuBERT met 78M parameters bereikte een WER van 15.77%.
• Dit is beter dan een Causal Transformer met 94.7M parameters (WER 16.66%).
• Conclusie: Mamba levert betere nauwkeurigheid met ongeveer 17% minder parameters in streaming scenario's.

3. Representatie-analyse

• Phone Purity: Mamba-modellen vertonen een hogere piek in "phone purity" (hoe goed de clusters overeenkomen met fonemen) dan Transformers, vooral in causale settings.
• CCA (Canonical Correlation Analysis): Mamba-modellen tonen een sterkere correlatie met spreker-embeddings (Resemblyzer, TitaNet) dan Transformers, wat aangeeft dat ze sprekerkenmerken duidelijker vastleggen.

4. Downstream Prestaties (SUPERB Benchmark)

• Causale Setting: Mamba-based modellen overtreffen Transformers in fonetische taken (PR) en overall score (SUPERBS), zelfs bij kleinere modelgroottes.
• Bidirectionele Setting:
  • Bij Base-size modellen (ca. 94M params) presteert de standaard Transformer beter dan ExtBiMamba.
  • Bij Small-size modellen (ca. 23M params) presteert ExtBiMamba echter beter dan de Transformer op de meeste taken.
  • Ablatie: Voor grotere modellen bleek InnBiMamba (interne bidirectionele structuur) beter te presteren dan ExtBiMamba, wat suggereert dat de schaalbaarheid van Mamba in bidirectionele settings nog verbeterd moet worden (training instabiliteit bij grote schaal).

Betekenis en Conclusie

Dit paper toont aan dat Mamba een veelbelovend en complementair alternatief is voor Transformers in de spraakverwerking, met name voor:

• Lange context modellering: Waar Transformers falen door geheugenbeperkingen.
• Real-time (streaming) toepassingen: Waar Mamba lagere latentie en betere prestaties biedt met minder parameters.
• Fonetische representaties: Mamba leert representaties die beter geschikt zijn voor taken zoals gesproken taalmodellen (SLM) die gebruikmaken van gekwantiseerde eenheden.

Hoewel de schaalbaarheid in bidirectionele settings (voor niet-streaming taken) nog uitdagingen kent, biedt deze studie empirisch bewijs en richtlijnen voor het ontwikkelen van efficiëntere, toekomstbestendige spraak-zelftoezichtende modellen. De code is openbaar beschikbaar gesteld.