在 6GB VRAM 上訓練生成式大鼓 (Kick Drum) 擴散模型

訓練高品質的生成式音訊模型不需要龐大的 GPU 集群或十億美元的預算。透過利用潛在擴散 (latent diffusion) 管線與激進的壓縮技術，在消費級硬體上訓練功能完備的生成式模型是可能的，例如使用一台擁有 6GB VRAM 的 7 年前 NVIDIA GeForce GTX 1660 SUPER。

生成管線：三個模型，一個工作流

為了讓音訊生成在有限的 VRAM 下在計算上可行，該系統避免了對原始音訊進行擴散。相反地，它使用一個三階段管線，將音訊壓縮到潛在空間 (latent space)，在該空間中生成新數據，然後再重建音訊。

1. 變分自編碼器 (VAE)

VAE 處理音訊表示的壓縮與解壓縮。它接收一個 log-mel spectrogram (128x173) 並將其壓縮成一個小的 3D 潛在張量 (4x8x11)。這代表了相對於 spectrogram 的約 63 倍，以及相對於原始音訊的約 250 倍的壓縮率。

為了確保潛在空間是連續且可用於生成的，VAE 採用了兩種機制：

• 均值與變異數輸出： 編碼器同時輸出均值與變異數，迫使解碼器學習從點的鄰域而非單一座標進行重建。
• KL Divergence Penalty： 應用懲罰項以將潛在分佈拉向標準常態分佈，防止模型產生「分散的孤島」數據。

2. 擴散 U-Net

擴散 U-Net 是核心生成引擎。它僅在 4x8x11 的潛在空間內運作。模型被訓練用以逆轉一個過程，即高斯雜訊在 1,000 個步驟中依序添加到潛在張量中。

為了引導生成，模型使用了文本條件 (text conditioning)（從檔名中提取的關鍵字）和無分類器引導 (classifier-free guidance, CFG)。透過使用 15% 的關鍵字丟棄率 (dropout) 進行訓練，模型可以學習到通用與特定的內容大鼓特性，讓使用者在推論時能夠放大特定關鍵字的影響力（例如「punchy」或「warm」）。

3. 聲碼器 (HiFi-GAN)

由於將音訊轉換為 mel spectrogram 是一個有損且不可逆的過程，因此需要一個聲碼器來將生成的 spectrogram 轉回波形。這裡使用了一個生成對抗網絡 (GAN)，由一個負責生成音訊的生成器 (generator) 和一個試圖將其與真實訓練數據區分開來的判別器 (discriminator) 組成。這種對抗過程確保了最終音訊聽起來「真實」，藉由填補在 spectrogram 轉換過程中丟失的細微細節。

數據準備與預處理

該模型是在從個人庫中提取的 13,615 個大鼓樣本數據集上進行訓練的。以下預處理步驟對於模型的穩定性至關重要：

• 過濾： 移除了包含 "loop" 或 "BPM" 的文件，並捨棄了小於 5KB 或大於 1MB 的文件。
• 正規化： 所有樣本皆重新採樣至 44.1kHz，填充或裁剪至精確的 2 秒，正規化至 -1dB peak，並給予 0.2 秒的淡出 (fade-out)。
• Spectrogram 轉換： 音訊使用 2048 個樣本的視窗 (window) 與 512 個樣本的步長 (hop length) 轉換為 log-mel spectrograms，產生 128x173 的表示形式。

部署與推論策略

為了避免 24/7 全天候運行 GPU 實例的成本，該專案使用了 Modal，一個無伺服器 GPU 平台。

• 基於類別的工作員 (Class-Based Workers)： 為了減輕「冷啟動」時間，推論引擎被封裝在一個類別中。這允許模型（約 300MB 的權重）在容器啟動時載入到 GPU 記憶體中一次，並在隨後的請求中重複使用。
• 性能： 熱請求 (warm requests) 約需 2.5 秒，而冷啟動 (cold boots) 則接近 10 秒。
• 快取： 模型權重託管於 HuggingFace 並快取在 Modal 機器上，以避免重複下載。

技術權衡與偽影 (Artifacts)

在有限的硬體上進行訓練必須做出某些妥協。作者注意到了一種「顆粒化 (granularization)」偽影——即大鼓聽起來像是由微小的顆粒而非連續波形組成，特別是在應用重度壓縮（透過 OTT）時。這歸因於潛在空間的高壓縮率 (63x) 以及聲碼器的上採樣過程。

未來模型品質的改進方向

為了減少這些偽影並提高條件引導能力，已確定了以下用於微調的超參數：

• 潛在尺寸 (Latent Size)： 增加潛在張量的大小以減少壓縮損失。
• FFT Window： 減少 2048 個樣本的視窗以銳化瞬態 (transients)。
• 關鍵字丟棄率 (Keyword Dropout)： 微調 15% 的丟棄率以提高對文本提示詞的響應性。
• 詞彙表 (Vocabulary)： 保留具有顯著電子音樂音效意義的數字標識符（例如 "808" 或 "909"）。

關鍵要點總結

組件	策略	優點
硬體	GTX 1660 SUPER (6GB VRAM)	證明了消費級硬體對於專門的生成式任務是足夠的。
表示形式	Log-Mel Spectrogram $\rightarrow$ Latent	透過在壓縮空間中進行擴散，降低了計算負載。
架構	VAE $\rightarrow$ U-Net $\rightarrow$ HiFi-GAN	將壓縮、生成與重建分開處理。
推論	無伺服器 GPU (Modal)	在維持可接受延遲的同時，消除了閒置 GPU 成本。