이미지와 텍스트의 표현 공간을 통합한 CLIP 논문 리뷰

04 Aug 2023 in Multi-Modal on Embedding, Transformer

Paper Review

Paper Review

Paper Link:

Learning Transferable Visual Models From Natural Language Supervision

CLIP의 주요 포인트 몇가지를 위주로 리뷰해보려고 합니다.

먼저 CLIP이 어떤 연구인지에 대하여 대략적으로 살펴본 후, 어떤 Encoder들을 사용했는지, 어떤 Loss를 사용했는지를 소개하겠습니다. 마무리로 간단하게 CLIP Official Repository에 있는 소스코드를 살펴보겠습니다.

Introduction

CLIP은 이미지와 텍스트를 같은 표현 공간(representation space)에 임베드(embed)하는 학습 기법입니다.

기존 연구의 한계

CLIP이 나오기 이전의 연구에서는 정해진 단어 집합(vocabulary)에 속한 단어들에 대하여 이미지들과의 관련도를 예측할 수 있도록 학습시켰습니다.

이를 위해서는 학습하고자 하는 모든 단어들에 대하여 라벨링이 된 데이터셋(labeled dataset)을 마련해야 한다는 문제점이 있었습니다.

CLIP 연구진들은 이로 인해 학습 가능한 데이터셋의 규모(scale)가 제한되며, 성능(performance)에도 제한이 있으리라고 보았습니다.

또한, 각 단어들을 단어 집합 내에서의 순서에 따라 0, 1, 2, … 과 같은 숫자로만 치환해놓고 학습하는 경우에는, 해당 단어의 의미와 관련된 지식을 이용하기에 어렵다는 문제점이 발생했습니다.

Why CLIP?

CLIP 연구진은 학습 가능한 데이터셋의 크기를 효과적으로 늘리기 위해서, 이미지와 그에 대응하는 캡션(caption) 또는 제목이 이미 매칭되어 있는 상태의 raw data를 수집했습니다. 예를 들어 인스타그램에서 이미지와 그에 대한 설명 또는 파일제목을 얻어 이를 데이터셋으로 사용했습니다.

이미지 캡션, 제목과 같은 텍스트를 인코딩(encode)할 때는 BPE Tokenizer를 사용했으며, Transformer를 이용하여 임베딩을 구했습니다. 이 덕분에 규모가 클 뿐더러 사전에 포함되지 않는 단어가 등장하기도 하는 raw data를 학습하기에 용이했을 것입니다.

CLIP은 기존보다 훨씬 큰 규모(scale)의 데이터셋(dataset)으로 학습할 수 있었으며, 따라서 학습에 사용한 데이터셋 (pretrained dataset)에 포함되지 않은 텍스트에 대해서도 더 잘 일반화할 수 있었다는 점에서 획기적인 연구였습니다.

CLIP의 이미지 및 텍스트 인코더

CLIP에서는 텍스트 인코더와 이미지 인코더를 따로 두되, 두 인코더가 같은 표현 공간 (representation space)에서 임베딩(embedding)을 만들어내도록 학습했습니다.

Untitled

텍스트 인코더로는 masked attention을 포함한 Transformer를 사용했습니다. 앞서 언급했듯이 토크나이징(tokenizing)에는 BPE를 사용했습니다.

이미지 인코더로는 두 가지 옵션을 두었습니다. 수정된 ResNet (modified ResNet)과 Vision-Transformer (ViT)입니다.

CLIP의 학습 방법; “Contrastive Learning”

Contrastive Learning은 유사한 두 오브젝트의 쌍(pair)들로 이루어진 데이터셋(paired datset)을 필요로 합니다. 유사도(similarity)를 계산하기 위한 metric을 정해 두고, 쌍을 이루는 유사한 두 오브젝트 간의 유사도가 최대가 되도록 학습(maximize)합니다.

CLIP은 이러한 Contastive Loss를 사용하는 Contrastive Learning입니다. CLIP에서 사용한 Loss는 그 중에서도 ‘Multi-class N-pair Loss’라고 불립니다.

