Text-to-Image Generation with Deep Image Prior

로보틱스, 특히 필드 로보틱스를 위한 AI에서 중요한 키워드를 꼽으라면 단연 "리소스 효율적"이라고 생각한다. LLM 시대가 도래하면서 로봇을 위한 AI 연구는 크게 양극단으로 분리된 것 같다. 하나는 "뭐든지 라지모델에게 물어봐~"이고 하나는 "더 가볍고 더 효율적으로!"이다. 

"뭐든지 라지모델에게 물어봐~" 쪽이 좀 더 요새 메인 기조이긴 하지만, 나는 <저렴한 보급형 농사로봇>에 관심이 있는 나는 후자를 열심히 보고 있다. 그래서 이번 학기에 고른 것은 <리소스 효율적 AI 리서치>라는 전산과 수업이었다. 

<리소스>라고 하면 GPU 메모리 뿐 아니라, AI 학습에 필요한 모든 단계에서의 리소스를 의미하며 이 수업에서는 학습 데이터 취득이나 학습 과정 자체 등 다양한 관점에서의 리소스 효율적 기법을 다룬다. 그 중에 하나가 Deep Image Prior (https://arxiv.org/abs/1711.10925)라는 페이퍼였다. 

이 페이퍼는  <훈련되지 않은 ConvNet 구조만으로도 이미지 문제를 해결할 수 있고, 이는 신경망 구조 자체가 이미지에 대한 Prior(사전 정보)이기 때문> 이라고 주장하고 있다. 한 문장으로 정의하면 다음과 같다. 

이미지 자체를 네트워크의 구조로 표현하고, 그 네트워크의 파라미터를 단일 이미지에 맞춰 최적화하는 기법

혹은

이미지를 네트워크로 Parametrize하고, 네트워크 구조를 implicit prior로 사용하는 방식

 

컨볼루션 필터는 이미지 전체에 적용되므로, 이미지의 저수준 통계를 포착할 수 있는 능력이 있으며 이런 연산자들을 순차적으로 사용하면 다중 스케일에서 픽셀들간의 관계를 이해할 수 있기 때문이라고 논문에서는 설명한다. 따라서 컨볼루션 연산을 통해 이미지를 생성하면 생성된 이미지에 self-similarity가 부여되는 경향이 있다. 

여기서 Not trained의 의미는, 다른 큰 데이터셋에 최적화되지 않고 이미지 한 장에 대해 네트워크 파라미터를 최적화 했음을 의미한다. 주의해야 할 점은, 이미지의 pixel을 최적화하는 게 아니라 네트워크의 weight를 최적화한다는 점이다.

Deep Image Prior는 그래서 Denoising, Super Resolution, Inpainting등 이미지 내부에서 픽셀들간의 관계를 잘 봐야 하는 태스크에서 DIP가 잘 작동한다고 주장한다. 

Text-to-Image Generation with DIP

그래서 이 글을 쓰려는 생각이 든 이유는, 이 DIP 논문에 연결해서 나온 과제가 흥미로웠기 때문이다. 

In this assignment, we utilize the Deep Image Prior mechanism for text-to-image generation, guided by visual-text alignment losses such as CLIP loss or SDS loss. DIP performs as an implicit image generator, producing an image that gradually aligns with the target text promt through optimization. 

SDS Loss

SDS Loss(Score Distillation Sampling)는 DreamFusion(https://dreamfusion3d.github.io/)에서 소개된 Loss인데, Pretrained 2D Text-to-Image Diffusion 모델을 활용해서 3D Object를 생성하는 알고리즘에 사용되었다. 이는 기존의 Diffusion 모델이 훈련되었던 Pixel space가 아니라 Parameter space에서 샘플링을 가능하게 하기 위해 도입되었다. SDS Loss는 Probability density distillation을 기반으로 한다. 이 손실은 Pretrained Diffusion 모델이 학습한 score function과, Diffusion의 forward process를 기반으로 하는 가우시안 분포 간의 KL-Divergence를 최소화하는 것이 목표이다. 

그래서 DreamFusion에서는
1) 너프가 무작위로 모델을 만들고,
2) 이 모델을 여러 각도에서 렌더링한 후
3) Pretrained Diffusion 모델에게 렌더링된 이미지와 텍스트를 넣어서 advice를 받고
4) SDS Loss에 기반해 너프 매개변수를 수정하는 방식으로 
Text 기반 3D 모델을 만든다. 

 그래서 과제에서 요구하는 CLIP Loss와 SDS Loss는 텍스트-이미지 정렬을 유도한다는 점에서는 비슷하지만 다음과 같은 차이가 있다. 

CLIP Loss는 모델이 만든 이미지 임베딩과 텍스트 임베딩의 코사인 유사도를 높이는 방식이다. 그래서 CLIP이 알려주는 것은 색감 조절, 전체적인 대비, 텍스트와 연관된 큰 스케일의 특징 등 다소 단순하고 전역적이다. (Low-frequency)
반면 SDS Loss는

• zₜ: latent에 noise를 넣은 샘플
• ε: 실제 noise
• ε_pred: UNet이 예측한 noise

	CLIP Loss	SDS Loss
사용 모델	CLIP encoder	SD UNet + scheduler
Gradient	정적 cosine 유사도	diffusion-based 확률적 gradient
표현 능력	global information	shape, 구조, 디테일
계산 비용	가벼움	무거움

CLIP Loss만 주고 optimize를 시키면, 전체 평균 색 바꿈 / 전체 대비 바꿈 등 전역적인 부분에서만 gradient 방향이 형성되어서 loss는 줄어들어도 노이즈에서 이미지로 denoise를 시킬 수는 없다. 

이미지 CLIP Loss 감소가 형태 복원으로 이어지지 않는다

좌) w/ DIP, 우) w/o DIP, promt : "A beautiful sunset"

이런 현상이 왜 나타나냐면, DIP 출력은 low-frequency smooth한 구조를 먼저 만들고 세부 형상은 늦게 학습되는데, CLIP loss는 Low-frequency 방향으로만 gradient를 제공하고 local structure 유도 능력이 부족하기 때문이다. 그래서 형상을 띠는 물체가 생기기 어렵다. (Only blob) 

여기서 SDS Loss에 대해 다시 말하자면, SDS는 "Denoising 과정을 역으로 이용해서 텍스트와 일치하는 방향으로 latent를 업데이트하는 loss"이다. 원래 Stable diffusion은 다음과 같이 동작한다. 

z (latent) → noise 추가 → noisy latent → UNet → noise 예측 → clean latent 복원

즉 Stable diffusion은 "노이즈에서 이미지를 복원하는 모델"인데, SDS는 "UNet이 예측한 방향대로 latent를 조정하면 텍스트와 일치하는 이미지가 된다"는 것이다. 그런데 Stable diffusion에서는 이미지를 직접 다루지 않고 VAE를 통과한 4x64x64 차원의 latent를 사용한다. 

실제 연구들은 VAE latent를 최적화하여 큰 구조를 생성하도록 한다. (이미지의 shape, composition, edge 정보들이 latent에 저장되어 있음)

Code Implementation

그래서 텍스트 프롬프트를 사용해서 이에 적합한 이미지를 생성하도록 하려면, 다음과 같은 프로세스로 진행해야 한다. 

    latent z (4×64×64) ← 학습되는 변수
                        ↓
SDS Loss Stable Diffusion UNet으로 gradient
                         ↓
            latent 업데이트
                         ↓
              VAE decode
                         ↓
             RGB 이미지 생성

내가 만든 코기🐶