Introduction

StaleyNET is named in memory of Layne Staley (1967 to 2002), vocalist of Alice in Chains.

Abstract

StaleyNET is a text-conditioned generative model of music audio, built and trained entirely from random initialization on two embedded NVIDIA Jetson devices. Every stage is implemented in Rust: feature extraction, a residual vector quantizer, a decoder-only Transformer with a hand-derived backward pass, the training loop and the sampler. All numerical kernels are compiled to GPU code with cuda-oxide, an experimental Rust-to-PTX compiler. No pretrained network, codec, vocoder or text encoder is used.

Sound is described by its log-mel spectrogram, and every 21 milliseconds of spectrogram is reduced to four whole numbers, each chosen from a list of 1 024 prototypes learned by k-means clustering. A Transformer of about seven million parameters learns to predict these numbers one after another, in the way a text model predicts words, given a text description of the music and, in an extended variant, the lyrics to be sung. The model is trained on the Song Describer Dataset (23 hours of Creative Commons music with human-written captions) and, for lyrics, on JamendoLyrics MultiLang (4.7 hours with line-level alignment).

The report documents the method, the verification of every kernel and of the full gradient, the performance of the GPU implementation, and the training results. It also decomposes the error of the audio pipeline. This analysis attributes the limited fidelity of the generated audio mainly to the tokenizer rather than to the language model or the vocoder. Finally, it describes a protocol for testing whether a model actually uses its lyrics, together with the results obtained so far.

Motivation

Recent text-to-music systems, such as MusicLM and MusicGen, combine large pretrained components: a neural audio codec, a text encoder and a language model of hundreds of millions to billions of parameters, trained on tens of thousands of hours of audio with datacenter accelerators. Their results are impressive, but the pretrained components conceal much of what the system has to learn, and the computational requirements place the complete pipeline out of reach of most individual researchers.

This project takes the opposite position. Each component is built from first principles, trained from scratch on public data, and made small enough to train on an edge device. The aim is not to match the quality of large systems but to make every part of the pipeline inspectable and measurable, and to identify which part limits the result.

Research questions

1. Can a small decoder-only Transformer over discrete log-mel tokens learn text-conditioned musical structure from roughly 23 hours of captioned audio?
2. What limits the fidelity of the generated audio, and how much of the limitation is attributable to the language model as opposed to the audio representation and the vocoder?
3. Can such a model learn to use lyrics when trained on a small corpus with line-level alignment, and how can this be tested objectively?
4. What do shared-memory tiling, register blocking, key/value caching and the delay pattern contribute to training and inference cost on an edge GPU?

Contributions

• A complete text-to-audio pipeline in Rust whose GPU kernels are written in Rust and compiled with cuda-oxide, including a strided, batched GEMM in three implementations and fused kernels for normalization, attention and cross-entropy.
• A testing approach in which each GPU kernel is checked against a CPU reference and the gradients of the whole network against numerical estimates.
• A lyrics-conditioning scheme that extends a trained model with a lyrics field without discarding what it has learned, and a paired evaluation protocol that measures whether the field is used.
• A decomposition of the reconstruction error of the audio pipeline into tokenizer and vocoder contributions.
• Measurements on two devices, the Jetson Orin NX 16 GB and the Jetson AGX Thor, with the commands needed to reproduce each of them.

How the results were verified

A system written from scratch has no established library to vouch for its arithmetic, so its correctness has to be demonstrated. Two kinds of test carry most of that weight. The GPU kernels are compared, one by one, with simple CPU implementations of the same formulas, on deliberately awkward matrix sizes where indexing mistakes tend to appear. The hand-written gradients of the whole network are compared with gradients estimated numerically, by nudging each parameter and observing how the error changes.

The figures in this report come from the project’s own commands, which are listed in Reproducibility: staleynet bench for speed, staleynet train for training curves, and staleynet reconstruct and staleynet evaluate for the analyses of the audio pipeline and of the lyrics. Where a statement rests on listening rather than on measurement, the text says so. The known weaknesses of the system are collected in Limitations.

Organization

Related Work places the project among existing systems. The method chapters then follow the data through the pipeline: the system overview, the audio representation, tokenization, the model, lyrics conditioning, training and generation. GPU Implementation describes the kernels. The evaluation part reports the experiments, the reproduction procedure and the limitations.

The figures are animated SVG and play in any modern browser without scripts.

Related Work

Language models of audio

AudioLM (Borsos et al., 2023) showed that audio can be generated by language-modeling discrete tokens, using a hierarchy of semantic and acoustic tokens. MusicLM (Agostinelli et al., 2023) conditioned this approach on text through a joint music-text embedding. MusicGen (Copet et al., 2023) simplified the hierarchy to a single Transformer over the codebooks of a neural codec, interleaved by a delay pattern so that all codebooks of a frame are predicted in one pass. StaleyNET adopts MusicGen’s delay pattern, top-k sampling and classifier-free guidance, but replaces the pretrained codec with a k-means quantizer over log-mel features and the pretrained text encoder with raw UTF-8 bytes, an approach studied for text by ByT5 (Xue et al., 2022).

Singing and lyrics

Jukebox (Dhariwal et al., 2020) generated singing conditioned on artist, genre and lyrics, using lyrics-to-audio attention trained on 1.2 million songs. Singing voice synthesis systems such as DiffSinger (Liu et al., 2022) instead condition on a musical score with phoneme durations, which gives precise control but requires annotated scores. StaleyNET follows the Jukebox direction at a much smaller scale: it conditions on lyrics as text and learns the alignment from data. The training signal for the alignment comes from JamendoLyrics (Stoller, Durand and Ewert, 2019; Durand, Stoller and Ewert, 2023), a lyrics-alignment benchmark whose line-level time stamps allow each training excerpt to be paired with the words sung in it.

Audio representations

Neural codecs (SoundStream, Zeghidour et al., 2021; EnCodec, Défossez et al., 2022; DAC, Kumar et al., 2023) learn an encoder, a residual vector quantizer and a decoder jointly, and reach high fidelity at a few kilobits per second. They require large training sets and substantial compute. StaleyNET retains the residual quantization principle but applies it, by k-means, directly to normalized log-mel frames, so that the tokenizer can be trained in minutes on an edge device. Waveforms are recovered with the Griffin-Lim algorithm (Griffin and Lim, 1984) in its accelerated form (Perraudin, Balazs and Søndergaard, 2013), rather than with a neural vocoder such as HiFi-GAN (Kong, Kim and Bae, 2020) or Vocos (Siuzdak, 2023). The consequences of both choices are quantified in Experiments.

Efficient Transformer computation

Key/value caching for autoregressive inference is standard practice (Pope et al., 2023). Register-blocked matrix multiplication on GPUs was analyzed by Volkov and Demmel (2008); FlashAttention (Dao et al., 2022) fuses attention to avoid materializing the score matrix. StaleyNET implements the first two in Rust through cuda-oxide and leaves fused attention to future work.

System Overview

This chapter describes the system as a whole before the following chapters examine its parts. It covers what the system does during training and during generation, the probability model it learns, and the way the code is organized.

Training and generation

During training, the system reads recordings together with their text descriptions (captions) and, where available, the words that are sung (lyrics). The text is kept as a plain sequence of bytes. The audio is turned into a sequence of tokens by the tokenizer of the previous chapters. A Transformer network then learns to predict each audio token from the text and from all the audio tokens that come before it. This is the same learning problem as predicting the next word of a sentence, applied to sound.

During generation, the trained network receives only the text. It produces the audio tokens one step at a time, each time feeding its own choice back in as context for the next. The finished token sequence is decoded into a spectrogram, and the spectrogram is converted into a waveform.

The probability model

The network learns a probability distribution over audio given text. To state it precisely, let \(x\) be the text and let the audio be a sequence of \(T\) steps, \(a_1, a_2, \dots, a_T\). Each step consists of the \(C = 4\) numbers chosen by the tokenizer’s four rounds, written \(a_t^0, a_t^1, a_t^2, a_t^3\), from coarse to fine. The network models the probability of a whole piece as a product of single predictions:

\[ p(\text{audio} \mid x) ;=; \prod_{t=1}^{T} \prod_{c=0}^{3} p!\left(a_t^{c} ,\middle|, x,; a_1, \dots, a_{t-1}\right). \]

In words, the probability of a piece is the probability of its first step given the text, times the probability of its second step given the text and the first, and so on. Training adjusts the network to make the recorded pieces as probable as possible. Generation follows the same product forward, drawing each step from the network’s predicted distribution.

The formula predicts the four numbers of a step together, each from the steps before it. The finer numbers of a step would benefit from knowing the coarser number of the same step, which this formula does not allow. The delay pattern recovers almost all of that benefit with a small shift of the sequence.

Organization of the code

The implementation is divided into six Rust crates, arranged so that each depends only on crates above it in the following table.

A separate program, studio/, serves a web page for generating and listening to samples and for recording new singing as training data.

Decisions that shape the implementation

Several decisions recur throughout the code, and each has a specific reason.

The GPU code is built from a small number of operations, the most important being one general matrix product. Spectrogram analysis, the tokenizer’s nearest-prototype search and every layer of the Transformer are expressed through it. Concentrating the arithmetic in one kernel means that one set of tests covers most of the numerical work, and that one optimization, such as the register-blocked kernel described in GPU Implementation, speeds up the whole system.

Each GPU array records its own dimensions, and the functions that launch GPU kernels check those dimensions first. A GPU kernel that receives an array of the wrong size does not fail; it silently reads or writes the wrong memory. Checking sizes on the CPU side turns such mistakes into immediate, readable errors.

The memory needed for a training step is allocated once, before training begins, and reused for every step. Allocating GPU memory is slow and can fragment the limited memory of an embedded device, and reusing buffers avoids both problems.

All randomness is drawn from seeded generators, and the one place where GPU threads would normally add into the same memory location in an unpredictable order (the gradient of the embedding table) is computed in a fixed order instead. As a result, a training run repeated on the same device produces exactly the same numbers, which makes differences between experiments attributable to the change being tested.

Finally, the trainer regularly saves not only the network’s weights but also the internal state of the optimizer. If a run is interrupted, it can continue exactly where it stopped instead of restarting its learning-rate schedule and losing the optimizer’s accumulated statistics.

Audio Representation

