Search
Duplicate

AI/ Memory Bank, Memory Attention

Background

Pretext Task

일반적으로 AI 모델들은 특정한 Task를 해결하기 위해 설계되고 학습되는데(예: 이미지 분류나 회귀 모델들) 특정 Task 자체를 해결하는데 목적을 두지 않고, 데이터셋의 좋은 representation(표현) 자체만 익히는 목적을 가진 작업을 Pretext Task라 한다.
이러한 작업은 모델에게 일반적인 표현을 학습시킨다음 다양한 downstream 응용에 전이(transfer) 시키기 위한 목적을 갖는다. 예컨대 이미지 모델에서 이미지 자체에 대한 좋은 representation을 학습하면, 이를 object detection이나 instance segmentation 같은 downstream task에서 fine-tuning하여 활용할 수 있다.

Contrastive Learning

Contrastive Learning은 일반적으로 Pretext Task를 위해 사용되는 방법으로 데이터에 명시적인 라벨링 없이 positive/negative 샘플만 활용하여 샘플의 representation을 학습하는 Self-Supervised learning 방법을 의미한다. 여기서 모델은 입력을 positive 샘플과 가깝게 하면서 동시에 negative 샘플과 멀리 떨어뜨리도록 학습된다.
Contrastive Learning은 라벨링 없이(unsupervised) 손 쉽게 positive/negative 샘플을 만들 수 있기 때문에(일반적으로 positive는 입력 데이터의 변형된 버전(random crop이나 flip, colorize 등)을 사용하고 negative 샘플은 임의의 샘플을 사용한다), 저렴한 비용으로 대량의 데이터를 학습할 수 있기 때문에 널리 쓰이게 되었다.
Contrastive Learning의 초기 모델인 Noise Contrastive Estimation(NCE)나 Triplet Loss에 명시적으로 존재하지 않았지만, Contrastive Learning의 학습 방식을 입력을 Query로 보고 positive/negative 샘플을 Dictionary의 Key로 보아 둘을 매칭하는 개념으로 생각할 수 있다.
Contrastive Learning의 Query, Key가 Attention에서의 Query, Key, Value와 다르다는 것에 주의. Contrastive Learning의 Dictionary에는 Query, Key만 존재하고 Value는 없으며, Query, Key, Value를 동일한 데이터를 사용하는 Self-Attention과 달리 Contrastive Learning의 Query는 입력이고, Key는 Positive나 Negative 샘플을 의미한다.
참고로 Negative Pair 없이 representation을 학습하는 방법도 존재하는데 이를 non-contrastive learning이라고 한다. BYOL이나 DINO 같은 방법이 이에 속한다.

Memory Bank

Memory Bank는 Contrastive Learning에 Query-Key 개념을 명시적으로 도입한다. 여기서 데이터셋의 모든 샘플을 Encoder를 통해 encoding 하여 CPU 메모리 공간(이곳을 Memory Bank라 한다)에 Key로 저장한 다음, Mini-Batch에서 입력 Query와 Memory Bank 내의 샘플링된 Key를 비교하여 loss를 계산하고 역전파 하여 Encoder를 업데이트한다.
새로운 Mini-Batch에서는 업데이트된 Encoder를 이용하여 Memory Bank에서 해당 Mini-Batch 내에 속하는 Key를 업데이트하고, 다시 Query와 샘플링된 Key를 비교하여 loss를 계산하고 역전파하여 Encoder를 업데이트하는 과정을 반복한다.
Memory Bank는 큰 Dictionary를 통해 다양한 샘플을 활용하여 모델이 더 일반화된 표현을 학습하는데 도움을 주지만, 무작위로 샘플링된 Key가 여러 단계의 encoder에 걸쳐 업데이트된 것이므로 일관성이 떨어진다는 단점을 갖는다.
이러한 문제를 개선한 방법으로 MoCo와 같은 방법이 등장했다.

참고