Neither did I. These kinds of documents are supposed to inform us of how our data is handled and the rights we have when using a website or app. Unfortunately, they’re usually either an unintelligible soup of legalese, or so long that they’d take a full afternoon to read. Companies have little incentive to provide summaries – they’re usually just filling a legal requirement.

Terms of Service; Didn’t Read (ToS;DR) has been trying to fix this. For each digital service they provide a collection of statements in yes/no format – for example, “This service tracks you on other websites” – that together offer a decent summary of the service’s policies. They even aggregate these statements into overall grades (Wikipedia gets a “B”), and offer browser extensions so you can see grades as you browse the Web.

Historically ToS;DR has relied solely on volunteers to do the tedious work of reading and annotating privacy documents. This is a heavy bottleneck on coverage. To date it has around 10,000 websites and apps catalogued, only ~8,600 of which are graded. It’s a great start, but it obviously falls short given how many websites and apps are out there handling your data.

The vast majority of websites on ToS;DR have yet to be graded. Data from December 2024.

Automation

To help, I teamed up with ToS;DR to scale their efforts of annotating and scoring privacy agreements with help from AI, a solution we call Docbot.

Specifically, I’ve used the human annotations collected over the years to train classification models that can estimate the strength of evidence for particular statements in privacy documents. If the models find evidence that scores above a certain threshold, the relevant excerpts are submitted to human curators. Once approved, the statement summaries are used automatically when calculating grades.

The output of Docbot – a queue of observations about privacy documents for human curators to later approve or reject. ‘Rating’ refers to whether the observation is a positive or negative thing for privacy, while ‘Score’ refers to the strength of evidence.

In the remainder of this post I’ll share my experience developing the solution and ultimately putting it to use within ToS;DR.

The Methodology

As mentioned above, ToS;DR has come up with a taxonomy of privacy-related statements that you might want to know about a digital service. We call these Cases, and here are some examples:

• Your personal data is not sold
• You must provide your legal name, pseudonyms are not allowed
• Tracking cookies refused will not limit your ability to use the service
• Your data is processed and stored in a country that is less friendly to user privacy protection
• The terms for this service are easy to read
• … and over 100 more

For each Case I trained a binary classification model that takes as input one or more sentences, and outputs a 0.0-1.0 score where 1.0 means there is extremely strong evidence that the Case statement is true. Like many other text classification problems, this was a good fit for fine-tuning.

Key to this approach was that not only did ToS;DR have data on which Cases are present in which documents, but specifically which excerpts provide the evidence. These can fit inside the context windows of transformer models during training, and most importantly, our solution can try different excerpts to find any that are high-scoring enough for submission to ToS;DR curators. Ideally, these AI submissions will look indistinguishable from those submitted by humans.

Sentence classification is a common problem in NLP, and this is almost that. There were two differences here that I had to figure out: 1) we wanted to apply the models to full terms of service documents and privacy policies, often with many thousands of words. We only wanted to extract the most convincing excerpt, even if multiple sentences in the document were relevant. And 2) sometimes policies are communicated in short paragraphs instead of single sentences, and we ideally wanted to preserve these spans of text as logical units, instead of only working at the single sentence level. Put another way: the problem actually sits somewhere between sentence classification and document classification.

My solution for inference (using the models) was:

• Split the document into sentences
• Apply the model to each sentence individually
• Take the highest scoring sentence, and see if appending additional neighboring sentences makes the score go up (until some limit, e.g. 5). If so, keep them.
• See if prepending prior neighboring sentences makes the score go up even more. If so, keep them.
• Use the final sentence span and score as the result for the document.

Forming Datasets

Much of the work for supervised classification approaches like this lies in designing effective training and evaluation datasets.

At the very least I needed a dataset of sentence spans to train each binary classifier, with positive and negative examples. These datasets could also be useful for evaluating models, but they didn’t seem fully adequate. Since our models were ultimately being applied to entire policy documents, I formed a companion dataset of full documents to evaluate against. For one, the accuracy metrics would be much more intuitive compared to the sentence span dataset with different class ratios. But also, it would evaluate my models in the same context in which they will be used in the real world, with the “highest score” inference procedure described above.

Sentence Span Dataset

Positive examples here were easy – each Case had a collection of verified quotes submitted by ToS;DR volunteers over the years, sometimes in the hundreds. I decided to only tackle Cases for which we had at least 40 positive examples, leading to 123 Case models in total. The remaining could some day be covered using other text classification methods like SetFit or few-shot prompting LLMs.

Negative examples

These were trickier. Obviously with such long documents, searching for proof of individual Case statements is like finding a needle in a haystack, and so there is plenty of hay – completely irrelevant sentence spans. One problem to be aware of is extreme class imbalance; sampling uniformly from entire documents would probably not end well.

The other big problem is that we don’t just want models to lazily learn to recognize the topic of a case. For example, we don’t want a model learning the case “Tracking cookies refused will not limit your ability to use the service” to just score highly any time it sees a statement about cookies. The solution here is to expose the model to plenty of “hard negatives”, which contain similar terms and force the model to more carefully discriminate.

For each case I put together a set of negative examples containing:

• Human submissions that were later rejected
• Sentence spans that surround positive examples. This approach is prone to accidental true negatives, but I believe it should help to sharpen the discriminative power of models.
• The ToS;DR taxonomy actually groups cases into topics, so I included approved submissions from other related cases
• Other random sentence spans from documents containing positive examples
• Random sentence spans from comprehensively reviewed documents that do not contain positive examples

Another helpful source of negative examples could come from manually constructing logical negations of positive examples. Assertions that data practices are not carried out are rare, but they exist. Using negations could lead to more robust models that can handle these correctly – instead of just learning phrases, they would have to learn to analyze the statement’s modality. I didn’t include these in our first model versions, but a similar project called MAPS[1] took this approach.

Another future direction is hard negative mining with vector similarity search, which would help surface examples right on the edge of what we’re trying to learn.

Sentence length bias

When constructing training examples, it’s important not to introduce any unintended biases. The models should learn to discriminate between positive and negative examples by way of learning the underlying task at hand, rather than learning irrelevant characteristics of the constructed examples that you introduced.

For example, here the human submitted quotations (positive training instances) could be any number of sentences. If our constructed negative instances were only one sentence, a model could theoretically learn the rule “If I’m shown multiple sentences, it’s positive.”

My solution here was to model each case’s typical number of sentences for submitted evidence as a multinomial probability distribution, and sample from it while constructing negative instances.

Another potential source of bias is that submissions don’t necessary have to begin and end at sentence boundaries – people can highlight partial sentences, even short phrases. Rather than try to replicate typical start and end times for each case, I decided to just expand positive instances to begin and end at their surrounding sentence boundaries. This restricts the resulting models to also only yield full sentence predictions, but that’s fine, and it avoids any related bias issue.

Document Dataset

The dataset of full documents (privacy policies, terms and conditions, cookie policies, etc.) was more straight-forward. Positive instances were those which had an approved submission of evidence from ToS;DR volunteers. Negative instances included documents that were comprehensively reviewed by volunteers without the case ever coming up, or rejected submissions if there were any.

This was utilized for evaluation only, using my custom “sentence expansion” inference procedure. It correlated with performance on the sentence span dataset used for training, but I think it’s more trustworthy for model selection because it’s closer to how we use the models downstream.

The only “gotcha” was to prevent data leakage by stratifying train/test splits the same way I did for the sentence span dataset: by website/service. I first split services into two groups 80-20, and then partitioned both datasets accordingly.

Training Docbot

I relied on the tried-and-tested fine-tuning of BERT family models using huggingface. bert-base-uncased, roberta-base, and a finetune of legal documents legal-bert-base-uncased all performed about the same. ModernBERT looks great, but wasn’t yet released at the time of training. Rather than fine-tuning the full base models I decided to train more parameter efficient LoRA adapters [2], via the huggingface peft library.

Interestingly the vast majority of time spent during training was actually the evaluation loop, when full documents were analyzed. This could surely be sped up by not re-running obviously unrelated sentences, but that’s an optimization for another day. This means using LoRAs didn’t actually have much of an effect on training speed, but they were still nice for disk space efficiency (each of the 123 models is just a 2.8MB adapter file), and there was virtually no reduction in accuracy compared to a full fine-tune.

As is typical when fine-tuning LLMs, accuracy can differ a decent amount depending on the random seed. This means it was beneficial to attempt training each Case multiple times, taking the highest test set accuracy observed across all attempts.

Training was carried out on vast.ai instances, using Weights & Biases to track experiments and visualize metrics.

Training curves for multiple restarts of one case model in particular – “You maintain ownership of your content”. The y axis shows F1 score on the test set fold of the full document dataset.