A model that generates music has to work with some description of sound. The most direct description, the waveform, lists the air pressure 24 000 times per second. Much of that detail is the exact phase of each vibration, which the ear barely notices, so a model trained on waveforms would spend most of its capacity on information that listeners do not hear. StaleyNET therefore describes sound with a log-mel spectrogram, a picture of how much energy is present in each frequency band during each short moment of time. This chapter explains how a recording becomes such a spectrogram and how a spectrogram is turned back into a waveform after generation.

From a recording to a clean signal

The training recordings arrive in different formats (MP3, FLAC, AAC, Ogg Vorbis and WAV), at different sample rates, and often in stereo. Before any analysis, every audio file in the dataset is decoded into a sequence of floating-point samples, and its stereo channels, where present, are averaged into a single mono channel.

The decoded audio is then resampled to a common rate of 24 000 samples per second, so that one second of music always corresponds to the same number of analysis frames. Resampling is a filtering problem: when the rate is reduced, any content above half the new rate (12 kHz here) must be removed first, or it folds back into the audible range as distortion, an effect known as aliasing. The resampler in staleynet-audio/src/resample.rs uses a windowed-sinc low-pass filter, following Julius O. Smith’s treatment of bandlimited interpolation at CCRMA. For a conversion between two fixed rates, such as 44.1 kHz to 24 kHz, the filter only ever needs to be evaluated at a small number of fractional positions (80 in this case), so its coefficients are computed once and each output sample is a single weighted sum. A unit test confirms the filter’s purpose: a 15 kHz tone, which lies above the new 12 kHz limit, is suppressed by more than 80 dB when 48 kHz audio is converted to 24 kHz.

Finally, the loudness of each recording is normalized to the same average level (an RMS of −20 dBFS). Commercial and amateur recordings differ widely in volume, and without this step the model would have to learn those differences as part of the music.

Measuring energy over time and frequency

A spectrogram is built by cutting the signal into short overlapping frames and measuring the strength of each frequency within each frame. The choices below fix how fine the picture is in time and in frequency.

The 128 bands are more than speech systems usually use (80 to 100), because music contains harmonic detail across a wider range than speech. The result is that one second of audio becomes a grid of about 94 frames by 128 bands, a far more compact description than 24 000 samples.

Computing the spectrogram on the GPU

The usual way to measure the frequencies in a frame is the fast Fourier transform (FFT). A Fourier transform can also be written as an ordinary matrix multiplication: each frequency is a fixed pattern of sines and cosines, and multiplying a frame by the matrix that holds these patterns measures how strongly each pattern is present. Written this way, the analysis of a whole recording becomes one large matrix product,

\[ X = F , W, \]

where each row of \(F\) is one frame of the recording (1 024 samples), each column of \(W\) is the windowed sine or cosine pattern of one frequency, and each row of the result \(X\) is the spectrum of one frame. A second matrix product then pools the 513 frequencies of each spectrum into the 128 mel bands.

The matrix form performs more arithmetic than an FFT, roughly 1 024 operations per frequency instead of about ten. It was nevertheless chosen because the GPU code for matrix products is the same code that runs the Transformer, which is tested thoroughly, so the spectrogram inherits that verification instead of requiring a second numerical kernel. In practice the cost is acceptable: on the Orin NX the GPU analyzes 60 seconds of audio in 47 ms, against 73 ms for a multithreaded CPU FFT. A test compares the GPU spectrogram with an FFT-based reference and finds them within 0.01 log units in every audible band.

Turning a spectrogram back into sound

After generation, the model’s output has to be heard, which means converting a spectrogram back into a waveform. This is harder than the forward direction, because the spectrogram records only how strong each frequency is, not the phase that says where in its cycle each vibration starts. The missing phase has to be estimated.

StaleyNET estimates it with the Griffin-Lim algorithm in its accelerated form (Griffin and Lim, 1984; Perraudin, Balazs and Søndergaard, 2013). The idea is to search for a waveform whose spectrogram matches the target. Starting from random phase, the algorithm converts the current estimate to a waveform and back, keeps the target strengths together with the newly obtained phase, and repeats. Each round makes the phase more consistent with the strengths, and 64 rounds are used. Before the search begins, the 128 mel bands are expanded back to the 513 linear frequencies with a regularized least-squares inverse of the pooling step.

Griffin-Lim needs no training, but its estimated phase leaves a characteristic metallic, slightly smeared quality. How much this matters, compared with the other losses in the pipeline, is measured in Experiments.

Why Griffin-Lim smears voices

A sustained harmonic, such as a sung vowel, needs a phase that advances smoothly from one frame to the next at the harmonic’s own frequency. Griffin-Lim never hears real audio; it only makes neighboring frames agree with each other. Its phases therefore drift, and where frames overlap they partly cancel, which the ear hears as a metallic, phasey smear. Voices suffer most, because their pitch and the vowels that make words recognizable live in exactly these harmonics, and sharp onsets such as drum hits spread into a faint pre-echo.

The cost is measurable. Even on a perfect spectrogram, Griffin-Lim adds an average error of 0.19 log units per band, and the audible damage is larger than that number suggests, because the measure compares magnitudes only and cannot see the phase errors that cause the smear. This is why the generated voices sound blurred even when the model predicts the right tokens.

The remedy is a neural vocoder: a network that learns, from real recordings, how to turn a spectrogram into a waveform, phase included (Kong et al., 2020; Siuzdak, 2023). StaleyNET’s vocoder follows the design of Vocos. For every frame it predicts the magnitude and the phase of each frequency from the mel spectrogram, and an inverse Fourier transform turns these into sound. It is trained on spectrograms decoded from StaleyNET’s own tokens, with the original recordings as targets, so that it learns not only the phase but also part of the detail that the tokens lose. The Neural Vocoder describes it in full.

Audio Tokenization

A language model predicts the next item in a sequence by choosing from a fixed list of possibilities, its vocabulary, in the same way that a text model chooses the next word. A spectrogram, however, consists of real numbers, which can take infinitely many values. Before a language model can learn from audio, each moment of the spectrogram has to be replaced by a choice from a finite list. This replacement is called tokenization, and the component that performs it is the tokenizer.

Large music systems use a pretrained neural codec for this purpose (for example SoundStream, EnCodec or DAC). This project builds every component from scratch, so it uses a simpler tokenizer that can be trained in minutes: k-means clustering applied in several rounds, known as residual vector quantization.

What is being tokenized

The spectrogram frames from the previous chapter are first standardized, band by band, using the average and spread of each band over the training set, so that loud and quiet bands contribute equally to the distances computed below. In the first tokenizer, two consecutive frames are joined into one unit, called a step. A step covers 21.3 ms of audio and is a list of 256 numbers (2 frames × 128 bands). A two-minute recording becomes about 5 600 steps.

The idea: describe each step by its nearest prototype

Suppose that a collection of 1 024 typical steps, called prototypes or centroids, has been found. Any new step can then be described by naming the prototype that most resembles it, which turns 256 real numbers into one whole number between 0 and 1 023. The prototypes are found with Lloyd’s k-means algorithm (Lloyd, 1982): starting from 1 024 randomly chosen training steps, it assigns every training step to its nearest prototype, moves each prototype to the average of the steps assigned to it, and repeats these two moves twenty times, by which point the prototypes barely change.

One list of 1 024 prototypes captures only the broad shape of a step. The error that remains, the difference between the step and its prototype, still contains useful detail. Residual quantization therefore repeats the procedure on these differences: a second list of 1 024 prototypes is trained on the errors left by the first, a third on the errors left by the second, and a fourth on those left by the third.

Written out, with \(x\) the step to be encoded, \(r_\ell\) the part of it that remains to be described before round \(\ell\), and \(e_\ell[k]\) prototype number \(k\) of round \(\ell\):

\[ r_0 = x, \qquad c_\ell = \arg\min_k \lVert r_\ell - e_\ell[k] \rVert^2, \qquad r_{\ell+1} = r_\ell - e_\ell[c_\ell]. \]

Each round chooses the prototype \(c_\ell\) closest to what remains and subtracts it.

A small example in two dimensions makes the procedure concrete. Suppose the first codebook holds the 4 prototypes (0, 0), (4, 0), (0, 4) and (4, 4), the second holds the 4 corrections (±0.8, ±0.8), and the point to encode is x = (3, 1). The squared distances from x to the first codebook are 10, 2, 18 and 10, so round 1 chooses (4, 0) and leaves the remainder (−1, 1). The correction closest to (−1, 1) is (−0.8, 0.8), so round 2 chooses it and leaves only (−0.2, 0.2). The point is now described by 2 small integers, and its reconstruction (3.2, 0.8) lies 0.28 from the original, against 1.41 after the first round alone.

Decoding reverses the process by adding up the four chosen prototypes:

\[ \hat{x} = e_0[c_0] + e_1[c_1] + e_2[c_2] + e_3[c_3]. \]

Every step of audio is thus described by four whole numbers, ordered from the coarse first choice to the finest correction. At 40 bits per 21.3 ms, this is about 1.9 kilobits per second, a very compact description; the price of that compactness is measured below. The coarse-to-fine order matters later, because the model predicts the coarse number before the corrections that depend on it (see the delay pattern).

Finding the nearest prototype quickly

Training and encoding repeatedly ask which of 1 024 prototypes is closest to a step. Comparing each step with each prototype directly is slow, but the squared distance can be rewritten as

\[ \lVert x - e_k\rVert^2 = \lVert x\rVert^2 - 2, x \cdot e_k + \lVert e_k\rVert^2 . \]

The first term is the same for every prototype and therefore does not affect which one is closest. The second term, the dot product of the step with each prototype, can be computed for 8 192 steps and all 1 024 prototypes at once as a single matrix product on the GPU, and the third term is computed once per prototype. A small GPU kernel then picks, for each step, the prototype with the smallest total. The averaging that moves the prototypes is done on the CPU. A prototype to which no step is assigned is restarted from a randomly chosen training step, so that none of the 1 024 entries is wasted.

Quantizing groups of bands

The measurements of Experiments showed that four rounds over a whole step lose too much detail: a step of 256 numbers is a great deal to describe with four choices from 1 024 prototypes. The current tokenizer makes two changes.

