[논문 리뷰] Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

Luyu Gao, Yunyi Zhang외 2명|arXiv (Cornell University)|2021. 01. 18.

Domain Adaptation and Few-Shot Learning참고 문헌 16인용 수 5

한 줄 요약

이 논문은 인코더 역전파를 배치 기반의 대비 손실에서 분리함으로써 메모리 제약이 있는 GPU에서 임의의 큰 배치 크기로 훈련할 수 있도록 해주는 Gradient Cache 기법을 소개한다. 이 기법은 표현의 기울기(gradient)를 캐시함으로써 구현되며, 단일 소비자용 GPU를 사용하여 밀도 높은 파assage 검색에서 최신 기술 수준의 성능을 달성한다. 이는 다수의 고성능 GPU로 훈련된 모델들과 동일한 성능을 내지만, 훈련 시간은 20% 더 오래 걸리지만 배치 크기에 관계없이 일정한 메모리 사용량을 유지한다.

ABSTRACT

Contrastive learning has been applied successfully to learn vector representations of text. Previous research demonstrated that learning high-quality representations benefits from batch-wise contrastive loss with a large number of negatives. In practice, the technique of in-batch negative is used, where for each example in a batch, other batch examples' positives will be taken as its negatives, avoiding encoding extra negatives. This, however, still conditions each example's loss on all batch examples and requires fitting the entire large batch into GPU memory. This paper introduces a gradient caching technique that decouples backpropagation between contrastive loss and the encoder, removing encoder backward pass data dependency along the batch dimension. As a result, gradients can be computed for one subset of the batch at a time, leading to almost constant memory usage.

연구 동기 및 목표

• 대조 학습에서 배치 크기를 제한하는 GPU 메모리 제약을 극복함으로써, 제한된 하드웨어를 가진 연구자들이도 활용할 수 있도록 하기 위해.
• 고품질의 표현 학습에 필수적인 대규모 배치로 훈련을 가능하게 하기 위해, 단일 소비자용 GPU에서 대규모 배치 훈련을 가능하게 하기 위해.
• 근사치 없이 대규모 배치 훈련과 동일한 정확한 기울기 업데이트를 유지함으로써, 다중 GPU 설정과 동일한 모델 성능을 확보하기 위해.
• 학술 및 자원이 제한된 환경에서도 대규모 배치 대비 학습에 접근 가능하도록 하기 위해, 이를 민주화하기 위해.

제안 방법

• 두 단계 역전파 프로세스를 도입함: 먼저 손실에서 표현으로의 기울기를 계산하고, 그 다음 표현에서 모델 파라미터로의 기울기를 계산함.
• 표현에 대한 손실의 기울기(표현 기울기 캐시)를 캐시하여 인코더 업데이트에서 배치 의존성을 분리함.
• 표현의 미니배치를 통해 서브 업데이트 방식으로 인코더 최적화를 수행함으로써, 고정된 GPU 메모리 내에서 큰 유효 배치 크기를 달성함.
• 파arameterized 유사도 함수를 위한 별도의 거리 기울기 캐시를 도입함으로써, 깊이 있는 유사도 함수에 동일한 원리를 적용함.
• 캐시된 기울기를 사용하여 전체 배치 기울기를 메모리 효율적으로 계산함으로써, 전체 배치의 활성화를 모두 저장할 필요 없이 정확한 기울기 업데이트를 수행함.
• 유사도 함수 Φ(·,·)에 대한 깊이 있는 유사도 함수를 지원하기 위해, 거리 기울기와 표현 기울기를 모두 캐시하여 공동 최적화를 가능하게 함.

실험 결과

연구 질문

• RQ1 메모리 제약이 있는 단일 GPU에서 대규모 배치 크기로 대조 모델을 훈련시킬 수 있는가?
• RQ2 대조 손실에서 인코더 역전파를 분리함으로써 전체 배치 훈련과 동일한 정확한 기울기 업데이트를 달성할 수 있는가?
• RQ3 제안된 방법이 단일 소비자용 GPU를 사용하여 밀도 높은 검색 벤치마크에서 최신 기술 수준의 성능을 재현할 수 있는가?
• RQ4 기울기 누적과 순차적 훈련과 비교해 본 결과, 이 방법의 훈련 속도 및 메모리 사용량은 어떻게 스케일링되는가?
• RQ5 간단한 내적 유사도를 초월한 깊이 있는 유사도 함수로도 이 기법을 일반화할 수 있는가?

주요 결과

• Gradient Cache 기법은 단일 RTX 2080ti에서 배치 크기 512로 훈련이 가능하며, 8장의 V100 GPU로 훈련된 모델과 동일한 성능을 달성한다.
• 20% 더 오래 걸리지만, MS MARCO 파assage 랭킹 벤치마크에서 표준 및 기울기 누적 훈련 기반 모델을 모두 초월하는 최신 기술 수준의 성능을 달성한다.
• 배치 크기에 관계없이 일정한 최고 메모리 사용량을 유지하므로, 하드웨어 한계를 초월한 배치 크기 확장을 가능하게 한다.
• 기울기 누적은 각 업데이트당 더 적은 인바이트 네거티브 샘플을 가지므로 성능이 열등하고, 순차적 훈련은 높은 분산과 낮은 정확도로 인해 성능이 열등하다.
• 이중 캐시 메커니즘을 통해 깊이 있는 유사도 함수로의 일반화에 성공하였으며, 유사도 및 표현 파라미터에 대해 정확한 기울기 계산을 유지한다.