What about just prompting LLMs?

Few-shot prompting is another way to perform text classification. Sometimes it’s the right tool for the job, especially when training data is hard to come by and you want a quick solution. For this problem I opted for fine-tuning because of it’s potential for higher accuracy at low cost, and because it provides interpretable and reliable confidence scores with each prediction.

LLMs could still be useful to help bootstrap Cases with very few training examples, and to audit training datasets for true negatives.

Model and threshold selection

I used early stopping to halt training, and selected the model checkpoint that maximized F1 on the full document dataset across all restarts.

One benefit of using regression models or neural networks for classification is that we can choose prediction thresholds at our preferred place on the precision-recall curve. Since we planned to send AI submissions to human curators before use in privacy grades, a lower threshold was preferable. This would mean more false positives crowding the approval queue (lower precision), but would ensure better coverage finding matches (higher recall).

I found that a single static post-softmax probability threshold would not work across Case models, as they were not very consistent in their relative implied confidence. And with 123 models in total, it was best to automate threshold selection by maximizing an objective criteria. F-score is a great choice because it incorporates both precision and recall and allows us to specify a tradeoff preference using the beta parameter. In my case I used a beta of 1.5 to prefer recall. The charts below show precision-recall curves and the corresponding f-scores as we vary the classification decision thresholds from 0.0 to 1.0.

On the left, precision-recall curves of five different case models. On the right, available f-scores for the same models, as a function of the decision threshold that could be used for classification. A black X marks the optimal point.

Results

Model accuracy varied quite a bit by Case.

Precision and recall for all 123 Case binary classification models.

Sorting by f-score shows us the best and worst performing models:

Prediction accuracy was higher for more strictly defined Cases that often use consistent language across privacy documents. Accuracy was lower for Cases with lots of edge cases, or ones with lower quality training sets.

When it comes to the quoted evidence submitted for approval, generally the Docbot submissions looked just like those from humans. They did tend to be a little longer on average, in part because they had to adhere to sentence boundaries, but also because too many sentences were appended during the inference procedure. The chart below shows differences in distributions of the number of sentences for a few random Cases, and you can see some are worse than others. This can be refined in future versions.

The Rollout

There was no need to have our models hosted as an API, since we just need to apply them once for each privacy document as a batch job. New services and documents are added to ToS;DR regularly, and so we can re-run the inference jobs daily or weekly.

To keep track of which documents have been already been analyzed by Docbot, we added a new database table that stores the results of inference: document ID, Case ID, Docbot version, the start and end character positions of the highest-scoring sentence span, and the probability score. This allows for audits and keeps us from running the same model on the same document twice. In the future we’ll likely also use a job queue to facilitate inference parallelization.

I also decided to have the inference job re-consider documents for which the human curators rejected a previous submission. If Docbot (or a human) thinks it found evidence for a Case but was wrong, that doesn’t mean evidence doesn’t exist elsewhere in the document. Previously rejected sentence spans act as “off limits” areas during inference.

After setting up APIs for Docbot to interact with the ToS;DR backend, we made a new frontend for the Docbot submission queue. Once again here’s what that looks like:

A preview of the queue shown to human curators who can approve or reject Docbot submissions.

Submissions are sorted by descending confidence score so that the most likely approvals are at the top. Curators can also filter by Case, if they just want to churn through the results of a single model.

The Impact

ToS;DR only assigns privacy grades to a service when it has enough knowledge about its policies, meaning it has enough approved evidence submissions. Thanks to Docbot, the bottleneck on grade coverage has now shifted from volunteers initially reading through and annotating documents, to volunteer curators that double-check submissions as a final sign-off.

The number of submissions of evidence for Cases from either humans or Docbot over time. The two big jumps to the right are from the first runs of Docbot. By the time remaining Cases are run, we expect there to be over 300,000.

In the coming months we’ll explore ideas for easing that bottleneck, such as:

• A Tinder-style swiping interface for approving or rejecting Docbot submissions more quickly on the go. We do have to be careful to make sure curators still apply enough scrutiny, for example by having redundent approvals and measuring inter-rater reliability.
• Rather than working through giant queues of submissions, encouraging curators to go service-by-service to increase grade coverage more quickly, starting with the most popular ungraded services.

Beyond optimizing the approval process workflow, we are also exploring options for how to incorporate Docbot more directly:

• Publishing fully automated grades for services, when we don’t yet have enough human-approved submissions
• For very high confidence predictions (say, a score >.99) we could auto-approve

They key challenges with both of these approaches are to 1) choose the right score thresholds to balance false-positives with coverage, 2) have an understanding of model biases, and their effect on end grades, and 3) adequately communicating everything to our users and API consumers. People have grown to trust ToS;DR, and ultimately any incorporation of AI has to be done carefully to maintain that trust.

Conclusion

This was a really fun side project. I got to work with a unique dataset, help solve a real world pain point, and contribute to free and open source software. Part way through this experience, I was invited to join the core dev team at ToS;DR.

Future directions

I previously mentioned how we’d like to apply Docbot more directly by assigning tentative grades using fully automated analyses. And, we’d like to build better interfaces for double checking Docbot submissions. Beyond that, there is still so much to do!

• More efficient inference, by doing a first pass that excludes paragraphs with no remote relevance to the Case topic.
• Use Docbot to audit for mistakenly approved points already in ToS;DR
• Automated retraining of models, CI/CD
• Expand coverage to non-English languages
• Expand coverage to Cases with very few training examples
• Release datasets for NLP and legal researchers, open source the LoRA adapters
• Put out a tool to help service owners craft their privacy docs to be more privacy friendly, by getting feedback in real time
• Beyond binary statements, extract targeted information from privacy docs, like a list of what data are collected, and how long the service keeps it

Other things that would help with digital privacy

• Crawling the web to automatically add services and privacy docs
• Regulation for companies to provide easier to understand summaries, as GDPR does, and protocols for machine-readable policies, as IEEE proposes
• Tracking data leaks and privacy scandals

If you found this project interesting you can help us by contributing analyses on edit.tosdr.org, getting in touch to contribute to development, or even donating.

References

• [1] S. Zimmeck. MAPS: Scaling Privacy Compliance Analysis to a Million Apps, 2019.
• [2] E. Hu. LoRA: Low-Rank Adaptation of Large Language Models, 2021.

In this situation, a really useful intermediate goal is figuring out the best way to turn your dataset into vectors, essentially just lists of numbers, each denoting a dimension with meaning. Vectors represent a vital stepping stone in the standard data science process – done right, they immediately unlock several paradigms for understanding, and building on top of, your data.

There’s a science to designing the right vector space for your use case. Sometimes working directly with tabular features is fine, usually with some scaling/preprocessing. Other times you might want to employ a training paradigm like word2vec, neural networks, or manifold learning. Admittedly this post should really be called “The unreasonable power that comes with representing your dataset as well-formed vectors of mostly-independent normalized features”, that just didn’t sound as snazzy.

But this is not a blog post about how to design vector spaces. Rather, I’ll cover the things you can do once you’ve embedded your dataset as a bag of vectors. Whether these are as simple as 3-dimensional vectors describing demographics, or something complex like 4096-dimensional embeddings from the latest LLM, all of the methods below should be applicable.

Scatterplots

One of the most basic things you can do is visualize your dataset in 2D space. This is especially useful in the exploratory phase of a project to get your bearings, understand the dataset’s overall structure, and identify outliers.

Some datasets might have two features that work great as X and Y dimensions, however not all do. The key step that makes this a universally applicable approach is to automatically reduce your vectors’ dimensionality, all the way down to two, in a way that preserves their global structure. That is to say, the distances between vectors in the original high dimensional space should correlate with the distances in 2D. There are a ton of methods for doing this, each with their pros and cons. I recommend tSNE or UMAP as your go-tos.

Obviously a bunch of dots alone without context are not useful, but you can decorate them with a visualization framework. Point size (for numeric features), shape (for categorical features), and color (for either) offer three ways of highlighting features that help you navigate. You can also add hover tooltips, letting you see as many interpretable features as you’d like.

Throughout this post I’ll be demonstrating with a toy dataset of popular music. The interactive scatterplot below is derived from “genre embeddings”, 128-dimensional vectors that come from the last hidden layer of a genre classification neural network. On desktop, hover over a dot for track metadata (I couldn’t get clicking on mobile to work). Code available here.

Each dot represents a track from a random selection of popular music artists.

Generally we see the tracks of artists end up near each other, which is a good sign. (As a side note, it’s also satisfying to see the dots representing Girl Talk right in between the cluster of hip-hop and other genres, because he makes mash-up music featuring hip-hop vocals over a backdrop of other genres)