First, a step is a single frame, 128 numbers covering 10.7 ms, so that each code describes a shorter moment. Second, the 128 bands are divided into \(G = 4\) groups of 32 neighboring bands, from the lowest frequencies to the highest, and each group receives its own residual codebooks. A group of 32 numbers is far easier to describe with 1 024 prototypes than a whole step of 256: the same number of prototypes now covers a much smaller space, so each one lies closer to the vectors it stands for.

Each group is quantized in 4 rounds, so a step is described by \(4 \times 4 = 16\) codes. They are ordered round by round: first the coarse code of every group, then their first corrections, and so on. This keeps the coarse-to-fine order on which the delay pattern relies, since every code is predicted after the coarser codes it refines.

The change halved the error: on the 39 held-out songs it fell from 0.776 to 0.401 log units per band, from about 3.4 dB to 1.7 dB. A control with 16 rounds and no groups reached 0.427, so most of the gain comes from describing each step with more codes, and a smaller part from the grouping (Experiments lists the whole study). The price is a larger vocabulary, 16 codebooks instead of 4, and 16 output heads in the model.

One vocabulary for several datasets

The model can only learn from datasets that share its vocabulary. When a new dataset is added after the tokenizer has been trained, such as the lyrics corpus or newly recorded singing, it is encoded with the existing prototypes (staleynet tokenize --reuse) rather than with newly trained ones. The prototypes transfer well: on held-out songs of the lyrics corpus, which the tokenizer never saw, the remaining error of the first tokenizer after four rounds is 0.821 log units, compared with 0.776 on held-out songs of the original corpus; with the current tokenizer it is 0.423.

How much detail is lost

staleynet tokenize decodes every track again and reports the average difference between the original and the decoded spectrogram after one, two, three and four rounds. This number matters for the whole project, since no model trained on these tokens can produce sound more faithful than the tokenizer’s own reconstruction. The measurements, and a comparison with the losses of the other stages, are reported in Experiments.

The Model

The network at the center of StaleyNET is a Transformer of the kind used for language models such as GPT-2 (Radford et al., 2019). It reads one long sequence that begins with the text and continues with the audio, and at every position it predicts what comes next. This chapter describes how text and audio share that sequence, how the four numbers of each audio step are arranged, what happens inside the network, and how its gradients are computed.

Putting text and audio in one sequence

A Transformer reads a sequence of symbols, each drawn from one shared vocabulary. StaleyNET’s vocabulary combines three kinds of symbol. The first 256 entries are the possible values of a byte, so any text can be written directly as its UTF-8 bytes without a separate text tokenizer (a choice also studied for text models by Xue et al., 2022). Three further entries mark the structure of the sequence: BEGIN opens it, SEPARATOR ends the text, and PADDING fills unused text positions. The remaining 4 096 entries are the audio numbers: 1 024 for each of the tokenizer’s four rounds.

The text part of the sequence always has the same length, 128 positions: the BEGIN mark, 126 bytes of caption, and the SEPARATOR. Shorter captions are filled with PADDING and longer ones are cut. Keeping the length fixed means that the first audio step always appears at the same position, which simplifies both training and generation.

BEGIN  b₀ b₁ … b₁₂₅  SEPARATOR  audio step 1  audio step 2  …

Each audio position carries all four numbers of one step. The network’s input at such a position is the sum of the learned vectors (embeddings) of the four numbers, plus a learned vector for the position itself, which tells the network where in the sequence it is. Text positions use the same scheme with a single byte. Lyrics Conditioning describes how a second text field for lyrics is added to this layout.

The delay pattern

At every position the network makes four predictions at once, one for each tokenizer round. This creates a problem. The finer rounds are corrections to the coarse one; to predict a correction well, the network should already know which coarse prototype it is correcting. If all four numbers of a step were predicted at the same position, the network would have to guess the coarse number and its corrections simultaneously.

MusicGen (Copet et al., 2023) solves this by shifting each round by one more position than the previous one. The coarse number of step \(t\) is predicted at one position, the first correction of step \(t\) one position later, the second correction two positions later, and the finest one three positions later. By the time a correction is predicted, the coarse number it refines is already part of the input. The cost is only three extra positions per sequence, whereas placing the four rounds in four separate positions would make every sequence four times as long.

Inside the network

The network is a stack of six identical blocks. Each block lets every position gather information from the earlier positions (attention) and then transforms each position on its own (a small feed-forward network). Each of the two parts adds its result to its input rather than replacing it, which makes deep networks much easier to train, and each part first rescales its input to a standard size (RMS normalization; Zhang and Sennrich, 2019). Written compactly, a block computes

\[ h = x + \mathrm{Attention}(\mathrm{Norm}(x)), \qquad y = h + \mathrm{FeedForward}(\mathrm{Norm}(h)), \]

where \(x\) is the sequence entering the block and \(y\) the sequence leaving it. Attention is causal: a position may only look at itself and at earlier positions, since during generation the later positions do not exist yet.

The network has 7.1 million adjustable parameters. Almost a third of them sit in two places: the table of input embeddings (1.1 million) and the four output layers that score the 1 024 possibilities of each round (1.05 million).

Before training, the parameters are set to small random values, drawn with a standard deviation of 0.02 as in GPT-2. The two layers in each block that add their output into the running sum are scaled down further, by \(1/\sqrt{12}\) for six blocks, so that the sum does not grow larger with every block it passes through.

Computing gradients by hand

Training requires, for every parameter, the gradient: how much the error would change if the parameter changed slightly. Machine-learning frameworks derive gradients automatically. No such framework exists for Rust on the GPU, so every layer of StaleyNET implements its gradient explicitly, derived on paper and documented next to the code.

Attention illustrates what this involves. In the forward direction, the network compares each position’s query \(Q\) with the keys \(K\) of earlier positions, turns the comparison scores into weights \(P\) that sum to one, and uses the weights to average the values \(V\):

\[ P = \mathrm{softmax}!\left(\frac{QK^\top}{\sqrt{d}} + M\right), \qquad O = PV, \]

where \(d = 32\) is the width of one head and \(M\) blocks later positions. Given the gradient of the error with respect to the output \(O\), written \(\partial O\), the chain rule yields the gradients of the inputs:

\[ \partial V = P^\top \partial O, \qquad \partial P = \partial O, V^\top, \qquad \partial S = P \odot \bigl(\partial P - \mathrm{rowsum}(\partial P \odot P)\bigr), \]

\[ \partial Q = \tfrac{1}{\sqrt{d}}, \partial S, K, \qquad \partial K = \tfrac{1}{\sqrt{d}}, \partial S^\top Q, \]

where \(\partial S\) is the gradient with respect to the scores before the softmax and \(\odot\) denotes element-wise multiplication. Each line is a matrix product, so the gradient runs on the same GPU kernel as everything else.

Hand-written gradients are easy to get subtly wrong, so they are checked numerically. The test in crates/staleynet-model/tests/model_is_trainable.rs nudges individual parameters up and down by a small amount, measures how the error actually changes, and requires the result to agree with the hand-computed gradient within 5 %. The test covers every parameter tensor of the network.

Lyrics Conditioning

A caption such as “slow grunge song with a raw male voice” tells the model what a piece should sound like, but not which words should be sung. To let a user supply the words, the model needs a second text input for lyrics, and it needs training examples in which it can observe which words are sung during which stretch of audio. This chapter describes how the lyrics enter the sequence, how an existing model is extended to read them, how training examples are paired with the right lines, how long lyrics are handled during generation, and how one can test whether a trained model actually uses them.

Where the lyrics go in the sequence

The lyrics occupy 256 additional byte positions between the caption and the SEPARATOR mark. A sequence of a model with lyrics therefore has this layout:

part           BEGIN  caption (126 bytes)   lyrics (256 bytes)    SEPARATOR  audio …
position       0      1 to 126              127 to 382            383        384 …
position row   0      1 to 126              768 to 1023           127        128 …

The last line needs explanation. The network learns one position vector for each row of a position table, and the vector added at a position is taken from the row shown. The caption, the separator and the audio use exactly the rows they used before lyrics were introduced. The lyrics use 256 new rows placed after the 768 that already existed.

Extending a model that was trained without lyrics

The reason for this arrangement is that it lets an already trained model acquire a lyrics input without losing what it has learned. If the lyrics had been inserted using the existing rows, every audio position would have moved to a different row, and the model’s knowledge of where it is in a piece would no longer apply. With the arrangement above, extending a model amounts to appending 256 new, randomly initialized rows to its position table. All other parameters keep their trained values, and training only has to teach the model what the new positions mean.

The function checkpoint::load_weights performs this extension when a new run is started from an older checkpoint: it copies every parameter by name and, where the new model’s table is longer, copies the old rows into its beginning. Two tests guard the scheme. One checks that the copied parameters are identical to the originals; the other confirms that generation with the incremental decoder gives the same predictions as the full network for a model with lyrics, which would fail if the position rows were mapped inconsistently.

Pairing audio with the words sung in it

Each training example covers a window of about 13.6 seconds taken from a song. For the model to learn which sounds correspond to which words, the lyrics given with that window must be the lines actually sung during it. Two sources provide this information in different ways.

JamendoLyrics MultiLang (Durand, Stoller and Ewert, 2023) is a set of 78 Creative Commons songs in English, French, German and Spanish, created for research on aligning lyrics with audio. For every line of lyrics it records the time at which the line starts and ends. A training window is given every line whose time span overlaps it.

Recordings made with the StaleyNET studio come with the text that was sung but without timings. For these, the lines are assumed to be spread evenly over the recording: in a 20-second take of four lines, the second line is taken to occupy seconds 5 to 10. This assumption is reasonable for short takes of continuous singing, which is what the recording page asks the singer to provide.

As with captions, the lyrics are occasionally removed during training (in 10 % of examples, together with the caption). This teaches the model to generate without any text as well, which the guidance technique described in Generation requires.

Mixing small and large datasets

The lyrics corpus is much smaller than the caption corpus (4.7 hours against 23 hours). If training examples were drawn from all songs with equal probability, examples with lyrics would be rare and the model would have little opportunity to learn from them. Each dataset is therefore given a weight: a dataset with weight 8, for example, makes each of its songs eight times as likely to be drawn as a song of the main corpus. The weights used in each experiment are reported with its results.

Generating with long lyrics

A song’s lyrics are usually far longer than the 256 bytes that the model reads at once. During generation the lyrics are therefore divided into consecutive windows, each holding as many complete lines as fit, and the requested length of audio is divided equally among them. When the generator reaches the end of one window’s share, it replaces the lyrics in its input with the next window, keeps the most recent audio as context, and continues, so the voice carries on into the next lines instead of starting again.

