DC-AE : AutoEncoder used at SANA

DEEP COMPRESSION AUTOENCODER FOR EFFICIENT HIGH-RESOLUTION DIFFUSION MODELS 논문 리뷰

김호진

Jun 08, 2025

DEEP COMPRESSION AUTOENCODER FOR EFFICIENT HIGH-RESOLUTION DIFFUSION MODELS

Nvidia image generation model인 SANA에 사용된 오토인코더이다.

Sana에서 줄어든 속도를 강조했지만 70%가 이 오토인코더의 압축력에 의해서 구현된 만큼 궁금했다.

We address this challenge by introducing two key techniques 1) Residual Autoencoding, 2) Decoupled High-Resolution Adaptation

기존에 많이 쓰는 SD-VAE는 width/height를 8x 낮췄다.

diffusion transformer model 단에서 ae를 새로 학습하긴 어려우니, patchify를 두고 내부적으로 작게 만들어서 처리하는 방법들도 썼다.

하지만 이걸로는 부족하다.

압축이 더 가능하다면 당연히 training, inference 둘다 매우 속도가 개선된다. 17.9x, 19.1x 빨라지고 퀄리티는 3.55 → 3.01로 좋아졌다.

Encodec에는 LSTM 모듈이 들어있다. Mimi에는 transformer.

Sound Stream, Dac, Stable Audio Open은 Fully convolution이다.

이전의 연구들과 다르게(latent dim을 높여서 recon 퀄리티를 높이려는), 본 연구에서는 공간 압축률을 높이려고 시도한다.

few-step samplers, distillation, quantization, leveraging sparsity, etc, data quality(PixelArt-a)

왜 기존의 오토인코더는 8보다 높은 압축률에서 퀄리티가 심하게 낮아졌을까?

이걸 위해서 f8, f16, f64로 실험을 해봤다.

f64는 우선 8배 압축 후 추가 압축 태스크로 보고, f8로 학습한 모듈 위에 레이어를 추가하는 식으로 구현했다. 그리고 spatial이 압축된 만큼 latent dim을 늘렸다.(= p^2, width/height가 각각 p만큼 압축되니까)

→ 공간 압축률을 높일수록 결국 latent dim이 늘어나니 총 압축률은 일정하고 대신 모델이 커지는데, 그래도 퀄리티가 낮아졌다(동일하지도 않음)

즉, 문제는 공간 압축률이 높을수록 학습이 어려워지기 때문이다. 모델의 capacity 문제라고 하기에는 동일한 압축률에 모델은 더 키웠으니까

우선 전체 Detail architecture

Residual

VQ를 쓴다는 뜻이 아니었구나?

원래 ResNet같은 곳에서의 residual은 input x + module output z를 더하는 식이다(x는 그대로 identity하게 더함)

근데 여기서는 한 Downsample block마다 channel은 2배, 해상도는 각 1/2이 되는 만큼 거기에 맞춰서 space-to-channel 매핑, channel averaging을 통한 연산 후 더해준다.

이건 학습하지 않는(non-parametric) 연산이다.

Decoupled High-Resolution Adaptation

high resolution이 될수록 어려워지는 부분이 있는데, 이게 주로 뭔가 지각적인 문제 같음 내 생각엔. 인간이 특정 resolution 이하에서는 차이를 크게 못느끼지만 어느 이후로는 크게 느껴서?

고해상도 이미지 학습은 비용이 크기 때문에 보통 오토인코더 자체는 저해상도로 학습하고, 그걸 그대로 고해상도에 쓰는 경우도 많다. 왜냐하면 어짜피 Fully convolution으로.. 1024를 256 16개로 볼 수 있으니?

저해상도 이미지의 압축률을 높이는건 residual만으로 좋았다. 그런데 동일하게 해도 high resolution의 압축률이 커졌을 때는 결과가 튀게 안좋았다. 그럼 고해상도로 학습해야할텐데 그건 1) 비싸고 2) GAN 최적화가 unstable하다.

generalization penalty?

따라서 기존의 VAE 학습 방법(SD-VAE)과 비교하여 두가지 변화를 주었다.

GAN 손실 학습을 전체 모델 학습에서 분리하여, 이를 위한 별도의 "국소 정제 단계(local refinement phase)"를 도입한다 (Figure 6의 3단계).

이 단계에서는 디코더의 head layer만 조정하고, 나머지 레이어는 모두 고정시킨다. 이 설계는 다음과 같은 인사이트에 기반한다:

재구성 손실(reconstruction loss)만으로도 콘텐츠와 의미 복원에는 충분하다.

반면, GAN 손실은 지역적 디테일 개선과 아티팩트 제거에 주로 기여한다 (Figure 5).

따라서 전체 모델을 학습하는 것보다, 디코더의 상위 레이어만 조정하는 것이 더 효율적이고 정확도도 더 높다.

또한, 이 방식은 GAN 손실 학습이 잠재 공간(latent space)을 변형시키는 것을 방지하므로, 낮은 해상도 이미지에서도 안정적인 학습이 가능하다. 이는 고해상도 GAN 손실 학습의 불안정성을 피하고, 3단계의 학습 비용도 줄여준다.

추가적인 "고해상도 latent 적응 단계(latent adaptation phase)"를 도입한다 (Figure 6의 2단계).

이 단계에서는 **중간 레이어(인코더의 초기 레이어 및 디코더의 입력 레이어)**만 조정하여, 잠재 공간을 고해상도에 맞게 적응시킨다.

실험 결과, 전체를 학습하지 않고 중간 레이어만 조정해도 일반화 문제를 충분히 해결할 수 있었으며 (Figure 3 b), 학습 비용도 전체 학습보다 훨씬 낮았다 (예: 메모리 비용 153.98 GB → 67.81 GB) (Cai et al., 2020).

저해상도 학습 → 고해상도에 적용할 때의 generalization 성능에 문제가 있다고 보고, 그걸 개선하는데 집중했다. 고해상도를 따로 학습하고 모델링했다기보다.