Workflow · Production

Training the Mamba architecture on speech and music data using Determined AI

The problem

The author wanted to learn the Mamba architecture through hands-on practice by reproducing an open-source speech synthesis script, and encountered repeated dataset failures that made it hard to produce usable model output.

First attempt

Three of four candidate datasets failed: Schmidt dialogues were too small causing the model to overfit; SpeechTokenizer discards music making the Taylor Swift dataset unusable; and AI Morgan Freeman audio contained excessive pauses that produced mostly empty model outputs.

Workflow diagram · grounded in source
1
Download and split audio data
trigger
“Download data in .mp4 and split into 10s .wav files”
2
Normalize and tokenize dataset
ai_action
“uses the SpeechTokenizer library to capture both semantic and acoustic tokens, while other audio tokenizers might only capture one or the other”
3
Train Mamba model
ai_action
“a project that initializes a 12M parameter Mamba model, establishes normalization + tokenization steps and a training scheme, and trains that model”
4
Track experiments with Determined AI
integration
“these graphs were generated using Determined AI's Detached Mode, which allows you to stick a few lines of code in your training loop to have your experiment appear on the Determined web UI”
5
Test generated audio
output
“This consists of downloading checkpoints and trying them out. You can either: generate audio unconditionally (without a seed sequence), or have the model complete the audio given a test sample from the dataset”
Reported outcome

Using the Alice in Wonderland audiobook dataset with 4 quantizers, the model produced audio that sounds like the input and is not memorized, with 4 quantizers achieving better training results than 8 quantizers.

Reported metrics
Mamba model parameters12M
Training loss with 8 quantizersnever goes below 2
Training loss with 4 quantizersfinally approaches 1
Optimal quantizer count4
Reported stack
MambaSpeechTokenizerDetermined AIffmpegWandB
Source
https://mlops.community/blog/audio-generation-with-mamba-using-determined-ai
Read source ↗

Frequently asked questions

What did this team achieve with this AI workflow?

Using the Alice in Wonderland audiobook dataset with 4 quantizers, the model produced audio that sounds like the input and is not memorized, with 4 quantizers achieving better training results than 8 quantizers.

What tools did this team use?

Mamba, SpeechTokenizer, Determined AI, ffmpeg, WandB.

What results were reported?

Mamba model parameters: 12M; Training loss with 8 quantizers: never goes below 2; Training loss with 4 quantizers: finally approaches 1; Optimal quantizer count: 4 (source-reported, not independently verified).

What failed first in this deployment?

Three of four candidate datasets failed: Schmidt dialogues were too small causing the model to overfit; SpeechTokenizer discards music making the Taylor Swift dataset unusable; and AI Morgan Freeman audio contained ex…

How is this workflow AI workflow structured?

Download and split audio data → Normalize and tokenize dataset → Train Mamba model → Track experiments with Determined AI → Test generated audio.