Testing whether the model uses the lyrics

A model’s training loss can improve on the lyrics corpus without the model reading the lyrics at all, because it can simply learn what the songs in that corpus sound like. Whether the words influence the predictions has to be tested directly. The test, run by staleynet evaluate, takes the same windows of held-out songs and scores the model’s predictions three times:

1. with the correct lyrics for each window;
2. with no lyrics;
3. with lyrics taken from a different song (always a sung passage, so that this input is never empty).

Because all three runs use exactly the same audio windows, the scores can be compared window by window, which removes most of the variation between windows. If the model ignores the lyrics, the three scores are equal. If it uses them, the correct lyrics should help it predict the audio, and the mismatched lyrics should not, so the error with mismatched lyrics should be higher. The reported quantity is the average of that difference,

\[ \Delta = \text{average over batches of } \bigl(\text{error with mismatched lyrics} - \text{error with correct lyrics}\bigr), \]

measured in nats per predicted number, together with its standard error. The model is said to use the lyrics only when \(\Delta\) is larger than twice its standard error, so that a difference produced by chance is not mistaken for an effect. Comparing the second and third runs adds a finer distinction: a model that reacts to the presence of any text, but not to which words it is given, scores the same with correct and with mismatched lyrics.

Training

Training adjusts the network’s parameters so that the recorded music becomes as probable as possible under the model. This chapter describes the quantity being minimized, how training examples are assembled, the data and how it is divided between training and evaluation, the optimizer, and the way progress is saved.

What is minimized

At every position of a training sequence, the network predicts a probability for each of the 1 024 possible values of each tokenizer round at the next position. The prediction is scored by the cross-entropy, which is the negative logarithm of the probability the network gave to the value that actually occurs:

\[ \text{error at one position} = -\log p(\text{actual value}). \]

A confident correct prediction (probability close to one) has an error close to zero, and a confident wrong one has a large error. The training loss is the average of this error over all positions, computed separately for each tokenizer round and then averaged over the four rounds so that each round counts equally. Errors are measured in nats, the unit of the natural logarithm. For reference, a network that assigns equal probability to all 1 024 values has an error of \(\ln 1024 \approx 6.93\) nats, which is the level any useful model must improve on.

The GPU kernel that computes the error also computes its gradient in the same pass, because for this loss the gradient has a simple closed form: the predicted probabilities minus one at the actual value.

How training examples are assembled

Each training step processes a batch of sequences, for example eight sequences of 768 positions, of which 128 hold the text and 640 hold about 13.6 seconds of audio. Each sequence is assembled as follows.

A song is drawn at random, with each dataset’s songs weighted as described in Lyrics Conditioning. A window of the song’s tokens is then chosen. In 20 % of cases the window begins at the start of the song, so that the model also learns how pieces open, and otherwise it begins at a random point. One of the song’s captions is chosen, together with the lyrics sung during the window. In 10 % of cases both are replaced by empty text, which teaches the model to generate without a prompt as well. Finally the inputs and the expected next values are written into a buffer that is reused for every batch.

The random choices for training step number \(s\) are derived from the run’s seed and \(s\) alone. A run that is stopped and resumed therefore sees exactly the batches that it would have seen without the interruption.

The data and its division

The main dataset is the Song Describer Dataset (Manco et al., 2023): 706 Creative Commons songs from the Jamendo platform, each a two-minute excerpt, described by 1 106 captions written by listeners. It spans pop, rock, electronic, folk, classical and other styles. Of the 706 songs, 298 have a caption that mentions singing or a voice (about 9.9 hours), and 131 are described as instrumental.

To measure whether the model generalizes, about 10 % of the songs are held out and never used for training. The division is made by artist rather than by song, so that no artist appears on both sides. Otherwise the model could appear to generalize while merely recognizing a familiar voice or production style. The assignment of an artist to one side is determined by a fixed hash of the artist’s identifier, which keeps it identical on every machine; it leaves 667 songs for training and 39 for evaluation. The lyrics corpus and the studio recordings are divided in the same way.

Some experiments train on a part of the main dataset selected by the words in the captions. Words are matched whole, so that “singer” does not match “singers”.

• The vocal subset keeps songs whose captions mention a voice (for example “vocal”, “singer”, “sung”, “choir” or “rap”) and are not described as instrumental. It contains 281 training songs (1.57 million steps) and 19 evaluation songs.
• The rock subset keeps songs whose captions mention rock or a related style (grunge, alternative, indie, punk, metal, garage, psychedelic, stoner, shoegaze, post-rock, progressive) or its typical sound (distorted or electric guitar, riffs). It contains 188 training songs, or 6.2 hours. Some of these styles are barely present in the dataset, with two songs described as grunge, one as psychedelic and none as soft rock, so the model cannot learn them as separate styles from this data.

Overfitting, and how often each song is seen

A model overfits when it learns the particular songs it trains on instead of what they have in common. The loss on the training songs keeps falling, while the loss on held-out songs, which the model never sees, stops falling. The difference between the two, the generalization gap, measures how much of what the model has learned applies only to its training songs. A model that has overfitted generates by recalling: it plays back stretches of a training song, and it can return the same song whatever the prompt asks for.

How much a model can memorize depends on how often it sees each song. Training draws every sequence from a song chosen at random, and a song \(i\) with sampling weight \(w_i\) is chosen with probability \(w_i / W\), where \(W = \sum_j w_j\) is the total weight of all training songs. A run of \(S\) steps with \(B\) sequences per step therefore draws song \(i\)

\[ E_i = S \cdot B \cdot \frac{w_i}{W} \]

times on average. Each draw is a randomly placed window of the song, so many draws of one song show the model most of its content many times over.

Weighting a small dataset more heavily makes it count for more in training, and it also multiplies how often each of its songs is seen. The table gives \(E_i\) for the three configurations used for the 42.9-million-parameter model of Experiments. The training songs are those of the Song Describer Dataset (667), the JamendoLyrics songs (73), the author’s two recorded takes, and, in the later runs, the 29 314 rock tracks of the Free Music Archive (Defferrard et al., 2017).

In pq-large, JamendoLyrics was weighted 4-fold and the two takes 10-fold, so each lyric-bearing song was drawn about 650 times and each take more than 1 600 times. These are also the only songs with lyrics, so whenever the model was given lyrics during training, the answer was one of 75 songs it had heard hundreds of times. In informal listening, the model returned the same song for different prompts and repeated a line of lyrics from its training data.

pq-rock-balanced keeps every song within about 16 draws. It does so by adding the 29 314 rock tracks at half the weight of a Song Describer song, giving JamendoLyrics the same half weight, and leaving out the two takes, since no weighting of two recordings would keep them from being memorized. The Song Describer songs keep twice the weight of the others because their captions are full sentences written by listeners, while the rock tracks carry only genre tags such as “rock, punk, garage”. The price is that lyrics become rare: 73 of 30 054 songs carry them, so this model will learn little about words. Singing recognizable words requires many more songs paired with their lyrics.

Two further measures limit the damage of overfitting. A random 10 % of the sequences are trained without their caption and lyrics, which teaches the model the unconditional distribution that guidance compares against (see Generation). And a run keeps, in best/, the checkpoint with the lowest held-out loss, so a model that begins to overfit late in a run is not the one that is used.

The optimizer

The parameters are updated with AdamW (Loshchilov and Hutter, 2019), which adapts the step size of each parameter to the recent size of its gradients and shrinks the weights slightly at every step (weight decay, here 0.1), a form of regularization. Weight decay is not applied to the normalization gains, the biases or the embedding tables, whose values carry meaning that shrinking would distort.

The learning rate first rises linearly from zero over a few hundred steps, which prevents the large, poorly directed updates of the first steps from damaging the randomly initialized network. It then falls along a cosine curve to one tenth of its peak by the end of the run, so that the final updates are small refinements. If the combined size of all gradients exceeds 1.0, they are scaled down to that size (Pascanu, Mikolov and Bengio, 2013), which protects training from rare, very large gradients. All arithmetic is carried out in 32-bit floating point. The peak learning rate and the schedule of each run are given with its results.

Saving progress

Training on an edge device runs for hours and can be interrupted. The trainer keeps two checkpoints in each run’s directory.

The latest/ checkpoint holds the weights together with the optimizer’s internal statistics and the step count, and is rewritten every 250 steps by default. It allows a stopped run to resume with --resume exactly where it was, without restarting the learning-rate schedule or losing the optimizer’s accumulated statistics; at most 250 steps of work can be lost. Its files are first written under temporary names and then renamed, so an interruption during saving cannot leave a half-written file in place.

The best/ checkpoint holds the weights at the point where the error on the held-out songs was lowest, which is usually the version worth listening to. When a run resumes, the trainer reads the best held-out error recorded so far from the run’s log, so that best/ is only replaced by a genuinely better model.

A new run can also start from the weights of an earlier one with --fine-tune-from. It then begins with a fresh optimizer and schedule, and the earlier model may lack the lyrics input, which is added as described in Lyrics Conditioning. Every logged event is appended to metrics.jsonl in the run directory.

A test that the training loop can learn at all

Before any training on real data, a test checks the training machinery as a whole. A small network is trained repeatedly on one fixed batch, which it should be able to memorize. Its error must fall from the level of random guessing (\(\ln 7 \approx 1.95\) nats for seven possible values) to below 0.1 within 300 steps. A wrong sign in a gradient, a missing term or a broken optimizer state prevents this, even when each individual gradient passes the numerical check described in The Model.

Generation

To generate music, the trained network is given the text and asked to continue the sequence. It predicts a probability for every possible value of the next audio step, one value is drawn from that prediction, the drawn step is appended to the sequence, and the process repeats until the requested length is reached. This chapter describes how this loop is made fast, how the values are drawn, how the output is pushed to follow the text more closely, and how the result becomes sound.

Reusing earlier work: the key/value cache

In attention, each new position compares itself with every earlier position through two vectors that those positions computed, called keys and values. These vectors do not change once computed, because a position never looks forward. A naive generator would nonetheless recompute them for the whole sequence at every step, which repeats the same work hundreds of times.

