AudioLDM 2 最適化ガイド: Hugging Face Diffusers を使った推論時間の短縮

Hugging Face は、AudioLDM 2 の最適化ガイドをリリースしました。AudioLDM 2 は、リアルなサウンドエフェクト、人間の音声、音楽を生成できるテキストからオーディオへの潜在拡散モデルです。diffusers ライブラリ内でコードとモデルの最適化を組み合わせることで、10 秒のオーディオサンプルの推論時間は、元の実装では 30 秒以上だったものが、1 秒未満に短縮できます。

AudioLDM 2 モデルアーキテクチャ

AudioLDM 2 は、テキスト埋め込みから連続表現を学習することでオーディオを生成するテキストからオーディオへの潜在拡散モデル（LDM）です。生成プロセスは多段階パイプラインに従います：

  1. テキストエンコーディング: モデルは、CLAP のテキストブランチ（オーディオに整合した埋め込み用）と Flan-T5 のテキストエンコーダー（意味表現用）の 2 つのエンコーダーを使用して、テキスト埋め込みを計算します。
  2. 投影: これらの埋め込みは、AudioLDM2ProjectionModel を通じた線形投影により共有埋め込み空間に渡されます。
  3. 自己回帰生成: GPT2 言語モデルは、投影された CLAP と Flan-T5 埋め込みに条件付けられた $N$ 個の埋め込みベクトルのシーケンスを生成します。
  4. 潜在拡散: UNet ベースの LDM は、$T$ ステップの推論ステップでランダムな潜在変数をデノイズします。ほとんどの LDM と異なり、AudioLDM 2 の UNet は GPT2 からのクロスアテンション埋め込みと Flan-T5 からのクロスアテンション埋め込みの 2 セットを使用します。
  5. デコード: 最終的なデノイズされた潜在変数は VAE デコーダーを通過して Mel スペクトログラムを復元し、これが vocoder によってオーディオ波形に変換されます。

利用可能なモデルチェックポイント

Checkpoint Task Model Size Training Data (hours)
cvssp/audioldm2 テキストからオーディオ 1.1B 1150k
cvssp/audioldm2-music テキストからミュージック 1.1B 665k
cvssp/audioldm2-large テキストからオーディオ 1.5B 1150k

推論の高速化のための最適化

元の実装の遅い推論速度に対処するため、Hugging Face は diffusers ライブラリで利用可能な 4 つの主要な最適化技術を特定しています：

1. Flash Attention (SDPA)

PyTorch 2.0 以上を使用すると、diffusers ライブラリは自動的に torch.nn.functional.scaled_dot_product_attention (SDPA) を有効にします。これにより、出力品質を変更せずに計算時間を削減する、Flash Attention に似たメモリ効率の良い注意操作が提供されます。

2. ハーフプレシジョン (float16)

モデルの重みと計算を float32 から float16 (ハーフプレシジョン) に変換すると、GPU メモリ使用量が大幅に削減され、音質にほとんど影響を与えずに推論速度が向上します。これは、.from_pretrained 呼び出し時に torch_dtype=torch.float16 を渡すことで実現されます。

3. Torch Compile

パイプラインで最も計算コストが高い UNet を torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) でラップすると、大幅な速度向上が得られます。最初の推論実行はコンパイルオーバーヘッドにより遅くなる可能性があります（最大 2 分程度）、しかしその後のすべての生成は大幅に高速化されます。

4. 効率的なスケジューラー

デフォルトの DDIMScheduler （通常 200 ステップ必要）を、DPMSolverMultistepScheduler などのより高性能なスケジューラーに置き換えることで、モデルは 20-25 の推論ステップで同様の品質を達成できます。この最適化を他の最適化と組み合わせることで、10 秒のサンプルの生成時間を 1 秒未満に短縮できます。

メモリ管理と長形式オーディオ

長いオーディオサンプル（例：150 秒）を生成したり、大きなチェックポイント（例えば audioldm2-large）を使用したりすると、潜在変数の幅が増加し、クロスアテンションメモリがシーケンス長の二乗に比例して増加するため、CUDA のメモリ不足（OOM）エラーが発生する可能性があります。

これを緩和するために、Hugging Face は pipe.enable_model_cpu_offload() を介した CPU オフロード を推奨します。この技術は、現在アクティブなモデルコンポーネントのみを GPU に残し、残りを CPU にオフロードすることで、限られた RAM の GPU でも長形式オーディオの生成や大きなモデルの使用を可能にし、推論時間へのペナルティを最小限に抑えます。

Sources