CLIP 구현을 이해하는 방법: 대칭 Contrastive Loss와 Zero-shot 분류

CLIP의 구현은 이미지와 텍스트를 인코딩(encode)하여 내적(inner product)을 취하는 것만으로 보입니다. 하지만 그 내적을 무엇에 대해 분류 손실(classification loss)로 변환할지, 학습 시의 로짓(logits)과 zero-shot 추론 시의 로짓이 어떻게 대응하는지를 추적하지 않으면, logit_scale나 프롬프트(prompt)의 역할을 오해하기 쉽습니다.

먼저 결론을 말씀드리자면, CLIP은 다음 세 가지 계산으로 나누어 읽는 것이 이해하기 쉽습니다.

• 이미지와 텍스트를 각각 다른 인코더로 처리하여 같은 차원의 정규화 벡터(normalized vector)로 변환합니다.
• 배치(batch) 내의 모든 조합으로부터 로짓을 만들고, 대각 요소(diagonal element)를 정답으로 삼습니다. (N×N)
• 추론 시에는 클래스 프롬프트(class prompt)의 텍스트 임베딩을 분류기 가중치(classifier weight)로 사용합니다.

이 글에서는 CLIP 원 논문과 OpenAI 공식 repository를 근거로, 텐서 형태(tensor shape), 손실 함수(loss), 레퍼런스 구현(reference implementation), zero-shot 분류, 프로덕션 환경에서의 주의점들을 연결하여 설명합니다.

두 개의 인코더는 파라미터가 아닌 좌표계를 공유한다

CLIP은 Dual Encoder(두 가지 입력을 별도의 네트워크로 인코딩하는 구조)입니다. 이미지와 텍스트는 입력 형식이 다르기 때문에, Image Encoder와 Text Encoder의 파라미터를 공유하지 않습니다. 두자의 출력을 linear projection을 통해 같은 차원으로 정규화한 후에는 내적을 코사인 유사도(cosine similarity)로 취급할 수 있습니다.

공유되는 것은 최종적으로 비교 가능한 임베딩 공간(Embedding space)입니다. 두 개의 인코더를 같은 네트워크에 넣는 것이 아닙니다. 원 논문에서는 Image Encoder에 수정된 ResNet 또는 Vision Transformer를, Text Encoder에 Transformer를 사용했습니다.

학습은 N×N의 대응 분류가 된다

온도(temperature)를 적용한 각 이미지에서 올바른 텍스트를 선택하는 손실은 행 방향의 교차 엔트로피(cross entropy)입니다. 각 텍스트에서 올바른 이미지를 선택하는 손실은 열 방향, 즉 전치된 로짓(transposed logits)에 대한 교차 엔트로피입니다.

CLIP은 이 두 방향을 평균합니다.

한 방향의 손실만으로도 대응 분류는 가능하지만, 이미지 검색과 텍스트 검색 양방향 모두 정돈된 공간을 만들기 위해서는 양쪽을 목적 함수에 포함하는 것이 자연스럽습니다.

수식을 최소 PyTorch 구현에 맞춘다

다음 코드는 인코더 이후의 피처(feature)를 받는 교육용 레퍼런스입니다. 분산 학습 시의 all-gather, mixed precision, gradient-preserving gather는 생략했습니다.

import torch
import torch.nn.functional as F
from torch import Tensor
...

구현을 읽을 때 대응 관계가 명확합니다.

프로덕션 환경의 분산 학습에서는, 로컬 디바이스 내에서만 negative를 처리할지, 모든 디바이스의 임베딩을 모을지에 따라 실효 배치(effective batch)가 달라집니다. all-gather한 텐서에 그래디언트(gradient)를 흐르게 하는 설계인지도 확인이 필요합니다.

Temperature와 batch negative를 분리해서 생각한다

temperature는 유사도 순위를 바꾸지 않으면서, softmax에 들어가는 차이를 확대하거나 축소합니다. 예를 들어 유사도가 0.8, 0.7인 두 후보가 있다면, 각각 52.5%, 47.5% 또는 73.1%, 26.9%가 됩니다.

원 논문에서는 temperature에 해당하는 logit scale을 학습 대상으로 삼았습니다. 구현에서 logit_scale을 로그 파라미터(log parameter)로 유지하는 경우, forward 시점에 exp()를 사용하여 사용하는 점을 놓치면 스케일의 의미가 달라집니다.

batch size

다만, 비지원 pair가 의미적으로도 negative인 것은 아닙니다. 강아지 이미지가 여러 장 있는 batch에서는, 다른 이미지에 붙은 설명문도 내용상으로는 올바를 가능성이 있습니다. 이것은 false negative입니다. batch를 크게 할수록 후보는 늘어나지만, data 중복이나 의미가 가까운 pair를 어떻게 다룰 것인가라는 문제도 커집니다.

