MERaLiON-SpeechEncoder-2

<h1 align="center">🎧 MERaLiON-SpeechEncoder-2 🎧</h1> <p align="center"> <a href="https://meralion.org/demo/asr">💻 ASR Web Demo </a> </p>

We introduce MERaLiON-SpeechEncoder-2, our next-generation multilingual speech encoder that was pre-trained from scratch on a greatly expanded corpus of 1.4 million hours of unlabeled audio, with a strong focus on Southeast Asian (SEA) languages and accents. As a speech foundation model, it encodes speech into a general-purpose, multilingual acoustic representation that can serve as a high-performance backbone for a wide range of downstream tasks — including automatic speech recognition (ASR), speech translation, speaker and language identification, and emotion recognition. This model can be finetuned on custom datasets, allowing developers to build speech systems tailored to their specific needs.

Unlike many existing models optimized for high-resource, Western languages, MERaLiON-SpeechEncoder-2 is designed from the ground up to reflect the linguistic diversity and complexity of Southeast Asia. Our training data was curated to contain a substantial amount originating from Singapore and SEA, including 60,000 hours of Singapore-accented speech, with a further 160,000 hours covering Singapore’s official languages Chinese, Malay and Tamil, along with a smaller portion of dialects like Hokkien and Cantonese. SEA data amounts to 200,000 hours, including significant proportions of Malay, Thai, Indonesian, Vietnamese, with smaller amounts of Tagalog, Burmese, Javanese, Sundanese, Khmer and Lao. See below for a regional breakdown of the language coverage of our pre-training data.

<p align="center"> <img src="data2.svg" width="580"/> </p>

Model Highlights

Small model size

With only 630M parameters (≈2.5 GB in memory), the model is easily deployable on most commercial GPUs, eliminating the need for distributed or large-scale compute setups.

Natively multilingual

Building on MERaLiON-SpeechEncoder-v1 (which focused on English and Singlish), this version expands to include English, Chinese, Malay, Tamil, Thai, Indonesian, and Vietnamese, along with codeswitching support across these languages. Given the wide coverage of languages in the training corpus, it may also be applicable beyond the officially supported languages.

Competitive performance on downstream speech tasks

The model retains near state-of-the-art results on the SUPERB benchmark for English, and showcases strong multilingual capabilities demonstrated through its integration into a high-performance ASR system.

Innovative pre-training techniques

MERaLiON-SpeechEncoder-2 was trained from scratch with a novel extension of the BEST-RQ self-supervised objective, by using more informative latent targets. We also adopted the Muon optimizer, which has previously only been shown to outperform the widely-used AdamW optimizer for LLM training. We find its advantages also carry over to speech-based models.

Model Summary

• Developed by: I²R, A*STAR
• Model type: Speech Encoder
• Language(s): English (Global & Singapore), Chinese, Malay, Tamil, Thai, Indonesian, and Vietnamese.
• License: MERaLiON Public License

For details on background, pre-training, tuning experiments and evaluation, please refer to our technical report.

Benchmarks

SUPERB

SUPERB is an English-based benchmark for speech encoders covering a wide range of downstream speech tasks across domains such as recognition, detection, semantics, speaker, and paralinguistics, where each task is finetuned separately with a frozen encoder.

MERaLiON-SpeechEncoder-2 is competitive to state-of-the-art, improving slightly against our own v1 model on speaker and paralinguistic tasks.

Automatic Speech Recognition (ASR)

<p align="center"> <img src="overall_wer.svg" width="720"/> <img src="audiobench_wer.svg" width="720"/> <img src="fleurs_wer.svg" width="720"/> </p>

Leveraging on the multilingual capabilities of MERaLiON-SpeechEncoder-2, we further finetuned the model for ASR on supervised speech data to produce a lightweight MERaLiON-SpeechEncoder-2-ASR-CTC, which is competitive to models many times its size in transcribing the target languages, while offering much faster inference speeds. It outperforms the popular Whisper large v3 across most languages in Audiobench and maintains close performance on FLEURS. Our comprehensive internal benchmarking, shown in the 'Overall ASR Performance', also contains several private datasets in addition to Audiobench and FLEURS.

Try out the ASR model here.

Direct Use

The following code snippet can be used to directly obtain latent features i.e. encoded speech by forwarding through the model. Inputs into the model are expected to be 80-dimensional Mel-spectrogram features transformed from audio sampled at 16kHz. The AutoFeatureExtractor method can carry out this conversion.

import torch
from datasets import load_dataset
from transformers import AutoModel, AutoFeatureExtractor

repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-2'
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# load model and feature extractor
model = AutoModel.from_pretrained(
    repo_id,
    trust_remote_code=True,
)
model = model.to(device)

feature_extractor = AutoFeatureExtractor.from_pretrained(
    repo_id,
    trust_remote_code=True
)

# prepare data
data = load_dataset("distil-whisper/librispeech_long", "clean",
                split="validation")

