Ok….so this is a bit of a series. Part of my side research on AI. This time, as recommended by Mr Hassan Rom, I would like to research more on AI for speech/sound/voice. I have explored text, images but not sound. So, this will be a good opportunity to take my time to research this, starting with Soundstream.
The paper for soundstream can be downloaded here: https://arxiv.org/pdf/2107.03312
As usual, let’s understand the abstract:
We present SoundStream, a novel neural audioccodec that can efficiently compress speech, music and general audio at bitrates normally targeted by speech-tailored codecs. SoundStream relies on a model architecture composed by a fully convolutional encoder/decoder network and a residual vector quantizer, which are trained jointly end-to-end. Training leverages recent advances in text-to-speech and speech enhancement,which combine adversarial and reconstruction losses to allow the generation of high-quality audio content from quantized embeddings. By training with structured dropout applied to quantizer layers, a single model can operate across variable bitrates from 3 kbps to 18 kbps, with a negligible quality loss when compared with models trained at fixed bitrates. In addition, the model is amenable to a low latency implementation, which supports streamable inference and runs in real time on a smartphone CPU. In subjective evaluations using audio at 24 kHz sampling rate, SoundStream at 3 kbps outperforms Opus at 12 kbps and approaches EVS at 9.6 kbps. Moreover, we are able to perform joint compression and enhancement either at the encoder or at the decoder side with no additional latency, which we demonstrate through background noise suppression for speech.
So, the first part of the abstract as below is just the introduction and functions of soundstream which is
efficiently compress speech, music and general audio at bitrates normally targeted by speech-tailored codecs.
Normally, speech are compressed at lower bit rates resulting in lower quality. However, Soundstream was able to compress music and general audio at this lower bitrate without losing it’s perceptual quality. Something that maybe we will explore further in the resulting parts of the paper
SoundStream relies on a model architecture composed by a fully convolutional encoder/decoder network and a residual vector quantizer, which are trained jointly end-to-end. Training leverages recent advances in text-to-speech and speech enhancement,which combine adversarial and reconstruction losses to allow the generation of high-quality audio content from quantized embeddings.
Two parts are mentioned here. The first one is the architecture. They mentioned that Soundstream use fully convolutional encoder/decoder network architecture. One interesting things that they mentioned is the residual vector quantizer (RVQ). That is an interesting concept that we need to explored in depth. The second part is the the training which utlizing the then advances in text-to-speech AI training which combine both adversarial and reconstruction losses. This is what allow them to generate high quality audio content even at lower bitrate. We will try to explore this losses in the next part.
By training with structured dropout applied to quantizer layers, a single model can operate across variable bitrates from 3 kbps to 18 kbps, with a negligible quality loss when compared with models trained at fixed bitrates. In addition, the model is amenable to a low latency implementation, which supports streamable inference and runs in real time on a smartphone CPU.
This is basically it’s results of training. One interesting things to comment is that this model can operate at lower bitrates causing it to be suitable to runs in real time on a smartphone CPU.
In subjective evaluations using audio at 24 kHz sampling rate, SoundStream at 3 kbps outperforms Opus at 12 kbps and approaches EVS at 9.6 kbps. Moreover, we are able to perform joint compression and enhancement either at the encoder or at the decoder side with no additional latency, which we demonstrate through background noise suppression for speech.
Another evaluation results which we will not explore much.
So basically, from the abstract, we know that we have to explore on 3 main concepts.
Architecture — the convolutional encoder/decoder
Architecture- the RVQ
The training losses (adversarial and reconstruction losses)
Architecture — Fully Convolutional Endocer/Decoder
Understanding what is Codec first
Ok, before we start understanding the arthitecture, I would like to made some comments on what is actually a codec. For so long, I’ve only been a user, not fully understands what codec is and what it’s actually doing.
For info, audio that comes from music has naturally large sizes. For example, one second of audio contains 44,100 samples of audio snapshot. Each of this samples is represented by 16 bits of depth and contains 2 channels (for steroe). If we’re multiplying all of this, we got around 10 MB of data for only 1 minutes seconds of audio.
That’s why codec is needed to compressed this audio into smaller size and decompress it backs when played in audio player. Hence the name (compression, decompression). The problem with codec is that, sometime during the decompression process, the quality of the audio deteriorated. So that’s why there’s a lot of research going on in Neural Codec where they want to perfom decompression without losing the quality of the audio. That’s the basis gist of it.
Alring, enough background. Let’s move in to the architecture.
Understanding the Architecture -Encoder/Decoder Blocks
Based on the reading and the image above, we can see that soundstream is a fully convolutional encoder-decoder neural network architecture. The convolutional layer is actually the main building blocks of this network. All encoder decoder blocks are actually consists of convolutional layer. I have extensively cover what is convolution layer all about in my CNN tutorial, you may read that here:
However, audio convolution is not the same as image convolution. The main difference is the kernel (filter) that they used.
Kernel in Audio vs Images
If we remember correctly, in images they used Sobel-like filter such as below:
This filter functions as a edge detection in the image data. That’s why it was designed like that. But how do want to use the same concept for audio? We used a different filter. There’s two major filters used in audio (there’s more but I don’t have time to cover them all). One is called as smoothing kernel/small symmetric kernel, where it functions as gaussian blur for the audio. It smooths any rapid changes in the audio waveform and capturing the trends in the audio waveform.
The second kernel is the high-pass kernel where it detects any rapid changes in the waveform.
So let’s test this out in code, and see how it functions.
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as signal
import scipy.io.wavfile as wav
import IPython.display as ipd
from pydub import AudioSegment
# Step 1: Convert MP3 to WAV
mp3_path = "audio-test.mp3" # Replace with your MP3 file
wav_path = "converted_audio.wav"
# Convert MP3 to WAV using pydub
audio = AudioSegment.from_mp3(mp3_path)
audio.export(wav_path, format="wav")
# Step 2: Load the converted WAV file
fs, waveform = wav.read(wav_path)
# Convert stereo to mono if needed
if len(waveform.shape) > 1:
waveform = np.mean(waveform, axis=1)
# Normalize waveform to range [-1, 1]
waveform = waveform / np.max(np.abs(waveform))
# Define 1D convolution kernels
high_pass = np.array([1, -1]) # High-pass filter (sharp transients)
low_pass = np.array([0.25, 0.5, 0.25]) # Low-pass filter (smoothing)
# Apply convolutions
filtered_high = signal.convolve(waveform, high_pass, mode='same')
filtered_low = signal.convolve(waveform, low_pass, mode='same')
# Plot results
plt.figure(figsize=(12, 8))
plt.subplot(4, 1, 1)
plt.plot(waveform, label="Original Audio", color='black')
plt.legend()
plt.subplot(4, 1, 2)
plt.plot(filtered_high, label="High-Pass Filtered", color='r')
plt.legend()
plt.subplot(4, 1, 3)
plt.plot(filtered_low, label="Low-Pass Filtered", color='g')
plt.legend()
plt.tight_layout()
plt.show()
# Play back the filtered audio to hear the differences
print("Original Audio:")
ipd.display(ipd.Audio(waveform, rate=fs))
print("High-Pass Filtered Audio:")
ipd.display(ipd.Audio(filtered_high, rate=fs))
print("Low-Pass Filtered Audio:")
ipd.display(ipd.Audio(filtered_low, rate=fs))
If everything works correctly, you will have the output as below:
As you can see above, the high pass filtered highlights any sudden increase in the temporal waveform while the low pass filtered seems like smoothing out the waveform. Although the pattern it’s not very apparent, but the operation did produce some results there. So at least we got the basic rights. And if we hear the audio properly, we can hear that the audio frequency seems to dip when pass through the high-pass filter (no bass can be heard) and seems smother when we hear the low-pass filtered.
So, the question some of you might have is, why is the need for the filter? Same with image convolutional, this filter functions as feature extraction that highlights important traits/context for these models to learn. By extracting this features, the model will have a more apparent pattern to learn from rather than scrambling from the pattern.
Ok, now we know how these filter in the convolutional layer works. Now, we have to understands the next part in the architecture, which is the FILM conditioning
FILM Conditioning
FILM stands for Feature-wise Linear Modulation. A technique used to modulate feature maps in neural network based on some external conditioning signal. Remember our feature that we has exteacted using convolutional layer previously? We will use the FILM conditioning to modulate this feature maps based on conditioning signal that may arise during training or inference.
This allows context-aware feature modulation. In SoundStream, it helps adjust encoding/decoding based on external conditions, such as:
Noise levels (for denoising).
Speaker identity (for voice adaptation).
Bitrate control (for efficient encoding at different bitrates).
Basically what I understand about FILM conditioning is that it allows for soundstream to dynamically process the sound in real time allowing for efficient sound modulation on the fly. We can basically perform simple operation of FILM conditioning in the code below
import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display
import scipy.signal as signal
import IPython.display as ipd
from pydub import AudioSegment
# Convert MP3 to WAV if needed
mp3_path = "audio-test.mp3" # Replace with your file
wav_path = "converted_audio.wav"
audio = AudioSegment.from_mp3(mp3_path)
audio.export(wav_path, format="wav")
# Load audio file
y, sr = librosa.load(wav_path, sr=None)
# Compute spectrogram
D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
# Estimate noise level dynamically (Simple: use signal energy)
noise_level = np.mean(np.abs(y)) # Higher = more noise
# Dynamic FiLM scaling based on noise level
gamma_dynamic = 1 + (0.5 * (1 - noise_level)) # Reduce effect in noisy conditions
beta_dynamic = -5 * noise_level # Lower brightness if noisy
# Apply dynamic FiLM transformation
modulated_D = gamma_dynamic * D + beta_dynamic
# Plot original vs FiLM-modulated spectrogram
plt.figure(figsize=(10, 4))
librosa.display.specshow(modulated_D, sr=sr, cmap='magma', x_axis='time', y_axis='log')
plt.title(f"FiLM Conditioned Spectrogram (Noise Level: {noise_level:.3f})")
plt.colorbar(label="dB")
plt.show()
# Convert modified spectrogram back to waveform
modulated_audio = librosa.istft(librosa.db_to_amplitude(modulated_D))
# Play original vs FiLM-modulated audio
print("Original Audio:")
ipd.display(ipd.Audio(y, rate=sr))
print(f"FiLM Conditioned Audio (Dynamically Adjusted to Noise Level: {noise_level:.3f}):")
ipd.display(ipd.Audio(modulated_audio, rate=sr))If it works ok, we will get the output as below:
Understanding the RVQ
Alright. Now is the time to understand RVQ. It stands for Residual Vector Quantization. As the name suggests, it is the process to compress the audio waveform for storage and transmission. But before we dig deeper, let’s try to understand what Quantization really means, especially in audio.
What is Audio Quantization
For people who are familiar with AI, the term quantization is already well understood. It is the process of reducing the bit-precision from higher to lower in order to save space. In basic computer science, we already now that computers store informations in bits. Bits stands for binary digits. Which is 0 and 1.
And we have 4 levels of bits to store all our data. Which is:
8 bit ~ 2⁸ =256 bits
16 bit ~ 2¹⁶ = 65,536 bits
32 bit ~ 2³² = 4,294,967,296 bits
64 bit ~ 2⁶⁴ =1.845×1⁰¹⁹ (which is almost near infinite).
But it’s a different story for sounds. Sound only uses up to 24-bits. Because each bits level already carried the maximum decibel that it can store. So for sound, we have a different bits level, which is:
8 bit — used for old games (like Pokemon), telehony and low quality audio. It has maximum decibel of around 50 decibel.
16 bit — CD quality and used for most music. Maximum decibel around 96 decibel.
24 bit — Studio quality and used for audio production. Maximum decibel is around 144 decibel.
All humans have hearing up to 120 decibel. So it doesn’t makes sense to go above 24 bit.
In terms of storage, we frequently the word bytes. Which is actually stands for 8-bit. So one byte contains 8 bit. So you can kind of guess where this goes. The higher the bits representation, the more bytes you have to store, and more space is needed. And we all know memory space in computer is not limitless.
Therefore, quantization is a process to convert high bits (maybe 24-bits) to lower bits quality like 16-bit or 8-bit for efficient storage and transmission.
What is Vector Quantization
Before we go into residual vector quantization, let’s first understand what is vector quantization. So, imagine that we have a set of data in a high dimensional space. the data is spread across the multiple dimensions. since the data is large, storing all this data in this high dimensional space will results in high memory usage.
So one way to tackle this issues is by mapping this data to sections/pointers. each of this data will have a leader that will group all these data into nearest cluster. Instead of storing all the data, we just store the leaders/pointers. By doing that, we will haveh huge memory savings.
This leader is what we called as codebook. and our goal during vector quantization training is to reduce the distance between the leader (codebook) with its peers in the high dimensional space.
Alright, let’s demonstrate the vector quantization in python with simple synthetic data as below:
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
# Set random seed for reproducibility
torch.manual_seed(42)
# 1. Generate synthetic 2D data
num_data = 300 # Number of points
data = torch.randn(num_data, 2) * 2 # Spread out the points
# 2. Initialize codebook (random cluster centers)
num_codewords = 5 # We create 5 representative codebook vectors or leaders
codebook = torch.randn(num_codewords, 2, requires_grad=True)
# Function to compute distances between data points and codebook vectors
def compute_distances(data, codebook):
return torch.cdist(data, codebook) # Shape: (num_data, num_codewords)
# Function to assign each data point to the nearest codebook vector
def quantize(data, codebook):
distances = compute_distances(data, codebook)
closest_indices = torch.argmin(distances, dim=1) # Find closest vector index
quantized_data = codebook[closest_indices] # Replace data with codebook vector
return quantized_data, closest_indices, distances
# 3. Train the codebook using gradient descent
optimizer = torch.optim.Adam([codebook], lr=0.05) # Optimize codebook vectors
num_epochs = 100
losses = []
for epoch in range(num_epochs):
optimizer.zero_grad()
quantized_data, _, _ = quantize(data, codebook)
loss = torch.mean((data - quantized_data) ** 2) # Mean squared quantization error
loss.backward()
optimizer.step()
losses.append(loss.item())
# 4. Plot Training Loss (Quantization Error)
plt.plot(losses, label="Quantization Error")
plt.xlabel("Epoch")
plt.ylabel("Mean Squared Error")
plt.title("Training Progress of Vector Quantization")
plt.legend()
plt.show()
# 5. Visualize Data and Codebook (Clustering Effect)
quantized_data, closest_indices, _ = quantize(data, codebook)
plt.scatter(data[:, 0], data[:, 1], c=closest_indices.numpy(), cmap="viridis", alpha=0.5, label="Original Data")
plt.scatter(codebook[:, 0].detach(), codebook[:, 1].detach(), c="red", marker="X", s=200, label="Codebook Vectors/Leaders")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.title("Vector Quantization in 2D Space")
plt.legend()
plt.show()Then we will get the output as below:
For anyone who is familiar with K-means clustering, this graph does look familiar right? Yeah, that’s the basic principle behind vector quantization. Since we have a cluster of data next to each other, we can use that cluster as our memory storage instead of overall data. This bring the memory usage lower.
So, what is Residual Vector Quantization?
So, Residual Vector Quantization is acttually using the residual error from the quantization to determine the new set of codebook. Instead of one time quantization, we perform this quantization multiple times. But, we did not intialize with larger codebook, we initialize with smaller codebook. And after each quantization, we will use the residual error to determine new codebook. By doing this multiple times, we will increase the accuracy of our quantization process, while at the same time, ensure that we limit the amount of memory needed.
Let’s demonstrate this via python as below:
import torch
import numpy as np
import matplotlib.pyplot as plt
# Set random seed for reproducibility
torch.manual_seed(42)
# 1. Generate synthetic 2D data
num_data = 300 # Number of points
data = torch.randn(num_data, 2) * 2 # Spread out the points
# 2. Initialize multiple codebooks (Residual Vector Quantization)
num_codebooks = 2 # Number of quantization stages (residual correction)
num_codewords = 3 # Number of vectors per codebook
# Create multiple codebooks (one per quantization stage)
codebooks = [torch.randn(num_codewords, 2, requires_grad=True) for _ in range(num_codebooks)]
# Function to compute distances
def compute_distances(data, codebook):
return torch.cdist(data, codebook) # Euclidean distance
# Residual Vector Quantization function
def residual_vector_quantization(data, codebooks):
residual = data.clone() # Start with original data
quantized_data = torch.zeros_like(data) # Accumulate quantized results
closest_indices_per_stage = [] # Store selected codeword indices
for stage, codebook in enumerate(codebooks):
distances = compute_distances(residual, codebook) # Compute distance to codebook
closest_indices = torch.argmin(distances, dim=1) # Find closest codeword index
closest_indices_per_stage.append(closest_indices)
stage_quantized = codebook[closest_indices] # Get selected codeword
quantized_data += stage_quantized # Accumulate quantized vectors
residual -= stage_quantized # Compute residual error for next stage
return quantized_data, closest_indices_per_stage, residual
# 3. Train RVQ using gradient descent
optimizer = torch.optim.Adam(codebooks, lr=0.05) # Optimize all codebooks
num_epochs = 100
losses = []
for epoch in range(num_epochs):
optimizer.zero_grad()
quantized_data, _, residual = residual_vector_quantization(data, codebooks)
loss = torch.mean(residual ** 2) # Mean squared residual error
loss.backward()
optimizer.step()
losses.append(loss.item())
# 4. Plot Training Loss (Residual Quantization Error)
plt.plot(losses, label="Residual Quantization Error")
plt.xlabel("Epoch")
plt.ylabel("Mean Squared Error")
plt.title("Training Progress of Residual Vector Quantization (RVQ)")
plt.legend()
plt.show()
# 5. Visualize Data and Codebook (Clustering Effect)
quantized_data, closest_indices_per_stage, _ = residual_vector_quantization(data, codebooks)
# Plot original data points
plt.scatter(data[:, 0], data[:, 1], c="gray", alpha=0.5, label="Original Data")
# Plot each codebook stage with different colors
colors = ["red", "blue", "green"]
for stage in range(num_codebooks):
plt.scatter(codebooks[stage][:, 0].detach(), codebooks[stage][:, 1].detach(),
marker="X", s=200, color=colors[stage], label=f"Codebook {stage+1}")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.title("Residual Vector Quantization in 2D Space")
plt.legend()
plt.show()Which outputted below graphs:
As you can see above, even with same codebook numbers (6), we still get lower residual loss with RVQ compared to standard VQ. So it’s a much better way to perform quantization.
Alright, now we understand RVQ. Let’s go to the training losses
Understanding the Training Loss
So in the abstract above it mentions that
Training leverages recent advances in text-to-speech and speech enhancement,which combine adversarial and reconstruction losses to allow the generation of high-quality audio content from quantized embeddings.
So there are two losses here being used as part of our training losses, which is the adversarial loss and reconstruction loss. For adversarial I think it refers to GAN (Generative Adversarial Network). It is a network that consists of two, the discriminator and the genrerator that is pitted against each other. The discriminator will act as a police that will determine whether the audio that we process/train is not fake. This is good, since we want to ensure that the sound that we has compressed and decompressed does not lose it’s quality.
But the other losses, the reconstruction losses, is the one that we have to focus on. And before we do that, we need to understand how losses are determined for audio.
Audio Losses : STFT (Short-Time Fourier Transform)
As per text, audio also has it’s own loss function. Mostly for audio it uses the same loss function we frequently used, such as mean-squared-error or mean-absolute-erorr. But, at the same time it also has one specialized loss funciton, which is the Short-Time Fourier Transform loss.
STFT is a mathematical technique used to analyze how the frequency content of a signal changes over time. The mathematical equation for STFT are as below:
This divides the signal into small overlapping windows and applies the Fourier Transform to each window separately, giving a spectrogram.
By doing this, we can apply those MSE/MAE to the windows (instead of the amplitude), to ensure that the reconsruction loss that we got preserved the quality and features of the audio.
We can demonstrate this below:
import numpy as np
import librosa
import librosa.display
import matplotlib.pyplot as plt
import IPython.display as ipd
# Load an audio file
audio_path = "converted_audio.wav" # Replace with your file
y, sr = librosa.load(audio_path, sr=None)
# Compute the STFT (Spectrogram)
D = librosa.stft(y) # Short-Time Fourier Transform
D_db = librosa.amplitude_to_db(np.abs(D), ref=np.max) # Convert to dB scale for visualization
# Plot the Spectrogram
plt.figure(figsize=(10, 4))
librosa.display.specshow(D_db, sr=sr, cmap='magma', x_axis='time', y_axis='log')
plt.colorbar(label="dB")
plt.title("STFT Spectrogram of Audio")
plt.show()
# Play original audio
print("Original Audio:")
ipd.display(ipd.Audio(y, rate=sr))Which shows the spectorgram of the audio wave as below:
Conclusion
Alright, now we have covered most of the aspects of SoundStream. It does uses a lot of techniques to ensure that compression and decompression happens efficiently while preserving audio quality. We saw that the features of the audio has been efficiently extracted by convolutional layer, and modulated by the FILM Conditioning. The RVQ process reduces the size of the memory while at the same time minimize the loss, and it uses both adversarial loss and reconstruction loss to ensure quality audio decompression happens.
So as I said above, this is just one of the series. We will cover more on audio AI model next.