Fine-Tuning XLS-R for Low-Resource Automatic Speech Recognition

Hugging Face has detailed a process for fine-tuning XLS-R, a successor to XLSR, for Automatic Speech Recognition (ASR) in low-resource languages. By leveraging self-supervised pre-training on 128 languages, XLS-R allows developers to achieve functional speech-to-text capabilities even with limited labeled training data.

XLS-R Model Architecture and Pre-training

XLS-R is designed to learn cross-lingual speech representations that are effective across multiple languages. It was pre-trained using nearly 500,000 hours of audio data across 128 languages. The model is available in three sizes: 300 million, 1 billion, and 2 billion parameters.

Self-Supervised Learning Objective

Similar to BERT's masked language modeling, XLS-R learns contextualized speech representations by randomly masking feature vectors before passing them through a transformer network during its self-supervised pre-training phase.

Fine-Tuning Mechanism

To adapt the pre-trained network for downstream tasks such as speech recognition, speech translation, or audio classification, a single linear layer is added on top of the transformer block. This layer maps the context representations to specific token classes based on the vocabulary of the target labeled dataset.

Technical Implementation for ASR

Fine-tuning XLS-R for ASR involves several critical preprocessing and architectural steps to ensure the model correctly maps audio signals to text.

Data Preprocessing and Tokenization

For ASR, two primary components are required:

  1. Wav2Vec2FeatureExtractor: Processes the raw speech signal into the model's input format. XLS-R requires audio sampled at 16kHz. If the source data (such as Common Voice) is sampled at a higher rate (e.g., 48kHz), it must be downsampled.
  2. Wav2Vec2CTCTokenizer: Maps the model's output to text. The vocabulary is built from the distinct characters found in the training and test datasets, including a word delimiter token (represented as |) and a special "blank token" required by the CTC algorithm.

Connectionist Temporal Classification (CTC)

XLS-R is fine-tuned using Connectionist Temporal Classification (CTC). This algorithm is essential for sequence-to-sequence problems where the input (audio signal) is significantly longer than the output (text transcription). CTC allows the model to predict characters without requiring a precise alignment between the audio frames and the text characters.

Training Configuration and Optimization

Fine-tuning on low-resource datasets, such as the Turkish subset of Common Voice (approximately 4 hours of validated data), requires specific optimization strategies to maintain stability.

Model Setup

To optimize GPU memory and training stability, the following configurations are recommended:

  • Freezing the Feature Extractor: The CNN layers used to extract acoustic features are frozen (model.freeze_feature_extractor()) because they are sufficiently trained during pre-training.
  • Gradient Checkpointing: Enabled to reduce memory consumption.
  • Loss Reduction: Set to "mean" for the CTC loss.

Hyper-parameter Tuning

Due to the noise inherent in crowd-sourced datasets like Common Voice, hyper-parameters such as dropout, SpecAugment masking dropout rate, and learning rate must be carefully tuned. In the provided demonstration, a learning rate of 3e-4 and 30 epochs were used.

Training Efficiency

Training efficiency is improved by using group_by_length=True, which groups samples of similar input lengths into the same batch to minimize the amount of padding required.

Evaluation and Performance

Model performance is measured using the Word Error Rate (WER), the standard metric for ASR. In the demonstration using the Wav2Vec2-XLS-R-300M checkpoint on Turkish data, the training loss and validation WER decreased steadily over 3,200 steps, reaching a WER of approximately 0.3195.

Sources