[Generative] Stable Audio Tools Autoencoder학습시 기록

<분류>
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/autoencoders/stable_audio_2_0_vae.json

[Pretransformer 동작 방식]

자세히보면 chunked=False로 되어있기 때문에 "샘플 사이즈" 차원으로 chunk를 만들어 iterate하며 동작하지 않고, iterate_batch=True이기 때문에 "배치" 차원으로 iterate하며 동작합니다.

 

그리고 mean과 stdev를 도출한 뒤, latent를 랜덤하게 얻어내고 KL Loss를 계산해줍니다.

 

이를 학습할 때 사용하는 Loss들에 대해 먼저 살펴보겠습니다.

 

Loss의 종류에 대한 설명에 앞서 Stable Audio Tools에서 사용하는 "Loss의 Wrapper 종류"를 살펴보겠습니다.

• MultiLoss : 여러개의 Loss를 묶어줍니다.
• ValueLoss : 해당 값에 weight만 곱해서 결과 얻음
• AuralossLoss :해당 로스로 구한 뒤에 weight를 곱해주기
• L1Loss : L1 Loss입니다

또한 아래와 같은 미리 object를 선언한 뒤 Loss에 활용합니다.

• sdstft : SumAndDifferenceSTFTLoss(“spectral”-“config”)
• lrstft : MultiResolutionSTFTLoss(“spectral”-“config”)
• discriminator: EncodecDiscrimnator(“discriminator”-“config”) → MultiScaleSTFTDiscriminator()
  ** https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/models/discriminators.py 

 

이제 본격적으로 위 Autoencoder를 학습할 때 사용하는 Loss를 정리하면 아래와 같습니다.

Generator의 Loss와 Discriminator의 Loss로 나뉩니다.

 

1. Generator : gen_loss_modules → (MultiLoss) losses_gen

• (ValueLoss) loss_adv("discriminator"-"weights" :0.1) : "loss_adv" (= adversarial loss)
• (ValueLoss) feature_matching_loss("discriminator"-"weights": 5.0) : "feature_matching_distance" 
• (AurolossLoss) mrstft_loss(sdstft, "spectral"-"weights" :1.0) : "reals"~"decoded" (= reconstruction error)
• (AurolossLoss) stft_loss_left(lrstft, "spectral"-"weights"/2 :1.0/2) : "reals_left" ~ "decoded_left"
• (AurolossLoss) stft_loss_right(lrstft, "spectral"-"weights"/2 :1.0/2) : "reals_right" ~ "decoded_right"
• (L1Loss) l1_time_loss("time"-"weights" :0.0) : "reals" ~ "decoded"
• +create_loss_modules_from_bottleneck(VAEBottleneck())  
  
  • (ValueLoss) kl_loss("bottleneck"-“weights” :1e-4) : "kl" (=bottleneck에서 구한 vae의 KL Loss)

2. Discriminator : disc_loss_modules → (MultiLoss) losses_disc

• (ValueLoss) discriminator_loss(weight=1.0) : "loss_dis" (=discriminator loss)

** 이외에 로깅되는 항목

• train/disc_lr : discriminator의 learning rate
• train/gen_lr : generator의 learning rate
• train/data_std : data 분포의 표준편차
• train/latent_std : latent 분포의 표준편차
• train/loss : 전체 loss (losses_gen 혹은 losses_disc)

 

근데 학습 중 Discriminator와 연관된 세개의 Loss가 아래와 같은 형태를 보입니다 : loss_adv, feature_matching_loss, discriminator_loss

[문제의 실험 현상]

 

이에 대한 원인과 해결책을 살펴보려고 합니다.

---

2. Why? (원인)

• X

---

3. How? (해결책)

 

위에서 살펴본 세개의 Loss가 무엇인지 자세하게 먼저 살펴보겠습니다.

설명에 앞서 아래의 real과 fake는 아시겠지만 각각 아래를 의미합니다.

