AnimateDiff

논문이 나오게 된 배경

해당 논문이 나오게 된 배경을 정리하면 아래와 같다.

1. 기존의 T2I 모델이 성능이 많이 올라왔다 (Diffusion Prior). 심지어 Personalization이 가능해지면서 (DreamBooth, LoRA 등의 방법) 보다 정교한 생성이라던지 원하는 이미지 생성 방법이 많이 공개되기 시작하였다.
2. 특히 personalize가 가능하다는 관점에서 큰 학습 리소스를 들이지 않고도 많은 개인화된 모델이 huggingface 등의 커뮤니티에 공개되기 시작하였다.
3. 그러나 이러한 객체나 스타일에 대한 personalize 이외에, T2I 모델을 기반으로 특정 모션을 주는 (Video prior, Motion prior) 연구는 여전히 부족하고, 특히 디퓨전의 경우 학습량이 방대하므로 이를 줄일 수 있는 방법이 절실하다.

T2I Model Personalization (개인화)란.

Text to Image model에서의 개인화라는 단어는 곧 특정 스타일이나 오브젝트에 미세 조정된 모델을 얻고자 하는 관점에서 처음으로 정의되었다.

DreamBooth

특정 오브젝트의 형태는 유지하면서, 해당 오브젝트를 주제로 한 다양한 Scene을 구성하고자 하는 문제가 있다.

기존의 T2I 모델은 방대한 데이터로부터 프롬프트 - 이미지 쌍을 학습하여 고퀄리티 이미지 생성의 prior를 학습할 수 있었다.

하지만 ‘강아지’라는 단어의 의미를 공유하는 다양한 형태의 prior를 생성할 수 있다는 장점은 있지만 특정 모습의 ‘강아지’만 집중적으로 생성할 수는 없는 문제가 있다.

아무리 텍스트 Description를 정밀하게 묘사하더라도 특정 객체의 모습이 계속 달라지는 문제가 생긴다.

이러한 문제는 Language model의 임베딩이 Vision model의 임베딩과 공유 공간에서 매핑되어 있더라도 사라지지 않고 발생한다. (아래와 같은 다양한 모델에서의 실패 사례)

따라서 이를 문제점으로 삼아, 특정 단어(강아지나 시계 등)과 생성될 이미지의 특징을 고정시키고자 한 연구가 dreambooth이다.

예상되는 문제

객체에 대한 prior를 diffusion model에 심고 (output domain에 원하는 물체 이미지를 추가하고) 싶다.

그러나 만약 Few-shot dataset으로 특정 객체에 대한 미세 조정을 진행한다면? 너무 뻔하게도 Mode collapse / overfitting 위험성이 증대되고, 원래 모델의 성능을 해칠 수 있다.

따라서 기존의 방법들은 보수적인 선택을 한다: 모델이 subject prior를 직접 학습하도록 하지 않고 target distribution에 따라가게끔만 task를 정의하였다. 즉, 객체를 완전히 personalize하진 못하는 문제가 있다.

Diffusion model에 대한 loss (아래 수식)만 잘 세팅하면 객체 정보를 학습시키면서도 기존의 prior를 유지할 수 있다.

$$\mathbb{E}_{\mathbf{x},\,\mathbf{c},\,\boldsymbol{\epsilon},\,t} \!\Bigl[ w_t \bigl\lVert \hat{\mathbf{x}}_{\theta}\!\bigl(\alpha_t \mathbf{x} + \sigma_t \boldsymbol{\epsilon},\,\mathbf{c}\bigr) - \mathbf{x} \bigr\rVert_2^{2} \Bigr]$$

Identifier 정의 방법

특정 객체 이미지를 학습시키려면 이를 특별하게 만들어줄 텍스트 묘사가 함께 필요하다.

(e.g.) [바로 그] [강아지] / [내가 개인화한] [고양이]

위의 예시에서 강아지 / 고양이는 일반적인 객체를 지칭하는 명사고, 그 앞에 특별한 identifier를 달아서 text로 컨디션을 줄 때 내가 원하는 객체가 나오게끔 학습하고자 한다.

그런데 여기서 문제가 발생하는데, 예컨대 unique / special 같이 이미 있는 영어 단어를 무작위로 사용하게 되면 원래 unique / special이 기존에 학습되었던 prior에서 disentangle을 수행하고(원래 해당 단어가 생성하던 이미지의 특정 property를 잊어야함) personalize하고자 하는 객체에 달라붙어야하는데, 이러한 구조는 학습이 어렵기 때문에 지양해야 한다.

