Training a Generative Kick Drum Diffusion Model on 6GB VRAM

Training a high-quality generative audio model does not require a massive GPU cluster or a billion-dollar budget. It is possible to train a functional latent diffusion model for specific sounds—such as kick drums—on consumer-grade hardware, specifically a 7-year-old NVIDIA GeForce GTX 1660 SUPER with 6GB of VRAM.

The Generative Pipeline: Three Models, One Workflow

Generating audio via diffusion is computationally expensive if performed on raw waveforms. To make this feasible on limited VRAM, the pipeline converts raw audio into a compressed latent space where the diffusion process occurs, then reconstructs the audio through a series of models.

1. Variational Autoencoder (VAE)

The VAE handles compression and decompression. The encoder compresses a 128x173 log-mel spectrogram (22,144 floats) into a 4x8x11 latent tensor (352 floats), achieving approximately 63x compression.

To ensure the latent space is continuous and navigable for the diffusion model, the VAE uses two specific mechanisms:

Mean and Variance Output: Instead of a single point, the encoder outputs a mean and variance, forcing the decoder to learn neighborhoods of sounds.
KL Divergence Penalty: This pulls the latent distributions toward a standard normal distribution (mean 0, variance 1), preventing the model from creating "scattered islands" of data.

2. Diffusion U-Net

The U-Net is the core generative engine. It learns to reverse a process where Gaussian noise is added to a latent tensor over 1,000 steps. By learning to denoise the tensor, the model can generate new, unique kick drum latents from pure noise.

Text Conditioning and CFG: The model is steered by keywords extracted from filenames (e.g., "warm", "hard"). To improve stability, the author implemented Classifier-Free Guidance (CFG), where the model is trained with keywords hidden 15% of the time. During generation, the difference between the conditional (with keywords) and unconditional (without keywords) outputs is used to amplify the influence of the prompt.

3. Vocoder (HiFi-GAN)

Because converting audio to a mel spectrogram is lossy and non-invertible, a vocoder is required to turn the generated spectrogram back into a waveform. A HiFi-GAN (Generative Adversarial Network) was used, consisting of a generator that upsamples the spectrogram and a discriminator that ensures the output sounds realistic. Due to the instability of GAN training, this model was trained at full precision (fp32).

Data Preparation and Processing

The model was trained on a dataset of 13,615 kick samples extracted from a personal library. The data pipeline involved several critical cleaning and normalization steps:

Filtering: Files containing "loop" or "BPM" were removed, and files outside the 5KB to 1MB size range were discarded.
Normalization: All samples were resampled to 44.1kHz, padded or trimmed to exactly 2 seconds, normalized to -1dB peak, and given a 0.2-second fade-out.
Log-Mel Spectrogram Conversion: Audio was converted to 128x173 log-mel spectrograms. The Mel scale was used to align frequency bins with human hearing, and the Log scale was used to handle amplitude (loudness) logarithmically.

Deployment and Inference Optimization

To avoid the cost of a 24/7 GPU instance, the project uses Modal, a serverless GPU platform. This architecture optimizes inference through several techniques:

Class-Based Workers: Models are loaded into GPU memory during the container's boot sequence rather than on every request, reducing "warm" generation time to approximately 2.5 seconds.
Weight Caching: Model weights (approx. 300MB) are cached on the Modal machine to avoid repeated downloads from HuggingFace during cold boots.
Rate Limiting: To manage costs, users are limited to 10 generations per day.

Technical Trade-offs and Artifacts

Training on limited hardware and using high compression ratios introduced specific artifacts:

Granularization: Heavy compression of the output reveals a "grainy" texture. This is attributed to the 63x latent compression and the vocoder's upsampling process.
DC Offset Clicks: The generated audio occasionally ended on a DC offset, causing an audible click. This was resolved by applying an exponential fade-out starting at the 1-second mark.
Keyword Limitations: Keywords with low frequency in the training set (e.g., "techno", which appeared only 39 times) produced poor results. The final application uses a hardcoded prompt ("hit house") that was found to be the most reliable.

Summary of Hardware and Training Efficiency

Component	Hardware	VRAM	Training Time
VAE	GTX 1660 SUPER	6GB	~1 Day
Diffusion U-Net	GTX 1660 SUPER	6GB	~1 Day
HiFi-GAN Vocoder	GTX 1660 SUPER	6GB	~1 Day