Open CLIP text embedding pooling

현재 인공지능의 대표적인 키워드를 하나 꼽으라면, 인간의 의사소통에서 사용하는 이미지, 텍스트, 음성 등의 다양한 매개체를 연결하는 multi-modality 구현일 겁니다. 특히 다양한 전달 매개체 중에 직관적이면서 다루기 쉬운 text 데이터와 이미지를 연결하려는 시도가 중심적으로 이루어졌고, OpenAI가 CLIP (Contrastive Language-Image Pretraining) 모델을 내놓으면서 multi-modality 연구의 중요한 분기점이 시작되었습니다. 이제는 CLIP 모델은 image-text 기반의 다양한 연구에 디폴트로 사용되고 있고 open_clip 라이브러리를 통해 매우 손쉽게 모델을 불러와 사용할 수 있습니다.

import open_clip

model, _, preprocess = 
	open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')

 

GitHub - mlfoundations/open_clip: An open source implementation of CLIP.

 

CLIP 모델의 디테일한 내용이나 훈련 방법, 사용한 데이터셋은 알려진 것이 워낙 많기에... 이번 포스트에서 다루고자 하는 것은 CLIP text embedder의 당연하지만 중요한 technique을 살펴보고자 합니다.

---

CLIP 사용하기

Single feature vector

CLIP 모델은 visual encoder와 text encoder로 구성되어 있고 각 encoder는 이미지와 그에 해당하는 caption을 받아 각각 single feature vector를 출력합니다. 따라서 각 encoder 출력은 [batch size, channel dim] 형태가 되고 channel dim은 encoder가 어떤 ViT 모델을 사용했느냐에 따라 768, 1024, 1280 중 하나로 정해집니다. open_clip 라이브러리의 사용법은 다음처럼 encode_image / encode text 함수를 써서 간단하게 image / text feature vector를 뽑아낼 수 있습니다.

import torch
from PIL import Image
import open_clip

model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
tokenizer = open_clip.get_tokenizer('ViT-B-32')

image = preprocess(Image.open("CLIP.png")).unsqueeze(0)
text = tokenizer(["a diagram", "a dog", "a cat"])

with torch.no_grad(), torch.cuda.amp.autocast():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)

    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)

print("Label probs:", text_probs)  # prints: [[1., 0., 0.]]

 

Diffusion 모델에서

Stable Diffusion 기준으로 일반적인 T2I (Text-to-Image) 생성 모델에서는 CLIP text encoder의 출력을 cross-attention 모듈에 넣어 text condition에 맞게 이미지를 생성하도록 학습합니다. 단, 이때 사용하는 text encoder의 출력은 single feature vector가 아닌 token 축을 가진 [batch size, token length, channel dim] 형태가 되죠. (보통 공개된 CLIP 모델을 사용하게 되면 token 길이는 77이 됩니다)

Cross-attention operation은 sequence 길이에 invariant 하니 모델은 token 길이에 상관없이 동작은 합니다. 하지만 궁금한 점은 본래 text가 tokenize 돼서 text encoder를 거쳐 나온 출력의 형태는 [batch size, token length, channel dim] 일 텐데 여기서부터 어떻게 text 입력을 대표하는 single feature vector를 뽑을까요??

 

open_clip 레포 탐험

open_clip 레포를 둘러보니 역시 답이 있었습니다. open_clip/src/open_clip/models.py 파일에서 encode_text 함수를 찾을 수 있었고 함수의 내용은 다음과 같습니다.

def encode_text(self, text, normalize: bool = False):
    cast_dtype = self.transformer.get_cast_dtype()

    x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]

    x = x + self.positional_embedding.to(cast_dtype)
    x = x.permute(1, 0, 2)  # NLD -> LND
    x = self.transformer(x, attn_mask=self.attn_mask)
    x = x.permute(1, 0, 2)  # LND -> NLD
    x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
    x, _ = text_global_pool(x, text, self.text_pool_type) <--------
    if self.text_projection is not None:
        if isinstance(self.text_projection, nn.Linear):
            x = self.text_projection(x)
        else:
            x = x @ self.text_projection

    return F.normalize(x, dim=-1) if normalize else x

• x=self.ln_final(x) 함수를 거치면 [batch size, token 길이, channel dim] 형태가 되고,
• 그 밑에 text_global_pool 함수를 거치면 이름에서 유추할 수 있듯 차원을 줄이는 pooling을 하는 것 같습니다.
• text_projection은 text feature를 pooling 한 이후에 적용되는데요, CLIP 모델 학습을 위해 image feature와 text feature의 차원을 맞춰주기 위해 사용되는 모듈로 아마 예전 크기가 작은 CLIP 모델에는 없는 경우도 있습니다.

 

text_global_pool

text_global_pool 함수는 open_clip/src/open_clip/transformers.py 파일에 다음과 같이 정의되어 있습니다.

def text_global_pool(x, text: Optional[torch.Tensor] = None, pool_type: str = 'argmax'):
    if pool_type == 'first':
        pooled, tokens = x[:, 0], x[:, 1:]
    elif pool_type == 'last':
        pooled, tokens = x[:, -1], x[:, :-1]
    elif pool_type == 'argmax':
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        assert text is not None
        pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x
    else:
        pooled = tokens = x

    return pooled, tokens

---

기본 pool_type은 argmax 형태이고 argmax 부분의 코드를 보면 "eot" embedding의 feature를 뽑는다라고 되어 있네요. "eot"라는 것은 ViT의 CLS처럼 CLIP을 훈련했을 때 text 마다 부여한 special token이고 추정상 "end of text"의 약자일 것 같습니다. (확실치는 않습니다) Text의 대표 벡터를 뽑아내기 위해 special token 부분의 feature를 가져오고 tokenize 될 때 "eot"는 제일 큰 임베딩 인덱스로 매핑되었을 가능성이 높을 것 같습니다. (이 부분 또한 제가 tokenizing 과정을 몰라 확실하지는 않습니다)

자 그러면 "eot" token 부분의 feature가 나머지를 과연 대표하느냐, 해당 sequence에 다른 feature는 무시되는 것 아니냐가 주요 이슈로 남을 듯합니다. 하지만 이 부분은 모든 text마다 "eot" token이 마지막에 붙은 채 text encoder가 "attention" 기반으로 학습이 되기 때문에 "eot" feature에는 해당 text의 알맹이 정보가 담기게 될 겁니다.

• 실제 open_clip 라이브러리의 tokenizer를 사용해 보면 모든 tokenized text마다 0이 되기 마지막에 가장 큰 수가 고정적으로 달려있음을 확인할 수 있습니다.

자세한 부분은 open_clip의 이 이슈를 보시면 좋습니다.

Question about clip.encode_text · Issue #217 · openai/CLIP

 

---

References

• open_clip 라이브러리
• encode_text 이슈