Multi-class N-pair Loss

CLIP에서는 각 Batch마다 서로 유사한 텍스트-이미지 쌍을 N쌍씩 학습시킵니다. 즉, 텍스트-이미지 쌍 N개가 나오도록 텍스트 N개와 이미지 N개씩을 준비해야 합니다.

이러한 텍스트 N개와 이미지 N개로부터 나올 수 있는 모든 조합의 텍스트-이미지 쌍을 만들면 N×N개의 쌍이 생깁니다. 이는 서로 유사한 텍스트-이미지 N쌍과 서로 유사하지 않은 텍스트-이미지 N×(N-1)쌍으로 이루어질 것입니다.

아래 그림을 보는 것이 이해를 도울 것입니다. 그림에서 파란색으로 표시된 쌍들은 서로 유사한 Positive Pair이며, 표시되지 않은 쌍들은 서로 유사하지 않은 Negative Pair입니다.

Untitled

유사도 메트릭 (similarity metric)으로는 코사인 유사도 (cosine similarity)를 사용합니다.

N개의 Positive Pair에 대해서는 코사인 유사도를 Maximize하고 N×(N-1)개의 Negative Pair에 대해서는 코사인 유사도를 Minimize하도록 학습합니다.

구체적으로, Symmetric Cross Entropy Loss를 사용합니다. 위 그림 상에서 Row별로, Column별로 각각 Binary Cross Entropy Loss를 구한 후, 두 Loss 값의 산술평균한 것을 Final Loss로 삼아 학습시킵니다.

논문에 첨부된 아래 Pseudo Code를 참고해서 이해할 수 있습니다.

Untitled

Code Review

Official Code Link:

GitHub - openai/CLIP: CLIP (Contrastive Language-Image Pretraining), Predict the most relevant text snippet given an image

마지막으로 위 Repository의 clip/model.py 소스코드를 리뷰해보겠습니다.

CLIP의 Main Code

CLIP 클래스

forward 메서드을 보면, 아래 코드와 같이, 먼저 이미지와 텍스트를 각자의 인코더를 이용해 임베딩합니다.

이 임베딩에 Normalization을 거친 후, 이미지와 텍스트를 각각 기준으로 잡고 코사인 유사도를 구합니다.

  def forward(self, image, text):
  	**image_features = self.encode_image(image)**
  	**text_features = self.encode_text(text)**

  	# normalized features
  	image_features = image_features / image_features.norm(dim=1, keepdim=True)
  	text_features = text_features / text_features.norm(dim=1, keepdim=True)

  	# cosine similarity as logits
  	logit_scale = self.logit_scale.exp()
  	**logits_per_image = logit_scale * image_features @ text_features.t()
  	logits_per_text = logits_per_image.t()**

  	# shape = [global_batch_size, global_batch_size]
  	return logits_per_image, logits_per_text

encode_image와 encode_text는 각각 사용할 인코더를 호출하는 메서드입니다.

encode_image는 아래 코드와 같이, ModifiedResNet과 VisionTransformer 중 하나로 지정된 Visual Encoder를 호출합니다.

  if isinstance(vision_layers, (tuple, list)):
  	vision_heads = vision_width * 32 // 64
  	self.visual = **ModifiedResNet**(
  	    layers=vision_layers,
  	    output_dim=embed_dim,
  	    heads=vision_heads,
  	    input_resolution=image_resolution,
  	    width=vision_width
  	)
  else:
  	vision_heads = vision_width // 64
  	self.visual = **VisionTransformer**(
  	    input_resolution=image_resolution,
  	    patch_size=vision_patch_size,
  	    width=vision_width,
  	    layers=vision_layers,
  	    heads=vision_heads,
  	    output_dim=embed_dim
  	)

