VaeZip: VAE 인코딩된 잠재 데이터셋을 품질 손실 없이 5~10배 적은 공간에 저장하는 방법

안녕하세요! 제가 모델을 훈련할 때마다 항상 비효율적이라고 느꼈던 부분이 하나 있었습니다. 바로 VAE(Variational Autoencoder) 잠재 파일들입니다. 이 파일들은 원본 데이터보다 훨씬 많은 공간을 차지하는데, VAE는 데이터를 압축해야 하는데 오히려 더 무겁게 만드는 것 같습니다. 이는 주로 JPEG나 비디오 코덱이 제 역할을 하여 데이터를 매우 잘 압축하는 반면, 잠재(latents) 값은 거의 압축되지 않았기 때문입니다. 그래서 저는 잠재를 위한 일종의 JPEG 같은 것을 만들어보고자 했습니다.

우선 말씀드리자면, 이 도구는 사용 사례가 매우 좁습니다. 즉, 잠재 데이터셋이 로컬 디스크에 들어가지 않을 만큼 클 경우에만 유용하며, 이런 경우 S3 스토리지와 같은 서비스를 사용해야 할 수 있고 이는 데이터셋이 너무 클 경우 비용이 많이 들 수 있습니다.

어쨌든 VaeZip은 PyTorch 기반 라이브러리로, 여러 잠재 양자화 코덱(latent quantization codecs)을 구현하고 압축된 텐서들을 일반적인 Safetensors 파일에 저장합니다. 이 방식은 해당 텐서를 fp32로 저장하는 것보다 훨씬 작습니다. 이 코덱들은 VAE 잠재 값에 특화되어 설계되었습니다. 일부는 공간적 중복성(spatial redundancy)과 VAE 잠재 분포가 보통 정규분포(normal distribution)를 따른다는 점을 모두 활용합니다. 이를 통해 디코딩 후 픽셀 MSE(Mean Squared Error)가 0.0003인 상태를 유지하면서도, 잠재 텐서를 약 3비트/값의 정밀도로 저장할 수 있게 되어 파일 크기가 10배 작아집니다.

예를 들어, 가장 좋은 코덱 중 하나는 전역 범위 스케일링(global range scaling), 사각형 공간 블록을 사용한 국소 범위 스케일링(local range scaling)과 소위 'Lloyd-Max 테이블'이라는 것을 조합하여 사용합니다. 이 테이블은 본질적으로 샘플 분포에 맞춰 조정된 코드북(codebook)입니다. 그 후 값들을 2~8비트로 양자화하고 패킹합니다.

이 코덱들은 재구성 품질, 충실도(fidelity), 공간 사용량 및 I/O 지연 시간을 기준으로 약 50개의 알고리즘을 8가지 이미지 및 비디오 VAE에 걸쳐 비교한 벤치마크에서 선정되었습니다. 저는 또한 다양한 코덱으로 압축된 잠재를 사용하여 작은 T2I(Text-to-Image) 플로우 매칭 DiT를 사전 학습시키는 데 미치는 영향도 측정했습니다. 한 테스트 결과에서는 단지 6비트 저장만으로 손실 곡선이 fp32 기준선과 구별할 수 없을 정도로 유사했습니다.

혹시 여기서 대규모 사전 학습이나 튜닝을 하시는 분들이 계시다면, 이 정보가 유용하기를 바랍니다! submitted by /u/TensorForger [link] [comments]