LD-Pruner: Task-Agnostic 인사이트를 활용한 Latent Diffusion Models (LDMs)의 효율적인 가지치기

Thibault Castells
Research Engineer, Nota AI

소개

Latent Diffusion Model(이하 LDM) [Rombach et al., 2021] 은 생성 작업에서 뛰어난 성능을 보였지만, 메모리 사용량이 많고 추론 속도가 느려서 자원이 제한된 장치에서의 배포가 어려웠습니다. 가지치기는 모델 크기를 줄이고 효율성을 높이는 일반적인 방법이지만 모델의 성능을 저하할 수 있으며, 대형 모델의 경우 재훈련이 필요해 비용과 시간이 많이 듭니다.

성능 기반 가지치기는 가지치기 과정에서 성능 손실을 최소화하는 것을 목표로 합니다. 하지만 LDM에 성능 기반 가지치기를 적용하는 것은 그 성능을 빠르고 범용적으로 평가할 방법이 부족했기 때문에 어려웠습니다. 이러한 문제를 해결하기 위해, 저희는 최근 연구 논문에서 LDM 압축에 특화된 새로운 성능 유지 가지치기 방법인 LD-Pruner를 제안했습니다. 잠재 공간을 활용해 가지치기 하는 과정에서 LD-Pruner는 추론 시간과 파라미터 수를 크게 줄이면서도 텍스트-이미지 생성, 무조건 이미지 생성, 무조건 오디오 생성 등 다양한 작업에서 성능 저하를 최소화합니다.

🍀 자세한 정보는 ArXiv를 참조하세요.

🍀 제1회 CVPR’24 Workshop on Efficient and On-Device Generation (EDGE)에서 채택되었습니다.

개요

LD-Pruner는 LDM의 Unet에서 각 연산자의 중요성을 체계적으로 수정하고, 그에 따른 생성된 잠재 표현의 변화를 추적하여 평가합니다. 또한, 맞춤화된 점수 공식으로 각 변동 사항이 어떤 영향을 미쳤는지를 정량화하여 가지치기 과정에서 덜 중요한 연산자를 제거하거나 대체하도록 안내합니다. 그림 1에서 보이듯, 이 접근법은 LD-Pruner가 가지치기 대상이 될 가장 적합한 연산자를 식별하면서 재훈련의 필요성을 최소화하여 비용을 절감할 수 있습니다. LD-Pruner는 잠재 공간에서 작동하기 때문에, 가지치기가 모델 성능에 미치는 영향을 빠르고 범용적인 방법으로 평가할 수 있습니다. 이를 통해 LDM의 성능 기반 가지치기를 효율적으로 수행할 수 있으며 전통적인 가지치기 방법의 한계를 극복할 수 있습니다. 압축된 모델은 메모리 소모가 줄어들고 추론 속도가 개선되어 자원이 제한된 장치에 더 적합해집니다.

Figure 1. LD-Pruner의 개요입니다. Unet에서 $k$ 개의 연산자가 있다고 할 때, 원본 Unet과 각 연산자가 하나씩 수정된Unet을 포함해 총 $k+1$ 개의 $N_{gen}$ 잠재 벡터 집합을 생성합니다. 그런 다음, 잠재 벡터를 비교하는 특별한 공식을 통해 각 연산자의 중요성 점수를 계산합니다. 이 공식은 잠재 벡터의 평균과 변동성을 고려하여 각 연산자의 중요성을 종합적으로 평가합니다.

방법

LD-Pruner의 과정은 LDM의 Unet에서 각 연산자를 하나씩 수정하는 것에서 시작됩니다. 이때 연산자는 완전히 제거되거나, 계산량이 적으면서도 원래 차원을 유지하는 연산으로 대체됩니다. 이후, 모델은 원본 Unet과 수정된 Unet을 사용해 잠재 표현을 생성하고, 다양한 연산자 수정에 따른 결과를 기록합니다.

각 연산자의 중요성을 평가하기 위해, 원본 잠재 표현과 수정된 잠재 표현 간의 차이를 측정하는 특별한 점수 계산 공식을 사용합니다. 이 공식은 잠재 표현의 평균 거리와 표준 편차 거리를 종합적으로 계산하여, 연산자 수정이 미치는 영향을 평가합니다.

$$ \text{score} = avg_{dist} + std_{dist} $$

With:

$$ avg_{dist} = |avg_{orig} - avg_{mod}|_2\\ std_{dist} = |std_{orig} - std_{mod}|_2 $$

여기서 $|\cdot|_2$ 는 유클리드 거리를 의미하며, $avg_{orig}$ 와 $avg_{mod}$ 는 각각 원본 모델과 수정된 모델(한 연산자가 변경된 상태)에서 생성된 잠재 벡터의 평균을 나타냅니다. $std_{orig}$ 와 $std_{mod}$ 는 해당 모델들에서 생성된 동일한 잠재 벡터 집합의 표준 편차를 의미합니다.