그래서 방법으로 아예 xy5syt00같은 랜덤한 텍스트를 넣어서 일종의 일련번호처럼 쓰는 방법을 고안해보았다.

근데 이렇게 랜덤하게 뽑아도 이상하게 tokenizer를 통과시켰을때, word 단위가 아니라 토큰 단위로 들어가다보면 strong prior를 가진 애들이 나왔다. 그래서 일반적인 단어를 쓰는 것과 별로 큰 차이가 없는 경우가 생겨서 이 또한 문제가 될 수 있다고 판단하였다.

결론은 직접 rare token을 찾아서 사용하였다. 이 부분은 일종의 리버스 엔지니어링으로 수행하게 되는데, token ID를 랜덤으로 샘플링해서 prior와의 entanglement를 측정하고, 이들 중 가장 낮은 entanglement를 가지는 단어를 역으로 decode해서 추리해낸다.

T5-XXL 토크나이저 기준으로 5000번대에서 10000번대 사이의 토큰을 디코딩했을 때의 단어가 가장 괜찮았다고 한다.

학습 시 주의할 점

Language drift : 언어 모델 함부로 학습시키면 이미 대량의 데이터로 학습한 언어 모델의 (Large Text Corpus로 학습된 결과) representation이 날아간다.

Reduced output diversity : 이미지 생성 모델을 잘못 학습시키면 이미지 다양성이 훼손될 수 있는 문제가 있다.

