[논문 리뷰] Effective post-training embedding compression via temperature control in contrastive training

Effective post-training embedding compression via temperature control in contrastive training
ICLR 2025 spotlight
Georgiana Dinu, Corey Barrett, Yi Xiang, Miguel Romero Calvo, Anna Currey, Xing Niu
Amazon, Oracle
[paper]

해당 논문 리뷰는 Contrastive Learning에 대한 기초적인 지식을 갖고 있다는 전제하에 작성되었음

1. Abstract & Introduction

기존 연구 문제점

• 별도의 문제점을 제시하지 않음

제안 방법

• 고정 크기의 학습된 표현(dense representation or embeddig)은 검색, RAG, 분류, 클러스터링 등 여러 응용프로그램에서 중요한 역할을 함
• 대부분 contrastive loss를 사용하여 훈련하며, infoNCE가 표준적임
• 온도 매개변수 \(\tau\)는 어려운 negative sample에 대한 모델의 민감도를 조절하는데, 작을수록 더 균일한 분포, 클수록 더 나은 정렬을 유도함
• \(\tau\)는 내재적 차원(intrinsic Dimension)에 직접적인 영향을 미치며, 작은 \(\tau\)는 더 큰 내재적 차원으로 이어짐. (내재적 차원: 데이터가 내재적으로 포함하고 있는 정보의 차원 수, 데이터를 설명하는데 필요한 독립적인 변수의 수, 즉 클수록 데이터의 복잡성이 높고, 다양한 정보고 포함되)
• 작은 내재적 차원은 사후 훈련 압축에서 품질 유지에 기여함
• 저자는 텍스트 임베딩을 위한 contrastive learning에서 \(\tau\)의 영향을 분석하고, \(\tau\)조정을 통해 압축과 품질 유지를 동시에 달성할 수 있는 방법을 제안

2. Related Works

Impact of temperature

• \(\tau\)가 0에 가까워질 때, infoNCE는 하나의 부정 샘플(가장 가까운)을 사용하여 마진 0의 triplet 손실로 변환됨

• 반대로, \(\tau\)가 무한대로 갈 경우, 소프트맥스 변환을 적용하지 않는 모든 부정 샘플을 사용하는 단순 소실로 수렴됨.

• Contrastive Learning with Hard Negative Samples (ICLR 2021)은 새로운 농도 매개 변수 \(\beta\)를 도입하여, 어려운 부정 샘플의 가중치를 높일 수 있음을 보여줌 -> 긍정 샘플의 온도를 고정하고, 부정 샘플의 온도를 변화시키는 것과 같음

• 그림 왼쪽 랜덤 벡터의 경우, \(\tau\)가 작을수록, negatve 샘플의 영향의 분포가 상이해짐을 확인할 수 있음
• 오른쪽 중간 단계의 벡터의 경우, 대부분의 샘플의 영향력은 낮지만, 일부 샘플의 영향력이 매우 큰 것을 확인할 수 있고, \(\tau\)가 작을수록 도드라짐.

Choosing the optimal temperature parameter

• \(\tau\) 값 간에는 트레이드오프가 존재하며, 작은 온도가 항상 우수한 것은 아님
• 큰 \(\tau\)는 모든 부정 샘플을 동일하게 가중치를 부여하여, 어려운 부정 샘플들을 사용할 때, 전반적으로 더 나은 성능을 낼 수 있음
• 작은 \(\tau\)가 tail 클래스를 큰 \(\tau\)가 head 클래스에 최적
• 즉, 작은 \(\tau\)는 세밀한 인스턴스 수준의 구별을, 큰 \(\tau\)는 더 일반적인 그룹 수준의 구별을 촉진하며, 기존 연구들은 훈련 중 온도를 변화시키는 것이 최적의 결과를 가져오고, 0.07에서 1.0까지의 범위에서 효과적이라고 주장함

• 부정 샘플의 수 M이 무한대일 때, 정규화된 손실은 alignment과 uniformity로 나누어짐
• alignment는 앵커와 긍정 샘플 간의 유사성을 측정하고, uniformity는 부정 샘플과의 유사성을 측정함 -> 목표는 임의의 점들 간의 거리를 최대하하는 동시에 정렬을 개선하는 것

3. Temperature in Contrastive Training of Text Embeddings

3.1. Experimental setup

• 1단계: 대량의 raw text를 사용하여, MLM 방식으로 학습
• 2단계: infoNCE contrastive learning을 사용

Model architecture

• CodeSage, 356M, 1024 embedding dimensions

Data

• 1단게에서 2T 토큰의 데이터(80% 영어, 20% 다른 100개 언어)를 사용
• 2단계에서 2M개의 데이터 사용, 배치 256

Evaluation

• MTEB 벤치마크 사용, 56개 데이터셋 평가, retrieval은 nDCG@10, clustering은 v-measure를 사용하여 평가

3.2 Performance When Varying Temperature