The generator therefore keeps the keys and values of all processed positions in a cache, one per layer (Pope et al., 2023). Each new step computes only its own query, key and value, appends its key and value to the cache, and attends over the cache. The cache has a fixed capacity, and positions not yet filled are excluded by the same mask that prevents looking ahead. Because the sizes of all arrays stay constant, generation allocates no memory after it starts.

A test confirms that this shortcut changes nothing: it feeds a sequence to the cached generator one position at a time and requires its predictions to match those of the full network within \(10^{-4}\) at every position, for models with and without lyrics.

When the cache is full, the generator rebuilds it from the text and the most recent half of the audio, and continues. Pieces longer than the network’s context can be generated in this way, with the most recent audio always in view.

Drawing values: top-k sampling

Always choosing the single most probable value tends to produce repetitive, lifeless output, while drawing from the full distribution occasionally picks very unlikely values that sound like errors. Top-k sampling (Fan, Lewis and Dauphin, 2018) takes a middle course: for each tokenizer round it keeps the \(k\) most probable values (250 by default, as in MusicGen), discards the rest, and draws among the survivors in proportion to their probabilities. A temperature setting adjusts how strongly the draw favors the most probable values; a temperature of zero always picks the most probable one.

Following the text more closely: classifier-free guidance

Because 10 % of the training examples had their text removed, the same network can predict audio both with and without the text. The difference between the two predictions indicates what the text contributes. Classifier-free guidance (Ho and Salimans, 2022) exaggerates that contribution. At each step the generator runs two streams, one with the text and one without, each with its own cache, and combines their scores as

\[ \text{guided} = \text{without text} + s \times (\text{with text} - \text{without text}), \]

where the scores are the network’s raw outputs before they are turned into probabilities and \(s\) is the guidance scale. With \(s = 1\) the result is the ordinary prediction with text; values above one, such as the default of 3, push the output further in the direction the text indicates. The cost is twice the computation per step and somewhat less variety between samples.

From numbers to sound

When all steps have been generated, the delay pattern is undone so that the four numbers of each step are back together. Each step is decoded by adding up its four prototypes, the standardization applied before tokenization is reversed, and the resulting spectrogram is converted into a waveform, either with 64 rounds of Griffin-Lim or, given --vocoder, with a trained neural vocoder. staleynet sing writes the waveform and a short report of the time spent in each stage, including the real-time factor: the time taken divided by the duration of the audio produced, so that a value below one means faster than real time.

The Neural Vocoder

A spectrogram says how strong each frequency is in each frame, but not where in its cycle each vibration is: the phase. Turning a spectrogram back into sound means supplying that phase. Griffin-Lim estimates it by making neighboring frames agree with each other, and the result has the metallic smear described in Audio Representation. A neural vocoder learns the phase instead, from real recordings. This chapter describes StaleyNET’s vocoder, how it is trained, and how its gradients are computed by hand like everything else in the project.

Staying at the frame rate

Early neural vocoders produced audio one sample at a time, and later ones grew the frame rate to the sample rate through a stack of upsampling layers. Both spend most of their computation at 24 000 steps per second. Vocos (Siuzdak, 2023) takes a shorter route, and StaleyNET follows it. The network runs at the frame rate, about 94 frames per second, and for each frame it predicts what the spectrogram leaves out: the magnitude and the phase of every one of the 513 frequencies. The last step, from a complex spectrum to samples, is the inverse short-time Fourier transform, which is exact and needs no learning.

This choice also fits the rest of StaleyNET. The analysis that turns audio into a spectrogram is already written as a matrix product (GPU Implementation), and the inverse transform is the same product read in the other direction.

The network

The input is the log-mel spectrogram that the tokens decode to: one row of 128 band energies per frame. Each band is first standardized with the tokenizer’s own statistics, so that every band enters on a similar scale and a spectrogram decoded from tokens looks like one computed from a recording.

A convolution over time then looks at each frame together with its three neighbors on either side and lifts those 7 × 128 values to a vector of 512 numbers. On the GPU this is one matrix product: each frame’s neighborhood is copied into a row of 896 numbers, and the rows are multiplied by the weights.

Eight ConvNeXt blocks (Liu et al., 2022) follow. Each block does two things in turn. A depthwise convolution filters each of the 512 channels over time with its own seven taps, mixing neighboring frames but not channels. Then, after RMS normalization, a small feed-forward network mixes the channels within each frame: it expands the 512 numbers to 1 536, applies the GELU activation, and projects back to 512. The result is scaled channel by channel and added to the block’s input:

\[ h \leftarrow h + \gamma \odot \bigl(\operatorname{GELU}(\operatorname{Norm}(\operatorname{DW}(h)),W_1 + b_1),W_2 + b_2\bigr). \]

The scales \(\gamma\) start at 1/8, so each block begins as a small correction to what it receives; this keeps an eight-block stack easy to train.

A final normalization and a linear head turn each frame’s 512 numbers into 1 026: a log-magnitude \(a_k\) and a phase \(\phi_k\) for each frequency \(k\). The complex spectrum is

\[ X_k = e^{\min(a_k,\ 4.6)},(\cos\phi_k + i\sin\phi_k). \]

Predicting the logarithm of the magnitude keeps loud and quiet frequencies on the same footing, and the cap of \(e^{4.6} \approx 100\) stops an untrained network from producing numbers that overflow. The phase needs no constraint: the cosine and sine wrap it around the circle by themselves. The whole network has 13.6 million parameters.

From spectrum to samples

Each frame’s spectrum becomes 1 024 samples through the inverse Fourier transform of a real signal. Only the 513 non-negative frequencies are stored, since the negative ones mirror them, so every frequency except the first and the last counts twice. For all frames at once this is two matrix products, one with the cosines and one with the sines, each of size \((F \times 513) \cdot (513 \times 1,024)\) for \(F\) frames, with the synthesis window folded into the constant matrices.

Overlap-add then lays the frames down 256 samples apart, so that every sample is covered by four frames, sums them, and divides each sample by the sum of the squared window values that covered it. That division is what makes analysis followed by synthesis return the original signal exactly.

Learning from recordings

The vocoder is trained to make its waveform sound like the recording that the spectrogram came from. Comparing waveforms sample by sample would punish a correct sound with a slightly different phase, so the comparison is made on spectra instead, at three resolutions:

Each term is the mean absolute difference between the logarithms of the estimate’s and the recording’s energies. Using several frame lengths keeps the network from satisfying one view of the sound at the expense of another (Yamamoto et al., 2020). Vocos also trains against learned discriminators, networks that try to tell generated audio from real audio. StaleyNET starts with the spectral terms alone, which are differentiable in closed form and stable, and adds discriminators as a later stage.

What it learns from

The language model produces tokens, so the vocoder trains on spectrograms decoded from tokens, not on spectrograms of the recordings. The recording supplies only the target waveform. The vocoder thus learns two things at once: the phase that the spectrogram never stored, and part of the detail that the tokenizer lost.

staleynet vocoder prepare stores every clip’s samples beside its tokens, resampled and level-normalized exactly as they were for analysis. A training example is then a window of 64 consecutive frames decoded from a clip’s tokens, together with the samples those frames cover. Frame \(f\) is centered on sample \(256 f\), so 64 frames cover \(63 \cdot 256 + 1,024 = 17,152\) samples, about 0.71 seconds. A batch holds 16 windows from random clips of the FMA rock corpus, 273 hours in all.

Training uses AdamW with \(\beta_1 = 0.8\) and \(\beta_2 = 0.9\), the values Vocos uses, a learning rate that warms up to \(5 \cdot 10^{-4}\) over 1 000 steps and then follows a cosine down to \(5 \cdot 10^{-5}\) at 200 000 steps, and gradients clipped to norm 1. Every step draws its windows from a random generator seeded by the step number, so a run that stops and resumes sees the same batches it would have seen without stopping. Checkpoints keep the optimizer’s moments for the same reason.

Gradients by hand

As with the language model, there is no automatic differentiation; each stage computes its own gradient. One observation saves most of the work. Cutting a signal into overlapping frames and overlap-adding frames back into a signal are adjoint linear maps: the gradient of one is the other. A single pair of GPU kernels therefore serves four purposes: the inverse transform, its gradient, the framing inside every loss, and that framing’s gradient.

The other pieces are short. The gradient of an absolute difference is the sign of the difference. The gradient of a logarithm with a floor is zero below the floor. The gradient of \(e^{a}(\cos\phi + i\sin\phi)\) follows from the chain rule, and a depthwise convolution’s gradient with respect to its input is the same convolution with the taps reversed.

A test checks the whole chain. It perturbs parameters of every kind, one at a time, measures how the loss changes, and compares that with the computed gradient. Because absolute differences have kinks, a single perturbation size can mislead, so the test tries three sizes and requires agreement to within 5 % at one of them; a wrong derivation agrees at none.

Using it

staleynet vocoder prepare --cache data/cache-fma-rock
staleynet vocoder train   --cache data/cache-fma-rock --run runs/vocoder-rock --config configs/vocoder-rock-thor.json
staleynet vocoder render  --cache data/cache-fma-rock --run runs/vocoder-rock --out samples/vocoder
staleynet sing --run runs/pq-large --vocoder runs/vocoder-rock --prompt "gritty rock vocal, distorted guitars"

A training run writes its metrics, a best/ and a latest/ checkpoint, and, every 5 000 steps, the same held-out clips rendered through the vocoder, next to the recordings and their Griffin-Lim rendering. In the studio, trained vocoders appear in a picker beside Griffin-Lim.

GPU Implementation

All of StaleyNET’s numerical work runs on the GPU, and all of its GPU code is written in Rust. This is made possible by cuda-oxide, an experimental compiler backend from NVIDIA that translates Rust functions marked #[kernel] into PTX, the GPU’s instruction format. Host code (running on the CPU) and kernels (running on the GPU) live in the same crate, and cuda-oxide generates a typed Rust function for launching each kernel. The roughly twenty kernels of the project live in crates/staleynet-cuda/src/kernels/, one file per topic: matrix products, attention, the loss, the optimizer, element-wise operations, normalization, embeddings, and signal processing. A kernel that needs more than a few numbers receives them as one small structure, such as Product for a matrix product or AdamWStep for one optimizer step, so that each value arrives with its name and its documentation. This chapter explains how the kernels divide their work among GPU threads, describes the matrix product in detail, and explains how the code is kept safe and correct.