• real 샘플($x$): 실제 학습 데이터가 Discriminator를 통과한 결과를 의미합니다. 
  • feature_real : Discriminator를 통과한 실제 feature입니다. 
  • score_real : Discriminator를 통과한 feature가 정답일 확률입니다. 클수록 Discriminator가 정답이라고 생각하는 것이며, 작을수록 정답이 아니라고 생각합니다.
• fake 샘플($D(G(x))$): Generator를 활용해 생성한 샘플이 Discriminator를 통과한 결과를 의미합니다.
  • Generator가 생성한 샘플은 실제로 fake이므로, Discriminator가 fake라고 인식해야하지만 Generator가 학습을 잘하면 real이라고 Discriminator를 속이게 될 것입니다.
  • feature_fake : 위와 같습니다.
  • score_fake : 위와 같습니다.

 

그럼 문제의 세개의 Loss를 살펴보겠습니다. 기존의 Minimax Loss와는 다릅니다.

• 1. Adversarial Loss (loss_adv): Generator의 Loss를 의미하며, Generator가 Discriminator를 잘 속일 수 있도록 유도하는 값을 사용합니다.
  ** Generator 학습 step에 Generator 파라미터에만 활용
  
  • $D(G(z))$가 크게 나오도록 학습 : $-\mathbb{E}[D(G(z))]$

adv_loss = -score_fake.mean()

• 2. Feature Matching Distance Loss : real샘플과 fake샘플이 비슷해지도록 만들어주는 추가적인 Loss입니다.
  ** Generator 학습 step에 Generator 파라미터에만 활용

feature_matching_distance = abs(feature_real - feature_fake).mean()

