토큰 병합을 확산 모델에 적용
Abstract
Token Merging을 Stable Diffusion에 적용하여 속도 향상
Introduction
대부분의 이미지는 중복성이 높고, 따라서 모든 토큰에 대해 계산을 수행하는 것은 리소스의 낭비이다.
추가적인 훈련 없이 평가 속도를 높일 수 있는 token merging(ToMe)을 SD에 적용.
ToMe는 SD의 속도를 크게 향상할 수 있었지만 이미지 품질이 크게 저하되었다.
속도 향상과 함께 이미지 품질을 유지하기 위한 새로운 기술을 도입한다.
Background
LDM(Stable Diffusion)
Token Merging for Stable Diffusion
확산 모델에서는 모든 토큰에 대해 제거할 노이즈를 알아야 하기 때문에 vanilla ViT에 적용하는 것처럼 간단하지 않다.
따라서 병합 해제라는 개념을 도입한다.
Defining Unmerging
c채널을 가진 유사한 토큰 둘을 병합하는 경우 :
병합 해제 :
두 토큰은 병합 이전에 유사했기 때문에 위와 같은 단순한 방법도 오차가 크지 않다.
An Initial Na¨ıve Approach
각 구성요소 전에 ToMe를 수행하고 병합 해제
또한 기존 ToMe보다 병합 횟수가 훨씬 더 많기 때문에 r개가 아닌 r%의 토큰을 병합하고 토큰 유사성 계산은 병합 당이 아닌 블록 당 한번 계산.
ToMe에서 제안한 proportional attention은 사용하지 않으며 유사성의 계산에도 key가 아닌 feature를 그대로 사용한다.
(Stable Diffusion의 feature 공간에서 정보가 이미 요약되어 있어서 그런듯?)
Further Exploration
ToMe는 추가적인 훈련 없이 즉시 사용할 수 있다.
하지만 많은 비율의 토큰을 병합할수록 콘텐츠가 손실되는데, 따라서 추가적인 개선을 수행한다.
A New Partitioning Method
기존 ToMe의 토큰 선택 방법을 적용하면 두 세트가 열을 형성하게 된다.(a)
dst 토큰의 Stride를 늘리면 조금 개선 (b)
랜덤성을 추가하면 FID가 엄청나게 떨어진다.
또한, classifier-free guidance를 사용할 때 컨디셔닝 된 샘플과 그렇지 않은 샘플의 dst 토큰의 위치가 동일해야 하기 때문에 배치 전체에서 랜덤성을 고정한다.(c)
b와 c를 결합하면 더욱더 향상 (d)
Design Experiments
ToMe 적용을 위한 더 나은 설계 탐색
What should we apply ToMe to?
FID는 프롬프트 반영도를 고려하지 않기 때문에 self attn에만 적용했을 때 FID가 가장 좋다.
Where should we apply ToMe?
깊은 U-Net에서는 토큰이 적은 계층도 있기 때문에 정보의 손실이 클 수 있어 토큰이 많은 계층에만 적용하는 것이 좋다고 한다.
When should we apply ToMe?
위 개념의 연장선으로 토큰이 적은 계층일수록 병합률을 점진적으로 줄일 수 있다.
Putting It All Together