def batch_collater(data):
    tensors = []
    for idx, sample in enumerate(data):
        tensors.append(sample['audio']['array'])
    return tensors

audio_array = batch_collater(data)
inputs = feature_extractor(audio_array, sampling_rate=16_000,
                        return_attention_mask=True,
                        return_tensors='pt', do_normalize=False)
inputs = inputs.to(device)

# model inference to obtain features
with torch.no_grad():
    model.eval()
    output = model(input_values=inputs['input_values'],
                attention_mask=inputs['attention_mask'],
                output_hidden_states=True)

# output is a Wav2Vec2BaseModelOutput or tuple containing:
# last_hidden_state: torch.FloatTensor containing hidden states of the last layer of the model
# extract_features: torch.FloatTensor containing extracted features from the convolution downsampling layers
# hidden_states: tuple(torch.FloatTensor) containing hidden states of each layer of the model
# attentions: tuple(torch.FloatTensor) containing attention states of each layer of the model

Downstream Use

<p align="center"> <img src="downstream.svg" width="580"/> </p>

Speech encoders are normally used in finetuning setups to provide the frontend to downstream speech applications. We provide an example below of an ASR finetuning setup with Huggingface. Please refer to this blog for the full ASR finetuning recipe using Huggingface Trainer. Alternatively, the Huggingface model can be loaded to any other frameworks such as Pytorch or ESPnet for custom finetuning loops.

import torch
import json
from datasets import load_dataset
from transformers import AutoModelForCTC, AutoFeatureExtractor, Wav2Vec2CTCTokenizer

repo_id = 'MERaLiON/MERaLiON-SpeechEncoder-2'
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# prepare data
def pre_processing(batch):
    batch["text"] = batch["text"].lower()
    return batch

def extract_all_chars(batch):
  all_text = " ".join(batch["text"])
  vocab = list(set(all_text))
  return {"vocab": [vocab], "all_text": [all_text]}

librispeech100h_train = load_dataset("openslr/librispeech_asr", split="train.clean.100")
librispeech100h_test = load_dataset("openslr/librispeech_asr", split="validation.clean")
librispeech100h_train = librispeech100h_train.remove_columns(
                                    ['file', 'speaker_id', 'chapter_id', 'id'])
librispeech100h_test = librispeech100h_test.remove_columns(
                                    ['file', 'speaker_id', 'chapter_id', 'id'])

librispeech100h_train = librispeech100h_train.map(pre_processing)
librispeech100h_test = librispeech100h_test.map(pre_processing)

vocab_train = librispeech100h_train.map(extract_all_chars, batched=True,
                                    batch_size=-1, keep_in_memory=True,
                                    remove_columns=librispeech100h_train.column_names)
vocab_test = librispeech100h_test.map(extract_all_chars, batched=True,
                                    batch_size=-1, keep_in_memory=True,
                                    remove_columns=librispeech100h_test.column_names)
vocab_list = list(set(vocab_train["vocab"][0]) | set(vocab_test["vocab"][0]))
vocab_dict = {v: k for k, v in enumerate(sorted(vocab_list))}

vocab_dict["|"] = vocab_dict[" "]
del vocab_dict[" "]
vocab_dict["[UNK]"] = len(vocab_dict)
vocab_dict["[PAD]"] = len(vocab_dict)

with open('ls_vocab.json', 'w') as vocab_file:
    json.dump(vocab_dict, vocab_file)

# load model, feature extractor and tokenizer
feature_extractor = AutoFeatureExtractor.from_pretrained(
    repo_id,
    trust_remote_code = True,
)

tokenizer = Wav2Vec2CTCTokenizer("./ls_vocab.json",
                            unk_token="[UNK]", pad_token="[PAD]",
                            word_delimiter_token="|")

model = AutoModelForCTC.from_pretrained(
    repo_id,
    trust_remote_code=True,
    vocab_size=len(vocab_dict),
    feat_proj_dropout=0.1,
    activation_dropout=0.1,
    hidden_dropout=0.1,
    conformer_conv_dropout=0.1,
    ctc_loss_reduction="mean",
    pad_token_id=tokenizer.pad_token_id,
    attention_dropout=0.1,
)
model = model.to(device)

# Continue with training loop...

Compute and Infrastructure

MERaLiON-SpeechEncoder-2 was trained on the ASPIRE 2A+ Supercomputer Cluster, provided by National Supercomputing Centre (NSCC), Singapore.

The model was trained with 64 H100 GPUs across 8 nodes for collectively around 3.5 million steps. The SEA data split was upsampled in the last phase of training. Training time took approximately 15 days.

Citation

If you find our work useful, please cite our technical report:

@misc{huzaifah2024speechfoundationmodelsingapore,
      title={MERaLiON-SpeechEncoder: Towards a Speech Foundation Model for Singapore and Beyond}, 
      author={{MERaLiON Team}},
      year={2024},
      eprint={2412.11538},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2412.11538}, 
}