• 3. Discriminator Loss (Minimax Loss) : Discriminator의 Loss를 의미하며, Hinge Loss를 활용했습니다.
  ** Discriminator 학습 step에 Discriminator 파라미터에만 활용
  
  • 아래 식을 통하면 Discriminator의 score가 real은 1보다 크게 fake는 -1보다 작게 만드는 것이 목표입니다.
  • score_real : $D(x)\geq1$이 되도록 학습 : $max(0,1-D(x))$
  • score_fake : $D(G(z))\leq-1$이 되도록 학습 : $max(0,1+D(G(z))$

dis_loss = torch.relu(1 - score_real).mean() + torch.relu(1 + score_fake).mean()

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

<실제 Default Discriminator 구현>

 

2025.08.28 기준 실제 default 셋팅에서 활용되는 Discriminator의 구현은 아래와 같습니다.

class EncodecDiscriminator(nn.Module):
    def __init__(self, normalize_losses=False, loss_type: tp.Literal["hinge", "rpgan"]="hinge", *args, **kwargs):
        super().__init__()
        from .encodec import MultiScaleSTFTDiscriminator
        self.discriminators = MultiScaleSTFTDiscriminator(*args, **kwargs)
        self.normalize_losses = normalize_losses
        self.fm_reduction = (lambda x, y: abs(x - y).mean()/(abs(x).mean() + 1e-3)) if normalize_losses else (lambda x, y: abs(x - y).mean())
        self.loss_type = loss_type

    def forward(self, x):
        logits, features = self.discriminators(x)
        return logits, features

    def loss(self, reals, fakes):
        feature_matching_distance = torch.tensor(0., device=reals.device)
        dis_loss = torch.tensor(0., device=reals.device)
        adv_loss = torch.tensor(0., device=reals.device)

        logits_true, feature_true = self.forward(reals)
        logits_fake, feature_fake = self.forward(fakes)

        # Compute per-scale losses
        for i, (scale_true, scale_fake) in enumerate(zip(feature_true, feature_fake)):
            feature_matching_distance = feature_matching_distance + sum(
                map(
                    self.fm_reduction,
                    scale_true,
                    scale_fake,
                )) / len(scale_true)

            if self.loss_type == "hinge":
                _dis, _adv = get_hinge_losses(logits_true[i], logits_fake[i])
            else:  # rpgan
                _dis, _adv = get_relativistic_losses(logits_true[i], logits_fake[i])

            dis_loss = dis_loss + _dis 
            adv_loss = adv_loss + _adv

        num_scales = len(logits_true)

        return dis_loss / num_scales, adv_loss / num_scales, feature_matching_distance / num_scales

 

위에서 1. Adversarial Loss와 3. Discriminator Loss를 구현하기 위한 함수는 아래와 같습니다.

def get_hinge_losses(score_real, score_fake):
    gen_loss = -score_fake.mean()
    dis_loss = torch.relu(1 - score_real).mean() + torch.relu(1 + score_fake).mean()
    return dis_loss, gen_loss
    
def get_relativistic_losses(score_real, score_fake):
    # Compute difference between real and fake scores
    diff = score_real - score_fake
    dis_loss = F.softplus(-diff).mean()
    gen_loss = F.softplus(diff).mean()
    return dis_loss, gen_loss

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

 

어떤 Loss들이 있는지 먼저 살펴보았는데, 미리 알아야하는 조건들을 먼저 살펴보겠습니다.

1. 우리는 VAE를 학습하는 것이므로, Generator의 학습이 우선적이며 Discriminator가 완벽하게 학습되는 것은 부차적인 목표입니다.
2. G-D학습의 경우 Loss를 두가지 방법으로 해석해야하며, 보통 절대적인 값으로 원인론적 해석이 가능하고 상대적인 값의 변화로 원인론적 & 결과론적 해석이 모두 가능합니다.
   위가 어려운 이유는, 상대적인 값의 변화는 "Loss를 높여 반대 효과를 억압하는 중"처럼 원인론적으로 해석이 가능하면서도 "Loss가 원하는 방향으로 학습되지 않는 중"처럼 두가지로 해석되기 때문입니다.
   ** 원인론적 해석 : Loss가 높으면 반대되는 효과를 누르기 위한 Loss가 강압적으로 동작하고 있다는 뜻이고, 낮으면 별로 힘이 없다.
   ** 결과론적 해석 : 원하는 방향과 학습되고 있지 않을때 Loss가 높아지고, 원하는 방향일 때 Loss가 낮아진다.
3. 일반적으로 Discriminator 학습이 어렵고, Generator가 Discriminator보다 더 강한 경우가 많습니다.
   따라서 1. Adversarial Loss가 3. Discriminator Loss보다 절대적인 값이 작아지는 것을 확인하는 것도 중요합니다.
4. 보통 Discriminator에서 나오는 Loss중 Generator를 위한 1. Adversarial Loss와 2. Feature Matching Distance 중에 Feature Matching Distance가 먼저 가까워지도록 학습이 되어야, Adversarial Loss가 따라옵니다.

 

이제, 위와 같은 결과는 단순히 Loss만 바라봤을때, 아래와 같은 현상을 의미합니다.

• 1. Adversarial Loss  증가 (Bad) : Generator가 Discriminator를 속이지 못하고 있습니다.
• 2. Feature Matching Distance 증가 (Bad) : Generator가 데이터와 점점 덜 비슷한 Discriminator의 feature를 만들고 있습니다.
• 3. Discriminator Loss 감소 (Good): Discriminator는 잘 학습되고 있습니다.

근데, 이는 Discriminator의 Loss 외 Generator의 Reconstruction Loss의 형태에 따라 다르게 해석이 가능합니다.

• CaseA. 만약 Generator는 수렴하고 있다면...
  • Generator가 Discriminator과 정상적인 상호작용으로 하며 학습하고 있지만, Discriminator로 오는 1. Adversarial Loss와 2. Feature Matching Distance를 Generator가 감당하기에는 너무 벅찹니다.
• CaseB. 만약 Generator가 발산하고 있다면...
  • Discriminator가 학습이 잘되고 있는 것처럼 보이지만, Generator와 상호작용할 수 있는 방향으로 학습이 되고 있지 않다고 해석 할 수 있습니다.

이 해석 외에도 다양한 해석이 가능하겠지만, 제가 실험한 결과는 위와 같았습니다.

 

그럼 Generator와 Discriminator 각각의 학습 속도를 조절할 수 있는 방법들을 살펴보겠습니다.

 

참고로 이와 같은 Autoencoder 형태의 학습 방법에서는 Loss의 변화만으로 명확하게 문제를 정의하기 하기는 사실 어렵습니다.

따라서 아래와 같은 방법으로 실제 학습이 잘되는지를 확인하면서 검토하는 것도 좋지만 이 글에서는 Loss의 변화로만 살펴봤습니다.

• Validation Metric을 활용
• Validation Demo Sample을 들으면서 검토

 

 

해결책 : Generator 다루기

 

 

방법1. Generator의 learning rate 조정

 

autoencoder의 learning_rate를 조정해 Generator의 학습속도를 조절할 수 있습니다.

"training": {
    "optimizer_configs": {
        "autoencoder": {
            "optimizer": {
                "type": "AdamW",
                "config": {
                    "betas": [0.8, 0.99],
                    "lr": 1.5e-4,
                    "weight_decay": 1e-3
                }
            },
        },
    },
}

 

 

방법2. Adversarial Loss나 Feature Matching Distance Loss의 가중치 조정

 

Generator에 영향을 주는 두개의 Loss 가중치인 weight를 조정해 얼마나 강하게 학습할지를 조절해줄 수 있습니다.

아래 보이는 config의 "discriminator"관련된 weights를 보면 feature_matching과 adversarial이 존재하는데 이를 올려서 조절합니다.

"training": {
    "loss_configs": {
        "discriminator": {
            "weights": {
                "adversarial": 0.1,
                "feature_matching": 5.0
            }
        },
        "bottleneck": {
            "type": "kl",
            "weights": {
                "kl": 1e-4
            }
        }
    },
}

 

 

방법3. Feature Matching Distance 방법 수정

 

Feature Matching Loss은 Generator의 학습을 도와 Discriminator를 빠르게 따라잡도록 만들기 위해, real의 Discriminator feature결과와 Generator의 결과가 비슷해지도록 도와주는 Loss입니다.

 

이 때, 기존 Stable Audio Tools에서는 Feature Matching Distance를 단순히 L1 Distance로 구하고 있습니다.

lambda x, y: abs(x - y).mean()

 

하지만 GAN에서는 Feature Matching을 강조하기 위해 L1 Distance가 아닌 다른 방법을 활용하거나, Perceptual Loss와 같은 Loss들을 추가해 안정성을 추구합니다.

** Discriminator에서 사용되었던 Loss들의 종류를 보려면 더보기를 참조하세요

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

<GAN에서 사용하는 추가적인 Loss의 종류>

 

1. Feature Matching Loss : L1 Distance, MAE (Mean Absolute Error)

abs(x - y).mean()

L1은 L2보다 노이즈에 덜 민감하고, 선명한 차이를 잡아내는 데 유리합니다.

 

이는 Sparse한 차이에 더욱 민감하기 때문에 noise에 강한것이며, 일반적으로는 안정적이고 효과적입니다. 따라서 많은 GAN에서 기본적으로 사용합니다.

 

2. Feature Matching Loss : L2 Distance, MSE(Mean Sqaure Error)

((x - y) ** 2).mean()

 

큰 차이를 더욱 많이 반영하기 위한 방법이며, 전체적으로 부드럽게 조정이 가능합니다.

 

따라서 더 부드러운 reconstruction을 유도할 수 있습니다.

 

3. Distribution loss: KL Divergence, JS Divergence, Wasserstein Distance 등

 

학습하는데 어려움은 있지만 fine-grained한 distribution의 차이를 줄이는 것에 유리합니다. 

 

예를 들어 StyleGAN과 같은 GAN에서는 아래 소개할 Perceptual이나 Distribution 기반의 loss를 같이 쓰기도 합니다.

 

4. Perceptual Loss

 

VGG 같은 pretrained 모델의 중간 feature들의 차이를 반영하는 방법입니다.

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

 

 

방법4. 일정 학습 이후에 encoder를 freeze하기 

 

VAE와 같은 모델을 학습할 때 Discriminator와 Generator를 모두 학습시키기 어려울 수도 있습니다. 

 

따라서 어느정도 학습이 되고 있는 상황이라면, 이후에는 Generator의 Encoder는 freeze하고 학습을 할 수 있습니다. Stable Audio Tools에는 encoder_freeze_on_warmup라는 옵션이 있는데 아래와 같이 진행됩니다.

if self.warmed_up and self.encoder_freeze_on_warmup:
    with torch.no_grad():
        latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True)