zero-shot 분류에서는 text embedding이 weight가 된다

학습된 CLIP을 이미지 분류에 사용할 때는, 대상 class별로 prompt를 만듭니다.

dog

, cat

, car

을 A photo of a {label}.

에 임베딩한다 - Text Encoder로 각 prompt를 encode하고 정규화한다

• 입력 이미지를 Image Encoder로 encode하고 정규화한다
• 이미지 vector와 모든 class vector의 similarity를 계산한다
• temperature scaling과 softmax를 적용한다

class

여기서 말하는 zero-shot은, '그 개념을 pre-training data에서 한 번도 보지 못했다'는 보장은 아닙니다. 대상 dataset의 label이 붙은 example로 classifier를 추가 학습하지 않고, text로 task를 지정하는 평가 방법입니다.

Prompt는 입력문이 아니라 classifier 설계의 일부

boxer는 견종과 운동선수, crane은 새와 건설 기계를 가리킬 수 있습니다. class 이름만 encode하면, 어떤 의미의 vector를 만들어야 할지 모호합니다.

| Prompt | 지정 가능한 문맥 | Trade-off |
|---|---|
| boxer | 거의 없음 | 짧지만 다의어에 취약 |
| a photo of a boxer dog | 견종, 사진 | 의미를 한정할 수 있다 |
| a photo of a boxer athlete | 인물, 운동 | 같은 label 단어를 다른 개념으로 분리할 수 있다 |
| a satellite photo of farmland | 촬영 domain | domain에 맞지 않으면 역효과가 날 수 있다 |

원 논문에서는, class 이름만 사용하는 것보다 A photo of a {label}.를 사용하는 것이 ImageNet top-1 accuracy를 1.3 point 개선했습니다. 80종류의 context prompt를 사용하는 ensemble은 default prompt에서 추가로 3.5 point 개선하고 있습니다.

이 수치를 다른 domain에 그대로 일반화할 수는 없습니다. 공통 template, domain 고유 template, 여러 template의 embedding 평균, 이용 언어를 나누어 평가해야 합니다. original CLIP은 영어 중심의 text data로 학습되었기 때문에, 일본어 prompt의 자연스러움만으로 판단하는 것은 안전하지 않습니다.

CLIP의 출력과 주변 model의 출력을 혼동하지 말 것

CLIP 단독의 핵심적인 출력은 Embedding과 similarity입니다. 자유문, Bounding Box, 이미지 pixel을 직접 생성하지 않습니다.

| 처리 | 주된 출력 | CLIP의 위치 설정 |
|---|---|
| text-to-image retrieval | 이미지 순위(rank) | query와 후보의 Embedding을 비교한다 |
| ... |
「CLIP을 사용하는 이미지 생성 system」과 「CLIP이 이미지를 생성하는 것」은 별개의 주장입니다. component의 입출력으로 책임을 분리합니다.

Production에서는 Embedding 외에도 version 관리를 할 것

검색이나 고정된 taxonomy 분류에서는, text 후보를 사전에 encode하여 cache할 수 있습니다. 대량 이미지를 검색 대상으로 하는 경우에는 image embedding도 사전 계산할 수 있습니다. 다만, 다음을 하나의 index version으로 관리해야 합니다.

• Image Encoder와 Text Encoder의 model version
• tokenizer와 이미지 preprocessing
• normalization과 logit scale 처리
• prompt template과 class taxonomy
• embedding dimension과 저장 dtype

encoder만 업데이트하고 오래된 index를 사용하면, 같은 좌표계에 있다는 보장이 없어집니다. 재encode 필요 여부와 rollback 단위를 미리 결정해야 합니다.

평가도 평균 accuracy만으로는 충분하지 않습니다.

요약

CLIP의 학습은 별개의 encoder가 출력한 정규화 벡터로부터

구현에서는 normalize,
행렬 곱셈(matrix multiplication), temperature, 대각 레이블(diagonal label), 전치된 양방향 손실(transposeした二方向loss)의 대응을 추적하면 전체가 연결됩니다. 운영 시에는 prompt뿐만 아니라 encoder, preprocessing, taxonomy, cache를 동일한 버전 경계에서 관리할 필요가 있습니다.

배경, 구체적인 4×4 계산, 논문의 데이터셋별 결과, CLIP 단독으로는 할 수 없는 것까지 순서대로 확인하고 싶은 분은 개인 블로그의 'CLIP 완전판'도 참고해 주세요.

참고 자료

토론(Discussion)