Boosting Wav2Vec2 with n-grams in Hugging Face Transformers

Hugging Face has integrated the pyctcdecode library into the 🤗 Transformers library, allowing Wav2Vec2 speech recognition models to be combined with n-gram language models (LMs). This integration significantly reduces spelling errors and improves the Word Error Rate (WER), particularly for models trained on limited data.

Improving Transcription Accuracy with n-gram LMs

Combining Wav2Vec2 with an n-gram language model corrects common acoustic-only transcription errors where words sound correct but are spelled incorrectly. While Wav2Vec2 can produce acceptable transcriptions without an external LM due to its transformer architecture and Connectionist Temporal Classification (CTC) fine-tuning, an LM provides essential linguistic context that prevents the model from predicting non-existent words.

In a demonstration using facebook/wav2vec2-base-100h, the addition of a 4-gram language model corrected errors such as "christmaus" to "christmas" and "simalyis" to "similes". However, some errors may persist if the incorrect transcription is a valid English word (e.g., "rose" instead of "roast"), as the n-gram model may still assign it a non-negligible probability.

Technical Implementation of LM-Boosted Decoding

Decoding with a language model differs from standard decoding in how it handles model outputs. Instead of using the argmax of the logits to find the most likely character, Wav2Vec2ProcessorWithLM utilizes the full probability distribution (logits) for all possible output characters at each time step.

This process applies a beam search through the probability matrix, leveraging the n-gram language model to weight the likelihood of the next letters based on linguistic patterns. This requires the pyctcdecode and kenlm libraries for efficient decoding and model storage.

Building a Custom n-gram Language Model

To create an n-gram LM for a specific domain or language, the following workflow is used:

1. Data Collection and Preprocessing

An effective LM requires text data that matches the target transcriptions of the speech recognition system. For a Swedish model based on facebook/wav2vec2-xls-r-300m, the europarl_bilingual dataset was used because its clean, read-out nature corresponds well to spoken audio. Preprocessing involves:

  • Extracting text in the target language.
  • Converting text to lowercase.
  • Removing specific characters (e.g., punctuation) to match the alphabet of the fine-tuned acoustic model.

2. Model Construction with KenLM

KenLM is used to build the n-gram model due to its low computational cost compared to Transformer-based LMs. While Transformer LMs can yield better results, n-grams offer a significant performance boost over no LM with much faster retrieval times (essentially a look-up table query).

Key steps in the KenLM pipeline include:

  • Building the n-gram (e.g., a 5-gram) using the lmplz command.
  • Manually adding the end-of-sentence (</s>) token to the .arpa file to ensure compatibility with 🤗 Transformers.
  • Converting the .arpa file to a binary .bin format using build_binary to reduce file size and improve loading speed.

Integration and Performance Gains

To integrate the n-gram model, a Wav2Vec2ProcessorWithLM object is created by combining the acoustic model's feature extractor, tokenizer, and a pyctcdecode beam search decoder initialized with the KenLM binary model.

Performance Impact

According to the official Wav2Vec2 paper, n-gram LMs provide a substantial reduction in Word Error Rate (WER), especially for models trained on very small datasets (e.g., 10 minutes of audio), where an n-gram can reduce WER by approximately 80% compared to no LM. In the Swedish xls-r-300m-sv example, the 5-gram LM-boosted decoder achieved a WER of 18.85% on the Common Voice 7 test set, representing a relative performance improvement of approximately 30%.

Sources