[11] DALL-E : Zero-Shot Text-to-Image Generation

long8v commented 2 years ago

TL;DR

task : zero-shot text-to-image generation
problem : generation 모델들은 모델 사이즈가 커지는데 text-to-image는 데이터가 한정적임
idea : 인터넷 크롤링으로 겁나 큰 text-image pair(MS-COCO)를 만들고 모델 사이즈도 키우자. 효율적으로 이미지를 인코딩하기 위해 discreteVAE를 사용하자.
architecture : discreteVAE으로 이미지를 코드로 인코딩 + 텍스트 인코딩해서 auto-regressive 하게
objective : cross-entropy loss
baseline : AttnGAN, DM-GAN, DF-GAN
data : 3.3M MS-COCO
result : FID(=Frechet Inception Distance)에서 SOTA. human evaluation에서도 DF-GAN보다 압도적 우위.
contribution : 엄청 큰 이미지 생성모델.

Details

120억 파라미터, 데이터셋은 2억 5천만 text-image pair 데이터

Method
이미지 픽셀 단위로 학습을 하면 1) 연산이 너무 늘고 2) likelihood를 적용했을 때, short-range 관계에만 집중해 high-frequency detail만 모델링하고, 인간이 식별 가능한 low-frequency를 모델링하지 못한다
이를 해결하기 위해 두단계로 나누었다.
- 첫번째 단계에서는 discrete VAE(dVAE)를 사용하여 256 x 256 RGB 이미지를 32 x 32 이미지 토큰이 있는 그리드로 나누고, 각각의 요소들이 8192개의 값을 가질 수 있게 한다. 이는 트랜스포머의 이미지 퀄리티의 큰 저하 없이 context size를 192배 줄였다.
- dVAE로 압축한 뒤 reconstruct한 결과
- 두번째 단계에서는 256개의 BPE-encoded text를 32 x 32 = 1024 이미지 토큰과 concat하여 텍스트와 이미지를 autoregressive하게 학습하였다.
- 이러한 전체적인 단계는 이미지 x와 캡션 y, dVAE로 encode된 token z간 joint probability의 evidence lower bound(ELB)를 최대화하는 것으로 볼 수 있다.
- 즉, 위와 같이 factorization을 할 수 있고 이 값의 ELB값을 구하면 아래와 같다.
식을 해석하면, dVAE에서 캡션과 이미지 토큰이 주어졌을 때 이미지의 likelihood를 최대화하면서, dVAE에서 이미지가 주어졌을 때 나오는 캡션 텍스트와 이미지 토큰의 확률분포와 트랜스포머를 통해 나오는 캡션과 이미지 토큰의 joint 확률 사이의 KL divergence를 최소화해야 한다.
이때 beta는 1이어야 이론적인 ELB이지만 실험 결과 크게 할 때 성능이 좋았다.

Step One: Learning the Visual Codebook

$\phi$와 $\theta$에 대해 ELB를 최대화하는 식으로 dVAE를 학습한다. code의 크기는 32 x 32이며 $K$=8192고 $p_\psi$는 uniform분포이다. code가 discrete 해서 미분 불가한 부분은 gumbel softmax를 사용해서 gradient를 흘려줬다.

$p\theta$는 log-laplace(정규 분포 지수 부분에 제곱대신 절대값) 분포로 평가됐다.

Stage Two: Learning the Prior

텍스트는 BPE encode해서 최대 256길이로 만들었고, 이미지는 dVAE encoder logit에서 argmax해서 1024개의 토큰을 얻었다. 두개의 인코딩을 concat해서 트랜스포머 디코더에 넣어줬고, 텍스트가 256보다 작을 경우에는 256개의 position에 따라 각각의 [PAD]토큰을 학습시켜줬다. -> OOD catpion에 더 강건했다. cross-entropy loss를 사용했고 텍스트와 이미지의 loss는 1/8, 7/8로 곱해줬다.

long8v commented 2 years ago

ELBO

어떤 random variable Z가 있고, Z와 X가 theta로 표현되는 p(Z, X | theta) 분포를 따른다고 하자. 보통 우리가 계산하는 건 1) theta와 X가 주어졌을 때 posterior Z 구하기 p(Z|X, theta) 2) P(theta | X) : likelihood를 최대화하는 theta 찾기이다. evidence란 theta가 주어졌을 때 likelihood를 부르는 다른 말이다.

이때 우리가 Z가 q라는 distribution을 따른다는 것을 알고 있을 때, p(X, Z| theta)는 p(X| Z, theta) * q(Z)로 표현할 수 있다. 이 때, log p(X | theta)의 최소값은 아래와 같이 표현할 수 있다.

이때 오른쪽 항을 ELBO라고 부른다

이 ELBO를 구하는 식은 아래와 같다.

마지막의 부등호는 jensen's inequality를 사용했다.

ELBO와 evidence의 차이는 KL divergence와 일치한다.

https://mbernste.github.io/posts/elbo/ https://www.youtube.com/watch?v=A8wf7QmmlUM

long8v commented 2 years ago

ELBO 2

https://www.youtube.com/watch?v=GbCAwVVKaHY https://long8v.notion.site/ELB-f7267ffc5301422c9965e3e6f0619958

long8v commented 2 years ago

encoder.py

attr.s 라는 decorator로 init에 대한 validate등을 하는듯

@attr.s(eq=False)
class Conv2d(nn.Module):
    n_in:  int = attr.ib(validator=lambda i, a, x: x >= 1)
    n_out: int = attr.ib(validator=lambda i, a, x: x >= 1)
    kw:    int = attr.ib(validator=lambda i, a, x: x >= 1 and x % 2 == 1)

    use_float16:   bool         = attr.ib(default=True)
    device:        torch.device = attr.ib(default=torch.device('cpu'))
    requires_grad: bool         = attr.ib(default=False)

    def __attrs_post_init__(self) -> None:
        super().__init__()

        w = torch.empty((self.n_out, self.n_in, self.kw, self.kw), dtype=torch.float32,
            device=self.device, requires_grad=self.requires_grad)
        w.normal_(std=1 / math.sqrt(self.n_in * self.kw ** 2))

        b = torch.zeros((self.n_out,), dtype=torch.float32, device=self.device,
            requires_grad=self.requires_grad)
        self.w, self.b = nn.Parameter(w), nn.Parameter(b)

long8v / PTIR