For a much more impressive scatterplot, check out the entirety of English Wikipedia on Atlas. Enjoy the rabbit holes.

For those working in the Python ecosystem like myself, I can recommend streamlit as a frontend platform with which to plot. Jupyter notebooks will work fine too, and there’s something to be said for code living right next to the plots it generates, but I find streamlit easier to work with for many use cases, including this blog post. Several plotting libraries are supported: Matplotlib, Altair, Plotly, Bokeh, and more.

Clustering

Once your dataset is in a vector space, you can always compute the distance between any two points. This could be Euclidean distance, cosine distance, or something else, but either way this simple ability unlocks a few go-to data science paradigms for free.

One is clustering – using an algorithm to find logical groupings. This can help you understand segments of your data, and can even work as a classifier for new data points. The scikit-learn library documentation offers a nice overview of common approaches. You’ll notice most of the APIs take a distance measure of your choosing as an input.

The shapes of the points below indicate membership in unsupervised clusters. This means they were not designed to delineate any existing groupings, like genre labels; they were assigned according to the hands-off approach HDBSCAN. And yet, the groups do resemble genres and would be useful for downstream analysis (I’ll admit I’m cheating a little in this example.. the embeddings come from a model optimized to recognize genre, so they have a head start in ending up that way.)

Like before each point represents a track, and the shape of the point is assigned according to the unsupervised clustering algorithm HDBSCAN.

Similarity Search

Another thing you can do after choosing a distance metric is find the most similar entities to some query – a paradigm called kNN (k nearest neighbors), or similarity search. This can even power user-facing search features or recommendation engines. See the example below.

The most similar tracks, according to the smallest Euclidian distance between 128-dimensional genre embeddings.
Click the dropdown to query a different track.

Aggregating

Hierarchies come up all the time with structured data. Documents are made up of paragraphs, which are made up of sentences, which are made up of words. Countries are made up of states. A customer’s activity is made up of individual actions they took.

Often you’ll find you want to navigate up these hierarchies, working with higher-level entities even though you have features/vectors for their components. Vectors allow for an elegant solution: just average each dimension, independently, across all components within a higher-level entity. A big advantage here is that it works with any number of components. For example, let’s say you’d like to measure the similarity between two documents, one with three paragraphs and one much longer with ten. Assuming you had an embedding vector for each paragraph, simply averaging the embeddings of document A’s paragraphs, and separately averaging those of document B’s paragraphs, would give you two vectors of equal length that you can measure the distance between.

Picking up on our example of music from above, we can average together the vectors of tracks within each album to get album vectors, seen below.

After aggregation, each dot now represents an album.

Training ML models

There are many “black box” ML paradigms that do well with inputs of arbitrary tabular data, and learn statistical patterns as needed for downstream tasks. If you’ve already done the work to represent your dataset as well-formed vectors, these downstream tasks can generally be prototyped very quickly. For example, a decent enough supervised classifier might be trainable with just an hour’s worth of labeling from a domain expert.

Another example of where this ease-of-application could come in handy is automated data imputation, or filling in missing values with substitutes. Imagine if a few dozen dimensions had missing values, and your goal was to generate reasonable guesses. Coming up with a unique process for each dimension could be a lot of work. However, if the vectors are in a good enough shape, you can automatically train a supervised regression model to predict each dimension of interest, iteratively holding them out as target labels and using the remaining dimensions as inputs.

Bridging data sources

It’s not uncommon to construct datasets of entities from multiple data sources. For example, maybe each entity is a customer, and you’d like to combine demographic data with purchase history. You might have already finished an analysis, only to discover there is more data about these customers on the way.

The easiest way to combine such data sources is to simply concatenate their vectors. This is a good choice if your downstream application is a ML model, which ideally will learn patterns across each source.

In theory, all of the other tools I’ve described can also work after concatenation. One thing to be mindful of, however, is unintentionally allowing data sources with a large dimensionality to have outsized influence. Many of the methods work by computing pairwise distances between entities, and depending on your distance metric of choice, each dimension treated equally means the number of dimensions has to be considered. One way around this is instead of concatenating vectors, computing the pairwise distance matrix of each data source separately, then averaging them together so that each source has equal weight. Scikit-learn APIs generally accept distance matrices as inputs instead of raw vectors by specifying metric='precomputed', as does umap-learn. This trick also offers an opportunity to customize the influence of each source with a weight factor of your choosing.

The scatterplot projections below are based on the genre-describing vectors used above, in addition to high-level audio features that describe mood, provided by the AcousticBrainz API, which itself uses the essentia library.

Each dot represents a track, as represented by two different data sources – 128-dimensional genre embeddings, and 14-dimensional acoustic mood descriptors provided by AcousticBrainz/essentia.

Bridging modalities

Going beyond data sources, high-level entities might be represented by multiple modes of data entirely. Text, images, videos, audio, etc. might all look very different from each other, but as long as they can be boiled down to vectors of (mostly independent) features, you can work with them together. Connecting modalities is a hot topic in modern ML (see CLIP, used by DALL-E). But the bag-of-vectors paradigm is, I think, the easiest way to do it for limited use cases.

So far the example music vectors above have been made up of features extracted from raw audio. Looking for complementary data, I came across a nice dataset of descriptors from rateyourmusic.com – albums were labeled by human listeners according to a taxonomy of atmosphere, form, lyrics, mood, style, and technique. For example, Nirvana’s Nevermind is labeled [energetic, rebellious, angry, malevocals, apathetic, sarcastic, alienation, passionate, anxious, self-hatred]. First, to get the labels into a more useful semantic vector space, I used UMAP to project the term-document matrix into 16 dimensions (I also tried Truncated SVD and PCA). These are then combined with the audio-based features from before, either at the distance matrix stage, or via concatenation in the case of PCA, while being projected down even further to two dimensions for the scatterplots below.

Each dot represents an album, as represented by up to three different data sources – aggregated 128-dimensional genre embeddings, aggregated 14-dimensional acoustic mood descriptors, and 16-dimensional vectors based on human labels from rateyourmusic.com.

Conclusion

If you employ these paradigms enough times for enough datasets, you’ll start to think in vectors. “My company needs to do X with a bunch of Ys… how would I design vectors to represent Ys?” Or even outside of work, “I should buy a new car soon.. hmm, I wonder what car model vectors would look like?”

Again, the title of this post is a bit facetious. Vectors themselves are not the magic sauce, but rather a mental and practical bridge between worlds. In one world, we obsess over how to represent Things as numbers – deciding their most important features and scrutinizing data quality. In the other world, we’re able to place these numbers inside a black box, and use them to learn and build.

But what does it even mean for two sounds to be similar? Ultimately, if the generated audio is intended for human ears, what matters most is human perception of sound similarity and quality – any mathematical functions we come up with are only as good as they correlate with those. This is obvious if we’re trying to use objective evaluation metrics to compare models, but it also tends to be true for the very loss functions used to optimize them in the first place.

This blog post is a survey of loss functions used in modern audio ML research, including trends and takeaways. I mostly focus on those that approximate the closeness of two sounds, so for example adversarial loss is not covered.

You might find this post useful if you’re working in:

• Speech synthesis, including text-to-speech (TTS)
• Speech denoising/enhancement
• Speech separation
• Voice conversion
• Music source separation
• Music synthesis
• Effect pedal simulation
• Phase reconstruction
• Audio super-resolution
• Representation learning for audio

A refresher on waveforms and spectrograms

If you’ve already worked with audio data you can skip this section. Otherwise, it helps to start with the very basics: sound is just, like, the air vibrating, man.

By sampling (measuring) these vibrations over time, we can represent audio digitally. This is typically done tens of thousands of times per second, resulting in the primary representation of digital audio: the waveform. Take a look at this helpful gif from Jan Van Balen, showing a waveform from a few seconds of a cello recording:

Zooming in, you can see it’s really just a very long one-dimensional array of floating point sample values, ranging from -1.0 to 1.0 and discretized according to the encoding’s precision. In the case of stereo audio, you would have two channels of samples (representing “left” and “right”), similarly to how there are RGB channels in computer vision. Easy, right?

Notice how the waveform above is not fully random – it goes up and down over time in a semi-repeated way, almost like a sine wave. Measuring the frequency (time between repeats) and amplitude (deviation from zero) of such oscillations is very useful for audio analysis tasks. The mathematician Joseph Fourier showed how any continuous function can be represented as a series of sine and cosine functions. This is the idea behind the Fourier transform, which decomposes a signal into two components – a magnitude spectrum and a phase spectrum – together called a spectrogram. The magnitude spectrum is usually visualized with frequency as the x-axis, and magnitude (how far from zero) as the y-axis, shown on the right in the figure below.

