Contrastive learning의 대표적인 논문 중 하나인 MoCo에 대해서 소개하도록 하겠습니다.
FaceBook, 현재 Meta인 Facebook AI research에서 Kaimming He께서 작성한 논문입니다. (ResNet 만드신 분)
MoCo에 대한 개념 설명 및 동작원리에 조금 더 초점을 맞추고 설명을 하도록 하겠습니다.
MoCo의 sudo코드와 실제 동작원리가 와닫지 않을 수 있어서 그림 및 차원을 통해 설명하도록 하겠습니다.
+ MoCo그림에서 x^key_i, k0 과 같은 부분이 와닿지 않을 수 있어서 그부분을 조금 더 풀어보았습니다.
BackGround.
배경지식으론, Representation Learning과, Contrastive Learning에 대해서 알고 넘어가야 합니다.
Representation Learning.
Representation Learning이란 Unsupervised-Learning 혹은 Self Supervised Learning에서 자주 등장하는 용어로, 단순하게 정확도가 높은 모델을 만드는 것이 아닌, 학습 데이터의 특징을 잘 학습하는, 데이터의 정보가 잘 구축된 모델을 만드는 것이 목표입니다.
Contrastive Learning
Contrastive Learning이란, Unsupervised, self-supervised 에서 데이터의 표현을 잘 학습하는 것을 목표로 합니다.
따라서 아래의 이미지와 같이 공간에다 이미지에 대한 정보를 mapping 하는 것과 같습니다. Anchor를 기준으로, Positive(Anchor와 같은 성질)일 경우 Anchor와 가깝게(Similar 라고 불립니다.), Negative(Anchor와 다른 성질)일 경우 Anchor와 멀게 (dissimilar라고 불립니다.) Mapping시키는 방법입니다. Mapping 거리에 기준은 loss function을 사용하며, Positive, Netative에 대한 기준은 Label을 사용합니다.
similar할 경우 Pull, dissimilar일 경우 push라는 용어를 사용하기도 합니다.
Contrastive Learning이란 개념이 어디서 왔는지 궁금하신 분은 아래 링크에서 가볍게 보신 후 돌아오시면 더욱 도움이 되실 것입니다. MoCo Abstract에서 언급되는 유일한 논문인 DrLIM입니다, 본문에도 자주 언급되기에 읽으시면 좋지만 핵심 내용은 간단하기에, Contrastive Learning이 와닿지 않으신다면 도움이 될 것이라 생각합니다.
Contrastive Learning 알아보기 (1) - DrLIM (Dimensionality Reduction by Learning an Invariant Mapping)
MoCo Review를 하기 전, MoCo에서 많이 언급되며, Contrastive Loss를 처음으로 사용한 논문으로 소개되는 Hadsell - Dimensionality Reduction by Learning an Invariant Mapping (DrLIM) in 2006 CVPR 논문에 대해서 먼저 정리하고
187cm.tistory.com
MoCo에서의 Contrastive Learing
Moco에서는 Query와 key라는 개념을 사용합니다. 입력 샘플에 대한 정보를 Query, Key로 제공하는데, Query는 Anchor sample을 의미하며, Key는 Similar(Positive sample) 혹은 Dissimilar(Negative sample)을 의미합니다.
또한 Dictionary(사전) 개념을 사용하는데, Dictionary는 전체 데이터 샘플에 대한 정보 (전체 데이터 샘플)를 의미합니다.
실제 사전이 단어의 뜻을 가지고 있는 것처럼, 데이터 샘플에 대한 정보, 표현을 가지고 있다고 생각하시면 될 것 같습니다.
Query: Anchor Sample
Key: Similar(Positive sample) 혹은 Dissimilar(Negative sample)
Dictionary: 전체 데이터 샘플의 정보를 가지고 있는 저장소 역할.
따라서 위의 그림을 통해 Contrastive Learing을 설명하자면, 2개의 입력 샘플 x_query와 x_key가 주어지는데,
encoder q, k를 통과한 값을 Contrastive Loss를 사용해 거리를 측정합니다.
x_query와 x_key가 Similar(positive)관계라면 가까운 값이, dissimilar(Negative)라면 큰 값이 나올 것입니다. 여기서 encoder를 Dictionary라고 부르며, x_query, x_key가 들어왔을 때 이 입력 query, key에 대한 정보를 제공하는 사전 역할을 합니다.
Related Work.
Related work로는 아래 3가지 Case가 존재합니다.
(a) end-to-end case는 위에서 잠깐 언급한 초창기 Contrastive Learning Method입니다. encoder q와 k는 동일한 모델을 사용합니다. (다른 모델을 쓰기도 함) encoder q는 샘플 하나에 대한 표현 정보를 제공하며, encoder k는 모든 key sample에 대한 representation정보를 저장한 후, 입력 Key에 대한 Representation 정보를 제공해야 합니다. 모델의 크기가 크면 클수록 다양한 정보를 표현할 수 있으므로, encoder의 크기가 커야한다는 문제가 있었으며, 이는 메모리 문제로 이어지는 경우가 많았습니다.
(b) Memory bank방식은 기존의 end-to-end case에서 메모리 문제를 해결하기 위해 나온 방식입니다. k는 memory bank에서 Random probability로 제공이 됩니다. query encoder에만 gradient를 통과시켜주면 되므로 메모리를 효과적으로 사용할 수 있습니다. 하지만 query encoder만 학습이 되므로 최신 정보가 반영되지 않는다는 점과, Random probability로 제공이 되므로 positive sample이 학습되기 어렵다는 단점이 있었습니다.
Moco (Momentum Contrast)
Purpose:
1. 좋은 표현을 학습한 모델을 만들어 Downstream task에 적용이 가능한 모델을 만드는 것.
2. Dynamic Dictionary를 만들어 지속적이고((b)단점 보완), 큰 규모의 학습((a)의 메모리 문제 해결)이 가능하게 만드는 것.
- Dictionary를 Queue로 만들어 큰 정보를 저장할 수 있게 하였습니다.
- Momentum encoder를 만들어 지속적인 정보를 가지는 encoder를 제공. (b)에서 query encoder만 학습되는 문제 해결
Method:
- InfoNCE Loss 사용하여 similar한 k는 q와 가깝게, dissimilar한 k는 q와 멀게 만듭니다.
Moco의 동작 원리
그림으로 설명하면 아래와 같습니다. 자세한 흐름은 바로 밑의 Algorithm과 같이 보시면 됩니다.
우선 초기화, for문 시작하는 부분을 다 제외하고 핵심 부분만 설명 드리면 다음과 같습니다
1. 입력 샘플 x_query는 encoder를 통과하며, x_key0, x_key1 ... 은 queue의 들어있는 압축된 k0, k1 .. 과 같은 Latent vector를 만드는 샘플들입니다. 이미 momentum encoder를 통과한 샘플들 입니다. 그림을 보면 화살표 방향으로 통과한 모습 대신, 중괄호와 같이 묶은 모습을 볼 수 있습니다. (그림의 모호성)
2. k0, k1, k2 ..은 momentum encoder를 통과하여 queue의 들어있는 값 즉 Latent vector들입니다.
3. 아래 코드와 같이 mini-batch 단위의 입력 sample x는 encoder 및 momentum encoder를 통과합니다. 따라서 그림을 위와 같이 수정하였습니다. (알고리즘과 그림의 차이)
+ 따라서 positive sample의 수는 batch의 단위인 N, Negative sample의 수는 queue의 수인 K=65536이 되어야 합니다.
+ sudo코드에서 l_pos가 1로 만들지지만 bmm을 통해 feature dimension을 압축하여 1로 만든 것이지 Positive sample의 수가 1이 되는 것이 아닙니다.
+ queue의 크기가 65536이고 (ImageNet) IN을 학습 데이터로 쓸 때 140만 장의 학습 이미지 중 6만장만 학습으로 쓰이고, 나머지는 안 쓰이는 것이 아닌, enqueue, dequeue를 통해 쓰고난 데이터는 뒤로 보내버립니다.
4. Positive 계산은 항상 q와 k_0의 batch multiplicaiton으로 진행됩니다. (알고리즘 참조 0번 idx는 positive)
5. Negative 계산은 q와 queue에 남아있는 원소들간의 matrix multiplication을 통해 진행하게 됩니다.
6. contrastive Loss를 구합니다.
7. query encoder에 대해 gradient update 후, momentum contrast 방식을 통해 momentum encoder를 업데이트 합니다.
8. Queue에 대해 enqueue, dequeue를 진행합니다. 새로 들어가는 원소는 Positive key의 latent vector입니다.
9. 그 다음 queue의 원소와 새로운 q에 대해 positive연산을 나머지 queue의 원소들로 negative 연산을 반복합니다.
+ 반복을 통해 queue의 원소는 각각 다른 시점에 만들어진 momentum encoder의 값으로 구성되게 됩니다. consistency
구체적인 알고리즘은 다음과 같습니다.
f_k.params = f_q.params # 같은 모델을 사용하므로, weight를 똑같이 맞춰줍니다.
for x in loader: # queue에서 나온 x를 가지고 positive 연산을 진행, 나머지 원소로는 Negative 연산을 합니다
# x의 이미지 사이즈는 224x224.
# augmentaiton은 "color jitter", "random resize", "horizontal flip", "grayscale"을 사용
x_q = aug(x) # random augmentaiton을 진행합니다.
x_k = aug(x) # 또 다른 Augmentaiton을 진행합니다.
q = f_q.forward(x_q) # queries: NxC
k = f_k.forwawrd(x_k) #keys: NxC
k = k.detach() # gradient 업데이트 X
l_pos = bmm(q.view(N, 1, C) , k.view(N, C, 1)) # Positive sample에 대해 연산.
l_neg = mm(q.view(N,c), queue.view(C,K)) # negative sample에 대한 연산
logits = cat([l_pos, l_neg], dim = 1) #logits: N x (1+k)
# Contrastive Loss! 왜 InfoLoss가 아닌 CE Loss를 사용하였을 까.
labels = zeros(N)
loss = CrossEntropyLoss(logits/t, labels)
# SGD를 사용하여 update // Only encoder
loss.backward()
update(f_q.params)
# momentum update.
f_k.params = m*f_k.params + (1-m) * f_q.params
dequeue(queue, k)
enqueue(queue)
알고리즘은 다음과 같습니다. 특이점이 있다면, InfoNCE Loss를 사용했다고 말하였지만, CrossEntropyLoss를 사용하는 트릭을 사용했다는 점입니다. 그 이유는 아래의 CrossEntropy 수식에서 생략되어있는 Softmax를 적용하면 됩니다. CrossEntropy에 공통으로 들어가는 log는 밖으로 빠지게 되며, label = zeros(N) 이라는 트릭을 사용하는 부분에서
Positive sample만 1, 나머지는 Negative sample, 즉 0으로 만들어 분자에는 Positive sample만 들어가게 되며, 분모는 Softmax의 분모와 같으므로 InfoNCE Loss의 분모와 같게 됩니다.
InfoLoss와 CrossEntropy Loss, Softmax는 다음과 같습니다.
성능
MoCo의 성능이 기존의 End-to-End 방식과 비교하였을 때, 더 좋은 것을 볼 수 있으며, 다양한 Downstream Task에 적용했을 때, 다양한 Task에서 높은 성능을 보인 것을 보였습니다. 또한 메모리 측면에서 queue를 사용하여 효율적인 메모리를 사용했다 라고 볼 수 있습니다.