How kernels divide their work

A GPU runs thousands of threads at once, grouped into blocks, and a kernel decides which part of the problem each thread handles. StaleyNET’s kernels follow four patterns.

In the simplest pattern, used for operations such as adding two arrays, each thread computes one element of the output. The last block usually has more threads than elements remain; cuda-oxide’s DisjointSlice type returns no element for such surplus threads, so they cannot write outside the array.

Operations that need a whole row of a matrix at once, such as normalization, the softmax inside attention and the cross-entropy, assign one warp (a group of 32 threads that execute together) to each row. The 32 threads each take every 32nd element of the row and then combine their partial sums or maxima by exchanging values directly between threads, without going through memory.

Matrix products assign one block of threads to each rectangular tile of the output, as described below. Finally, the sums of squares needed to measure the size of all gradients use a fixed number of blocks that sweep over the data, each producing a partial sum that the CPU adds up.

The matrix product

The matrix product is by far the most important kernel: the spectrogram, the tokenizer’s distances, every projection in the network and both products inside attention go through it. A single kernel handles all of these cases. It computes, for a batch of matrix pairs,

\[ C \leftarrow \alpha \cdot A \cdot B \quad (\text{optionally } + ; C), \]

where either input can be read as stored or as transposed. The kernel does not need a separate flag for this: each input is described by the distance in memory between neighboring rows and between neighboring columns (its strides), and a transposed matrix is the same memory read with the two distances exchanged. The batch allows the attention of all heads and all sequences to be computed in one launch, and reading an input transposed avoids creating transposed copies in memory.

Three versions of this kernel exist, and comparing them shows where the time goes. The naive version gives each output element its own thread, which reads a full row of \(A\) and a full column of \(B\) from the GPU’s main memory. Main memory is slow compared with the arithmetic units, and every value is read many times by different threads, so this version spends most of its time waiting for data.

The tiled version lets a block of 16 × 16 threads cooperate. The block copies a 16 × 16 tile of each input into shared memory, a small, fast memory that the threads of a block share, and computes with the tiles before moving on to the next pair. Each value is now fetched from main memory once per block instead of once per thread, sixteen times less traffic. On the Orin NX this raises throughput from about 35 to about 250 billion operations per second.

The register-blocked version (Volkov and Demmel, 2008) goes one step further. Shared memory is fast, but reading it still costs time, and in the tiled version each multiplication needs two shared-memory reads. In the register-blocked version each thread computes a 4 × 4 group of outputs instead of one. It reads four values of each input and combines them in all sixteen pairings, keeping the running sums in its registers, the fastest storage on the chip. Four times as much arithmetic is done for each value read. The threads are also arranged so that neighboring threads read neighboring memory words, which lets the hardware serve a whole group of threads in one access. On the AGX Thor this version reduced the time of a training step from 238 ms to 147 ms, and it is now the default.

Making the embedding gradient repeatable

During the backward pass, the gradient of the embedding table is obtained by adding the gradient at every position into the table row of the token found there. Many positions share a token, so many additions target the same row. The usual GPU solution is atomic addition, which lets threads add into the same location safely but in an unpredictable order. Because floating-point addition is not exactly associative, the result then differs slightly from run to run.

StaleyNET avoids this by first building, on the CPU, an index that lists for every token the positions where it occurs (a counting sort, stored in compressed sparse row form). One GPU thread per token and column then adds up exactly those positions, always in the same order. The result is identical on every run, which keeps experiments repeatable.

Safety

Launching a kernel is inherently risky: if the launch dimensions do not match what the kernel expects, threads read or write memory they should not, and nothing reports the error. cuda-oxide therefore marks such launches unsafe in Rust’s sense, meaning the programmer must guarantee their validity. In StaleyNET every launch is wrapped in a small, safe function in staleynet-cuda/src/ops/. The function first checks each size the kernel depends on and only then launches it, with a comment stating why the launch is valid. Inside the kernels, writes go through two small helpers in kernels/device.rs that check the index against the length of the output before writing, so the kernels themselves contain no unchecked memory access. The model and its layers call only the safe wrappers.

Testing the kernels

The file crates/staleynet-cuda/tests/kernels_match_reference.rs runs every kernel on random data and compares the result with a straightforward CPU implementation of the same formula. The test sizes are chosen to be awkward, for example 17 × 33 × 9 and 5 × 129 × 31, so that tiles and warps are only partly filled, which is where indexing errors usually hide. The matrix product is tested in every combination of transposed inputs, with and without batching and accumulation, for all three versions. Results must agree within a relative difference of \(10^{-5}\), or \(10^{-4}\) for reductions whose additions happen in a different order.

Practical notes

Several details of cuda-oxide may help others using it on Jetson devices. The libm::sqrtf function does not compile for the GPU on ARM hosts because its implementation is written in inline assembly, whereas f32::sqrt works. A CUDA context belongs to one CPU thread at a time, so multi-threaded test runners must attach it in each test (Gpu::bind_to_current_thread). A crate’s tests are run through the backend with cargo oxide test -- -p <crate>. The project follows cuda-oxide’s main branch, and the revision it was last tested with is recorded in Cargo.lock (see the README).

Experiments

Experimental setup

Hardware. Two devices were used, often concurrently with other workloads:

• NVIDIA Jetson Orin NX 16 GB (Super developer kit), CUDA 12.6;
• NVIDIA Jetson AGX Thor with 128 GB of memory shared by its CPU and GPU. In the early runs other GPU services shared the device; the larger models were trained with most of the device to themselves.

Unless stated otherwise, figures were measured on the Orin NX.

Data. The Song Describer Dataset and JamendoLyrics MultiLang, partitioned by artist as described in Training.

Metrics. Losses are mean cross-entropies in nats per code (chance level 6.93). Spectral errors are mean absolute differences of natural-log mel energies; one unit is a factor of e in energy, about 4.3 dB. Throughput is given in sequence positions per second.

Procedures. Every number in this chapter is produced by a command listed in Reproducibility. Raw results are stored in benchmarks/, runs/*/metrics.jsonl, samples/reconstruction/report.json and runs/*/lyrics-evaluation-*.json.

1. Correctness

All 169 tests pass on the Orin NX, and the kernel tests also pass on the Thor.

2. GPU kernels

Matrix products at the shapes used by the first model (a batch of 4 × 512 positions), measured with staleynet bench on the Orin NX. Each timing is the mean of five runs after one warm-up run, and each entry is the median of three such measurements, because the clock frequency of the device varies noticeably from run to run. These are the figures of the paper’s appendix.

A forward and backward training step on 4 × 512 positions takes 2 930 ms with the direct kernel, 576 ms with the tiled kernel and 315 ms with the register-blocked kernel: 1.8 times faster than tiling and 9.3 times faster than the direct kernel.

On the Thor. Measured during training runs of identical shape, register blocking reduced a training step of the first model on 4 × 768 positions from 238 ms to 147 ms, raising throughput from 12 900 to 20 800 positions per second. A step of the 42.9-million-parameter model on 8 × 1 408 positions takes about 2.0 seconds.

Feature extraction. The log-mel spectrogram of 60 s of audio takes 47 ms on the GPU and 73 ms with the multithreaded rustfft reference on the Orin’s eight Cortex-A78AE cores. The \(O(N^2)\) cost of the DFT-as-GEMM formulation, about 100 times the arithmetic of an FFT, consumes much of the GPU’s advantage on this device.

3. Key/value caching

Sampling 200 audio frames (203 delayed steps) with a randomly initialized default-size model:

The cache is 2.8 times faster at this context length, and the gap grows with the context, since recomputation repeats the work of every earlier position.

4. The tokenizer

The tokenizer was trained on 3.69 million steps from the 667 training tracks (subsampled to 200 000 vectors) with 20 Lloyd iterations per level. Mean absolute log-mel error, averaged over tracks, after decoding with the first n codebooks:

The validation error lies within 0.7 % of the training error, so the codebooks generalize to unseen artists, and they transfer to the lyrics corpus with a modest increase. The absolute error is nevertheless large: 0.77 natural-log units correspond to about 3.4 dB per mel band (one unit is a factor of e in energy, 10 log₁₀ e ≈ 4.3 dB), and each additional level removes less error than the previous one (−0.13, −0.08, −0.06).

5. The audio pipeline in isolation

Listening to generated samples suggested pitched, melodic content with a blurred and unintelligible voice. Two experiments were designed to locate the source of the blur without involving the generative behavior of the model.

Reconstruction without a model

staleynet reconstruct takes 12-second excerpts of held-out vocal tracks and renders them through the pipeline in stages: the exact log-mel through Griffin-Lim, the tokens decoded with one codebook, and the tokens decoded with all four codebooks. The last rendering is the best output that any language model trained on these tokens could produce. Errors are measured against the exact log-mel of the excerpt.

The final tokenizer, described below, is included for comparison.

The tokenizer introduces between three and four times the spectral error of Griffin-Lim, and the end-to-end error of a perfect language model is almost entirely that of the tokenizer. An error of 0.7 units, about 3 dB per mel band on average, is consistent with the blurred quality heard in generated voices, which depend on fine spectral detail (harmonics and formants). The measure does not capture the phase artifacts of Griffin-Lim, which are audible in the second rendering although its spectral error is small; the comparison therefore establishes the tokenizer as the dominant spectral limitation, not as the only audible one. The renderings are in samples/reconstruction/.

Memorization of a single excerpt

To test the language model independently of data scarcity, a fresh default-size model was trained on a single 12-second excerpt (sdd-103089, 562 frames) with its caption, using batch 4, peak learning rate 10⁻³ and no weight decay (configs/overfit-one-take.json).

The model was then asked to generate 12 seconds from the caption without guidance, and the sampled codes were compared with the excerpt’s codes:

Greedy free-running generation reproduces the excerpt token for token, and top-k sampling departs from it only at isolated tokens without accumulating error. The architecture, the delay pattern, the key/value-cached decoder and the sampler are therefore capable of reproducing audio exactly, and a model that has learned an excerpt produces precisely the tokenizer reconstruction measured above. Together, the two experiments indicate that improving the audio representation is a precondition for intelligible singing: further training of the language model on the present tokens cannot exceed the fidelity of the reconstruction. The renderings are in samples/overfit/.

A finer tokenizer