encode_text는 아래 코드와 같이, token_embedding과 positional_embedding을 더하고 Transformer에 넣은 후 Laynorm과 Linear Projection을 거칩니다.

  def encode_text(self, text):
    x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]

    x = x + self.positional_embedding.type(self.dtype)
    x = x.permute(1, 0, 2)  # NLD -> LND
    x = self.transformer(x)
    x = x.permute(1, 0, 2)  # LND -> NLD
    x = self.ln_final(x).type(self.dtype)

    # x.shape = [batch_size, n_ctx, transformer.width]
    # take features from the eot embedding (eot_token is the highest number in each sequence)
    x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection

    return x

Encoders

ModifiedResNet 클래스

기존 ResNet과 다른 점이 3가지 있습니다. 기존 ResNet 구현은 아래 이미지를 참고하면 됩니다.

Untitled

Residual Layer들 이전 stem layer가 다릅니다.

아래 코드의 모듈들을 순서대로 적용합니다.

 # __init__ 내부

 self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
 self.bn1 = nn.BatchNorm2d(width // 2)
 self.relu1 = nn.ReLU(inplace=True)
 self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
 self.bn2 = nn.BatchNorm2d(width // 2)
 self.relu2 = nn.ReLU(inplace=True)
 self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
 self.bn3 = nn.BatchNorm2d(width)
 self.relu3 = nn.ReLU(inplace=True)
 self.avgpool = nn.AvgPool2d(2)

 # ------------

 def forward(self, x):
 	  def stem(x):
 	      x = self.relu1(self.bn1(self.conv1(x)))
 	      x = self.relu2(self.bn2(self.conv2(x)))
 	      x = self.relu3(self.bn3(self.conv3(x)))
 	      x = self.avgpool(x)
 	      return x

Average Pooling을 사용합니다.
기존 ResNet과 달리, stem layer의 마지막과 각 residual layer의 마지막에서 Average Pooling을 사용합니다. (BottleNeck 클래스에서 확인 가능)

Attention Pooling을 사용합니다.

ModifiedResNet의 마지막 layer는 Attention Pooling입니다. 그 구현은 아래 코드를 참고할 수 있습니다. 이는 multi-head attention을 포함하고 있습니다.

 class AttentionPool2d(nn.Module):
     def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
         super().__init__()
         self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
         self.k_proj = nn.Linear(embed_dim, embed_dim)
         self.q_proj = nn.Linear(embed_dim, embed_dim)
         self.v_proj = nn.Linear(embed_dim, embed_dim)
         self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
         self.num_heads = num_heads

     def forward(self, x):
         x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
         x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
         x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
         x, _ = F.multi_head_attention_forward(
             query=x[:1], key=x, value=x,
             embed_dim_to_check=x.shape[-1],
             num_heads=self.num_heads,
             q_proj_weight=self.q_proj.weight,
             k_proj_weight=self.k_proj.weight,
             v_proj_weight=self.v_proj.weight,
             in_proj_weight=None,
             in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
             bias_k=None,
             bias_v=None,
             add_zero_attn=False,
             dropout_p=0,
             out_proj_weight=self.c_proj.weight,
             out_proj_bias=self.c_proj.bias,
             use_separate_proj_weight=True,
             training=self.training,
             need_weights=False
         )
         return x.squeeze(0)

Transformer 클래스

Transformer는 ResidualAttentionBlock 여러 층을 쌓은 것입니다.

  class Transformer(nn.Module):
      def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
          super().__init__()
          self.width = width
          self.layers = layers
          self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])

      def forward(self, x: torch.Tensor):
          return self.resblocks(x)

ResidualAttentionBlock은 아래 코드와 같이 먼저LayerNorm → (Masked) Multi-Head Attention을 하면서 Residual Connection을 한 번 넣은 후, LayerNorm → MLP를 하면서 Residual Connection을 한 번 넣는 연산입니다.

  class ResidualAttentionBlock(nn.Module):
      def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
          super().__init__()

          self.attn = nn.MultiheadAttention(d_model, n_head)
          self.ln_1 = LayerNorm(d_model)
          self.mlp = nn.Sequential(OrderedDict([
              ("c_fc", nn.Linear(d_model, d_model * 4)),
              ("gelu", QuickGELU()),
              ("c_proj", nn.Linear(d_model * 4, d_model))
          ]))
          self.ln_2 = LayerNorm(d_model)
          self.attn_mask = attn_mask

      def attention(self, x: torch.Tensor):
          self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
          return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]

      def forward(self, x: torch.Tensor):
          x = x + self.attention(self.ln_1(x))
          x = x + self.mlp(self.ln_2(x))
          return x