On the left we see a simple waveform (the blue line), which is the same as the sum of the two overlapping grey lines. After a Fourier transform, on the right we see the resulting frequency magnitude spectrum. Source.

Together, the magnitude and phase of a sine wave at a certain frequency form a complex number. The phase is the angle, describing at what point in the cycle a sine wave begins.

Taking the Fourier transform of long portions of audio isn’t particularly useful. Instead, it’s commonly applied to small windows across time, known as the Short-Term Fourier Transform (STFT). Stacking these next to each other lets us easily visualize how different frequencies are activated in a signal across time. See below the magnitude and phase spectrogram of a piano recording, borrowed from Sander Dieleman’s fantastic blog post on generating audio in the waveform domain:

Magnitude (top) and phase (bottom) spectrums from a piano recording. X-axis is time, Y-axis is frequency. Source.

Notice that the magnitude spectrogram is easily interpretable, but phase looks essentially random. Unsurprisingly, it follows that magnitude spectrograms carry most of the perceptually important information about audio signals and so for analysis, phase can often be discarded altogether. However, phase is still quite important for generating high-quality outputs. To see what I mean, listen to the following piano piece with its original phase information and with a random phase.

Left: piano recording with original phase. Right: the same with random phase. Source.

You can still hear the melody with a random phase, but it doesn’t sound like something you’d want to come out of a generative audio model.

Generating spectrograms

If you have both magnitude and phase components for a signal, you can do an inverse STFT operation to get back to a waveform. However, given that magnitude spectrums alone offer an effective and more condensed way to model audio signals… can you just drop the phase component altogether?

As it turns out – yes this is feasible, and commonly done! The only issue is getting the result back into the time domain as a waveform. Here we have several options.

For sequence-to-sequence problems, one strategy is to use the original phase component of the input signal.

Sequence-to-sequence audio modeling in the magnitude spectral domain by re-using the original phase.

This can work, but it might not sound great for longer sequences that were drastically transformed by the model.

When that’s not feasible, another option is to come up with a phase component from scratch. This is the field of research called phase reconstruction. One time-tested approach is the Griffin-Lim algorithm, which uses the intuition that the STFT usually results in the same frequencies being activated in neighboring frames to iteratively come up with a reasonable guess for a phase.

More recently, generative models like the WaveNet[1] and WaveGlow[2] have been used as “vocoders”, or front-end components that probabilistically generate waveforms given magnitude spectrograms. You can either add these to your neural net architecture and train it all end-to-end, or you can use pretrained versions and just focus on training your custom components.

In recent years it has seemed like generative audio research was heading in the direction of always modeling the waveform domain end-to-end. However, high quality vocoders like WaveNet have enabled researchers to continue utilizing the more compact and in some cases better performing magnitude spectrograms as primary modeling domains.

With that out of the way, onto the loss functions.

L1 and L2 of waveforms

These are the bread and butter of learning objectives. If you need a reminder: L1, also known as Mean Absolute Error (MAE), involves simply aligning the two waveforms, finding the absolute difference between each pair of sample points, and averaging them. L2, or Mean Squared Error (MSE) is the exact same except you square each difference before averaging, which has the effect of penalizing larger errors and being more forgiving of smaller ones.

L1 and L2 are used all the time for audio models in the waveform domain, including many of the top music source separation approaches [3, 4, 5, 6], and are always worth a try. However, there are some potential weaknesses to be aware of.

For one, they don’t reflect natural biases in human hearing. Interestingly, we perceive certain frequencies to be louder or quieter than others, even when they played with an equal amount of energy. In the 1930’s, researchers actually measured this, resulting in the first of many “curves” that try to capture human loudness bias as a function of frequency.

Source

If training a model for human ears, with L1/L2 you risk overweighting the importance of low frequency sounds. One way around this is to use pre-emphasis filters, which will adjust the energy levels of different frequency channels according to an equal-loudness curve, such as the one above. Librosa and SpeechPy offer this functionality.

They also fail to correlate well with human judgement of speech quality and intelligibility [7]. This makes sense from an evolutionary standpoint – humans had plenty of reasons to favor our ability to perceive speech, so we evolved neural pathways adept at doing that, which simpler metrics might not capture.

Shift Invariance

Another issue of note is that L1 and L2 in the waveform domain are not shift invariant. If your models comes up with the exact correct output, but it’s shifted just 1 millisecond too early or late, a human would consider it a perfect match but error would likely spike. This can especially manifest as phase shift invariance, where two waves have similar frequencies but out-of-sync phases, leading to a large gap in waveform distance.

This might not be a problem for you depending on your goals, or your model of choice. Recent approaches for sequence-to-sequence problems are converging on the U-Net architecture, or similar. These models have skip connections, which are more robust to shift invariance because they have access to the original input waveform late in the forward pass.

One trick for incorporating shift invariance is Dynamic Time Warping (DTW), whereby dynamic programming is used to find the minimal-cost alignment path between two time series. This means they can be considered a close match even with mutations or temporal shifts. A great example can be found in the paper “End-to-End Adversarial Text-to-Speech”[31], where DTW is used to give the model flexibility on the timing of speech utterances. Note that for loss functions, you will need to use the differentiable variant soft-DTW.

Aside from shifts, there are other transformations for which L1/L2 will unfairly penalize. When it comes to loss functions and evaluation metrics, it’s important to keep in mind which invariances you care about as a researcher.

Waveform L1/L2 make for great baselines, and it is also common to see them added to other losses.

Spectral Losses

Earlier I mentioned how magnitude spectrograms resulting from Fourier transforms offer a rich data representation, great for both analysis and generation. What I didn’t mention is that STFT is a differentiable operation! It should be unsurprising then that L1 and L2 work well as loss functions in the spectral domain, too. Sometimes you will see this referred to as spectral loss, though there are enough variations that I consider it a family of loss functions.

If your model is already working with spectrograms, then applying L1/L2 is trivial. But even if it stays in the waveform domain, both tensorflow and pytorch offer differentiable tensor versions of the STFT that you can apply during training only.

Usually the log magnitude spectrogram is used, sometimes with a small number epsilon added first, which controls the trade-off between accurately representing low energy and high energy spectral components [8].

$$ \big|\big| log(|STFT(y)| + \epsilon) - log(|STFT(\hat{y})| + \epsilon)\big|\big|_{\ell} $$ Log magnitude loss. Here \(y\) is the target signal, \(\hat{y}\) is the predicted, \(||\cdotp||_{\ell}\) is the \(\ell\) norm (like L1 or L2), and \(\epsilon\) is a small number. The absolute value around STFT is a way to discard the phase component.

Log magnitude spectrogram of an example sound, a prediction from a model, and their squared difference.

As mentioned previously, STFT involves sliding a window across a waveform, taking Fourier transforms as you go. This hints at a couple of parameters that will effect the resulting spectrogram, and thus any downstream learning. For one, you have to choose a window size. Larger windows will be able to capture smaller frequencies, but will take up more space. You also need a hop length – it’s common to have successive windows of samples overlap each other, and hop length in relation to window size determines the extent of this overlap. In order to minimize the bias of any one choice of parameter, sometimes people use a multi-resolution spectral loss by summing the results of several runs with different window parameters.

Although logarithmic magnitude spectral losses are more common, you do sometimes see magnitude squared. In signal processing terminology, this turns amplitude into power, and so you will see this called power loss. As you can see below, a power spectrum concentrates around the strongest frequency bands, drowning out the others. It follows that power loss is great at ensuring your predictions have an accurate distribution across the frequency spectrum, in line with your target domain. For example, it could help a text-to-speech (TTS) system match the highs and lows of believable human speech [9].

Power spectrogram of an example sound, a prediction from a model, and their squared difference.

Another related loss you might see is spectral convergence loss:

$$ \frac{\big|\big| |STFT(y)| - |STFT(\hat{y})| \big|\big|_E}{\big|\big| |STFT(\hat{y})| \big|\big|_E} $$

\(||\cdotp||_E\) above \(x\) is the Euclidean matrix norm (AKA Frobenius norm). In words: you take the Euclidean distance of the difference between two magnitude spectrograms, and normalize it by the Euclidean “length” of the original signal. Similarly to the power spectrum, since we’re squaring each error component, it highly emphasizes when any one frequency bucket prediction is way off, and is more forgiving of several frequency bucket predictions being just a little bit off. According to one paper I read [10], this can especially help in early phases of training.

Mel spectrogram loss