$$\mathbb{E}_{\mathbf{x},\,\mathbf{c},\,\boldsymbol{\epsilon},\,\boldsymbol{\epsilon}',\,t} \!\Bigl[ w_t\, \bigl\lVert \hat{\mathbf{x}}_{\theta}\!\bigl(\alpha_t \mathbf{x} + \sigma_t \boldsymbol{\epsilon},\,\mathbf{c}\bigr) - \mathbf{x} \bigr\rVert_2^{2} \;+\; \lambda\,w_{t'}\, \bigl\lVert \hat{\mathbf{x}}_{\theta}\!\bigl(\alpha_{t'} \mathbf{x}_{\mathrm{pr}} + \sigma_{t'} \boldsymbol{\epsilon}',\,\mathbf{c}_{\mathrm{pr}}\bigr) - \mathbf{x}_{\mathrm{pr}} \bigr\rVert_2^{2} \Bigr]$$

뒤쪽에 있는 텀이 정규화 텀이다. 사전 학습 모델에서 미리 샘플링해두고 학습할 때 source dataset 형태로 사용한다. 흔한 domain adaptation 방법과 동일하다.

LoRA (Low-Rank Adaptation)

원래 모델은 건들지 않고 기존 weight보다 낮은 차원에서 conditioning을 주는 방법을 통해 적은 파라미터 수로 기존 모델의 성능을 바꾸지 않고 plug-and play가 가능한 학습법. Low Rank 영역에서의 implicit space를 학습하는 것이 성능에 악영향을 미치지 않는다는 관점에서 나온 논문이다.

$$h = W_0 x + \Delta W x = W_0 x + BAx$$

위의 수식에서 $BA$만 효율적으로 학습하려고 한다. 결국 잔차 학습을 통해 fine-tuning의 효과를 얻고자 한다.

학습 프레임워크

AnimateDiff는 기본적으로 Text to Image Diffusion Model을 사용하고, 이러한 T2I 모델 아키텍처로 대표적으로 잘 알려진 것들은 GLIDE, DALL-E 등 많지만 그 중 가장 대표적으로 많이 쓰이면서 퀄리티가 높은 Stable Diffusion (LDM) 모델에 대해서 보면 다음과 같다.

기존 DDPM에서는 ResNet 중간에 $16 \times 16$ 피쳐맵에서 어텐션을 사용하여 타임코드에 대한 정보를 넘겨주는 방식을 사용하는데, ADM에서는 이를 $32 \times 32$ 부터 $8 \times 8$로 확장시키고, 어텐션 헤드 수나 dimension 등을 조절하여 최적의 아키텍처를 찾았다. ADM부터 본격적으로 GAN의 성능을 넘어서게 되었다.

이후 후속으로 나온 연구들은 이러한 ADM 구조의 어텐션에 text embedding 정보를 잘 넣는 방식을 연구했는데, 이러한 연구들이 대표적으로는 GLIDE, DALL-E 그리고 Imagen 이라고 할 수 있겠다.

LDM은 이런 아키텍처와는 조금 다른 독자노선을 선택했는데, 이는 이미지 단위가 아니라 VQ-GAN과 같은 모델 기반으로 latent 단위로 인코딩을 해서, 노이즈로부터 이미지가 아닌 디코딩 이전의 이미지 임베딩을 예측하도록 하여 연산량을 줄였다. 또한 conditioning 파트에서도, 기존 방법들은 대부분 text embedding을 pooling해서 기존 아키텍처에 있는 어텐션에 타임 코드($t$)로부터 온 positional embedding과 함께 잘 추가해주는 방식을 선택하여 학습하는데, 이 모델은 self-attention 구조를 self-attention 과 cross-attention이 혼합된 말 그대로 트랜스포머 구조를 차용하였다. Latent 연산이 기본적으로 연산량을 줄여주면서, 텍스트 임베딩에 대한 conditioning에 더욱 집중할 수 있던 흐름이 되는 것이다.

AnimateDiff에서도 대부분의 T2I 레퍼런스는 personalized model에서 나오게 되었는데, 사실상 거의 모든 T2I 모델은 LDM을 LoRA 튜닝하는 방식으로 생성되는 것이 한때 인기였기 때문에 이 논문에서 적용한 방법론은 ablated UNet이 아닌 transformer architecture로 생각하는 것이 맞는 것 같다. 그 중 Stable Diffusion Model은 이를 보다 확장시켜 모델링을 더욱 진행한 형태다.

AnimateDiff는 3개의 파이프라인으로 구성되어서 학습된다.

1. 이미지 도메인과 비디오 도메인 사이의 GAP 없애기 (추가 모듈)
2. 움직임에 대한 prior를 학습시키기 (추가 모듈)
3. 원하는 움직임에 대한 personalize 진행하기 (추가 모듈)

여기서 움직임에 대한 정보를 학습시키는 구조는 2번과 3번이 유사하지만, 2번은 일반적인 비디오의 특성 (temporal information)에 대한 학습을 진행하는 한편, 3번은 생성 시 특정 움직임을 위한 모델을 만들고자 한다. 그렇기 때문에 1번과 3번은 미세 조정의 관점으로 LoRA를 적용하며, 2번은 transformer 구조를 차용하여 pretraining하게 된다.

Alleviate Negative Effects

사실 이미지와 비디오는 프레임별로 보게 되면 크게 다르지는 않다. 어찌 보면 이미지는 정적인 프레임 하나만 있는 비디오라고 생각하면 된다.

하지만 비디오 프레임은 일반적인 이미지와 다르게 각 프레임별로 보게 되면 객체가 이동하거나 카메라가 이동하게 되면 중간중간 모션 블러가 발생하거나 몇몇 인코딩/디코딩 방식에 따른 compression이 발생해서 일부 픽셀이 뭉개진다던지, 워터마크가 있다던지 등등 깔끔함과는 차이가 있다.

이는 곧 이미지와 비디오 도메인 간의 distribution misalignment 문제를 발생시키는데, 이를 조정하지 않고 motion prior를 학습하려하면 다음과 같은 문제가 발생한다.

1. 사전 학습된 DM의 text-image conditioning space가 깨진다.
2. 사전 학습된 DM의 분포에서 벗어난 이미지에 대해 학습을 수행하지 못한다.

따라서 좌측과 같이 ResNet 블록 간에 conditioning을 위한 transformer block (Attention Modules)에 Adapter를 달아서 학습하게 된다. 수식을 보면 알 수 있듯이 LoRA로 학습시킨 것을 알 수 있다.

$$Q = \mathcal{W}^Q z + \operatorname{AdapterLayer}(z) = \mathcal{W}^Q z + \alpha \cdot AB^\top z$$

Learn Motion Priors

Adapter로 DM Block의 출력값에 대한 분포를 맞춰줬기 때문에 이제는 temporal 정보를 학습시키고자 한다.

앞에 학습한 파이프라인은 비디오 데이터셋을 일부 프레임을 샘플링해서 Domain adapter를 비디오 프레임 형태의 이미지에 익숙하게 한거라, 이미지 단위에서 도메인을 맞춘거지 비디오 도메인으로 끌고왔다고 할 수 없기 때문에 따라서 여전히 DM(Diffusion Model)은 단일 이미지에 대해서만 적용이 가능하다는 문제가 있다.

움직임에 대한 특성만 학습하기 위해 해당 학습 파이프라인을 추가하였다.

DM을 비디오 영역으로 확장시키는 개념은 Make a video 논문에서 대부분이 구현되었다. 가장 간단한 방식으로는 video frame을 이미지 단위로 처리하는 방식이 있다. 예를 들어 배치 $b$에 프레임 $f$라면, 이미지를 $b \times f$ 배치만큼 처리한다고 생각하면 된다.

그런데 이 구조에서는 video motion에 대한 prior를 따로 학습할 수 없기 때문에, motion 부분만 학습하고자 temporal Transformer 아키텍처를 추가한다. 이때, 모션 prior를 학습하는 모듈은 일반 영상을 처리하는 트랜스포머와 다르게 spatial dimension을 무시해서 처리하게끔 한다 ($b \times h \times w$). Conditioning이 되는 파트이므로 zero initialization + residual 연결을 통해 motion에 대한 잔차 정보만 학습시키는 방법을 사용하였고, 전반적으로 모두 Make a video 논문 아이디어에서 차용된 것으로 보인다. 아래는 깃허브 소스에서 Versatile Attention 부분을 참고한 것이다.

def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None):
        batch_size, sequence_length, _ = hidden_states.shape

        if self.attention_mode == "Temporal":
            d = hidden_states.shape[1]
            hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length)
            
            if self.pos_encoder is not None:
                hidden_states = self.pos_encoder(hidden_states)
            
            encoder_hidden_states = repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d) if encoder_hidden_states is not None else encoder_hidden_states
        else:
            raise NotImplementedError

        # ... hidden states 구하는 부분 생략

        # linear proj
        hidden_states = self.to_out[0](hidden_states)

        # dropout
        hidden_states = self.to_out[1](hidden_states)

        if self.attention_mode == "Temporal":
            hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)

        return hidden_states

