[Generative] Stable Audio Tools 동작 방식 몇가지 분석해보기

<분류>
A. 수단
- OS/Platform/Tool : Linux, Kubernetes(k8s), Docker, AWS
- Package Manager : node.js, yarn, brew, 
- Compiler/Transpillar : React, Nvcc, gcc/g++, Babel, Flutter
- Module Bundler  : React, Webpack, Parcel

B. 언어
- C/C++, python, Javacsript, Typescript, Go-Lang, CUDA, Dart, HTML/CSS

C. 라이브러리 및 프레임워크 및 SDK
- OpenCV, OpenCL, FastAPI, PyTorch, Tensorflow, Nsight

 

---

1. What? (현상)

이번 글에서는 Stable Audio Tools에서 아래와 같은 몇가지 특징들을 기록해두려고 합니다.

 

코드는 계속 업데이트 되기 때문에, 이 글이 등록된 시기를 기점으로 구현된 코드를 참조 부탁드리고, 하기의 모든셋팅은 아래 링크의 default셋팅을 기준으로 설명합니다.

** https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0.json

 

 

---

2. Why? (원인)

• X

---

3. How? (해결책)

 

아래와 같은 특징들을 기록해두려고 합니다.

• 1. Pretransformer 동작 방식
• 2. Attention 차원 분석
• 3. 활용되는 모든 Masking

---

1. Pretransformer 동작 방식

 

Pretransformer는 학습된 Autoencoder를 감싸고 있는 wrapper입니다.

 

즉, 따로 학습된 AudioAutoencoder는 AutoencoderPretransform이라는 구조로 wrapping되어서 Conditioned DM의 VAE Encoder-Decoder로 활용됩니다.

 

이때 Pretransformer인 AutoencoderPretransform가 동작하는 방식은 아래와 같습니다.

[Pretransformer 동작 방식]

 

(학습시) input audio 차원이 [배치, 채널수=2, 샘플사이즈] 일 때,

• OobleckEncoder : [배치, Latent 채널수=64, 샘플사이즈/downsampling_ratio] 형태로 각각 mean과 variance를 만들고,
• latents : 해당 mean과 variance를 통한 random sampling으로 [배치, Latent 채널수=64, 샘플사이즈/downsampling_ratio] 형태의 Latent를 만들어냅니다.
• 즉, 샘플사이즈 개수만큼 Latent 채널길이의 샘플링 벡터(Latent Vector)를 만들어내는 것이죠.

 

(실행시) output audio 차원이 [배치, 채널수=2, 샘플사이즈] 일 때,

• latents : 랜덤하게 [배치, Latent 채널수=64, 샘플사이즈/downsampling_ratio] 형태의 gaussian noise를 만들어줍니다.
• OobleckDecoder : 앞의 gaussian noise를 latent로 생각하고 [배치, 채널수=2, 샘플사이즈] 형태의 output을 만들어줍니다.

 

위 AudioAutoencoder의 아래 두가지 옵션을 좀 알아둘 필요가 있습니다.

• "chunked" : Encoder-Decoder는 1D Conv Layer로 구성되어 있어, 샘플사이즈의 길이에 상관없이 동작합니다. 따라서 샘플 사이즈가 크면 여러번 동작하기 때문에 메모리를 너무 많이 먹을 수 있기 때문에 "샘플 사이즈" 차원으로 chunk를 만들어 iterate하며 동작하는 방식입니다.
• "iterate_batch" : 배치가 큰 경우에도 메모리를 많이 먹을 수 있기 떄문에 "배치" 차원으로 iterate하며 동작하는 방식입니다.

 

또한 VAEBottleneck 내의 아래 두가지가 그림에 표현이 안되어 있어 따로 설명합니다

• softplus : 기존의 ReLU 함수의 smooth approximation버전으로, output이 항상 양수가 될 수 있도록 해줍니다. 이 함수와 그림은 아래와 같습니다.
  ** https://pytorch.org/docs/stable/generated/torch.nn.Softplus.html#torch.nn.Softplus
  $$Softplux(x)=\frac{1}{\beta}*log(1+exp(\beta*x))$$

[softplus 함수]

 