else:
    latents, encoder_info = self.autoencoder.encode(encoder_input, return_info=True)

 

또한 보통 Encoder를 freeze하고 Decoder만 학습할 때는, 다른 Loss로 학습하게 되니 멈추고 학습하는 것도 방법입니다. 

** Decoder 학습시에는 Empty Audio를 배제하고 학습하지 않으면 normalize에서 폭주하는 경우가 있습니다. 따라서 아래와 같이 Empty Audio를 배제하는 것을 추천합니다.

import random

if not torch.any(audio):
	return self[random.randrange(len(self))]

** Dataset.__getitem__ 안에서 위와 같이 return하면 다른 무작위 샘플 하나를 대신 반환하고, 다시 해당 샘플로 Dataset.__getitem__를 부릅니다. 

 

방법5. Gradient Clipping활용하기

 

Generator가 감당가능한 Loss만큼만 학습하도록 Gradient를 Clipping할 수 있습니다.

 

Stable Audio Tools에는 clip_grad_norm라는 옵션으로 적용할 수 있는데, 아래와 같이 적용됩니다.

if self.clip_grad_norm > 0.0:
    torch.nn.utils.clip_grad_norm_(self.autoencoder.parameters(), self.clip_grad_norm)

 

해결책 : Discriminator 다루기

 