이 부분을 보면 cross attention은 사용되지 않는 것으로 보아, Temporal Attention 파트는 오로지 앞단에서 생성되었을 latent 간의 temporal consistency를 위한 장치로 생각된다.

즉, 이미 앞부분에서 어느 정도 프레임별 이미지가 완성되었다고 가정하고, temporal attention 파트에서는 프레임 간의 temporal information만 학습하고자 한 것.

Customize motion prior

2번까지만 해도 T2I 모델에 T2V prior를 넣을 수 있는데, 만약에 특정 모션 정보에 대해 특화된 모델을 만들고자 한다면(카메라 모션에 따라서) 해당 모션에 대한 레퍼런스 영상 20~50개만 가지고 튜닝할 수 있는 방법을 제안하였다. 앞서 했던 내용은 personalize T2I 모델을 personalize T2V 모델로 바꾸는 거였고, 3번은 personalize T2V를 personalize + motion personalize T2V로 바꾸는 작업이 되겠다. 특별한 작업은 없고 stage2에 추가했던 motion module에 LoRA를 붙여서 학습시킨다.

실험 과정

학습 파트

Domain Adapter는 기존의 Stable Diffusion Model 학습 objective로 튜닝이 가능하고, Motion module과 LoRA는 비슷하기는 한데 차이가 있다면 motion을 학습시킬때는 학습 범위 내에서 frame에 대한 정보가 필요하기 때문에

$$\mathcal{L} = \mathbb{E}_{\mathcal{E}\!\left(x_{0}^{1:f}\right),\,y,\,\epsilon^{1:f}\sim\mathcal{N}(0,I),\,t} \!\left[ \left\| \epsilon - \epsilon_{\theta}\!\left(z_{t}^{1:f},\,t,\,\tau_{\theta}(y)\right) \right\|_{2}^{2} \right].$$

Objective가 관측하게 되는 empirical도 frame 1 ~ f 를 기준으로 업데이트된다. 물론 파라미터는 학습되는 각 파트를 제외하고는 모두 freeze하고 학습한다.

정량 평가

텍스트랑 얼마나 잘 align 되는가에 대한 CLIP metric 그리고 실제 사용감 테스트를 위한 User Study

추론 결과

대부분의 모듈은 LoRA 형태로 학습되었기 때문에 (Motion prior 학습한 temporal Transformer 제외). 이에 대한 $\alpha$값으로 ablation 수행이 가능하다. 상대적으로 video domain이 이미지 퀄리티가 image domain보다는 확연히 떨어지기 때문에 점진적으로 떨어지는 것을 확인할 수 있다.