• \(\tau\)을 0.04에서 0.4까지 변화시키며 MTEB 성능 측정
• \(\tau\)가 증가함에 따라 retrieval은 일관되게 감소, clustering은 반대로 증가하다 수렴하는 경향이 있음
• \(\tau\)가 group-wise, instance-wise 구별에 영향을 미치며, traid-off가 존재한다는 가설을 검증

• t-SNE 투영 결과, 이전 연구의 이미지 공간과 유사하게 \(\tau\)가 증가함에 따라 잘 정의된 클러스터가 나타나는 것을 확인할 수 있음.
• uniformity는 \(\tau\)가 증가할 때, 증가하지만, 그 이후에는 느린 감소 추세를 보임

• PCA를 기반으로, 데이터의 분산을 설명하는데 필요한 주성분의 수를 계산하여, 내재 차원성(intrinsic dimensionality)의 척도로 사용, 이떄 95%의 임계값을 사용
• 더 큰 \(\tau\)를 사용할수록 내재 차원이 감소하는 경향이 있음을 보여줌, 이는 \(\tau\)가 클러스터링 가능성을 높이는 방향으로 작용한다는 것을 의미함

4. Post-Training Embedding Compression

Compressing embeddings

• Random feature selection: 내재 차원성이 낮은 임베딩 공간은 차원 축소 방법을 적용할 때, 품질을 유지될 것으로 예상되며, 벡터를 잘라내어 크기를 줄이는 방법을 사용
• Binarization: 부호함수를 사용하여, 이진 양자화를 수행, 이는 저장 공간을 32배를 줄일 수 있으며, 이진 임베딩에서는 코사인 유사성을 해밍 유사성으로 대체함

• 왼쪽은 차원 갯수 별 실제 성능을, 오른쪽은 원래 차원(1024)대비 성능이 얼마나 보존되는지를 나타낸 것
• 온도가 높을수록, retrieval 성능은 떨어지지만, 임베딩 품질 보존은 더 잘함

• Binarization 또한 유사한 경향을 보이며, 압축률 대비 성능 보존은 Binarization이 더욱 뛰어남, 또한 두 가지를 같이 사용하였을 때는 128배 축소되었지만, 품질 유지율은 87%로 감소

Matryoshka Representation Learning (MRL)

• MRL은 원본 벡터를 잘라낸 후 더 작은 벡터를 최적화하는 로스를 사용하는 방법
• \(k=3\), \(d_i = [256, 512, 1024]\)를 사용

• MRL은 원래 임베딩(1024)에 대해서도 개선된 성능을 보여줌
• MRL은 더 작은 \(\tau\)에서 성능 개선이 두드러지지만, 전체 \(\tau\)에서 성능 개선을 보여줌
• MRL은 모든 \(\tau\)에서 내재 차원성을 낮추는 효과가 있음

5. Multiple Temperatures in Training

• 큰 \(\tau\)가 임베딩을 효율적으로 압축할 수 있으나, 이 특성과 성능 간의 트레이드 오프는 여전히 존재하며, 이러한 더 나은 트레이드 오프를 얻는 방법을 조사

• Plain temperature aggregation: 개별 \(\tau\)를 사용하는 infoNCE 로스의 합으로 대체
• 3개의(0.03, 0.06, 0.1) \(\tau\)를 사용

• MRL을 통해 성능을 개선 할 때, 식 6번의 로스 식을 사용, 여기선 각 가중치는 동일

• 이어서 낮은 \(\tau\)가 retrieval에 더 유리한 점을 활용하기 위해 256 차원에는 더 작은 \(\tau\)를 1024에는 더 큰 \(\tau\)를 사용하여 학습

• bin re-rnk는 이진 표현으로 상위 100개를 retrieved한 후, 전체 정밀도 쿼리로 re-ranking하는 방식
• full 기준으로 retrieval 성능은 최고 성능에 비해 조금 감소허나, clustering 성능은 크게 증가
• 압축 후 성능의 경우는 TempAggMRL이 가장 우수하며, 압축 후 성능 보존율의 경우 TempSpecMRL이 가장 우수함
• 다중 온도 사용이 임베딩의 압축 효율성과 성능 간의 균형을 잘 맞출 수 있음을 보여줌

6. Future Work

• 임베딩 공간의 내재 차원성을 더 깊이 탐구하고, 다양한 학습 방식에서의 압축과의 관계를 조사
• 이진화가 높은 압축 비율과 성능을 보였는데, 이 관찰이 유연한 임베딩을 얻는 네만 적용되는 것이 아니라, 학습에서의 over-parameterization 형태로 전반적인 성능 개선에 기여할 수 있는지를 탐구

리뷰

• 논문의 구성이 조금 독특함
• 대조학습의 시작은 vision 쪽인데, NLP쪽에서의 이러한 임베딩 압축 경향이 vision쪽, multimodal 측면에서는 어떻게 관찰 될 지 궁금함, 여기서 제사힌 관련 연구들을 훑어봐야할 것 같음
• 유사한 압축 경향을 보인다면, 임베딩 압축이 더 필요한 쪽은 vision & multimodal 쪽이 아닐지? 또 다른 방향의 경량화 연구가 될 수 있을 것 같음