방법1. Discriminator의 learning rate 조정

 

Discriminator의 learning_rate를 조정해 Discriminator의 학습속도를 조절할 수 있습니다.

"training": {
    "optimizer_configs": {
        "discriminator": {
            "optimizer": {
                "type": "AdamW",
                "config": {
                    "betas": [0.8, 0.99],
                    "lr": 3e-4,
                    "weight_decay": 1e-3
                }
            },
        }
    },
}

 

 

방법2. Discriminator 학습 빈도 조정

 

보통 학습시에 step%2로 나누어 Generator-Discriminator를 1:1로 번갈아가면서 학습합니다.

 

하지만, 예를 들어 학습 빈도 Generator:Discriminator를 2:1 등으로 바꿔주면 Generator를 더욱더 학습하는 방향으로 학습할 수 있습니다.

 

 

방법3. Gradient Penalty 또는 다른 regularization 기법 도입

 

Gradient Penalty 또는 다른 regularization 기법 도입해 Discriminator가 너무 강해지는 것 방지할 수 있습니다.

** Gradient Penalty : WGAN과 같은 생성모델은 weight clipping을 사용해 Lipschitz 제약을 적용했는데, 이는 gradient가 소실하거나 폭발하는 문제를 초래합니다. 따라서 Discriminator의 gradient norm을 1 근처로 유지해 Lipschitz 제약을 만족시키는 Regularization Term입니다.

 

예를 들어 feature의 길이만큼의 normalization하는 간단한 방법을 활용할 수도 있습니다.

 

방법4. Discriminator Warmup

 

Discriminator의 Loss를 Generator가 감당하지 못한다면, Generator를 먼저 어느정도 학습한 후에 Discriminator와 함께 학습하는 warmup을 적용할 수 있습니다.

 

Stable Audio Tools에 구현된 내용은 아래와 같은데, 두가지 모드가 있습니다.

• full : warmup 전에 Discriminator의 Loss와 Discriminator의 학습을 전부 차단하기
• adv : warmup 전에 Discriminator의 Loss중 Generator에 영향을 주지 않는 loss_dis만을 이용해 Discriminator를 학습하기

if self.global_step >= self.warmup_steps:
    self.warmed_up = True

if self.warmed_up:
    loss_dis, loss_adv, feature_matching_distance = self.discriminator.loss(reals=reals, fakes=decoded)