Guided by these two experiments, tokenizers were compared that differ in the number of frames per step, the number of band groups (the 128 mel bands are divided into groups of neighboring bands, each quantized with its own codebooks) and the number of residual levels in each group. All use 1 024 prototypes per codebook and were trained on the same songs; the error is the mean over all 39 held-out songs, in log units per mel band.

Quantizing one frame per step in 4 band groups with 4 levels each lowered the error from 0.776 to 0.401, about 1.7 dB per band. At 8 codes per step, dividing the bands into groups did not help; with 16 codes it helped a little (0.427 without groups against 0.401 with them). Most of the improvement therefore comes from describing each step with more codes, and a smaller part, about 6 %, from quantizing the band groups separately. On the four vocal excerpts above, the error of the tokens fell from 0.70 to 0.36 and the end-to-end error from 0.71 to 0.41. The price is a vocabulary of 16 codebooks instead of 4.

6. Training

Initial run (Orin NX)

configs/smoke.json: the default 7.1-million-parameter model, batch 8 × 512 positions, 1 500 steps (about 6.1 million positions, or 1.7 passes over the 3.69 million training steps), warm-up of 100 steps and cosine decay. Wall-clock time was 25 minutes at 4 100 positions per second. Validation loss on the 39 held-out tracks:

The validation loss decreased at every evaluation and remained close to the training loss, indicating that the run ended under-fitted. Codebook 0 is the most predictable (3.1 nats, a perplexity of about 23 of 1 024); the finest codebooks are the least predictable (about 5.9 nats), since they encode residual detail that is close to noise.

Extended run (AGX Thor)

configs/thor.json (batch 4 × 768, planned for 30 000 steps) ran at 12 900 positions per second, 3.3 times the throughput of the Orin NX, and was stopped at step 15 625 to free the device for the subsequent experiments. The best validation loss on the 39 held-out tracks was 4.02 nats at step 15 000, against 4.93 for the initial run; with the model and data unchanged, the improvement is attributable to tenfold longer training.

Vocal subset: fine-tuning against training from scratch

Two runs were trained on the vocal subset (281 training and 19 validation tracks):

These losses are measured on the vocal validation subset and are not comparable with the 4.02 of the extended run, which was measured on all 39 validation tracks.

Initialization from the extended run is worth about 0.5 nats per code. The fine-tuned model reached its best validation loss after only 250 steps; thereafter its validation loss rose to about 4.30 and remained there while its training loss continued to fall, which indicates that 281 tracks do not support prolonged fine-tuning at this learning rate. The model trained from scratch improved until near its final step and ended with a training loss of 3.53 against a validation loss of 4.72, a gap that likewise points to insufficient data. On informal listening, samples from vocal-scratch were perceived as blurred but melodic voices; this is a single listener’s impression rather than a controlled evaluation.

A larger model on the finer tokens (AGX Thor)

With the finer tokens, a larger model was trained on the Thor (configs/pq-large-thor.json): 8 blocks of width 512 with 8 heads, 42.9 million parameters, with the lyrics field enabled from the start. It was trained on all 742 training songs of the combined corpora, with JamendoLyrics weighted 4-fold and the author’s two takes weighted 10-fold, in batches of 8 sequences. A step took 2.0 seconds and about 11 GB of memory. Its best held-out loss was 5.062 nats per code, averaged over its 16 codebooks, at step 18 500 of 20 000; because the vocabulary differs, this figure is not comparable with the losses above. Its curves are published on Hugging Face.

The per-codebook panel shows the coarse-to-fine structure directly: the four codebooks of the coarsest level are predicted at about 3.0 nats, each finer level is harder, and the finest approach the uniform guess of 6.93, because they encode residual detail that is close to noise.

The model overfitted

The table compares the training loss, averaged over the 1 000 steps before each point, with the held-out loss at that point. The training loss is measured with the prompt dropped from 10 % of the sequences, which raises it slightly, so the true gap is, if anything, larger.

From step 15 000 the held-out loss no longer improved, while the training loss kept falling: everything the model learned in its last quarter applied only to its training songs. The gap grew to almost one nat per code, which means the model assigned to each code of its training songs, as a geometric mean, \(e^{0.975} \approx 2.7\) times the probability it assigned to a code of a song it had not heard. Listening agreed: different prompts produced nearly the same song, and a line of lyrics from a training song recurred. The cause was the sampling weights, under which each lyric-bearing song was drawn about 650 times and each of the author’s takes more than 1 600 times (Training).

A balanced run

A second run, pq-large-rock, added the 29 314 rock tracks of the Free Music Archive but kept the heavy weights on the small datasets, and was stopped after 1 000 steps for that reason. Its replacement, pq-rock-balanced (configs/pq-rock-balanced-thor.json), trains the same 42.9-million-parameter model on 30 054 songs for 30 000 steps, with weights chosen so that no song is drawn more than about 16 times. It is in progress; its curves are published alongside those of pq-large, and the quantity to watch is the gap between the two losses.

7. Lyrics conditioning

Models with a 256-byte lyrics field were warm-started from the two vocal models and trained on a mixture of the vocal subset, JamendoLyrics and, after the first 750 (Orin) and 4 750 (Thor) steps, the rock subset with weight 3. The use of the lyrics was measured with the paired protocol of Lyrics Conditioning on 48 batches of four crops.

First runs. With JamendoLyrics at weight 2 (73 training songs against 281 vocal-subset tracks), about one third of training sequences were drawn from the lyrics corpus. The Thor run (8 000 steps, warm-started from vocal-finetune) gave the following losses at its best checkpoint (step 4 750):

The difference between mismatched and aligned lyrics is negligible on held-out songs and very small on training songs. The model responds weakly to the presence of text (the removed condition is slightly worse) but not to its content: it has not learned to use the lyrics. The likely causes are the small share of lyric-bearing sequences, the short schedule, and the difficulty of learning a text-to-audio alignment from 78 songs in four languages.

Second run. A Thor run of 20 000 steps with JamendoLyrics at weight 8 and the register-blocked GEMM (configs/lyrics-thor-long.json) overfitted quickly: its training loss fell to 2.66 while its held-out loss stayed near 4.34, and its best held-out checkpoint is that of step 250. The paired test on that checkpoint, now reported with the standard error of the paired difference, gave:

On held-out songs the correct and the wrong lyrics give the same loss, so the model does not use the words. On training songs the difference is statistically clear but small, which points to memorization of those songs rather than a general relation between words and sound. In both cases the model scores slightly worse with no lyrics than with any lyrics: it reacts to the presence of text, not to its content.

8. Generation

Three prompts of 12 seconds each, generated from the initial checkpoint with top-k 250, temperature 1.0, guidance scale 3 and seed 1 on the Orin NX:

Generation runs about three times faster than real time, including guidance, which doubles the number of forward passes.

The spectrograms show rhythmic onsets and, for the pop prompt, sustained harmonic lines between 300 and 600 Hz. The prompt influences the output: the share of energy below 250 Hz is 0.57 for the pop prompt, 0.88 for folk and 0.90 for dance. Broadband clicks, visible as vertical stripes, are present in all three samples; in the light of Section 5 they are attributed in part to the tokenizer and Griffin-Lim, although the contribution of each factor has not been isolated. No formal listening test has been conducted, so these results do not establish perceived quality, musicality or intelligibility.

Reproducibility

This chapter lists what is needed to rebuild the system and the commands that produce each result reported in this book. Because sampling and initialization are seeded and the GPU code adds numbers in a fixed order, a command repeated on the same device produces the same numbers.

Requirements

The code runs on an NVIDIA Jetson (Orin or Thor) or any other CUDA GPU supported by cuda-oxide. Building it requires the cuda-oxide toolchain: the nightly Rust compiler named in rust-toolchain.toml and the cargo-oxide tool described in the cuda-oxide README. The datasets occupy about 5 GB on disk and the prepared features about 3 GB.

Building and testing

cargo oxide build -- --release                    # builds target/release/staleynet
cargo test --release -p staleynet-audio           # signal processing on the CPU
cargo oxide test -- --release -p staleynet-cuda   # every GPU kernel against its CPU reference
cargo oxide test -- --release -p staleynet-model  # gradients, cached decoding, memorization, checkpoints, warm start
cargo oxide test -- --release -p staleynet-data   # GPU spectrogram, tokenizer, batches, lyrics alignment

Obtaining the data

The Song Describer Dataset is published on Zenodo. Its audio carries per-track Creative Commons licenses, listed in audio_licenses.txt, and its captions are licensed CC BY-SA 4.0.

mkdir -p data/raw/song-describer && cd data/raw/song-describer
for file in song_describer.csv audio.zip; do
  curl -L -o "$file" "https://zenodo.org/records/10072001/files/$file?download=1"
done
unzip -q audio.zip -d audio && cd -

JamendoLyrics MultiLang is distributed as a Git repository whose songs and lyrics carry the Creative Commons license stated for each song in JamendoLyrics.csv.

git clone --depth 1 https://github.com/f90/jamendolyrics data/raw/jamendolyrics

Recordings made with the studio are stored in data/recordings/ in the layout that staleynet prepare --folder reads: each audio file is accompanied by a .txt file holding its description and a .lyrics.txt file holding the words sung.

The pipeline, step by step

# Spectrograms and tokens of the main dataset, and the tokenizer itself
staleynet prepare  --song-describer data/raw/song-describer --cache data/cache
staleynet tokenize --cache data/cache --listen samples/tokenizer

# The lyrics corpus, encoded with the same tokenizer
staleynet prepare  --jamendo-lyrics data/raw/jamendolyrics --cache data/cache-jamendolyrics
staleynet tokenize --cache data/cache-jamendolyrics --reuse data/cache/tokenizer.json

# Training runs (see configs/ for the settings of each experiment)
staleynet train --cache data/cache --run runs/vocal-scratch --config configs/vocal-scratch.json
staleynet train --run runs/lyrics-orin --config configs/lyrics-orin.json --fine-tune-from runs/vocal-scratch/best

# Analyses
staleynet reconstruct --cache data/cache --out samples/reconstruction
staleynet evaluate    --run runs/lyrics-orin --cache data/cache-jamendolyrics
staleynet bench       --out benchmarks/results.json

# Listening
staleynet sing --run runs/lyrics-orin --prompt "slow grunge song with a raw male voice" \
               --lyrics-file song.txt --seconds 15 --out samples/song.wav