Earlier I mentioned how humans have natural biases in how we hear loudness at different frequencies. We also have biases in how we hear pitch. Low frequencies sound “low” to us, and higher ones sound “high”, but the relationship between frequency and perception is non-linear – a 100 Hz and 200 Hz sine wave will sound much further apart than a 10,000 Hz and 10,100 Hz wave. There have been several attempts at codifying this relationship through listening tests; the most widely used formulation is the Mel scale. Obtaining the Mel spectrogram is similar to the STFT operation, but frequency energies are placed into different buckets of equal perceived pitch difference, according to the filterbanks of the Mel scale.

The Mel filterbank, showing windows of frequencies that are aggregated into a Mel spectrogram. Source.

Any of the previously mentioned spectral losses can also be carried out using the Mel spectrogram.

A quick side note: the Mel spectrogram is actually a common general representation for modeling audio. Along with being a better correlate with human pitch perception, it’s also more compact then normal spectrograms. How much so depends on how precise you make the Mel filterbanks (e.g. the n_mels parameter in librosa), but depending on your problem, you can potentially reduce your space footprint by an order of magnitude without sacrificing performance. You will, however, need a specialized vocoder model if you want to get back to waveforms from a Mel spectrogram.

Spectral losses come up in many areas, but especially speech enhancement, TTS, and music synthesis.

Source Separation

Audio source separation is the problem of separating an audio signal into the components (sources) that were mixed together to create it. One application is in music, where source separation can be used to “demix” a song (usually into predefined categories, like bass, drums, vocals, and other). Another application is voice separation, which attempts to tackle the classic cocktail party problem by isolating individual voices from a conversation, or from the background of a recording.

During model training, each predicted source is compared to a ground truth target using an objective function. While the losses covered above can work just fine, another is more common. For the following, suppose you are trying to predict the source signal \(s\), and your model comes up with an imperfect \(\hat{s}\).

SDR (Source to Distortion Ratio):

This metric is measured in decibels, where higher is better. Any mistake in the prediction will cause the denominator to grow, which will cause the overall metric to go down.

SDR is equivalent to the classic signal-to-noise ratio, so sometimes you will see it reported as SNR. However, there is a totally different SNR metric (Sources to Noise Ratio) also sometimes used to evaluate source separation results, so be careful to know which one you’re dealing with.

Scale invariance

A major downside of vanilla SDR is that it’s sensitive to the loudness of the predicted signal. If you were to scale the prediction up or down, the metric would rise or fall along a curve around some point of “ideal” scale. In the context of SDR as an evaluation metric, this means some researchers might optimize scale to achieve the highest score and others might not, leading to unfair comparisons. In the context of SDR as an objective function, I think this would lead to a slower and less-smooth learning curve.

Thankfully, some researchers surfaced the issue in 2018 and came up with a scale-invariant version of SDR where the optimal scaling factor is baked in.

SI-SDR (Scale-Invariant Source to Distortion Ratio):

Again, sometimes you will see this reported as SI-SNR. At this point, I see no reason why anyone would want to use plain SDR in favor of the scale-invariant version.

Permutation invariance

When you know the categories of sources ahead of time and your training set comes with labeled sources, you can just add the losses from each prediction. However, for problems with an unknown number of sources (called blind source separation), such as in the cocktail party problem, it’s not that simple. You need to know which predictions correspond to which ground truth targets. Further, you’ll want a one-to-one correspondence between the predictions and targets, so that you don’t end up with two predictions of the same source or not enough predicted sources. Using random assignments at every optimization step would unfairly penalize the model for a bad roll of the dice, leading to very unstable training. Rather, in these cases you should apply your loss function in a permutation invariant way. One such method is Permutation Invariant Training (PIT), wherein every combination of assignments is attempted, and the best set (lowest combined loss) is used. This incentivizes the model to settle into learning each source once, without needing to care in what order they are output.

The asteroid library offers a wrapper Python class that makes it easy to turn any loss function into a permutation-invariant version.

Perceptual Loss

Neural networks are great at discovering features that distinguish between inputs. And it’s often easy to imagine other problems for which the same features would be helpful. For example, a facial recognition model might have a neuron that fires when presented with a horizontal line, which might be used by a higher-level neuron to decide whether the picture contains a chin. If you were then designing a person-detection model, it could probably benefit from being able to detect both chins and faces.

So how could you utilize features from a network trained to solve problem \(A\) to help train a network to solve a similar problem \(B\)? This is the research area known as transfer learning. The simplest way is to directly feed the feature values into your new network. In our example above, the person-detection network would see its normal inputs, but would also be given the neuron values from the facial detection network to work with. Another way is fine-tuning, where network \(A\) is first trained to solve problem \(A\), and then further trained to solve problem \(B\).

Neural features are not just numbers; they are connected to other neurons, that, once fully unrolled, are large differentiable functions. We should be able to utilize these function gradients without being forced to re-use the neural architecture of network \(A\) to solve problem \(B\). These are the insights that motivate perceptual loss, also known as deep feature loss. Perceptual loss was inspired by research on style transfer in computer vision.

Lets assume you have already trained an auxiliary task \(B\) (seen in the figure below in red). The basic idea is to use the features of \(B\) to compute a gradient that updates \(A\) in a way that makes audio predicted by \(A\) closer to ground-truth audio in \(B\)’s feature space. Usually in deep learning the parameters being updated during optimization are the same ones that were used to compute the gradient. However, here \(B’s\) parameters are frozen.

A depiction of perceptual loss learning in the audio denoising domain. First a batch of training inputs is put through the network being optimized (top). The output is then put through an auxiliary network (bottom), pre-trained on a similar problem. Finally, the original batch of signals is also put through the auxiliary network, and the two activations are compared to compute a gradient that updates the original network. Source.

The loss formula is usually something called feature reconstruction loss, where you sum the L1/L2 distances of neural activations of the auxiliary network’s first \(k\) layers. Using earlier layers favors reusable patterns over high-level features that tend to be specific to the auxiliary task.

Here are some examples from audio research:

• “Parallel wavenet: Fast high-fidelity speech synthesis”[9] – among other loss components, they include a perceptual loss term using a network trained to detect the phonemes of speech signals. They basically utilize speech-to-text features as a means of improving text-to-speech, which I find to be an interesting parrallel.
• “Voice Separation with an Unknown Number of Multiple Speakers”[11] – they attempt the cocktail party problem using perceptual loss against a networked trained on speaker identification. The intuition is that if a network is able to pick apart the nuances of different speakers’ voices, it should also provide useful features for separating voices from background audio.
• “Hierarchical Timbre-Painting and Articulation Generation”[12] – they attempt a kind of music synthesis whereby you take a source recording of some instrument playing a melody, and render it as a completely different instrument, while preserving the nuances of pitch and loudness. Their auxiliary network of choice is the pitch-detecting CREPE, with some pretty cool results here.
• “Deep Network Perceptual Losses For Speech Denoising”[13] – they utilize a deep feature loss based on speech detection as an auxiliary task, as well as one trained on the audio event detection dataset AudioSet.

Learned Perceptual Loss

In the search for a metric that highly correlates with human perception, the authors of “A Differentiable Perceptual Audio Metric Learned from Just Noticeable Differences”[14] take a unique approach.

They first crowdsource a dataset of human judgements that capture “just noticeable differences” (JND) – users are presented with two audio recordings, where one is a slightly perturbed version of the other, and asked whether they are exactly the same sounds or not. The amount of perturbation (e.g. reverb, compression, dropouts) is carefully chosen to be small enough so that the authors end up with lots of data right around the threshold of what’s barely noticeable.

They then fit a deep neural network to predict these human judgements, by minimizing binary cross-entropy. The resulting network can be used as a quality metric, and, since it is differentiable, as a learning objective (similarly to how you would use the perceptual losses above). They call their metric DPAM, available for both tensorflow and pytorch. More recently they released an improved metric called CDPAM[15], which is more robust to perturbations not used in the original JND dataset.

The released version of CDPAM is meant to be general-purpose, but their approach also seems quite adaptable. If you have particular requirements for your ideal similarity metric, or you’re working in a specialized domain with access to a pool of experts, you should be able to capture your requirements in a JND dataset and train a metric yourself.

Other

Speech Quality Metrics

If you work in speech enhancement or TTS, you may come across the metric PESQ. Originally a commercial telecomm specification, it is differentiable, but only defined up to frequencies of 8k so not ideal for learning [16]. It did inspire a newer trainable variant PMSQE[17], which looks promising.

You may also come across Short-Time Objective Intelligibility (STOI)[18], which is differentiable but usually just used as an evaluation metric.

Categorical cross entropy

You might not have expected to see a categorical loss for inherently continuous waveform data. Nonetheless, architectures like WaveNet[1] model waveform sample probabilities as discrete classes under a softmax distribution.