else:
    loss_adv = torch.tensor(0.).to(reals)
    feature_matching_distance = torch.tensor(0.).to(reals)

    if self.warmup_mode == "adv":
        loss_dis, _, _ = self.discriminator.loss(reals=reals, fakes=decoded)
    else:
        loss_dis = torch.tensor(0.0).to(reals)

 

방법5. Discriminator 종류 변경

 

Discriminator는 학습전반을 책임진다기 보다, "Reconstruction Error만으로 극복할 수 없는 부분을 Discriminator가 인지해서 Penalty를 주는" 구조로 학습하기 위한 도구입니다.

 

위에서는 기본적으로 "encodec" 옵션으로 제공되는 EncodecDiscriminator를 활용했는데, 아래와 같은 다양한 Discriminator들이 존재합니다.

** https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/models/arc.py

** https://github.com/Stability-AI/stable-audio-tools/blob/main/stable_audio_tools/models/discriminators.py

• EncodecDiscriminator : 직접 정의한 MultiScaleSTFTDiscriminator를 통해 판별
  • MultiScaleSTFTDiscriminator : 주파수 영역으로 바꾸되 hop size와 window size를 다양하게 여러가지 스케일로 포먼트/spectral envelope, 하모닉 라인, 잡음 밴드 등을 잘 잡습니다.
• OobleckDiscriminator : MultiScaleDiscriminator와 MultiPeriodDiscriminator를 활용해 디테일하게 판별
  • MultiDiscriminator : 아래 둘을 묶어주는 역할
  • MultiScaleDiscriminator : 1D의 SharedDiscriminatorConvNet로 이루어져 짧은 구간의 미세한 디테일(transient, 고주파)부터 긴 구간의 거시 구조(폼, Dynamic)까지 동시에 판별
    ** SharedDiscriminatorConvNet : 같은 형태의 convolutional layer 여러개로 이루어진 구조
  • MultiPeriodDiscriminator : 2D로 접어서 2D의 SharedDiscriminatorConvNet로 처리하며, 서로 다른 기본주파/배음 주기에 정렬된 패턴(Pitch, Harmonic, 주기성)을 잘 잡습니다.
    ** SharedDiscriminatorConvNet : 같은 형태의 convolutional layer 여러개로 이루어진 구조
• DACGANLoss : Descript Audio Codec을 활용해서 다양한 특징들을 판별
  • DACDiscriminator :  아래 세가지를 묶어주는 역할
    
  • MSD : Descript Audio Codec에 존재하는 WNConv1d를 활용한 Multi-Scale Discriminator로, Waveform 전체의 자연스러움과 Global Structure을 확인.
    
  • MPD : Descript Audio Codec에 존재하는 WNConv2d를 활용한 Multi-Period Discriminator로, 주기적 패턴(pitch 주기, harmonic structure)을 인지
  • MRD : Descript Audio Codec에 존재하는 WNConv2d와 STFT를 활용한 Multi-Resolution Discriminator로, 주파수 영역에서 음질의 Harmonic, 포먼트, 잡음 분포 등을 인지
• BigVGANDiscriminator : 위 DACGANLoss와, MultiScaleSubbandCQTDiscriminator를 함께 활용해 판별
  
  • MultiScaleSubbandCQTDiscriminator : CQT(Constant-Q Transform)는 STFT가 선형 주파수인 것과 다르게 로그 주파수이므로, 피치·화음·하모닉 구조를 더 직접적으로 볼 수 있습니다. 이를 다양한 hop size, 옥타브 세분화 정도를 통해 다양한 스케일로 분석하여 인지합니다.
• ConvDiscriminator, ConvNeXtDiscriminator : ARC(Adversarial Relativistic-Contrastive)라는 post-training학습 기법에 쓰이는 판별기로, Diffusion Model을 학습할 때 prompt 등의 조건을 더 잘 따르도록 학습하는 방법입니다.
  ** Fast Text-to-Audio Generation with Adversarial Post-Training (arxiv'25)

---