6GB VRAM에서 생성형 킥 드럼 디퓨전 모델 훈련하기

고품질 생성 오디오 모델을 훈련하기 위해 거대한 GPU 클러스터나 수십억 달러 규모의 예산이 필요하지 않습니다. 잠재 디퓨전 파이프라인과 적극적인 압축을 활용하면 6GB VRAM을 가진 7년 된 NVIDIA GeForce GTX 1660 SUPER와 같은 일반 소비자용 하드웨어에서도 실용적인 생성 모델을 훈련할 수 있습니다.

생성 파이프라인: 세 모델, 하나의 워크플로우

제한된 VRAM에서 오디오 생성을 계산 가능하게 만들기 위해 시스템은 원시 오디오를 직접 디퓨징하지 않습니다. 대신 오디오를 잠재 공간으로 압축하고, 그 공간에서 새로운 데이터를 생성한 뒤, 오디오를 복원하는 3단계 파이프라인을 사용합니다.

1. 변분 오토인코더 (VAE)

VAE는 오디오 표현의 압축 및 복원을 담당합니다. 로그‑멜 스펙트로그램(128×173)을 받아 작은 3D 잠재 텐서(4×8×11)로 압축합니다. 이는 스펙트로그램 기준 약 63배, 원시 오디오 기준 약 250배의 압축률에 해당합니다.

잠재 공간이 연속적이고 생성에 활용 가능하도록 VAE는 두 가지 메커니즘을 사용합니다:

• 평균 및 분산 출력: 인코더가 평균과 분산을 모두 출력함으로써 디코더가 단일 좌표가 아니라 점들의 이웃에서 복원하도록 강제합니다.
• KL 발산 패널티: 잠재 분포를 표준 정규 분포로 끌어당기는 패널티를 적용해 모델이 "흩어진 섬" 형태의 데이터를 만들지 않게 합니다.

2. 디퓨전 U‑Net

디퓨전 U‑Net은 핵심 생성 엔진입니다. 4×8×11 잠재 공간 내에서만 작동합니다. 모델은 가우시안 노이즈가 1,000 단계에 걸쳐 순차적으로 추가되는 과정을 역전하도록 훈련됩니다.

생성을 안내하기 위해 모델은 텍스트 조건화(파일명에서 추출한 키워드)와 **분류기‑프리 가이드(CFG)**를 사용합니다. 15% 키워드 드롭아웃으로 훈련함으로써 모델은 일반적인 킥 드럼 특성과 특정 특성(예: "punchy" 또는 "warm")을 모두 학습하고, 추론 시 특정 키워드의 영향을 증폭시킬 수 있습니다.

3. 보코더 (HiFi‑GAN)

오디오를 멜 스펙트로그램으로 변환하는 과정은 손실이 발생하고 역변환이 불가능하기 때문에, 생성된 스펙트로그램을 파형으로 되돌리기 위해 보코더가 필요합니다. 여기서는 생성자와 실제 훈련 데이터를 구분하려는 판별기로 구성된 생성적 적대 신경망(GAN)을 사용합니다. 이 적대 과정은 스펙트로그램 변환 중 손실된 미세 디테일을 채워 최종 오디오가 "실제"처럼 들리도록 합니다.

데이터 준비 및 전처리

모델은 개인 라이브러리에서 추출한 13,615개의 킥 샘플 데이터셋으로 훈련되었습니다. 다음 전처리 단계가 모델 안정성에 핵심이었습니다:

• 필터링: "loop" 또는 "BPM"이 포함된 파일을 제거하고, 5KB 이하 또는 1MB 이상인 파일을 삭제했습니다.
• 정규화: 모든 샘플을 44.1kHz로 재샘플링하고, 정확히 2초가 되도록 패딩·트리밍했으며, 피크를 -1dB로 정규화하고 0.2초 페이드‑아웃을 적용했습니다.
• 스펙트로그램 변환: 오디오를 2048 샘플 윈도우와 512 샘플 홉 길이로 로그‑멜 스펙트로그램으로 변환하여 128×173 형태로 만들었습니다.

배포 및 추론 전략

24/7 GPU 인스턴스 비용을 피하기 위해 프로젝트는 Modal이라는 서버리스 GPU 플랫폼을 사용합니다.

• 클래스 기반 워커: "콜드 부팅" 시간을 완화하기 위해 추론 엔진을 클래스에 래핑했습니다. 이를 통해 모델(약 300 MB 가중치)이 컨테이너 부팅 시 한 번 GPU 메모리에 로드되고 이후 요청에서 재사용됩니다.
• 성능: 워밍된 요청은 약 2.5초, 콜드 부팅은 약 10초가 소요됩니다.
• 캐싱: 모델 가중치는 HuggingFace에 호스팅되고 Modal 머신에 캐시되어 반복 다운로드를 방지합니다.

기술적 트레이드오프와 아티팩트

제한된 하드웨어에서 훈련하면 일부 타협이 필요합니다. 저자는 특히 강한 압축(OTT 사용) 시 킥이 연속 파형이 아니라 작은 입자들로 이루어진 듯 들리는 "granularization" 아티팩트를 발견했습니다. 이는 잠재 공간의 높은 압축 비율(63배)과 보코더의 업샘플링 과정 때문입니다.

모델 품질 향상을 위한 향후 개선점

아티팩트를 줄이고 조건화를 개선하기 위해 다음 하이퍼파라미터를 조정할 것을 제안합니다:

• 잠재 크기: 압축 손실을 줄이기 위해 잠재 텐서 크기 확대.
• FFT 윈도우: 2048‑샘플 윈도우를 줄여 트랜지언트를 선명하게.
• 키워드 드롭아웃: 15% 드롭아웃 비율을 조정해 텍스트 프롬프트에 대한 반응성 향상.
• 어휘: 전자음악에서 중요한 음향 의미를 가진 "808"·"909"와 같은 숫자 식별자를 유지.

핵심 내용 요약

구성 요소	전략	이점
하드웨어	GTX 1660 SUPER (6GB VRAM)	소비자 하드웨어가 특화된 생성 작업에 충분함을 증명
표현	Log‑Mel Spectrogram → Latent	압축된 공간에서 디퓨징함으로써 계산 부하 감소
아키텍처	VAE → U‑Net → HiFi‑GAN	압축, 생성, 복원을 분리
추론	서버리스 GPU (Modal)	유휴 GPU 비용을 없애고 허용 가능한 지연 시간 유지

Summary

Developer zhinit trained a latent diffusion model to generate kick drum samples using a 7-year-old NVIDIA GTX 1660 SUPER, demonstrating that high-quality generative audio can be achieved on consumer hardware through aggressive latent space compression.

Title

Training a Generative Kick Drum Diffusion Model on 6GB VRAM