One benefit of quantization is that it prevents mode collapse. Sometimes models predicting continuous-valued waveforms will cut their losses by mostly just outputting the safest option between -1 and 1: zero (or near it). This is especially true earlier in the training process. By instead predicting non-ordinal classes, no such safe option exists and the model is forced into trying something.

Conclusion

During my informal survey of research from the past five years or so, I noticed some trends:

These are not endorsements, just patterns. Other areas had no clear consensus. Here are some additional tips for practitioners:

• Start with L1 in whichever domain is easiest given your model architecture, then expand your search.
• Don’t be afraid to mix several loss functions. It’s not uncommon to see three or even four components.
• Strive to use as many evaluation metrics as you can. Be especially careful if you’re using the same metric for both an objective function and evaluation, as it may leave you thinking your model is performing better than it actually is.

If you’re using pyTorch and looking to try some of these out, I recommend the auraloss library, or asteroid if you’re working on source/voice separation. Sadly I could not find anything similar for tensorflow (any readers in need of a side project?…)

There are many objective functions available for your audio models, and given the number of potential combinations, you may not be able to try them all. By being informed about their pros, cons, and prevalences, I hope you can save precious GPU hours.

References

• [1] A. Oord. WaveNet: A generative model for raw audio, 2016.
• [2] R. Prenger. WaveGlow: a Flow-based Generative Network for Speech Synthesis, 2018.
• [3] A. Défossez. Music Source Separation in the Waveform Domain, 2019.
• [4] D. Stoller. Wave-U-Net: A Multi-Scale Neural Network for End-to-End Audio Source Separation, 2018.
• [5] R. Hennequin. Spleeter: a Fast and Efficient Music Source Separation Tool with Pre-Trained Models, 2020.
• [6] N. Takahashi. D3Net: Densely connected multidilated DenseNet for music source separation, 2020.
• [7] S. Fu. MetricGAN+: An Improved Version of MetricGAN for Speech Enhancement, 2021.
• [8] A. Défossez. SING: Symbol-to-Instrument Neural Generator, 2018.
• [9] A. Oord. Parallel WaveNet: Fast High-Fidelity Speech Synthesis, 2017.
• [10] S. Arik. Fast Spectrogram Inversion using Multi-head Convolutional Neural Networks, 2018.
• [11] E. Nachmani. Voice Separation with an Unknown Number of Multiple Speakers, 2020.
• [12] M. Michelashvili. Hierarchical Timbre-Painting and Articulation Generation, 2020.
• [13] M. Saddler. Deep Network Perceptual Losses For Speech Denoising, 2020.
• [14] P. Manocha. A Differentiable Perceptual Audio Metric Learned from Just Noticeable Differences, 2020.
• [15] P. Manocha. CDPAM: Contrastive learning for perceptual audio similarity, 2021.
• [16] POLQA Vs PESQ: Objective quality scoring explained.
• [17] J. Martín-Doñas. A Deep Learning Loss Function based on the Perceptual Evaluation of the Speech Quality, 2019.
• [18] C. Taal. A short-time objective intelligibility measure for time-frequency weighted noisy speech, 2010.
• [19] Y. Luo. Conv-TasNet: Surpassing Ideal Time-Frequency Magnitude Masking for Speech Separation, 2019.
• [20] C. Subakan. Attention is All You Need in Speech Separation, 2020.
• [21] M. Pariente. Filterbank design for end-to-end speech separation, 2019.
• [22] H. Kameoka. StarGAN-VC: Non-parallel many-to-many voice conversion with star generative adversarial networks, 2018.
• [23] S. Wang. NoiseVC: Towards High Quality Zero-Shot Voice Conversion, 2021.
• [24] J. Zhang. Sequence-to-Sequence Acoustic Modeling for Voice Conversion, 2018.
• [25] Y. Wang. Tacotron: Towards End-to-End Speech Synthesis, 2017.
• [26] J. Shen. Natural TTS Synthesis by Conditioning WaveNet on Mel Spectrogram Predictions, 2017.
• [27] Y. Jia. Transfer Learning from Speaker Verification to Multispeaker Text-To-Speech Synthesis, 2018.
• [28] A. Gritsenko. A Spectral Energy Distance for Parallel Speech Synthesis, 2020.
• [29] X. Wang. Neural source-filter-based waveform model for statistical parametric speech synthesis, 2018.
• [30] W. Ping. Deep Voice 3: Scaling Text-to-Speech with Convolutional Sequence Learning, 2017.
• [31] J. Donahue. End-to-End Adversarial Text-to-Speech, 2021.

As my own drum library has exploded in size (ok I’m a hoarder), it’s got me thinking of ways to take control and organize it better. One thing that has bugged me is the inconsistent naming of sound files:

• Abbreviations like CR.wav for a crash cymbal, that wouldn’t come up in a search
• Different classification schemes, e.g. Perc.wav for ‘percussion’, instead of something more specific like rimshot
• Labels I don’t agree with, or are just flat out wrong

It would help to have my own way of assigning drum sound labels, and in this post I will share several ways to achieve that with ML.

All code is open-source, so you can follow along and try this out yourself! You only need python and your own drum sounds.

What do drum sounds look like?

If you’ve ever worked with audio data you’re probably familiar with waveforms, which show the raw signals that might be sent to speakers for playback. Spectrograms are more insightful because they separate a signal into different frequency bands, allowing you to see the amount of energy up and down the spectrum. See examples of both below.

Snares

Kicks

Toms

Plots above made with librosa, a great python library for audio processing.

You’ll notice some patterns, like how the snare waveforms look “fuzzier” and have energy up and down the spectrum. Kicks are obviously strongest in the lower frequencies (you can see spikes around 64 Hz), and toms look similar to kicks in ways but have spikes anywhere from 100Hz-500Hz. Now, lets see what ML can do.

A taxonomy

I needed a set of drum type classes. A few options came up in research, the simplest being bass drum, snare drum, hi-hat (used in this thesis and elsewhere). A limited set like that would make prediction easy, but it’s not as useful. The authors of one paper used a scheme closer to that of a standard rock kit: kick, snare, low tom, medium tom, high tom, open hi-hat, closed hi-hat, ride, crash. There’s even the General MIDI spec with dozens of percussive types – though, with that many classes it becomes hard to build up a dataset with enough of each.

I wanted to strike the right balance of taxonomy size, while also choosing classes more reflective of modern hip-hop and electronic drum kits found on the web. So, I made my own: hat, tom, ride [cymbal], open [hi-hat], kick, bongo [and conga], clap, snare, rim, snap, crash [cymbal], shaker

Building a clean dataset

Drum libraries are usually scattered across different folders, so given the top-level folder, it shouldn’t be hard to recursively walk and find all eligible sounds. I’m a fan of using pandas to explore data, so I initialized a DataFrame with one sound per row.

def read_drum_library(input_dir_path):
    logger.info(f'Searching for audio files found in {input_dir_path}')

    dataframe_rows = []
    for input_file in input_dir_path.glob('**/*.*'):
        absolute_path_name = input_file.resolve().as_posix()
        if not can_load_audio(absolute_path_name):
            continue

        properties = {
            'audio_path': absolute_path_name,
            'store_path': file_store_path.as_posix(),
            'file_stem': Path(absolute_path_name).stem.lower(),
            'start_time': 0.0,
            'end_time': np.NaN
        }
        # Tack on the original file duration (will have to load audio)
        audio = read_audio.load_raw_audio(absolute_path_name, fast=True)
        properties['orig_duration'] = len(audio) / float(read_audio.DEFAULT_SR)

        dataframe_rows.append(properties)

    return pandas.DataFrame(dataframe_rows)

I got rid of really quiet sounds by thresholding RMS, and also excluded those over 5 seconds.

I wanted to go even further in isolating single percussive hits, because I had noticed some loops with multiple hits that might throw off a model. I also wanted some consistency around how much silence appeared at the begging of sounds. My solution for both of these issues was to use librosa’s onset detection API. An onset is the moment that marks the beginning of a rise in energy of a sound. All I had to do is set the start time of each sound to just before the first onset, and the end time to just before any second onset.

One downside of the supervised deep learning techniques we’ll see later is that they require a lot of data. It’s easy to find tons of drum kits on the web, but what about labels? Fortunately, we can build up a dataset without manual annotations by just trusting the original filenames. They won’t be perfect, but they’ll be good enough to start.

DRUM_TYPES = ['hat', 'tom', 'ride', 'open', 'kick', 'bongo', 'clap', 'snare', 'rim', 'snap', 'crash', 'shaker']
drums_df = read_drum_library(drum_lib_path)
for drum_type_class in DRUM_TYPES:
    drum_sounds.loc[drum_sounds.file_stem.str.contains(drum_type_class), 'file_drum_type'] = drum_type_class

