CDE(Contextual Document Embeddings) 모델 소개
문서 검색에서 문서 임베딩(Document Embedding)은 텍스트 데이터를 잠재 공간에 임베딩하여 문서와 질의 간의 유사도를 계산하는 중요한 요소입니다. 기존의 신경망 기반 임베딩 방식은 개별 문서만을 인코딩해 문맥을 고려하지 못하는 한계가 있습니다. Contextual Document Embeddings라는 제목의 이 연구는 문서 주변의 문맥 정보를 활용하여 문서 임베딩 성능을 개선하는 방법을 활용하여 문맥 기반 문서 임베딩(Contextual Document Embedding, CDE)을 제안하고 있습니다.
CDE 연구에 영향을 준 임베딩에 대한 기존의 연구들을 살펴보면, 크게 다음의 4가지로 나누어 볼 수 있습니다:
-
문서 검색(Text Retrieval): 전통적으로 단어 빈도 등을 활용한 통계적 방법이 사용되어 왔습니다. BM25와 같은 초기 텍스트 검색 시스템에서 활용한 통계적 모델은 **n-그램 빈도(n-gram TF; Term Frequency)**와 역문서 빈도(IDF; Inverse Document Frequency)를 통해 특정 단어의 중요도를 평가했습니다. 최근에는 신경망 기반 이중 인코더(Bi-encoder) 방식이 발전하면서, 문서와 사용자 질의(Query)를 밀집 벡터 공간으로 변환하여 더 나은 성능을 보이고 있습니다. DPR(Dense Passage Retrieval), GTR(Generalized Text Representations) 등과 같은 방법은 문서의 의미는 더 잘 반영하는 임베딩을 생성할 수 있지만, 각 문서를 독립적으로 처리하고 있기 때문에 전체적인 문맥을 반영하지는 못한다는 한계가 있습니다.
-
대조 학습(Contrastive Learning): 대조 학습은 **양성 샘플(예. 서로 유사한 항목들)**과 음성 샘플(예. 서로 다른 항목들) 간의 차이를 극대화하여 임베딩을 학습하는 방법입니다. 이러한 과정을 통해 비슷한 항목은 가까이에 모이고, 서로 다른 항목은 멀리 두는 방식으로 학습이 진행됩니다. 예를 들어, 문서 검색에서 대조 학습은 사용자 질의와 관련한 문서들을 긍정 샘플로, 그 외의 문서들을 부정 샘플로 취급하여 임베딩 공간 상에서 이를 구분하는 방식으로 사용됩니다. 이러한 방식은 다양한 도메인에서 효과적으로 동작하나, 성능 향상에 중요한 역할을 하는 음성 샘플을 선정하는 방법(Hard Negative Sampling)이 까다로운 것이 문제입니다.
-
시험 시간 적응(Test-time Adaptation): 모델이 새로운 데이터셋에 직면했을 때 이를 더 잘 처리할 수 있도록 사용자 질의와 문서 간의 유사도를 조정하(여 적응시키)는 방법입니다. PRF(Pseudo-Relevance Feedback)가 이러한 시험 시간 적응 방법의 대표적인 예시로, 사용자 질의와 관련한 문서들을 사용하여 사용자 질의 자체를 보완하는 방식입니다. 하지만 이러한 방식은 검색 도메인마다 추가적인 모델 학습이 필요하거나, 특정 도메인에 한정된 성능 개선만 연구되었습니다.
-
비-매개변수 모델링(Non-parametric Modeling): 매개변수의 수를 고정하지 않고, 데이터의 복잡성에 따라 모델이 유연하게 변화하는 방법입니다. 최근 딥러닝 연구에서 이러한 접근 방식이 많이 연구되었으며, NPT(Non-parametric Transformer)와 같은 모델이 대표적입니다.
CDE 연구에서는 기존의 이중 인코더(Bi-encoder)를 개선하기 위해 문맥 정보를 활용한 새로운 학습 방법과 아키텍처를 제안하고 있습니다. 특히, 문서 간의 문맥 정보를 활용하여 데이터를 여러 하위 데이터셋으로 분할하는 클러스터링 알고리즘을 사용하는 하드 네거티브 샘플링 방법을 제안합니다.
또한, PRF와 같은 문서 간의 문맥을 반영한 임베딩을 생성하여 문서와 사용자 쿼리 간의 연관성을 더 잘 학습하도록 합니다. 추가적인 성능 향상을 위해 검색 작업에 최적화한 Semi-Parametric Model을 구축하여 더 나은 검색 결과를 제공할 수 있도록 합니다.
CDE 연구가 제안하는 방법론 소개
기존의 신경망 기반의 임베딩 방법 \phi(d) 은 문서 자체만을 고려하는데, 이는 동일 도메인 내에서는 큰 문제가 되지 않지만, 서로 다른 도메인의 문서들에서 임베딩 성능이 저하될 수 있습니다. 특히 학습된 데이터셋의 도메인 \mathcal{D}_\mathcal{T} 과 테스트 시점의 도메인 \mathcal{D} 이 서로 다를 경우에 이러한 문제는 더욱 두드러집니다. CDE 연구에서는 문서 임베딩 함수인 \phi 와 \psi 에 문맥 정보를 통합한 문맥 기반 학습을 제안하고 있습니다.
문맥 기반의 대조 학습(Adversarial Contrastive Learning)
학습 데이터셋 \mathcal{D}_T 에서 자주 등장하지 않는 단어들이 테스트 시점의 도메인에서는 매우 흔한 용어가 될 수 있습니다. 예를 들어, "NFL"이라는 단어는 일반적인 텍스트에서는 희귀하지만, 스포츠 관련 도메인에서는 매우 흔한 용어가 됩니다. 이와 같은 통계적 차이는 모델이 학습한 데이터셋 \mathcal{D}_\mathcal{T} 과 테스트 도메인 \mathcal{D} 간에 대조적인 상황(Adversarial Setting)을 야기하며, 모델이 새로운 도메인에서 성능이 떨어지게 만드는 요인이 됩니다.
이를 해결하기 위해 CDE는 메타학습(Meta-Learning)과 유사한 방식을 활용하여, (독립적인 문서-질의 쌍을 곧바로 샘플링하지 않고) 먼저 도메인을 기준으로 샘플링한 뒤 해당 도메인 내에서 예시들을 선택해 학습하는 방법을 제안하고 있습니다. 이렇게 함으로써, 모델이 각 도메인 내에서 관련된 데이터를 더 많이 학습할 수 있게 되, 문맥에 맞는 임베딩을 학습할 수 있습니다.
대조 학습의 학습 목표 정의
구체적으로는 각 도메인에서의 문서와 질의 쌍의 유사성을 극대화하는 것으로, 다음과 같은 최적화 문제를 해결하는 것으로 볼 수 있습니다:
구체적으로는 학습 데이터셋 \mathcal{D}_\mathcal{T} 를 여러 의사 도메인(pseudo-domain)들로 분할하여, 각 그룹들(\mathcal{B}^1, \mathcal{B}^2, ...\mathcal{B}^\mathcal{B})에 속하는 문서들이 유사한 특징을 갖도록 합니다. 이를 통해 임베딩 함수 \phi 와 \psi 가 도메인별로 더 구체적이고 세밀한 힘베딩을 학습할 수 있게 됩니다.
연산 시에는 내부 항을 단일 배치로 구현하여 별도의 하드 네거티브 \mathcal{H} 없이 주어진 문서 d 와 쿼리 q 의 확률 p(d \mid q) 을 최대화하도록 합니다. 함수 f() 는 이러한 문서 d 와 질의 q 간의 유사성을 계산하는 함수입니다.
이러한 최적화 문제를 통해 각 도메인 내에서 문서와 쿼리 쌍을 대상으로 대조학습을 진행하는 방식으로 임베딩을 더욱 미세하게 조정할 수 있습니다.
대조학습을 위한 배치 구성: 클러스터링 활용
대조 학습에서 하드 네거티브(Hard Negative)를 선택하는 과정은 모델의 성능에 크게 영향을 미치지만, 정확하게 나누는 것은 매우 시간과 비용이 많이 드는 비현실적인 과정입니다. CDE 연구에서는 도메인 간의 유사성을 비대칭 K-Means 클러스터링(Asymmetric K-Means Clustering) 문제로 보고, 빠른 속도의 유클리디안 K-Means(Euclidean K-Means) 패키지를 사용하였습니다.
이 과정에서 각 데이터 포인트별로 2개의 개별적인 벡터 \phi(\mathcal{d})\oplus\psi(\mathcal{q}) 와 \psi(\mathcal{d})\oplus\phi(\mathcal{q}) 을 계산하여 사용합니다. (\oplus 는 연결(concatenation) 연산을 나타냅니다)
대조학습을 위한 배치 구성: FN(False Negative) 필터링
대조 학습에서 중요한 문제 중 하나는 바로 양성이어야 하는데 음성으로 잘못 표기된 거짓-음성(FN, False Negative, 오탐) 샘플이 학습에 부정적인 영향을 미칠 수 있다는 것입니다. 실제로 Qu 등의 연구(2021)에 따르면, MS Macro 데이터셋의 상위 검색 문서 70%가 FN으로 분류되는 문제가 있었으며, 이러한오탐 문제는 대규모 검색 데이터셋으로 갈수록 심해집니다.
이를 해결하기 위해 CDE는 동등 클래스(Equival Class)를 계산하는 방법을 사용하여 잘못된 FN 샘플들을 배치에서 제외합니다. 동등 클래스 S(q, d) 는 주어진 질의 q 와 문서 d 에 대해서 쿼리와 더 높은 점수를 가지는 다른 문서 d' 들을 원소로 갖는 집합입니다.
이 때, 쿼리와 문서의 유사성을 계산하는 함수 f 는 대체 스코어링 함수(Surrogate Scoring Function)입니다. 결과적으로 위에서 정의한 동등 클래스 S 는 쿼리와 문서의 유사성을 계산하는 함수 f 가 문서 d 보다 더 높은 점수를 부여한 모든 문서 d' 를 원소로 포함하고 있도록 합니다.
학습 과정에서 문서 d 에 대한 파티션 함수(Partition Function)를 변경하며 S(q, d) 에 속하는 FN 샘플이 없도록 합니다. 이를 통해 위 최적화 문제에서 확률을 계산하는 수식을 다음과 같이 다시 정의할 수 있습니다:
또한, 위에서 정의한 스코어링 함수 f 로 사전 학습된 임베딩 모델을 사용합니다. 이 때, 일부 정확하게 탐지한 경우(TN, True Negative)도 포함하여 과도하게 제거할 수가 있으나, 모델의 정확를 위해서는 이 방법을 사용하였습니다.
대조학습을 위한 배치 구성: 배치 크기 균등화
위 클러스터링 과정을 통해 나눈 각 클러스터 그룹의 크기는 서로 다르기 때문에 같은 크기의 배치(batch)로 묶는 과정이 필요합니다. 따라서 마지막에 무작위 파티셔닝(random partitioning)과 탐욕스러운 클러스터 레벨의 순회 세일즈맨을 통한 그룹화(grouping via greedy cluster-level traveling salesman) 방법을 모두 고려하였습니다.
두 경우 모두 큰 그룹을 작은 배치로 분할하고, 같은 도메인 내에서 가까운 작은 배치들을 균등한 크기의 배치로 병합합니다.이렇게 하면 여러 에폭(epoch)에 걸쳐 학습 시 그룹별 무작위성을 도입할 수 있다는 추가적인 이점이 있습니다.
문맥 기반의 임베딩(CDE, Contextual Document Embedding)
이제 신경망을 사용하여 문맥 기반의 문서 임베딩(CDE)을 진행합니다. 이 때 기존의 검색 방식이 말뭉치(corpus)의 통계적 정보를 활용하여 임베딩 형태를 결정하는 방식을 차용하여, 이중 인코더(Bi-Encoder) 모델을 수정하여 말뭉치에 접근할 수 있도록 하였습니다. 즉, 임베딩 함수 \phi(d; \mathcal{D}) 와 \phi(d; \mathcal{D}) 가 임베딩을 생성할 때 말뭉치의 정보를 반영하도록 합니다.
주요 도전 과제
가장 큰 도전 과제는 데이터셋의 문맥 정보(dataset contextualization)를 신경망 구조에 어떻게 반영할 것인가입니다. 한 가지 방법은 BM25와 같이 고정된 말뭉치(corpus) 통계 정보를 미리 계산한 후, 이를 문서 인코더에 입력하는 방식입니다. 이러한 방식의 반대편 극단에는 인코더가 교차 어텐션(cross attention)을 통해 전체 말뭉치에 대해 접근하여 연산하는 것입니다. 하지만 이러한 방식은 대규모 데이터셋에서는 확장성이 떨어집니다.
CDE는 연산 효율성을 유지하면서 말뭉치의 통계 정보를 학습할 수 있는 구조를 설계함으로써 이 두 가지 접근법 사이에서 중간 지점을 찾았습니다. 기존의 연구에 따르면, 문서 임베딩은 임베딩 후에도 상당한 양의 어휘적 정보(lexical information)를 유지한다고 합니다. 따라서 코퍼스의 일부를 미리 임베딩하면, 임베딩 과정 중에 중요한 데이터셋 정보(key dataset information)를 동적으로 계산할 수 있다고 가정할 수 있습니다.
문맥 기반의 임베딩 생성 과정은 다음의 두 단계로 이루어집니다:
문맥 기반의 임베딩 생성 과정: 1단계. 문맥 수집 및 문맥 임베딩 생성
문맥 수집 및 임베딩(Gather and embed context) 단계에서는 전체 문서의 일부 하위 집합들 d^1, d^2, ... d^J \in \mathcal{D} 을 미리 임베딩하여 코퍼스 정보를 수집합니다. 하나의 임베딩 모델(unique embedding model, 아래 M_1)을 사용하여 각 문서들을 임베딩하고, 이러한 임베딩 벡터들을 연결하여 하나의 시퀀스를 생성합니다:
M_1 은 첫번째 단계에서 사용하는 임베딩 모델을 뜻하며, 이러한 과정을 통해 문서 간의 관계를 효과적으로 파악할 수 있습니다.
문맥 기반의 임베딩 생성 과정: 2단계. 추가적 문맥 토큰을 결합한 임베딩 생성
문서 d' 에 추가적 문맥 토큰을 결합하여 임베딩(Embed document with additional context tokens)을 생성합니다.
이 때, 이전 단계에서 생성한 문맥 기반의 임베딩 시퀀스 M_1(d^1), ..., M_1(d^J) 를 또다른 임베딩 모델(M_2)의 입력으로 사용합니다:
앞에서 설명한 것과 같이, M_1 은 첫번째 단계에서 사용하는 임베딩 모델이며, M_2 는 두번째 단계에서 사용하는 임베딩 모델입니다. E 는 d' 의 각 토큰에 M_2 를 적용한 토큰 임베딩 행렬(token embedding matrix)로, 이 과정에서 문서와 문맥의 정보를 함께 사용하여 최종 임베딩을 생성합니다.
실제 구현 시에는 M_1 과 M_2 모델은 전통적인 양방향 트랜스포머(Bi-Directional Transformer)를 사용하여 2단계로 구성된 이중 인코더 모델이 됩니다.
또한, 쿼리 인코더 \psi 에도 동일하게 문서의 문맥을 참고하여 생성합니다. (학습 시점에는 쿼리에 대한 문맥이 없으므로)
문맥 기반의 임베딩: 구현 상의 세부 사항
지금까지 설명한 구조를 구현하며 몇 가지 주목할만한 속성들이 있습니다:
먼저, 학습 시 모든 문서마다 문맥 임베딩을 계산하면 J 에 비례하여 연산 비용이 매우 커지게 됩니다. 이 문제를 해결하기 위해, CDE는 배치 내의 문서들이 동일한 문맥을 공유하도록 하였습니다. 이렇게 하면 문맥 임베딩을 학습 단계별로 한 번씩만 연산하고, 연산 그래프(computational graph)를 활용하여 문서들에서 재사용할 수 있게 하여 효율성을 높였습니다. 이는 클러스터링을 통해 동일한 배치 내 문서가 동일하게 그룹화되어 있기 때문으로, 일반적으로 동일한 배치 내의 문서들은 많은 컨텍스트를 공유하고 있기 때문에 가능합니다.
또한, 첫번째 임베딩 과정에서 사용하는 M_1 모델의 경우 향후 검색 시점에도 고정되어 있으므로 이 과정의 결과물을 캐싱해둘 수 있습니다. 즉, 새로운 말뭉치 \mathcal{D} 를 인덱싱할 때, 첫 번째 단계의 임베딩 M_1(d^1), \dots, M_1(d^J) 을 미리 연산하여 캐싱합니다. 이렇게 하면 검색 과정에서 추가적인 계산 시간 없이 문맥 임베딩을 사용할 수 있습니다.
문맥이 없는 상황에서의 임베딩(Embedding without Context)
일부 개별 말뭉치(Individual Corpora)는 학습 단계에서 충분한 문맥 정보를 제공하지 않을 수 있습니다. 이러한 상황에서는 모델의 일반화 능력(Model's Generalization)을 향상시키기 위해 시퀀스 드롭아웃(Sequence Dropout) 기법을 사용합니다. 이는 무작위로 문맥 임베딩 M_1(d^*) 을 빈 토큰(null token) \mathcal{v}_ \varnothing 으로 대체하는 방식으로 문맥 임베딩을 제거하는 것으로, 모델이 특정한 문맥에 과도하게 의존하지 않도록 하는 방법입니다.
이후 실행 시점(test time)에서도 말뭉치 정보를 사용할 수 없는 경우에 모든 시퀀스 토큰 입력을 \mathcal{v}_ \varnothing 으로 대체하여 비문맥적 이중 인코더(non-contextual bi-encoder)로 동작하게 됩니다.
위치에 의존하지 않는 임베딩(Position-Agnostic Embedding)
말뭉치 \mathcal{D} 내의 문서들은 순서가 없는 데이터셋이기 때문에, 임베딩 과정에서 위치 정보(Positionality)를 제거해야 합니다. 전통적인 트랜스포머(Transformer) 모델은 입력에 위치 인코딩(Positional Encoding)을 통해 각 토큰의 위치 정보를 반영하지만, CDE에서는 이를 제거합니다.
실제로는 각 셀프 어텐션에서 RoPE(Rotary Positional Embedding)를 사용하는 FlashAttention에 기반한 트랜스포머 구현체를 사용합니다. 이 렇게 하면 문서 간의 순서나 위치 정보에 상관없이 임베딩을 생성할 수 있으며, 말뭉치 내의 문서들이 순서없이 배치된 경우에도 일관된 성능을 보장할 수 있습니다. (자세한 내용은 논문의 부록 섹션 10.4 참고해주세요.)
두 단계의 그래디언트 캐싱(Two-Stage Gradient Caching)
효율적인 학습을 위해 CDE에서는 GradCache 기법과 유사한 두 단계의 그래디언트 캐싱(Two-Stage Gradient Cache) 기법을 사용합니다. 이는 더 큰 배치 크기와 긴 시퀀스, 더 많은 문맥 샘플 수에서도 메모리 부족 문제를 최소화할 수 있습니다.
이 방법은 트랜스포머의 순전파 단계(forward pass)를 두 번 실행하는 것으로 추가적인 연산을 하게 되지만, 메모리 사용량을 줄일 수 있어 전체적으로 메모리 효율을 크게 향상시킬 수 있습니다.
문맥 기반 문서 임베딩(CDE) 모델 평가 및 활용
문맥 기반 문서 임베딩 모델의 성능 평가: MTEB
CDE의 저자들은 위와 같은 방법을 통해 생성한 모델 cde-small-v1
을 허깅페이스에 공개하였습니다. 이 모델은 MTEB(Massive Text Embedding Benchmark)의 여러 작업들에 대해서 평균 65.00을 기록하여 250M 이하의 크기 모델들 중에서는 가장 나은 성능을 보였습니다.
또한, 배치 크기별로 성능(Performance)과 (사전 학습 및 감독 학습 종료 시의 손실 값으로 측정한) 평균 배치 난이도(Average Batch Difficulty)를 비교해보았습니다. 주어진 배치 크기 내에서 개별 배치를 더 어렵게 만들면 성능이 뚜렷하게 향상되는 것을 관찰할 수 있습니다.
CDE 모델 불러오기 & 사용하기
CDE 연구의 저자들이 공개한 cde-small-v1
모델은 transformers
라이브러리를 통해 다음과 같이 불러올 수 있습니다:
import transformers
model = transformers.AutoModel.from_pretrained("jxm/cde-small-v1", trust_remote_code=True)
tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
cde-small-v1
모델은 특정 작업에 맞는 접두어(prefix)를 붙여 질의와 문서를 처리합니다. 질의와 문서에 각각 다음 프리픽스를 사용합니다:
query_prefix = "search_query: "
document_prefix = "search_document: "
첫 번째 단계: 말뭉치 임베딩
먼저 말뭉치(Corpus)의 일부를 임베딩하여 말뭉치 정보를 학습해야 합니다. 이를 위해 말뭉치의 샘플을 준비하고, 이를 모델에 입력하여 임베딩을 얻습니다:
# CDE 관련 말뭉치 샘플 및 크기 정의
minicorpus_size = model.config.transductive_corpus_size
minicorpus_docs = [ ... ] # 여기에 말뭉치 샘플 문장들을 입력합니다. 예를 들어, random.sample(corpus, k=minicorpus_size) 같은 함수를 호출하여 추출할 수 있습니다.
assert len(minicorpus_docs) == minicorpus_size # minicorpus_size에 해당하는 수의 문서를 사용해야 합니다. 말뭉치가 이보다 작으면 오버샘플링될 수 있습니다.
minicorpus_docs = tokenizer(
[document_prefix + doc for doc in minicorpus_docs],
truncation=True,
padding=True,
max_length=512,
return_tensors="pt"
)
# 말뭉치 임베딩 수행
import torch
from tqdm.autonotebook import tqdm
batch_size = 32
dataset_embeddings = []
for i in tqdm(range(0, len(minicorpus_docs["input_ids"]), batch_size)):
minicorpus_docs_batch = {k: v[i:i+batch_size] for k,v in minicorpus_docs.items()}
with torch.no_grad():
dataset_embeddings.append(
model.first_stage_model(**minicorpus_docs_batch)
)
dataset_embeddings = torch.cat(dataset_embeddings)
두 번째 단계: 문서 및 질의 임베딩
앞에서 말뭉치 임베딩을 얻은 후, 이를 활용하여 문서와 질의를 각각 임베딩할 수 있습니다. 문서에는 문서 접두어(document_prefix
), 질의에는 질의 접두어(query_prefix
)를 붙여 임베딩합니다.
다음은 문서를 임베딩하는 예시입니다:
docs = tokenizer([document_prefix + doc for doc in docs], truncation=True, padding=True, max_length=512, return_tensors="pt").to(device)
with torch.no_grad():
doc_embeddings = model.second_stage_model(input_ids=docs["input_ids"], attention_mask=docs["attention_mask"], dataset_embeddings=dataset_embeddings)
doc_embeddings /= doc_embeddings.norm(p=2, dim=1, keepdim=True)
질의를 임베딩하는 과정도 유사합니다:
queries = tokenizer([query_prefix + query for query in queries], truncation=True, padding=True, max_length=512, return_tensors="pt").to(device)
with torch.no_grad():
query_embeddings = model.second_stage_model(input_ids=queries["input_ids"], attention_mask=queries["attention_mask"], dataset_embeddings=dataset_embeddings)
query_embeddings /= query_embeddings.norm(p=2, dim=1, keepdim=True)
말뭉치 정보를 사전에 알 수 없는 경우
만약 사전에 말뭉치 정보를 알 수 없다면, 임의의 문자열을 입력해도 모델은 동작합니다. 다만, 이러한 경우 성능이 약간 떨어집니다. (MTEB 기준, 65.0에서 63.8로 감소)
저자들이 제공하는 임의의 문자열을 사용해보세요.
문맥 기반 문서 임베딩(CDE; Contextual Document Embedding) 논문
문맥 기반 문서 임베딩 모델 다운로드: cde-small-v1
문맥 기반 문서 임베딩 데모(Google Colab)
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.
파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~