• Kullback-Leibler (KL) divergence : KL은 VAE 기반 모델에서 Latent 공간이 표준 정규분포처럼 되도록 유도하기 위해 사용하는 Loss입니다. mean과 variance를 각각 활용하며 아래 코드와 식처럼 처럼 구해집니다. 
  ** https://github.com/Stability-AI/stable-audio-tools/blob/4fdc25f4f2cd84482af58ecee37959ece55c4a0a/stable_audio_tools/models/bottleneck.py#L57C1-L65C27
  $$\mathrm{KL}(q(z|x)\|p(z))=\frac{1}{2}\sum_i(\mu^2_i+\sigma^2_i-log\sigma^2_i-1)$$

def vae_sample(mean, scale):
    stdev = nn.functional.softplus(scale) + 1e-4
    var = stdev * stdev
    logvar = torch.log(var)
    latents = torch.randn_like(mean) * stdev + mean

    kl = (mean * mean + var - logvar - 1).sum(1).mean()

    return latents, kl

** KL Loss에 대한 추가정보가 궁금하시면 참조하세요

-----------------------------------------------------------------

<VAE의 KL Loss>

 

VAE에서 KL Loss는 Latent 공간이 표준 정규분포처럼 되도록 유도하기 위한 Loss입니다.

 

위 Loss가 아닌, 일반적인 KL(Kullback-Leibler) Divergence Loss는 아래와 같이 생겼습니다.

$$\mathrm{KL}\left( q(z) \,\|\, p(z) \right) = \int q(z) \log \left( \frac{q(z)}{p(z)} \right) \, dz$$

• $q(z)$: 실제 분포 (posterior or approximate distribution)
• $p(z) $: 목표 분포 (prior)
• $\log \left( \frac{q(z)}{p(z)} \right)$: 두 분포의 차이를 측정하는 로그 비율

 

여기에 일반화된 Gaussian vs Gaussian의 KL divergence를 적용하기 위해 $q(z)=\mathcal{N}(\mu,\Sigma)$, $p(z)=\mathcal{N}(\mu_0, \Sigma_0)$를 대입하면 아래와 같을 것입니다.

$$\mathrm{KL}\left( \mathcal{N}(\mu, \Sigma) \,\|\, \mathcal{N}(\mu_0, \Sigma_0) \right)
= \frac{1}{2} \left[ \log \frac{|\Sigma_0|}{|\Sigma|} - d + \mathrm{Tr}(\Sigma_0^{-1} \Sigma) + (\mu_0 - \mu)^\top \Sigma_0^{-1} (\mu_0 - \mu) \right]$$

• $d$: 차원수
• $\Sigma, \Sigma_0$: 공분산행렬
• $|\cdot|$ : 행렬식
• $Tr$ : 대각합 (Trace)

 

다음으로 위에서 표준 정규 분포로 만들어주기 위해 $\mu_0 = 0, \Sigma_0 = I$를 넣으면 아래와 같은 수식이 되고, 이는 실제 KL Loss로 활용되는 수식입니다.

$$\mathrm{KL}\left(q(z|x) \,\|\, \mathcal{N}(0, I)\right)
= \frac{1}{2} \sum_{i=1}^{d} \left( \mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1 \right)$$

• $q(z|x)$: 인코더가 추정한 posterior 분포
• $\mathcal{N}(0, I)$ : 표준 정규 분포 (평균 0, 단위 분산)
• $\mu_i^2,\sigma_i^2$ : 인코더가 출력한 평균과 분산
• $d$ : latent vector의 차원 수

 

근데 VAE에서는 정확히 어디의 분포를 표준정규분포와 비슷하게 만든다는 것일까요? 아래 코드는 실제 KL Loss를 구하는 식입니다.

kl = (mean * mean + var - logvar - 1).sum(1).mean()

 

각 상황에 대해 코드 상의 dimension을 살펴보면 아래와 같습니다.

 

1. Time-series (audio 등) 입력일 경우의 예

• input shape: [Batch, 2, Samplesize]
• latent shape: [Batch, C=64, T=Samplesize/D]
• KL loss : [Batch, T]

즉, 각 시간 프레임 t 별로 latent vector([C], 64차원)이 표준 정규분포에 가까워지도록 학습됩니다. 

 

2. Image 입력일 경우의 예

• input shape: [Batch, 3, H, W]
• latent shape: [Batch, C=64, H', W'] (downsampling된 spatial map)
• KL loss는 [Batch, H' × W']

즉, 각 위치 (pixel) 별로 latent vector([C], 64차원)이 표준 정규분포에 가까워지도록 학습됩니다.

  

위 두가지 예시 중 1. Time series일 경우를 예로 코드에서 다시 살펴보면 

