6GB VRAMでの生成型キックドラム拡散モデルの学習

6GB VRAMでの生成型キックドラム拡散モデルの学習

高品質な生成オーディオモデルの学習には、巨大なGPUクラスターや数十億ドルの予算は必要ありません。潜在拡散パイプラインと強力な圧縮を利用することで、6GBのVRAMを持つ7年前のNVIDIA GeForce GTX 1660 SUPERのようなコンシューマー向けハードウェアでも、機能的な生成モデルを学習させることが可能です。

生成パイプライン:3つのモデル、1つのワークフロー

限られたVRAMでオーディオ生成を計算可能にするため、システムは生のオーディオを拡散させることを避けています。代わりに、オーディオを潜在空間に圧縮し、その空間内で新しいデータを生成し、その後オーディオを再構成する3段階のパイプラインを使用します。

1. 変分オートエンコーダー (VAE)

VAEはオーディオ表現の圧縮と展開を行います。log-mel spectrogram (128x173) を受け取り、小さな3D潜在テンソル (4x8x11) に押しつぶします。これは、spectrogramから約63倍、元の生のオーディオから約250倍の圧縮率を表しています。

潜在空間が連続的で生成に使用可能であることを確実にするため、VAEは2つのメカニズムを採用しています:

  • Mean and Variance Output: エンコーダーは平均と分散の両方を出力し、デコーダーが単一の座標ではなく、点の近傍から再構成することを学習するように強制します。
  • KL Divergence Penalty: 潜在分布を標準正規分布に引き寄せるためのペナルティが適用され、モデルがデータの「散在する島々」を作成するのを防ぎます。

2. 拡散 U-Net

拡散 U-Netは生成エンジンの核となる部分です。これは4x8x11の潜在空間内でのみ動作します。モデルは、ガウスノイズが1,000ステップにわたって潜在テンソルに逐次的に追加されるプロセスを逆転させるように学習されます。

生成をガイドするために、モデルはtext conditioning (ファイル名から抽出されたキーワード) と classifier-free guidance (CFG) を使用します。15%のキーワードドロップアウトで学習させることにより、モデルは汎用的および特定のキックドラムの特徴を学習し、ユーザーが推論時に特定のキーワード(例:「punchy」や「warm」)の影響を増幅させることが可能になります。

3. ボコーダー (HiFi-GAN)

オーディオをmel spectrogramに変換するプロセスは、損失を伴う非可逆的なプロセスであるため、生成されたspectrogramを再び波形に戻すためにボコーダーが必要です。ここでは、オーディオを生成するジェネレーターと、それを実際の学習データと区別しようとするディスクリミネーターからなるGenerative Adversarial Network (GAN) が使用されます。この敵対的プロセスにより、spectrogram変換中に失われた微細な詳細を補完することで、最終的なオーディオが「リアル」に聞こえるようにします。

データ準備と前処理

モデルは、個人のライブラリから抽出された13,615個のキックサンプル・データセットで学習されました。モデルの安定性のために、以下の前処理ステップが極めて重要でした:

  • Filtering: 「loop」や「BPM」を含むファイルは削除され、5KB未満または1MBを超えるファイルは破棄されました。
  • Normalization: すべてのサンプルは44.1kHzにリサンプリングされ、正確に2秒間にパディングまたはトリミングされ、-1dB peakに正規化され、0.2秒のフェードアウトが与えられました。
  • Spectrogram Conversion: オーディオは、windowの2048サンプルとhop lengthの512サンプルを使用してlog-mel spectrogramに変換され、128x173の表現となりました。

デプロイメントと推論戦略

24時間365日のGPUインスタンスのコストを避けるため、プロジェクトはサーバーレスGPUプラットフォームであるModalを使用しています。

  • Class-Based Workers: 「コールドブート」時間を軽減するため、推論エンジンはクラスにラップされています。これにより、モデル(約300MBの重み)を1をコンテナの起動時にGPUメモリにロードし、再利用することが可能になります。
  • Performance: ウォームリクエストは、約2.5秒、コールドブートは約10秒かかります。
  • Caching: モデルの重みはHuggingFaceにホストされており、Modalマシン上でキャッシュされることで、繰り返しのダウンロードを回避します。

技術的なトレードオフとアーティファクト

限られたハードウェアで学習させることは、ある種の妥協を必要とします。著者は、特に重い圧縮(OTT経由)が適用された場合に、キックが連続的な波形ではなく、小さな粒子の集合体のように聞こえる「granularization」というアーティファクト(ノイズ)が発生することを指摘しています。これは、潜在空間の高い圧縮率(63x)とボコーダーのアップサンプリングプロセスに起因します。

モデルの品質向上のための将来の改善策

これらのアーティファクトを減らし、コンディショニングを改善するために、以下のハイパーパラメータのチューニングが検討されています:

  • Latent Size: 圧縮損失を減らすために、潜在テンソルのサイズを大きくします。
  • FFT Window: トランジェントを鋭くするために、2048サンプルのwindowを小さくします。
  • **Keyword Dropout: 15%のドロップアウト率をチューニングして、テキストプロンプトへの応答性を向上させます。
  • Vocabulary: 電子音楽において重要な音響的意味を持つ数値識別子(例:「808」や「909」)を保持します。

主要な要約のまとめ

コンポーネント 戦略 メリット
Hardware GTX 1660 SUPER (6GB VRAM) コンシューマー向けハードウェアが、特化型生成タスクにおいて十分であることを証明しています。
Representation Log-Mel Spectrogram $\rightarrow$ Latent 圧縮された空間で拡散させることで、計算負荷を軽減します。
Architecture VAE $\rightarrow$ U-Net $
ightarrow$ HiFi-GAN 圧縮、生成、再構成の役割を分離しています。
Inference Serverless GPU (Modal) 許容可能なレイテンシを維持しながら、アイドル状態のGPUコストを回避します。

Sources