VisionTransformer 클래스

Convolution으로 Image Feature를 얻고, positional embedding도 더하며, 이후 위에서 소개한 Transformer로 처리합니다.

  class VisionTransformer(nn.Module):
      def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
          super().__init__()
          self.input_resolution = input_resolution
          self.output_dim = output_dim
          self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)

          scale = width ** -0.5
          self.class_embedding = nn.Parameter(scale * torch.randn(width))
          self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
          self.ln_pre = LayerNorm(width)

          self.transformer = Transformer(width, layers, heads)

          self.ln_post = LayerNorm(width)
          self.proj = nn.Parameter(scale * torch.randn(width, output_dim))

      def forward(self, x: torch.Tensor):
          x = self.conv1(x)  # shape = [*, width, grid, grid]
          x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
          x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
          x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
          x = x + self.positional_embedding.to(x.dtype)
          x = self.ln_pre(x)

          x = x.permute(1, 0, 2)  # NLD -> LND
          x = self.transformer(x)
          x = x.permute(1, 0, 2)  # LND -> NLD

          x = self.ln_post(x[:, 0, :])

          if self.proj is not None:
              x = x @ self.proj

          return x

CLIP의 Training Code

CLIP의 학습에 관해서는 위에 링크를 걸어놓은 Official Repo에서 공개해놓은 코드가 따로 없어서 아래 링크에 누군가 올려놓은 코드를 잠시 살펴보겠습니다.

CLIP Training Code · Issue #83 · openai/CLIP

아래 코드에서는 CLIP에서 얻은 텍스트, 이미지 사이의 코사인 유사도에 대하여 Ground Truth와의 Cross-Entropy Loss를 계산했습니다. 그렇게 구한 두 가지 Loss 값을 평균낸 것을 최종 Loss로 사용하여 학습합니다.

loss_img = nn.CrossEntropyLoss()
loss_txt = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=5e-5,betas=(0.9,0.98),eps=1e-6,weight_decay=0.2) #Params from paper

for batch in train_dataloader :
    optimizer.zero_grad()

    list_image,list_txt = batch #list_images is list of image in numpy array(np.uint8)

    images= torch.stack([preprocess(Image.fromarray(img)) for img in list_image],dim=0)
    texts = clip.tokenize(list_txt)

    logits_per_image, logits_per_text = model(images, texts)

    ground_truth = torch.arange(BATCH_SIZE).to(device)
    total_loss = (loss_img(logits_per_image,ground_truth) + loss_txt(logits_per_text,ground_truth))/2
    total_loss.backward()

    convert_models_to_fp32(model)
    optimizer.step()
    clip.model.convert_weights(model)

Repo에 CLIP 사용 예제 cookbook도 포함되어 있던데, CLIP을 어떻게 이용할 수 있을지 궁금하다면 참고하면 되겠습니다 (Interesting 하더라구요 ~~)

이상, CLIP Paper Review를 마치겠습니다. 읽어주셔서 감사합니다~!

이미지와 텍스트의 표현 공간을 통합한 CLIP 논문 리뷰

Paper Review

Introduction

기존 연구의 한계

Why CLIP?

CLIP의 이미지 및 텍스트 인코더

CLIP의 학습 방법; “Contrastive Learning”

Multi-class N-pair Loss

Code Review

CLIP의 Main Code

Encoders

CLIP의 Training Code

301 Agency

Error

Paper Review

Introduction

기존 연구의 한계

Why CLIP?

CLIP의 이미지 및 텍스트 인코더

CLIP의 학습 방법; “Contrastive Learning”

Multi-class N-pair Loss

Code Review

CLIP의 Main Code

Encoders

CLIP의 Training Code

Templates (for web app):

Error