• (mean * mean + var - logvar - 1) : [Batch, C, T]
• (mean * mean + var - logvar - 1).sum(1) : [Batch, T] → KL Loss의 정의
• (mean * mean + var - logvar - 1).sum(1).mean() : 모든 배치와 time series에 대해 KL Loss의 평균 값 적용

kl = (mean * mean + var - logvar - 1).sum(1).mean()

 

정리하면 개별 위치에서의 "latent vector"의 분포가 각각 $\mathcal{N}(0, I)$를 따르도록 유도되고, 그 결과를 평균 또는 합산하여 전체 loss를 계산합니다.

-----------------------------------------------------------------

-----------------------------------------------------------------

<VAE의 default bottleneck>

 

특이한 것은 위 default 셋팅에서 활용하는 "vae" Bottleneck은 Continuous Latent를 활용하는 VAEBottleneck입니다.

 

근데 실제 논문에서는 RVQ(Residual VQ)를 활용한 VQ-GAN구조를 활용했다고 하고 있는데, 이는 Discrete(Quantized) Latent를 활용하는 RVQBottleneck이라고 따로 있습니다.

 

정확한 원인은 알 수 없지만, VAEBottleneck이 더욱 학습이 안정적이고, 품질이 일정하기 때문에 활용한 것으로 보입니다.

 

따라서 기존의 posterior collapse를 해결하려는 아래와 같은 시도가 되어있지 않음을 알 수 있습니다.

• RVQ(Residual VQ)
• KL annealing
• Skip-connection to encoder to decoder
• beta-vae
• Free-Bits
• info VAE의 mutual information 보존 목적 loss

 

참고로 Stable Diffusion에서도 VQModel이 아닌 AutoencoderKL을 활용합니다.

from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

# 1. Load the autoencoder model which will be used to decode the latents into image space. 
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

vae.to('cuda')

# 2. Load the tokenizer and text encoder to tokenize and encode the text. 
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")

text_encoder.to('cuda')

# 3. The UNet model for generating the latents.
unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet")

unet.to('cuda')

 

즉, Latent space를 구성할 때 아래와 같은 두가지 방법이 있는데, 두 방법 모두 KL-regularized VAE를 활용하는 것입니다.

• VQ-style Quantization : codebook을 사용
• KL-regularized Reparameterization : continuous Gaussian Latent를 활용 

-----------------------------------------------------------------

---

2. Attention 차원 분석

 

Attention은 일반적인 형태이지만, 내부적으로 config를 복잡하게 활용하기 때문에 그림으로 그려봤습니다.

** https://github.com/Stability-AI/stable-audio-tools/blob/4fdc25f4f2cd84482af58ecee37959ece55c4a0a/stable_audio_tools/models/transformer.py#L290

 

내부는 아래와 같이 동작합니다. 이 중 config로 제공되는 파라미터는 파란색으로 표시되어 있으며, 아래와 같습니다. 

• "sample_size" : input 오디오가 가지고 있는 샘플사이즈를 의미합니다. 이때 샘플 사이즈는 "sample_rate * seconds"로 구해집니다.
  ** bitrate : 음원의 품질을 이야기할 때 bitrate로 이야기하기도 하는데, 이는 "sample_rate * audio_channels * bits per sample(16bits)"로 구해냅니다.
• "model" > "pretransform" > "downsampling_ratio" : input 오디오가 위 설명한 pretransformer를 통해 downsample되는 크기 입니다.
• "model">"diffusion">"config">"cond_token_dim" : conditioner에서 토큰의 길이로 맞춰서 들어오는 길이이며, 따라서 "model" > "conditioning" > "cond_dim"과 같습니다.
• "model">"diffusion">"config">"embed_dim" : embedding의 사이즈로 선택된 값입니다. 이는 아래의 num_heads로 나눠져야할 것 같습니다.
• "model">"diffusion">"config">"num_heads" : 위 embed_dim과 num_heads를 통해 아래 그림에서와 같이 dim_heads의 크기가 결정됩니다.

[Attention 내의 차원 설명]

 

---

3. 활용되는 모든 Masking

 

Stable Audio Tools 내부에는 다양한 Masking들이 존재합니다.

---

 

A. cfg_dropout_prob (training default=0.1, inference default=0.0)

 

▶ 설명 : cfg를 위해 dropout 학습을 하기 위해 condition쪽에 masking을 합니다.

▶ Config : "training" > "cfg_dropout_prob"

▶ Mask 제작 : 모델 forward()에서 만들어집니다. (training/diffusion.py > models/dit.py)