If you’re following along at home, you can run:

python preprocess.py --drum_lib_path ~/Music/drums

With that, my dataset looked like this:

> import pickle
> drums = pickle.load(open('data/interim/dataset.pkl', 'rb'))
> drums = drums[~drums.file_drum_type.isna()]   # Restrict to sounds with recognizable file labels
> drums.info()
Data columns (total 7 columns):
 #   Column          Non-Null Count     Dtype  
---  ------          --------------     -----  
 0   audio_path      11672 non-null     object
 1   store_path      11672 non-null     object
 2   file_stem       11672 non-null     object
 3   start_time      11672 non-null     float64
 4   end_time        818   non-null     float64
 5   orig_duration   11672 non-null     float64
 6   file_drum_type  11672 non-null     object
dtypes: float64(3), object(4)
memory usage: 1.2+ MB

> drums[['file_stem', 'start_time', 'end_time', 'file_drum_type']].sample(5)
   file_stem  start_time  end_time file_drum_type
15   openhat    0.000000       NaN           open
5     kick18    0.000000       NaN           kick
38   schat13    0.000000       NaN            hat
2    lakick9    0.000000       NaN           kick
19    kick8     0.022653  0.092494           kick

> drums.drum_type.value_counts()
snare     2952
kick      2845
hat       1805
tom       1371
clap      1048
rim        436
open       337
ride       263
crash      257
snap       110
bongo      101
shaker      98

As you can see, I had a pretty large class imbalance. This can be bad for model performance on the lesser represented classes. One way to help with this is over-sampling via something like imbalanced-learn or pyTorch’s WeightedRandomSampler. Instead, during training procedures I simply capped class size at 2000 like so:

drums = drums.groupby('file_drum_type').head(2000)

Method 1: Hand-crafted Features + Random Forest

Analyzing percussion sounds is not new. Let’s look at some of the hand-crafted features that have been used in the past. First, a few from the MPEG-7 standards that reflect the change in a signal’s power over time:

• Log attack time - “attack” here means how quickly the signal reaches its peak.
• Temporal centroid - this measures how far along into the sound we’ve reached half of the signal’s power. A sound that starts loud and fades will have an early temporal centroid, while something like a crash cymbal will have a later one.
• Spectral centroid - similarly, this measures the center of gravity in the frequency domain. You can think of it as how low or high the sound is. Since it changes over the duration of the signal, we’ll take the average.

Since we’re using “attack”, we can also include “release”, which measures how long the signal lasts from its peak before dipping below a threshold (I use 2% of the peak, following the lead of this paper). There are also other spectral features we can utilize, many provided by librosa:

• Spectral bandwidth - Intuitively, how spread out the frequency spectrum is. Technically, the second central moment of the spectrum.
• Spectral flatness - measures how noise-like a sound is, as opposed to an isolated tone
• Spectral rolloff - as opposed to the spectral centroid, this measures at what frequency a certain percentage of the magnitude distribution is less than. To capture the low and high ends of the frequency spectrum, we’ll compute the rolloff at 15% and 85%.

Some additional features we can pull in:

• Duration
• Average, max, and standard deviation of log RMS - RMS is root mean squared energy, which you can just think of as volume for our purposes.
• Average change in RMS - since we compute RMS per frame of audio (each frame is ~23 ms in my implementation), we can also look at how it goes up and down between frames. This is also true for other features in this list.
• Crest factor - measures how intense the audio’s peaks are. Peak RMS divided by average RMS.
• Zero crossing rate (ZCR) - If you zoom in on a waveform, you’ll see the signal oscillating above and below zero. This measures how many times that happens in one frame of audio. We can take the average across all frames, the standard deviation, or the ZCR at the loudest frame since that’s probably revealing.

One last set of features worth consideration is the Mel Frequency Cepstral Coefficients (MFCCs). If you hear multiple instruments playing a note at the exact same pitch and volume, you will probably still be able to tell them apart. This is because they exhibit different “timbre” characteristics. MFCCs are a set of (typically 10-20) features great at capturing timbre. I found a great explanation in another blog.

After applying summary statistics (average, max, standard deviation, ZCR, and derivative) to frame-based features, I ended up with 72 features in total.

There were some missing values from when I couldn’t take the derivative of short single-frame sounds, but I used Scikit-learn’s IterativeImputer to come up with reasonable guesses. I also scaled features to have a mean of zero and std deviation of one.

# Turn class labels into numbers for scikit-learn
drum_type_labels, unique_labels = pandas.factorize(drums.file_drum_type)
drums = drums.assign(drum_type_labels=drum_type_labels)

# We'll train on 75% of the data
from sklearn.model_selection import train_test_split
train_clips_df, val_clips_df = train_test_split(drums, random_state=0, test_size=0.25)

# Get numpy arrays of our features (which all start with an underscore)
train_np = train_clips_df.filter(regex='^_', axis=1).to_numpy()
test_np = val_clips_df.filter(regex='^_', axis=1).to_numpy()

# Fill in missing values, normalize
imp = IterativeImputer(max_iter=25, random_state=0)
imp.fit(train_np)
train_np = imp.transform(train_np)
test_np = imp.transform(test_np)
scaler = preprocessing.StandardScaler().fit(train_np)
train_np = scaler.transform(train_np)
test_np = scaler.transform(test_np)

Finally, time to train!

from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report

model = RandomForestClassifier(n_estimators=400, min_samples_split=2)
model.fit(train_np, train_clips_df.drum_type_labels)
pred = model.predict(test_np)
print(classification_report(pred, val_clips_df.drum_type_labels,
                            target_names=drum_class_labels, zero_division=0))

To run this experiment yourself:

python drum_sound_classifier/models/train_sklearn.py --inputs descriptors --model random_forest --max_per_class 2000

optional arguments:
  -h, --help            show this help message and exit
  --inputs {cnn_embeddings,descriptors}
                        The source of features on which to build a model
  --model {lr,svc,random_forest,gb,knn,all}
  --max_per_class MAX_PER_CLASS
                        limit common drum types to lessen effects of class imbalance

Scikit-learn makes it easy to try several models using the same API.

An 82% accuracy and F1 score, not bad.

              precision    recall  f1-score   support

        open       0.54      0.38      0.44        56
       bongo       0.89      0.32      0.47        25
      shaker       0.70      0.42      0.52        38
         rim       0.72      0.60      0.66       108
        snap       0.79      0.63      0.70        30
        clap       0.79      0.74      0.76       215
         hat       0.77      0.86      0.81       395
       snare       0.79      0.86      0.83       484
         tom       0.87      0.84      0.86       327
       crash       0.93      0.81      0.87        68
        ride       0.95      0.80      0.87        71
        kick       0.90      0.94      0.92       509

    accuracy                           0.82      2326
   macro avg       0.80      0.68      0.73      2326
weighted avg       0.82      0.82      0.82      2326

As you can see, performance on the test set really depends on the drum class. Open hi-hats and bongos are difficult, while kicks are easily classified (not surprising, because their lower frequency presence is unique). For a better understanding of the model’s misclassifications, lets look at a confusion matrix. The following is normalized by row, so for example 8% of the time bongos are confused as claps.

One reason I like using Random Forest models is their ability to give insight into which features were most useful for differentiating between classes. In this case, 72 features is way too many to comprehend in one chart, but we can group most of them into feature sets:

Spectral rolloff features are very important because they help differentiate sounds at the extremes of the frequency spectrum. The next most important are the remaining spectral features, Zero Crossing Rate, and a small number of outlier MFCC features. The MFCC distribution makes sense because they represent coefficients of a series of terms that get smaller and smaller; while the first few MFCCs explain a lot of the data, as you include more they represent coefficients of smaller less explanatory terms.

Method 2: Convolutional Neural Networks

A major trend of ML in the past decade has been the abandonment of higher-level hand-crafted features in favor of deep neural networks that do their own feature engineering over lower-level signals. A particularly powerful family of these is the Convolutional Neural Network (CNN), which rose to prominence in the Computer Vision field but quickly found a home in other domains too, like language and music.

How do they work? A CNN uses a bunch of smaller filters to find patterns in spatial or temporal data. As an example, if you were working with images, a single filter might represent something like a horizontal line sitting in a 3x3 pixel window. A “convolution” is an operation that essentially moves a filter window across an image, scanning for matches and returning a score representing how close each part of the image is to the filter.

An example convolution. A 3x3 filter is scanned over the blue array, with each green square representing an activation score (Source)

In practice, the situation is more complicated:

• you will use dozens or even hundreds of filters that each specialize in a different pattern
• the filter patterns will be determined automatically by a backpropagation learning procedure
• after the convolution operation you typically add additional layers, such as: batch normalizations, nonlinearities to give the neural net more expressive power, and pooling layers such a max pool which essentially turns “where in the input is this pattern the strongest?” into “does this pattern appear?”
• often people stack multiple convolutional + pooling layers on top of each other. In this case, you can think of the higher-level ones as looking for patterns of patterns.
• for classification problems, we add standard fully-connected layers atop the convolutional layers, and finally a softmax function that outputs a probability for each drum class.

Example of a CNN architecture (Source)

All examples above operate on two dimensional inputs, but the same concept could apply in one dimension to audio data – instead of scanning up and down, the filters are also one dimensional, and are scanned beginning to end. That said, I decided to use two dimensional Mel Spectrograms following the lead of research papers I came across. These are similar to the spectrograms we looked at above, but they use an alternate frequency axis that more closely reflects the way humans hear sounds.

After some testing I settled on an architecture with two convolutional layers and two linear layers. Here’s what my PyTorch implementation looks like:

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 256, kernel_size=(12, 4), stride=2)
        self.conv1_batch = nn.BatchNorm2d(256)
        self.conv2 = nn.Conv2d(256, 256, kernel_size=(4, 4), stride=(1, 2))
        self.conv2_batch = nn.BatchNorm2d(256)
        self.fc1 = nn.Linear(512, 128)
        self.fc2 = nn.Linear(128, len(DRUM_TYPES))

    # x is a Tensor object of size 1x128x259 containing normalized mel spectrogram data
    def forward(self, tensor, softmax=True):
      # First convolution
      tensor = F.leaky_relu(
          F.max_pool2d(
              self.conv1_batch(self.conv1(tensor)),
              stride=(4, 4)
          )
      )
      # Second convolution
      tensor = F.leaky_relu(
          F.max_pool2d(
              self.conv2_batch(self.conv2(tensor)),
              stride=(4, 8)
          )
      )

      # Now two fully connected layers, and a softmax
      assert np.prod(tensor.shape[1:]) == 512
      tensor = tensor.view(-1, 512)
      tensor = F.leaky_relu(self.fc1(tensor))
      tensor = F.dropout(tensor, training=self.training)
      tensor = self.fc2(tensor)

      return F.log_softmax(tensor, dim=-1) if softmax else tensor

One downside of applying deep neural networks to low-level data is the extra processing required. It helps that due the parallel nature of convolutions we can use a GPU instead of a CPU, which makes this feasible. Still, it took me 3-4 hours to train on an RTX 2080 Ti.

To run this yourself:

python drum_sound_classifier/models/train_cnn.py

train_cnn.py [-h] [--batch_size BATCH_SIZE]
                    [--val_batch_size VAL_BATCH_SIZE]
                    [--max_epochs MAX_EPOCHS]
                    [--early_stopping EARLY_STOPPING] [--lr LR]
                    [--momentum MOMENTUM] [--max_per_class MAX_PER_CLASS]
                    [--log_interval LOG_INTERVAL]
                    [--continue_name CONTINUE_NAME] [--eval]

From these training curves, you can see my model generally got better after each training epoch:

Validation accuracy during training. The x-axis is the number of epochs, y-axis is accuracy

Negative log likelihood loss on the validation set

There is a spike in accuracy of 83.62% on the 50th epoch, a slight improvement over method 1. Let’s keep going!

Method 3: CNN features + SVM

Using method 1 I showed that a handful of hand-crafted features and a Random Forest classifier does pretty well. With method 2, I showed that CNNs can discover their own hierarchical feature representations, which make them even more effective.

The CNN architecture offers a few points of interception where instead of running inputs through the entire pipeline, one can stop to measure their values coming out of a particular layer. In my own architecture, the second convolutional layer results in 512 values, which gets fed into a linear layer that outputs 128 values, which then gets fed through another linear layer down to 12 (the number of drum types). At each of these points, I can intercept the values and use them as a general-purpose embedding representation of a drum input.

What would happen if I fed the 512 or 128 sized embeddings into a Random Forest or other standard ML classifier? On one hand, the pure CNN solution might have an advantage in that the model we’re putting in front of the embeddings (the final linear layer followed by softmax), is the exact same model setup that was used during the training procedure to optimize the embeddings themselves. It would be like if some engineers built a specialized engine for a particular high-end car – what are the chances some other car runs it better?

But on the other hand, the Random Forest frontend has already shown promise with less complex data, and it has the benefit of being an ensemble method which should make it more robust.

To get size-128 embeddings from the CNN all I had to do was add a method similar to forward(), but that stops after the first fully-connected layer.

def embed(self, tensor):
    tensor = F.leaky_relu(
        F.max_pool2d(
            self.conv1_batch(self.conv1(tensor)),
            (4, 4)
        )
    )
    tensor = F.leaky_relu(
        F.max_pool2d(
            self.conv2_batch(self.conv2(tensor)),
            (4, 8)
        )
    )
    assert np.prod(tensor.shape[1:]) == 512
    tensor = tensor.view(-1, 512)
    return F.leaky_relu(self.fc1(tensor)).detach().numpy()[0]

And now, I can feed these into scikit-learn like any other dataset.

To do this yourself, assuming you ran all the above commands:

python drum_sound_classifier/models/train_sklearn.py --inputs cnn_embeddings --model random_forest --max_per_class 2000

Interestingly, with a Support Vector Machine classifier the accuracy is now 85% – an improvement over the pure CNN strategy!

Inspecting model misses

Like with any classification or regression problem, it’s wise to take a closer look at instances the model struggles with.

Filenames said snare, but my model said otherwise:

Filenames said hat, but my model said otherwise:

Filenames said tom, but my model said otherwise:

There are a few different things happening here. Remember how I decided to trust the filenames of drum kits as the ground truth labels? Well, some of them are obviously wrong (like any of the toms above). Not only can this degrade my model’s performance, but it obfuscates it so the true accuracy is not known. Dataset cleaning doesn’t make for great blog content so I will spare you, but at this point before exploring any additional model architectures, it would be worth manually annotating such discrepancies.

Some of the examples above do reveal straight-up mistakes by the model, but others are less clear cut and reasonable people could disagree about the true class (the third snare above – snare or clap?). Ultimately, the boundaries between drum types are blurry so 100% accuracy is not a reasonable goal.

Conclusion

It’s possible to turn a disorganized pile of drum sounds into a dataset well-suited for ML. And, there are a few standard approaches for training a classifier of drum types that achieve good accuracy.

One takeaway that can be applied to other domains is that if you have hand-crafted features readily available that you think will capture the contours of your problem, throwing standard classification tools at it can be the best bang for your buck. You typically require fewer data, it’s easier to implement, and you can avoid high training costs.

Another takeaway is that sometimes CNNs (and deep neural networks in general) are best viewed as feature extractors, not merely as end-to-end models. This is particularly true when the neural net is trained for the problem of interest, or a similar problem. I also tried putting a Random Forest in front of features derived from a convolutional autoencoder, but didn’t see the same gains.

So what’s next? How can we do better? One way might be to use a WaveNet, a special type of CNN that applies dilated convolutions to a raw audio signal. WaveNets have proven themselves very effective as generative models for applications like speech synthesis or even drum synthesis. But they can also be applied in a classification setting. This may be the subject of a future blog post.

Data augmentation would make for another quick win. The idea is to systematically modify training examples to increase the diversity of data seen by the model, and thus increase its robustness. In the audio ML field, one way to do this is random pitch shifts or time stretches.

Here are some uses I’ve gotten out of my drum classification model:

• corrected mislabeled sounds by writing new filenames
• projected drum embeddings down to two dimensions and plotted them, for an interesting new method of browsing through sounds
• outlier detection to come up with “safe” random drum racks

References

• O. Gillet. Transcription des signaux percussifs. Application à l’analyse de scènes musicales audiovisuelles. PhD thesis, 2007.
• P. Herrera, A. Yeterian, R. Yeterian and F. Gouyon. Automatic classification of drum sounds: a comparison of feature selection methods and classification techniques” Perfecto Herrera, Alexandre Yeterian, Fabien Gouyon. 2002.
• E. Pampalk, P. Herrera, M. Goto. Computational Models of Similarity for Drum Samples. 2008
• X. Zhang, Y. Gao, Y. Yu, W. Li. Music Artist Classification with WaveNet Classifier for Raw Waveform Audio Data. 2020.
• Neural Drum Machine - An interactive system for real time synthesis of drum sounds
• https://musicinformationretrieval.com/