在 6GB VRAM 上训练生成式底鼓扩散模型

训练高质量的生成式音频模型并不需要庞大的 GPU 集群或数十亿美元的预算。通过利用潜空间扩散（latent diffusion）流水线和激进的压缩技术，在消费级硬件上训练功能完备的生成式模型是可能的，例如使用一台拥有 6GB VRAM 的、已有 7 年历史的 NVIDIA GeForce GTX 1660 SUPER。

生成式流水线：三个模型，一个工作流

为了使音频生成在有限的 VRAM 下在计算上可行，该系统避免了对原始音频进行扩散。相反，它使用了一个三阶段流水线，将音频压缩进潜空间，在该空间中生成新数据，然后重建音频。

1. 变分自编码器 (VAE)

VAE 处理音频表示的压缩和解压缩。它接收一个 log-mel 频谱图 (128x173) 并将其压缩成一个小的 3D 潜张量 (4x8x11)。这代表了相对于频谱图约 63 倍的压缩比，以及相对于原始音频约 250 倍的压缩比。

为了确保潜空间是连续且可用于生成的，VAE 采用了两种机制：

• 均值与方差输出： 编码器同时输出均值和方差，迫使解码器学习从点的邻域而非单个坐标进行重建。
• KL 散度惩罚： 应用惩罚项以将潜分布拉向标准正态分布，防止模型创建数据的“分散孤岛”。

2. 扩散 U-Net

扩散 U-Net 是核心生成引擎。它仅在 4x8x11 的潜空间内运行。该模型通过训练来逆转一个过程，即在高斯噪声在 1,000 个步骤中被顺序添加到潜张量中。

为了引导生成，模型使用了文本条件 (text conditioning)（从文件名中提取的关键词）和无分类器引导 (CFG)。通过使用 15% 的关键词丢弃率进行训练，模型学习了通用和特定的底鼓特征，允许用户在推理阶段放大特定关键词（例如，“punchy”或“warm”）的影响。

3. 声码器 (HiFi-GAN)

由于将音频转换为 mel 频谱图是一个有损且不可逆的过程，因此需要一个声码器将生成的频谱图转回波形。这里使用了一个生成对抗网络 (GAN)，由一个创建音频的生成器和一个试图将其与真实训练数据区分开来的判别器组成。这种对抗过程通过填补频谱图转换过程中丢失的精细细节，确保最终音频听起来“真实”。

数据准备与预处理

该模型是在从个人库中提取的 13,615 个底鼓样本的数据集上训练的。以下预处理步骤对于模型的稳定性至关重要：

• 过滤： 移除了包含 “loop” 或 “BPM” 的文件，并丢弃了小于 5KB 或大于 1MB 的文件。
• 归一化： 所有样本都被重采样为 44.1kHz，填充或裁剪至精确的 2 秒，归一化为 -1dB 峰值，并应用了 0.2 秒的淡出。
• 频谱图转换： 使用 2048 个样本的窗口和 512 个样本的步长将音频转换为 log-mel 频谱图，从而得到 128x173 的表示。

部署与推理策略

为了避免 24/7 全天候运行 GPU 实例的成本，该项目使用了 Modal，一个无服务器 GPU 平台。

• 基于类的工人 (Class-Based Workers)： 为了减轻“冷启动”时间，推理引擎被封装在一个类中。这允许模型（约 300MB 的权重）在容器启动期间加载到 GPU 显存中一次，并在随后的请求中重复使用。
• 性能： 热请求（Warm requests）大约需要 2.5 秒，而冷启动（Cold boots）接近 10 秒。
• 缓存： 模型权重托管在 HuggingFace 上，并在 Modal 机器上进行缓存，以避免重复下载。

技术权衡与伪影

在有限的硬件上进行训练必须做出一定的妥协。作者注意到了一种“颗粒化”伪影——即底鼓听起来像是由于微小的颗粒组成的，而不是连续的波形——特别是在应用重度压缩（通过 OTT）时。这归结于潜空间的高压缩比 (63x) 以及声码器的上采样过程。

未来改进以提升模型质量

为了减少这些伪影并改进条件控制，确定了以下用于调优的超参数：

• 潜空间大小 (Latent Size)： 增加潜张量的大小以减少压缩损失。
• FFT 窗口： 减小 2048 样本的窗口以使瞬态更加锐利。
• 关键词丢弃率 (Keyword Dropout)： 调整 15% 的丢弃率以提高对文本提示词的响应性。
• 词汇表： 保留具有显著电子音乐声学意义的数字标识符（如 “808” 或 “909”）。

关键要点总结

组件	策略	益处
硬件	GTX 1660 SUPER (6GB VRAM)	证明了消费级硬件对于专门的生成式任务是足够的。
表示	Log-Mel 频谱图 $
ightarrow$ 潜空间	通过在压缩空间中进行扩散来减轻计算负载。
架构	VAE $
ightarrow$ U-Net $
ightarrow$ HiFi-GAN	将压缩、生成和重建分离。
推理	无服务器 GPU (Modal)	在保持可接受延迟的同时，消除了闲置 GPU 的成本。