적응형 토큰 압축을 통한 확산 모델(Diffusion Model) 비용 절감

지연 시간(Latency)은 여전히 확산 모델(Diffusion Model)이 대화형 편집 도구에서 실용적으로 사용되는 것을 가로막고 있습니다. HiLo-Token은 입력 적응형 고주파-저주파(high-low frequency) 토큰 압축을 통해 생성 품질을 유지하면서도 해당 지연 시간을 절반으로 줄일 수 있음을 증명하며, 일반적인 편집 작업에서 최대 3.13배의 속도 향상을 제공합니다 \u2111.

현재의 파이프라인에서는 확산 트랜스포머(Diffusion Transformer, DiT)가 연산의 대부분을 소비하며, 50개의 타임스텝(timesteps)에서 8개의 타임스텝으로 증류(distilled)된 후에도 전체 모델 지연 시간의 약 73%를 차지합니다. 따라서 이전의 최적화 방식은 타임스텝 감소에 집중되어 있었으며, 토큰 단위의 비용은 그대로 유지되었습니다.

저자들은 확장된 편집 마스크(dilated edit mask)에 전체 해상도 토큰을 할당하고, 마스크 외부의 저주파 토큰을 가지치기(prune)하며, 나머지를 16배 다운샘플링된 표현으로 보충함으로써 작은 마스크 편집 시 최대 3.13배의 DiT 속도 향상을 달성했습니다. 마스크 크기에 따라 평균 토큰 비율은 6.38%(소형), 15.92%(중형), 35.36%(대형)로 감소하며, 이에 상응하는 DiT 속도 향상은 각각 3.13배, 2.59배, 1.67배입니다 \u2111.

보고된 이점은 마스크가 커짐에 따라 감소하여, 토큰의 35%가 유지되는 대형 마스크의 경우 속도 향상이 1.67배로 떨어지며, 이는 압축률과 실행 시간 사이의 트레이드오프(trade-off)를 보여줍니다. 또한, 평가는 마스크 가이드 편집(mask-guided editing)에 국한되어 있습니다. 동일한 토큰 예산 전략이 제한 없는 이미지-투-이미지(image-to-image) 또는 텍스트-투-이미지(text-to-image) 확산 모델에서 어떻게 작용할지는 여전히 미해결 과제로 남아 있습니다.

Photoshop의 'Remove' 기능을 구동하는 배포 환경에서는 줄어든 토큰 예산 덕분에 Amazon AWS p5.48xlarge 노드 수를 3분의 1만큼 절감할 수 있습니다. 따라서 실무자들은 제한된 하드웨어를 위한 확산 기반 비전 편집기를 구축할 때 적응형 토큰 할당(adaptive token allocation)을 기본 최적화 계층으로 취급해야 합니다.

References

1. HiLo-Token: Input-Adaptive High-Low Frequency Token Compression for Efficient Image Editing