A run interrupted at any point continues with --resume. The resumed run restores the optimizer state and draws the same batches it would have drawn without the interruption.

The single-excerpt memorization experiment of Experiments, Section 5 uses configs/overfit-one-take.json on a folder containing one excerpt rendered by staleynet reconstruct, and compares the generated tokens with the excerpt’s own tokens through staleynet sing --save-codes.

Where the results are stored

Limitations and Future Work

This chapter collects the known weaknesses of the system, beginning with those that most affect what a listener hears, and outlines the work that would address them.

Limits on what the output sounds like

The most important limitation is the fidelity of the audio representation. The experiments in Section 5 of Experiments show that the tokenizer loses about 0.7 log units of spectral detail, roughly 3 dB per frequency band, and that a model which reproduces its training audio perfectly still sounds only as good as this reconstruction. The detail that singing depends on, the harmonics and formants that make vowels recognizable, is precisely what is lost. Until the representation improves, longer training of the language model cannot make voices intelligible.

Griffin-Lim adds a further loss that the spectral measure understates: it estimates phase without ever having heard real audio, so harmonics lose their frame-to-frame coherence and voices take on a metallic smear (see Audio Representation). The neural vocoder, trained on spectrograms decoded from StaleyNET’s own tokens, is meant to replace it. It is trained with spectral losses only, and vocoders trained that way tend to sound slightly muffled in the highest frequencies; adversarial training, as in Vocos, would address this.

The lyrics input does not yet influence the output measurably. The first models trained with it produced the same predictions with correct lyrics as with lyrics from another song. Learning which sounds belong to which words from 78 songs in four languages may require more lyric-bearing data than is available, or a stronger form of conditioning.

The training data contains comparatively little singing. Of the 706 songs in the main dataset, 298 mention a voice in their captions, and several styles of interest are nearly absent: two songs are described as grunge, one as psychedelic, and none as soft rock.

Small datasets invite overfitting, and the first 42.9-million-parameter model overfitted: its training loss kept falling while its held-out loss stopped at 5.06 nats per code, and it returned nearly the same song for different prompts (Experiments). The cause was heavy weighting of the smallest datasets, under which some songs were drawn more than a thousand times. The balanced run that replaced it draws no song more than about 16 times, but it pays for this by making lyric-bearing songs rare. The model also has no dropout inside the network; the only regularization is weight decay, prompt dropout, and keeping the checkpoint with the lowest held-out loss.

The scale of the system is small by the standards of published music generators. Its largest model has 42.9 million parameters and is trained on about 300 hours of audio, most of it rock tracks described only by genre tags, whereas MusicGen trains between 300 million and 3.3 billion parameters on 20 000 hours. The output is mono at 24 kHz.

Limits of the implementation

All arithmetic is carried out in 32-bit floating point, and the GPU’s tensor cores, which accelerate lower-precision matrix products, are not used. Storing values in 16-bit form while accumulating in 32 bits would roughly halve the memory traffic.

Attention stores the full matrix of scores between all pairs of positions, which limits the context length that fits in memory. A fused attention kernel in the style of FlashAttention (Dao et al., 2022) would remove that matrix.

At every training step the loss of each tokenizer round and the size of every gradient are copied from the GPU to the CPU, which forces the CPU to wait for the GPU and limits the speed-up on the faster AGX Thor.

The spectrogram is computed as a matrix product, which costs about a hundred times more arithmetic than an FFT for the frame length used, and is therefore only 1.3 times faster than the CPU on the Orin NX.

Future work

The analysis above suggests an order of priorities.

1. A better audio representation. A learned tokenizer, consisting of a convolutional encoder and decoder trained together with residual quantization, as in SoundStream, would preserve far more detail at a similar number of tokens. A lighter intermediate step is to quantize single frames instead of pairs, or to use more and larger codebooks, and to measure the effect with staleynet reconstruct.
2. Adversarial training for the vocoder: discriminators that judge short stretches of waveform and of spectrogram, added to the spectral losses once those have converged.
3. More lyric-bearing singing data, in particular short takes recorded with the studio, in which each line can be aligned with the audio.
4. Cross-attention to the lyrics, so that every audio position can consult the words directly rather than through the shared sequence.
5. Ablation studies of the guidance scale, the rate at which text is removed during training, and the delay pattern.
6. Engineering improvements: mixed precision, tensor-core matrix products, fused attention, and fewer synchronisations between GPU and CPU.

References

Models and training

• A. Vaswani et al. Attention Is All You Need. NeurIPS 2017. https://arxiv.org/abs/1706.03762
• A. Radford et al. Language Models are Unsupervised Multitask Learners. OpenAI, 2019.
• T. Brown et al. Language Models are Few-Shot Learners. NeurIPS 2020. https://arxiv.org/abs/2005.14165
• B. Zhang, R. Sennrich. Root Mean Square Layer Normalization. NeurIPS 2019. https://arxiv.org/abs/1910.07467
• R. Xiong et al. On Layer Normalization in the Transformer Architecture. ICML 2020. https://arxiv.org/abs/2002.04745
• D. Hendrycks, K. Gimpel. Gaussian Error Linear Units (GELUs). 2016. https://arxiv.org/abs/1606.08415
• I. Loshchilov, F. Hutter. Decoupled Weight Decay Regularization. ICLR 2019. https://arxiv.org/abs/1711.05101
• I. Loshchilov, F. Hutter. SGDR: Stochastic Gradient Descent with Warm Restarts. ICLR 2017. https://arxiv.org/abs/1608.03983
• R. Pascanu, T. Mikolov, Y. Bengio. On the Difficulty of Training Recurrent Neural Networks. ICML 2013. https://arxiv.org/abs/1211.5063
• L. Xue et al. ByT5: Towards a Token-Free Future with Pre-trained Byte-to-Byte Models. TACL 2022. https://arxiv.org/abs/2105.13626

Audio generation

• A. Agostinelli et al. MusicLM: Generating Music From Text. 2023. https://arxiv.org/abs/2301.11325
• Z. Borsos et al. AudioLM: a Language Modeling Approach to Audio Generation. IEEE/ACM TASLP 2023. https://arxiv.org/abs/2209.03143
• P. Dhariwal et al. Jukebox: A Generative Model for Music. 2020. https://arxiv.org/abs/2005.00341
• J. Liu et al. DiffSinger: Singing Voice Synthesis via Shallow Diffusion Mechanism. AAAI 2022. https://arxiv.org/abs/2105.02446
• A. Défossez et al. High Fidelity Neural Audio Compression (EnCodec). TMLR 2023. https://arxiv.org/abs/2210.13438
• R. Kumar et al. High-Fidelity Audio Compression with Improved RVQGAN (DAC). NeurIPS 2023. https://arxiv.org/abs/2306.06546
• J. Copet et al. Simple and Controllable Music Generation (MusicGen). NeurIPS 2023. https://arxiv.org/abs/2306.05284
• N. Zeghidour et al. SoundStream: An End-to-End Neural Audio Codec. IEEE/ACM TASLP 2021. https://arxiv.org/abs/2107.03312
• J. Ho, T. Salimans. Classifier-Free Diffusion Guidance. 2022. https://arxiv.org/abs/2207.12598
• A. Fan, M. Lewis, Y. Dauphin. Hierarchical Neural Story Generation. ACL 2018. https://arxiv.org/abs/1805.04833
• R. Pope et al. Efficiently Scaling Transformer Inference. MLSys 2023. https://arxiv.org/abs/2211.05102
• J. Kong, J. Kim, J. Bae. HiFi-GAN. NeurIPS 2020. https://arxiv.org/abs/2010.05646
• H. Siuzdak. Vocos: Closing the Gap between Time-Domain and Fourier-Based Neural Vocoders. 2023. https://arxiv.org/abs/2306.00814

Data

• D. Stoller, S. Durand, S. Ewert. End-to-End Lyrics Alignment for Polyphonic Music Using an Audio-to-Character Recognition Model. ICASSP 2019. https://arxiv.org/abs/1902.06797
• S. Durand, D. Stoller, S. Ewert. Similarity-based Audio-Lyrics Alignment of Multiple Languages. ICASSP 2023. https://arxiv.org/abs/2306.07744
• I. Manco et al. The Song Describer Dataset: a Corpus of Audio Captions for Music-and-Language Evaluation. NeurIPS 2023 Workshop. https://arxiv.org/abs/2311.10057
• M. Defferrard, K. Benzi, P. Vandergheynst, X. Bresson. FMA: A Dataset for Music Analysis. ISMIR 2017. https://arxiv.org/abs/1612.01840

Signal processing

• D. Griffin, J. Lim. Signal Estimation from Modified Short-Time Fourier Transform. IEEE TASSP 1984. https://doi.org/10.1109/TASSP.1984.1164317
• N. Perraudin, P. Balazs, P. Søndergaard. A Fast Griffin-Lim Algorithm. WASPAA 2013. https://doi.org/10.1109/WASPAA.2013.6701851
• J. O. Smith. Digital Audio Resampling Home Page. CCRMA, Stanford. https://ccrma.stanford.edu/~jos/resample/
• J. O. Smith. Spectral Audio Signal Processing. CCRMA, Stanford. https://ccrma.stanford.edu/~jos/sasp/
• M. Slaney. Auditory Toolbox, Version 2. Interval Research TR 1998-010.
• S. Lloyd. Least Squares Quantization in PCM. IEEE Trans. Information Theory 1982. https://doi.org/10.1109/TIT.1982.1056489
• A. Oppenheim, R. Schafer. Discrete-Time Signal Processing. 3rd ed., Pearson, 2010.

GPU programming

• V. Volkov, J. Demmel. Benchmarking GPUs to Tune Dense Linear Algebra. SC 2008. https://doi.org/10.1109/SC.2008.5214359
• T. Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. https://arxiv.org/abs/2205.14135
• NVIDIA. cuda-oxide. https://github.com/NVlabs/cuda-oxide
• D. Kirk, W. Hwu, I. El Hajj. Programming Massively Parallel Processors. 4th ed., Morgan Kaufmann, 2022.
• NVIDIA. CUDA C++ Programming Guide. https://docs.nvidia.com/cuda/cuda-c-programming-guide/
• G. Steele, D. Lea, C. Flood. Fast Splittable Pseudorandom Number Generators. OOPSLA 2014. https://doi.org/10.1145/2714064.2660195