중요성 점수가 계산되면, LD-Pruner는 점수가 가장 낮은 연산자를 가지치기 후보 또는 대체할 후보로 식별합니다. 이러한 연산자들은 모델 출력에 가장 적게 기여하는 것으로 간주됩니다.

결과

LD-Pruner는 텍스트 기반 이미지 생성 (T2I), 비조건부 이미지 생성 (UIG), 비조건부 오디오 생성 (UAG)과 같은 여러 생성 작업에서 평가되었으며, LDM을 압축하면서도 성능을 유지하는 데 효과적임을 보여주었습니다.

T2I 작업에서 가지치기된 모델은 원래 모델과 유사한 성능을 보였으며, 34.89% 더 빠르게 작동했습니다(엔비디아 RTX 3090 GPU에서 측정됨). 이는 표1에서 확인할 수 있습니다.

표 1: MS-COCO 256 X 256 검증 세트에서 다양한 모델들이 생성한 T2I를 비교했습니다. 속도 향상 값은 SD-v1.4를 기준으로 상대적으로 측정되었습니다.

그림 2는 가지치기된 모델의 출력과 다른 T2I모델들을 정성적으로 비교하여 해당 모델의 경쟁력 있는 성능을 보여줍니다.

그림 2: T2I를 대상으로 한 제로샷 MS-COCO 벤치마크를 정성적으로 비교했습니다. 이전 연구 결과는 해당 연구에서 공식적으로 배포한 모델을 사용해 얻었습니다.

원래 모델이 410k 반복 후 프레셰 인셉션 거리 (FID) 13.84를 기록한 것에 비해, UIG 작업에서 압축된 모델(23.47% 속도 향상)은 빠르게 수렴하여 46k 반복 후 FID 15.03을 달성했습니다(그림 3 참조)

그림 3: CelebA-HQ 256 X 256 데이터셋에서 UIG 작업을 위한 훈련 과정 중 프레셰 인셉션 거리의 변화와 두 가지 다른 압축 비율에 대한 결과를 나타냅니다.

UAG 작업에서는 90개의 연산자를 수정하여 LD-Pruner가 19.2%의 속도 향상과 2.0의 사후 미세 조정 FAD를 달성했으며, 이는 원래 모델보다 더 우수한 성능을 보였습니다(표 2 참조). 사람의 귀로는 압축된 모델의 오디오 출력이 원래 모델과 정확히 동일하게 들립니다.

표 2: AudioDiffusion을 사용한 UAG 작업에서의 압축 성능을 보여줍니다. 12k 단계 동안 미세 조정을 진행했습니다.

LD-Pruner는 가지치기 과정에서 가중치 보존에 중점을 두었고, 이는 성능 유지에 중요한 역할을 했습니다. 가중치가 보존된 확산 모델은 동등한 압축 및 훈련 조건하에서 처음부터 훈련된 모델보다 일관되게 더 우수한 성능을 보였습니다(표 3 참조).

표 3: CelebA-HQ 256 X 256에서 UIG 작업을 위한  압축된 모델(31개의 연산자 수정)이 처음부터 훈련된 경우와 가중치 보존 후 훈련된 경우의 프레셰 인셉션 거리 점수를 나타냅니다. 두 경우 모두 동일한 훈련이 적용되었습니다. 원래 모델의 프레셰 인셉션 거리는 13.85입니다.

결론

LD-Pruner는 LDM 압축에 있어 중요한 발전으로, 작업에 구애받지 않고 성능을 유지할 수 있는 가지치기 접근 방식을 제공합니다. 잠재 공간을 활용하고 잠재 표현을 비교하기 위한 새로운 점수 산정 지표를 도입함으로써, LD-Pruner는 LDM에 적용될 때 전통적인 가지치기 방법의 한계를 극복합니다. 그 결과 압축된 모델들은 성능을 크게 저하시키지 않으면서 더 빠른 추론 속도와 적은 파라미터 수를 제공하여 자원이 제한된 장치에서의 배포를 보다 효율적이고 쉽게 만듭니다. LD-Pruner의 개발은 LDM이 다양한 응용 분야에서 광범위하게 채택될 수 있는 새로운 가능성을 열어주고, 강력한 생성 모델을 더 많은 사용자와 장치에 제공하려는 목표에 한 걸음 더 다가가게 합니다.

이 연구에 대해 추가로 궁금한 사항이 있으시면, 📧 contact@nota.ai.로 문의해 주세요.

또한, AI 최적화 기술에 관심이 있으시면 🔗 netspresso.ai를 방문해 보세요.