이 논문은 인-배치(in-batch) 대조 학습(contrastive learning)에서 미니 배치(mini-batch) 샘플링 문제를 다루며, 모델 성능 향상을 위한 BatchSampler라는 새로운 방법을 제안합니다.
인-배치 대조 학습의 맥락
대조 학습은 컴퓨터 비전(vision), 자연어 처리(language), 그래프(graphs) 등 다양한 데이터 영역에서 널리 사용되는 자기 지도 표현 학습(self-supervised representation learning) 전략입니다. 그 핵심은 의미적으로 유사한 인스턴스(instances)는 가깝게, 의미적으로 다른 인스턴스는 표현 공간에서 멀어지도록 학습하는 것입니다.
인-배치 대조 학습에서는 미니 배치 내의 모든 인스턴스가 다른 인스턴스에 대한 부정 샘플(negative sample)로 활용되는 '부정 공유 전략(negative sharing strategy)'이 사용됩니다. 대조 최적화 과정에서 부정 샘플은 매우 중요한 역할을 하며, 특히 '어려운 부정 샘플(hard negatives)'은 학습 성능을 개선하는 데 기여하는 것으로 알려져 있습니다.
문제점: 기존 미니 배치 샘플링 방법의 한계
기존의 인-배치 대조 학습 모델은 주로 데이터에서 각 미니 배치를 순차적으로 로드하는 방식에서 출발합니다. 로드된 미니 배치 내에서 어려운 부정 샘플을 샘플링하여 성능을 개선하려는 연구들이 있었지만, 미니 배치 크기는 전체 데이터셋 크기보다 훨씬 작기 때문에 이러한 '지역적인 샘플링 방법'으로는 모든 어려운 부정 샘플을 효과적으로 탐색하기 어렵다는 한계가 있습니다.
전통적으로 미니 배치를 구성하는 두 가지 주요 방법이 있습니다:
- Uniform Sampler (균일 샘플러): 각 학습 단계에서 배치를 무작위로 샘플링하거나 순차적으로 로드하는 가장 일반적인 방법입니다. 이는 단순하고 모델에 독립적이지만, 어려운 부정 샘플의 효과를 간과하며, 최적화에 거의 기여하지 않는 낮은 기울기의 '쉬운 부정 샘플(easy negatives)'을 포함하는 경향이 있습니다. 샘플링된 배치에 유사한 인스턴스 쌍의 비율이 낮습니다.
- kNN Sampler (kNN 샘플러): 표현 공간에서 서로 가장 가까운 인스턴스들을 클러스터링하여 배치를 구성합니다. 이는 자연스럽게 많은 어려운 부정 샘플을 포함하지만, 동일한 기본 레이블을 가진 인스턴스들도 함께 클러스터되는 경향이 있어, 배치 내에 거짓 부정 샘플(false negatives, FN)의 비율이 높아지는 문제가 있습니다. 이는 모델 학습을 오도할 수 있습니다.
따라서 인-배치 대조 학습을 위한 이상적인 부정 샘플러는 어려운 부정 샘플 활용과 거짓 부정 샘플 문제 완화 사이의 균형을 맞춰야 합니다.
BatchSampler 제안
BatchSampler는 인-배치 대조 학습을 위해 구분하기 어려운(hard-to-distinguish) 인스턴스로 미니 배치를 샘플링함으로써 대조 학습 성능을 개선하고자 합니다. 이는 기존의 지역적 샘플링 방법과 달리, 전체 데이터에서 인스턴스 간의 유사성을 기반으로 전역적으로(globally) 미니 배치를 샘플링합니다. BatchSampler는 비전, 언어, 그래프 등 기존의 인-배치 대조 학습 모델에 직접 적용할 수 있는 간단하고 일반적인 기술입니다.
BatchSampler의 기본 아이디어는 다음과 같습니다:
- 근접 그래프(Proximity Graph) 구축: 무작위로 선택된 인스턴스들의 근접 그래프를 구축합니다. 이 그래프의 각 엣지(edge)는 인스턴스 표현 간의 쌍별 유사성, 즉 부정 샘플의 난이도를 제어하는 데 사용됩니다. kNN Sampler의 거짓 부정 샘플 문제는 무작위 인스턴스를 선택하여 그래프를 구성함으로써 완화됩니다.
- 근접 그래프 샘플링: 근접 그래프에서 인스턴스(노드)를 샘플링하기 위해 재시작이 있는 랜덤 워크(Random Walk with Restart, RWR)를 활용합니다. 워커(walker)가 샘플링한 지역 이웃(local neighbors)은 서로 유사하여 구분하기 어렵다는 전제에 기반합니다.
근접 그래프 구축 상세
근접 그래프는 유사한 인스턴스를 연결하면서 거짓 부정 쌍을 줄이는 것을 목표로 합니다. 주어진 데이터셋에서 인코더로 생성된 인스턴스 표현 {eᵢ}을 사용하여 근접 그래프 G = (V, E)를 구성합니다. 여기서 V는 인스턴스 노드 집합이고 E는 노드 쌍의 집합입니다. 각 인스턴스 vᵢ에 대해, 먼저 전체 데이터셋에서 무작위로 M개의 후보 집합 Cᵢ를 선택합니다. 그 다음 Cᵢ에서 표현 eᵢ와 내적(inner product) 유사성이 가장 높은 K개의 인스턴스를 이웃 Nᵢ로 선택합니다. 후보 집합의 크기 M은 중심 노드와 즉각적인 이웃 노드 간의 유사성을 제어합니다. Proposition 1에 따르면 M이 클수록 두 인접 노드가 유사할 확률이 커지며, kNN으로 클러스터링된 노드 그래프와 유사해집니다. M이 작고 K에 가까우면 그래프의 인스턴스는 Uniform Sampler에 의해 무작위로 선택된 것으로 간주될 수 있습니다.
근접 그래프 샘플링 상세
BatchSampler는 근접 그래프에서 '워크(walk)'로서 미니 배치를 샘플링하며, 시작 노드에서 방문한 인스턴스를 수집합니다. Random Walk with Restart(RWR)가 사용되는데, 이는 워커의 행동을 제어하는 이론적으로 뒷받침되는 능력을 제공합니다. Algorithm 3에 설명된 대로, RWR은 시작 노드에서 출발하여 확률 α로 시작 노드로 순간 이동하거나(teleporting back) 엣지 가중치에 비례하여 이웃 노드로 이동합니다. 이 과정은 고정된 수(미니 배치 크기 B)의 노드가 수집될 때까지 계속됩니다.
Proposition 2는 RWR이 지역 클러스터에서 벗어날 확률이 그래프 컨덕턴스(graph conductance)와 재시작 확률 α에 의해 제한될 수 있음을 나타냅니다. 또한 RWR은 BFS(Breadth-first Sampling)와 DFS(Depth-first Sampling)라는 두 가지 샘플링 방법의 혼합을 보일 수 있습니다. α가 높을수록 워커는 BFS 행동에 근접하여 작은 지역 내에서 샘플링하는 경향이 있으며, α가 낮을수록 DFS처럼 더 먼 노드를 탐색하도록 장려합니다.
Pairwise Negative Sampling과의 차이점
미니 배치 샘플링 문제는 쌍별 부정 샘플링 문제와 다릅니다. Figure 1에 표시된 것처럼, 쌍별 부정 샘플링은 각 긍정 쌍에 대해 부정 인스턴스를 샘플링하는 반면, 미니 배치 샘플링 방법은 배치 내의 다른 인스턴스를 부정으로 재사용하는 미니 배치를 샘플링하는 데 중점을 둡니다. 이전 연구들은 주어진 앵커(anchor)에 대해 전체 데이터셋에서 유사한 부정 샘플을 전역적으로 선택하여 성능을 개선하려 했지만, BatchSampler는 각 쌍에 대해 하드 네거티브만 선택하는 것이 아니라, 많은 하드 네거티브를 포함하는 미니 배치를 전역적으로 샘플링하는 데 초점을 맞춥니다.
BatchSampler의 프레임워크 (Algorithm 1)
BatchSampler는 Algorithm 1에 요약된 대로 미니 배치 샘플러 역할을 하며 어떤 인-배치 대조 학습 방법에도 쉽게 적용할 수 있습니다. 학습 과정 동안 BatchSampler는 먼저 근접 그래프를 구축하고, 이 그래프는 t번의 학습 단계 후에 업데이트됩니다. 그런 다음 무작위로 시작 노드를 선택하고 RWR을 통해 근접 그래프에서 미니 배치를 샘플링합니다. 이 샘플링된 미니 배치는 표준 대조 학습 파이프라인에 공급되어 표현을 학습하고 모델 매개변수를 업데이트하는 데 사용됩니다.
실험 결과 및 분석
BatchSampler는 비전(SimCLR, MoCo v3), 언어(SimCSE), 그래프(GraphCL, MVGRL) 등 다양한 양식의 대표적인 대조 학습 알고리즘에 적용되어 평가되었습니다.
- 비전: SimCLR과 MoCo v3 기반 실험에서 BatchSampler는 ImageNet 데이터셋에서 일관되게 성능을 향상시켰습니다. CIFAR10, CIFAR100, STL10, ImageNet-100 등 다른 벤치마크 데이터셋에서도 SimCLR 및 그 변형 모델(DCL, HCL 포함)의 성능을 꾸준히 개선했습니다. 특히 중규모 데이터셋에서 개선 폭이 더 컸습니다.
- 언어: SimCSE 프레임워크에서 BERT 및 RoBERTa 백본을 사용하여 문장 표현 학습을 평가했습니다. BatchSampler는 7개의 STS(Semantic Textual Similarity) 작업에서 SimCSE-BERTbase의 평균 성능을 1.09% 개선하는 등 기준 모델의 성능을 일관되게 향상시켰습니다. 또한, DCL과 HCL이 사전 학습된 언어 모델의 분포 문제로 인해 성능 저하를 겪는 상황에서도, BatchSampler가 명시적으로 샘플링한 하드 네거티브가 이 문제를 완화하여 DCL 및 HCL의 성능을 크게 개선했습니다.
- 그래프: GraphCL과 MVGRL 프레임워크를 사용한 그래프 수준 분류 작업에서 BatchSampler는 7개의 벤치마크 데이터셋 모두에서 성능을 꾸준히 향상시켰습니다. 평균 0.4%에서 2.9%의 성능 개선을 보였으며, 분산(variance)을 줄여 더 견고한 표현 학습에 기여했습니다. BatchSampler를 사용한 MVGRL은 일부 데이터셋에서 최신 지도 및 자기 지도 방법보다 우수한 성능을 달성했습니다.
BatchSampler가 더 나은 성능을 보이는 이유 (Section 5.2)
Figure 4와 Figure 8의 분석 결과는 BatchSampler가 왜 더 나은 성능을 보이는지 보여줍니다.
- kNN Sampler는 명시적으로 유사한 쌍을 포함하는 배치를 생성하지만, 상당히 높은 거짓 부정 비율로 인해 성능 저하를 초래합니다.
- Uniform Sampler는 학습 중 거짓 부정 비율이 일정하게 유지되지만, 어려운 부정 샘플을 효과적으로 샘플링하는 데 실패합니다.
- BatchSampler는 후보 크기 M과 재시작 확률 α를 조절하여 샘플링된 배치의 난이도를 제어할 수 있습니다. 이를 통해 BatchSampler는 이 두 샘플링 방법 사이의 최상의 균형을 달성합니다. BatchSampler는 어려운 미니 배치를 샘플링하면서도 최적의 매개변수 설정 시 Uniform Sampler보다 약간 더 높은 거짓 부정 비율만 보입니다. 이러한 균형이 최상의 성능으로 이어집니다.
매개변수 분석 (Section 5.3, B.8)
- 후보 크기 M: M은 샘플링된 배치의 난이도를 조절하는 핵심 매개변수입니다. Table 4에 따르면, 대부분의 경우 M=1000일 때 성능이 최고조에 달하며 M이 증가함에 따라 성능이 빠르게 하락합니다. 이는 M이 클수록 유사한 인스턴스를 선택할 확률이 높아지지만, 동시에 거짓 부정 샘플을 포함할 가능성도 커져 성능 저하로 이어진다는 직관과 일치합니다. 제안된 기준은 소규모 데이터셋의 경우 M=500, 대규모 데이터셋의 경우 M=1000입니다.
- 재시작 확률 α: α는 샘플링된 배치의 난이도를 조절합니다. Figure 5와 Figure 9는 α가 증가함에 따라 코사인 유사도가 왼쪽으로 치우치지만(더 어려운 부정 샘플), 배치 내 거짓 부정 샘플도 더 많이 발생하여 트레이드오프가 존재함을 보여줍니다. 모델에 따라 최적의 α 값이 다를 수 있으며, 특히 사전 학습된 언어 모델의 경우 높은 α (예: 0.7)가 하드 네거티브를 통해 분포 문제를 완화하는 데 유리할 수 있습니다. 다른 방법론의 경우, α를 학습 진행에 따라 선형적으로 감소시키는 동적 전략(예: 0.2에서 0.05로)이 최상의 성능을 달성했습니다.
- 그래프 업데이트 간격 t: 근접 그래프는 t 훈련 반복마다 업데이트됩니다. Table 12에 따르면 t가 너무 짧거나 길면 성능이 저하될 수 있습니다. 너무 빈번한 업데이트는 학습 불안정을 유발할 수 있고, 너무 드문 업데이트는 학습 과정 중 임베딩 공간의 변화를 따라잡지 못해 하드 네거티브를 제대로 포착하지 못할 수 있습니다.
- 배치 크기 B: Table 11은 배치 크기 B의 영향을 보여줍니다. 비전에서는 B가 클수록 결과가 좋았고, 언어에서는 B=64에서 최적 성능을 보였으며, 그래프에서는 B 증가에 따라 성능이 소폭 개선되었습니다.
- 이웃 수 K: Figure 12는 이웃 수 K의 영향을 보여줍니다. K가 증가함에 따라 성능이 향상되지만, 특정 시점 이후에는 더 이상 큰 개선이 없거나 성능이 저하될 수 있습니다. K가 커질수록 근접 그래프의 규모가 커지고 BatchSampler가 더 작은 범위의 지역 클러스터(즉, 배치 내에서 더 어려운 부정 샘플)를 탐색하게 되지만, 너무 많은 K는 쉬운 부정 샘플을 도입할 수도 있습니다.
효율성 분석 (Section 5.4)
Table 5는 BatchSampler의 시간 비용을 분석합니다. 배치 샘플링 비용(Costₛ)은 대부분의 경우 배치를 사용한 훈련 비용(Costₜ)보다 훨씬 적습니다. 근접 그래프 구축 비용(Cost_G)은 ImageNet에서 100초 정도 걸리지만, 이는 t 훈련 단계에 걸쳐 분할되므로(Cost_G/t) 배치당 비용은 미미합니다.
추가 분석 (Section 5.5)
- 근접 그래프 샘플링 전략: RWR 외에 DFS, BFS, RW 등 다른 그래프 샘플링 방법과 비교했습니다. Table 6과 Figure 6은 RWR이 일관되게 더 나은 성능을 달성함을 보여줍니다. BFS는 가장 유사한 쌍을 가져오지만 많은 거짓 부정 샘플을 유발하여 성능이 저하됩니다. DFS와 RW는 거짓 부정 비율은 낮지만 하드 네거티브를 충분히 수집하지 못합니다. RWR의 재시작 특성은 DFS와 BFS의 혼합을 통해 샘플링된 배치의 난이도를 유연하게 조절하고 하드 네거티브와 거짓 부정 샘플 사이의 최적 균형을 찾을 수 있게 합니다.
- 근접 그래프 vs. kNN 그래프: 근접 그래프는 전체 데이터셋에서 단순히 가장 가까운 K개의 이웃을 선택하는 kNN 그래프보다 우수한 성능을 보입니다 (Table 7). kNN 그래프를 사용한 BatchSampler는 거짓 부정 샘플 문제로 인해 원래 대조 학습 방법보다도 성능이 떨어질 수 있습니다. Figure 7은 학습 과정에서 근접 그래프가 kNN 그래프에 비해 거짓 부정 비율을 현저히 감소시킴을 보여줍니다.
결론
이 논문은 인-배치 대조 학습을 위한 미니 배치 샘플링 문제를 해결하기 위해 BatchSampler를 제안했습니다. BatchSampler는 구축된 근접 그래프 위에서 워크로서 미니 배치 샘플링을 수행하며, 근접 그래프는 인스턴스 간의 쌍별 유사성을 제어하고, RWR은 배치를 형성하는 데 사용됩니다. 이론적 및 실험적으로 BatchSampler가 kNN Sampler와 Uniform Sampler 사이의 균형을 맞출 수 있음을 입증했습니다. 비전, 언어, 그래프 등 3가지 양식의 5가지 대표적인 대조 학습 알고리즘에 대한 실험을 통해 BatchSampler가 일관되게 성능을 향상시킴을 보여주었습니다.
BatchSampler의 코드는 https://github.com/THUDM/BatchSampler 에서 사용할 수 있습니다.