▷▷ 학습시 : cfg_dropout_prob가 10% 확률로 Mask를 만듭니다.

모두 cfg_dropout_prob 값을 가지는 [batch,1,1]의 형태로 만든 뒤,

a = torch.full((cross_attn_cond.shape[0], 1, 1), cfg_dropout_prob, device=cross_attn_cond.device)

[batch,1,1] 형태 그대로 1또는 0을 cfg_dropout_prob의 확률로 Mask를 만들어줍니다.

torch.bernoulli(a).to(torch.bool)

▷▷ Inference시 : 위에 나올 수 있는 것 처럼 default cfg_dropout_prob가 0.0이기 때문에 Mask를 만들지 않습니다.

 

▶ Mask 적용 : 모델 forward()에서 사용됩니다.

이후에 이 Mask는 cross_attn_cond와 prepend_cond의 batch 단위에 적용됩니다.

Mask 값이 1인 경우는 해당 batch를 모두 0.0으로 만들 것이고, Mask 값이 0인 경우는 해당 batch 값을 그대로 사용합니다.

 

---

 

B. padding_mask(default = False)

 

▶ 설명 : 오디오에서 비어있는 부분에 대한 정보를 주기 위한 mask입니다. 보통 학습에 활용되는데, default로는 사용하지 않습니다.

▶ Config : "training" > "mask_padding" 

▶ Mask 제작 : 데이터로더에서 만들어집니다. (data/dataset.py)

직접 데이터를 로드하는 과정에서 PadCrop_Normalized_T()함수로 만들어 "padding_mask"라는 값으로 저장합니다.

실제 전체 sample size중에 오디오가 있는 곳은 1, 없는곳 0인 Mask를 만들어줍니다.

 

▶ Mask 적용 : 학습 Wrapper인 pl.LightningModule에서 사용되거나, 모델안에서 사용됩니다.

▷▷  사용처 1. Loss로 활용 “padding_mask”가 mask_key로 활용되어 사용됩니다. (training/losses/losses.py)

학습시에 위 옵션이 켜져있는 경우, Loss 중에 mask가 1인 부분만 Loss로 활용하게 됩니다. 

▷▷  사용처2. 모델안에서 mask라는 파라미터로 들어옵니다. (models/dit.py)

안쓰입니다. 이유는 정확히 모르겠지만 사용하려고 해보니, flash_attn이 지원이 안됩니다.

따라서 flash_atten을 활용하기 위해 사용하지 않는 것 같습니다.

 

 

---

 

C. Conditioning Mask들

 

▶ 설명 : Condition을 제공하기 위한 다양한 mask들입니다. 

▶ Config : X

▶ Mask 제작 : 모델 Wrapper에 존재하는 model.conditioner() & get_conditioning_input() 함수로 만들어집니다. (models/conditioners.py)

conditioner들은 mask를 모두 return하게 되어있는데, 아래와 같이 사용되는 condition의 종류에 따라서 mask를 사용해서 넘기기도하고, 사용하지 않고 버리기도 합니다.

▷▷  cross_attn : mask 얻음 → cross_attn_mask

▷▷  global_cond : X

▷▷  input_concat : X

▷▷  prepend_cond : mask 얻음 → prepend_attn_mask

 

▶ Mask 적용 : 모델 forward()에서 사용됩니다.

▷▷  적용1 (global_cond, input_concat) : Mask가 없습니다.

▷▷  적용2(prepend_cond) : prepend_cond_mask로 활용됩니다. (models/dit.py)
안쓰입니다. 이유는 정확히 모르겠지만 사용하려고 해보니, flash_attn이 지원이 안됩니다.

따라서 flash_atten을 활용하기 위해 사용하지 않는 것 같습니다.

▷▷  적용3 (cross_attn) : cross_attn_cond_mask로 활용됩니다. (models/dit.py)

현재 코드 내에서 사용하지 않도록 항상 Mask를 None으로 바꿔주고, 코드 상에도 아래와 같이 기록되어있습니다.

Temporarily disabling conditioning masks due to kernel issue for flash attention

즉, 사용하지 않습니다.

 

빨간색으로 된 부분은 flash attention을 사용하기 위해 Disabled된 부분이 있습니다. 따라서 flash attention을 적극적으로 사용하는 것이 좋을 것 같습니다.

** 참고로 flash_attention은 cuda 12.0혹은 12.4이상이어야 사용할 수 있습니다.

 

---