Saxophone highlight
A soulful saxophone excerpt backed drums.
01
Current approach
Pretrained diffusion with LoRA finetuning
One of the key realizations that led to LoRA fine tuning was that training a diffusion model from scratch required an enormous scale of data.
I had originally expected large amounts of data would be needed, but I was willing to sacrifice some quality for less data. It turns out, before a certain threshold, the audio will simply sound like noise before you start to discern music (see below).
To get coherent music when training from scratch, one would need 100s of GB to TBs of audio data in order to turn random noise into somewhat pleasant music.
This was impractical for myself to implement and unnecessary - thinking cost wise, the amount of compute resources I had and the data I could effectively process.
On the other hand, using Low-Rank Adaptation (LoRA) fine tuning, only 3 GB to 6 GB is needed (this is about 10-20 hours of audio) to turn a general text-to-audio diffusion model to a model capable of generating jazz style music.
In essence, an adapter is attached to the output of the pretrained model, with pretrained model weights frozen, and only the adapter weights trainable.
The number of weights in the adapter is greatly smaller than the pretrained model (at only about 1-3% of the total weights).
This keeps the heavy lifting of what music sounds like handled by the pretrained model unaltered (things like pitch, instrument sounds, timbre), and having the adapter fine tuning focus only on capturing the style of jazz (swing feel, brushed drums, bending of notes on sax).
Although, LoRA fine tuning does not generally cover all the weights that full fine tuning achieves, the original LoRA paper research (Hu et al., 2021) shows that useful adaptations appear to have low dimensional structural changes than the total parameters in the original model.
Note that the data for training was hand selected to be around 10 hours of audio, all which resembled similar structures in style, phrasing, instrument sounds for training consistency.
Data cleansing was performed to pick out oddities that might cause the model to diverge.
Generated music
Good quality swing piano
A piano generation with a swing feel that transitions into straight time.
Saxophone — swing
A saxophone sample with a swing-oriented prompt.
Slow saxophone
A slower, more spacious saxophone generation.
Saxophone — blues
A blues-oriented saxophone generation.
Piano
A piano-focused generation.
The pipeline
Turn the source recordings into consistent clips, split them into training and holdout sets, and pair them with useful text captions.
Encode the training audio once with the pretrained variational autoencoder so training does not repeat this expensive step.
Optimize only the LoRA attention weights while keeping the pretrained AudioLDM2 model frozen.
Load an adapter checkpoint and synthesize music from jazz-focused text prompts.
Compare matched base and adapted generations through listening and measures such as Fréchet Audio Distance.
Training details
80training epochs
4samples per batch
500steps between checkpoints
Dataset scale and quality
One of the major challenges in fine tuning was getting the scale of the data right, as well as picking out a quality dataset to train on. As I noticed, and also mentioned in the scaling laws paper (Kaplan et al., 2020), the number of trainable model parameters must scale proportionally with dataset size, or else there will be diminishing returns. In this project, when I tried fine tuning with a small dataset (even when the style was more similar), the quality of music at the best checkpoint before overfitting was less than what I achieved with a decently sized dataset (10 hours of audio, about 7 GB).
This relationship can be expressed by holding the ratio \(N^{0.74}/D\) approximately constant.
Model and adapter architecture details
This project uses AudioLDM2 (Liu et al., 2024), a pretrained latent diffusion model designed for text-conditioned audio generation. It provides a strong general musical prior, which means the experiment can focus on adapting style and instrumentation rather than learning audio generation from scratch. LoRA adapters are attached to the attention layers of the diffusion network while the larger pretrained pipeline remains frozen.
What did AudioLDM2 produce as jazz music?
This original AudioLDM2 sample was used as the baseline for comparison. The fine-tuned model produced a significant improvement.
A LoRA rank of 16 was selected as a practical balance: it offers enough capacity to learn the timbre and phrasing of the jazz dataset while keeping the adapter compact, reducing memory use, and limiting the tendency to overfit a comparatively focused collection of recordings.
Learning rate tuning
The learning rate is tapered near the end of training. Early updates can make broader stylistic changes; the smaller late-stage updates help refine the adapter without abruptly disturbing musical structure that has already been learned. The final taper is also a useful safeguard against audible degradation or overfitting late in the run.
Suggested next steps
- Increase to rank 32 with more training data to see if more complex patterns and phrasing can be captured
- Chord quality struggles right now with generation - with more samples, this multi-pitch harmony can be better handled
- Curate a smaller second-stage dataset around phrasing, articulation, and instrument balance.
- Expand captions with tempo, instrumentation, articulation, and mood.
- Experiment with other techniques to get better layering of instrumentation - a bit out of sync in some parts
02
Learning from first principles
Previous trials with diffusion
Diffusion trial
This trial explores learning diffusion (Ho et al., 2020) by injecting random noise into an original audio sample.
All diffusion training and ML architecture was written from the ground up in PyTorch.
Original audio sample
gen_orig_4.wav is the original audio sample.
Noisy diffusion sample
gen_4.wav is the sample with random noise injected while the model is in the process of learning diffusion.
How noising and denoising work
DDPM defines a fixed forward process that gradually adds Gaussian noise, then learns a reverse process that removes that noise one step at a time (Ho et al., 2020). Here, x0 is the clean sample, xt is its noisy state at step t, and βt controls how much noise is added at that step.
Forward process: add noise
\[q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t; \sqrt{1-\beta_t}\,x_{t-1}, \beta_t I\right)\]
Each step keeps most of the previous signal and adds a small amount of fresh Gaussian noise. After many steps, the original structure is replaced by something close to pure noise.
Direct sampling: reach any noise level
\[x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I)\]
With αt = 1−βt and ᾱt equal to the product of all α values through step t, a noisy training example can be created at any chosen step without simulating every earlier step.
Reverse process: learn to denoise
\[x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\,\epsilon_\theta(x_t,t)\right) + \sigma_t z\]
The model predicts the noise component εθ(xt, t), scales it according to the noise schedule, and subtracts it from xt. The σtz term adds the variance required for reverse diffusion, where z ~ 𝒩(0, I) for intermediate steps and z = 0 for the final step. Repeating this update produces xt−1, then xt−2, until the generated sample reaches x0.
Training objective: predict the noise
First, separate the known noise from the model's predicted noise to form the residual:
\[r_\theta = \epsilon - \epsilon_\theta\!\left(\sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon,t\right)\]
Then calculate the expected squared size of that residual:
\[L_{\mathrm{simple}} = \mathbb{E}_{t,x_0,\epsilon}\!\left[\left\lVert r_\theta \right\rVert_2^2\right]\]
The expectation averages the loss over randomly selected timesteps, clean samples, and Gaussian noise. Training minimizes this prediction error; the actual removal of predicted noise from xt happens in the reverse-process update above.
03
Earlier and new experiments
Transformer autoregressive generation
The original idea was that music, specifically jazz, can be treated as a language, and if we could tokenize it, then like LLMs, we can perform next-token prediction to generate intelligent pieces of creation. Fortunately, there does exist a mechanism, audio to MIDI, that produces discrete pitches from wav audio, and through encoding the MIDI, I was able to train a autoregressive transformer to accomplish this task.
Samples from the 8 layer transformer model built from the ground up
A sample using Autumn Leaves as primer melody
(The generated melody starts at 0:27)
Previous experimentation with a Google Magenta pretrained model
Note that this was the very first initial experiement tinkering in trying to generate jazz music with ML models.
The result was with piano instrumentation was decent, however it was constrained to particular timbres enforced by MIDI (whereas diffusion on raw audio was more permissive).
Some background:
- An Recurrent Neural Network (RNN) with an attention mechanism was used to capture longer patterns within the music.
- The song was based on the Charlie Parker tune, Blues for Alice. Piano was selected to be the lead instrument.
- The model was trained on collected recordings on the web of pianists, as well as snippets which I personally recorded!
- Check out the GitHub for more detailed information: GitHub
Some samples to listen to below:
Sample 1
Notes & Notable Timestamps:
- Trained on a 2 layer RNN with 64 processing units each
- 0:51 - 1:03 - Repetition of a common starting phase
- 1:15 - 1:16 - Blend nicely with Cm7 - F7 harmonic structure
- Notable dissonance throughout (playing E)
Sample 2
Notes & Notable Timestamps:
- Trained on a 2 layer RNN with 128 processing units each
- Relatively more swing feel than others
- At 1:20, there is instance of overfitting
- 1:32 - interesting swing groove happening! (solid chrous of improv)
- 2:02 - fit the Abm7 -> Db7 harmony nicely
- Occasional jazz articulations/phrases appearing (one at 2:37 - 2:38)
Sample 3
Notes & Notable Timestamps:
- Trained on a 3 layer RNN with 64 processing units each
- 1:20 - Interesting stepping up sequence
- The stepping up sequence transitions from mid to high range at 1:30 - 1:33
- 1:33 - 1:34 - Melodical phrase generated!
- 1:40 - 2:36 - A period of exotic back and forths with the upper and lower ranges of pitch (effects of attention)
- 2:36 - Return to natural phrasing