However, this means that a short sentence and a document containing thousands of tokens are given the same fixed-size representation budget. When an information-dense document is compressed into one vector, fine-grained signals relevant to retrieval can be weakened or lost.

Late-interaction models such as ColBERT address this information bottleneck by retaining one embedding for each query and document token. At retrieval time, the MaxSim operator allows every query token to find its best-matching token in the document. This enables fine-grained token-level matching, but it also requires storing many more vectors per document, resulting in larger indexes and higher retrieval costs.

KURE-v2 is a Korean–English bilingual late-interaction retriever designed to address both retrieval effectiveness and serving efficiency. It uses a lightweight 154M-parameter encoder and represents each token with a 128-dimensional vector.

Our key results are as follows:

• KURE-v2 achieves an average nDCG@10 of 0.8162 across nine retrieval tasks in MTEB(kor, v2).
• At the time of evaluation, it outperformed every single-vector model we tested, including a 27B-parameter model approximately 175 times larger than KURE-v2.
• Combining token pooling with 1-bit quantization reduces the total index size across all nine corpora to 1.7GB, smaller than every single-vector HNSW index evaluated under the same setting.
• On a single A100 80GB GPU with a batch size of 1, KURE-v2 achieves an end-to-end throughput of 43–55 QPS, roughly twice the approximately 25 QPS achieved by the 4B–8B single-vector models.

We are releasing the model together with all training and serving code.

• 🐱 Github: https://github.com/nlpai-lab/KURE
• 🤗 Models: https://huggingface.co/nlpai-lab/KURE-v2

1. What is a Late Interaction model?

A typical Single-Vector retrieval model compresses the query and the document each into one vector, then computes a cosine similarity between the two. Search is fast, but everything a long text contains has to fit into a single vector, and that becomes a bottleneck.

A Late Interaction model postpones the interaction between query and document until after encoding. Like a Single-Vector model, it encodes the two inputs independently of each other, but instead of pooling everything into one vector at the end, it keeps a small vector for every token. The query tokens and document tokens interact for the first time at search time, which is why it is called ‘Late Interaction’.

The representative scoring function is ColBERT’s MaxSim.

s(q, d) = ∑ᵢ maxⱼ(qᵢᵀdⱼ)

Each query token qᵢ finds the document token dⱼ that matches it best, and those maxima are summed. For example, given the query “requirements for terminating an employment contract”, each query token individually locates the most relevant expression inside a long statute. Rather than summarizing the whole document into one vector, the model directly checks, token by token, the specific expressions that matter for retrieval.

The differences between the three retrieval architectures can be summarized as follows.

• Single-vector models: compare one query vector against one document vector, once. Fastest, with the smallest index, but the single vector can become an information bottleneck.
• Late interaction (multi-vector) models: encode queries and documents separately in advance while keeping token vectors, and compute MaxSim at search time. Fine-grained matching becomes possible, but there are many more vectors to store.
• Cross-encoder (reranker) models: encode the query and the document together as a single input. Full query-document interaction is possible, but every document requires a fresh model forward pass, making them too expensive as a first-stage retriever over an entire corpus.

In other words, late interaction sits between the efficiency of single-vector retrieval and the fine-grained interaction of a cross-encoder.

Despite these advantages, few open late-interaction retrievers explicitly support Korean. KURE-v2 addresses this gap while examining whether late interaction can remain practical under real-world serving constraints.

2. Model architecture

KURE-v2 is built on skt/A.X-Encoder-base, a Korean-English bilingual encoder supporting a 16K-token context.

Each token representation passes through a projection head into a 128-dimensional vector, and similarity is computed as MaxSim between normalized query and document token vectors. Queries use ColBERT-style [MASK] expansion padded to 64 tokens, and documents support up to 8,192 tokens at inference.

Instead of the usual single-layer projection hidden → 128, KURE-v2 uses a multi-layer projection head of the form hidden → 2 × hidden → hidden → 128.

The first two layers carry residual connections, and no layer uses a bias or a non-linear function. This deeply factorized parameterization together with the residual connections can affect the optimization path and act as implicit regularization. The ablation study later in this post shows that this difference translated into a real performance gain in SFT. The multi-layer head adds about 4.8M parameters, bringing KURE-v2 to 154M parameters in total.

3. Datasets

Training proceeded in two stages.

1. Stage 1 learns retrieval representations from large-scale weakly relevant pairs (weakly-supervised fine-tuning; Pre-FineTuning; PFT)

2. Stage 2 refines the ranking using hard negatives and reranker scores (Supervised Fine-Tuning; SFT)

In both stages we balanced the scale of the Korean and English data so that neither language dominates training.

Stage 1 PFT dataset: weakly-supervised pairs

The Pre-FineTuning (PFT) stage uses pairs that are related but carry no precise retrieval labels, such as title-body, title-abstract, and instruction-output pairs.

• Korean, 10M pairs: AIHUB, Modu Corpus, KULLM, mC4, Wikipedia, synthetic data
• English, 10M pairs: 12 public corpora including MS MARCO, NQ, and S2ORC

The Korean data consists of title-body pairs from AIHUB and the Modu Corpus, instruction-output pairs from KULLM, pairs extracted from Korean Wikipedia and mC4, and Magpie-style synthetic pairs. After deduplication, about 10.35M pairs remained.

For English, we selected 12 corpora centered on real queries from https://huggingface.co/datasets/lightonai/embeddings-pre-training-curated, then sampled them to nearly the same scale as the Korean data to keep the two sides balanced.

Stage 2 SFT dataset: supervised pairs

The Supervised Fine-Tuning (SFT) stage uses training samples containing 1 positive and 7 hard negatives per query.

• Korean, 1.4M samples: starting from the KURE-v1 training data, cleaned with an LLM
• English, 1.5M samples: 9 public retrieval datasets that come with hard negatives (https://huggingface.co/datasets/lightonai/embeddings-fine-tuning#scores)

The Korean data started from the query-positive pairs built for KURE-v1. Concretely, we inspected the data by hand and observed pairs mapped as positives that were not actually positive (false positives). We therefore re-judged relevance with Qwen3–32B and removed unrelated pairs, leaving about 1.46M cleaner samples.

We then mined hard negatives using the percpos method from NV-Retriever, and assigned Qwen3-Reranker-4B scores to every query-document pair, positives and negatives alike. These scores served two purposes.

1. Distillation: the student model (KURE-v2) learns the teacher’s (Qwen3-Reranker-4B) relative score distribution over the documents of each query.
2. A false-negative filter: if the teacher scored some negative higher than the mapped positive, that negative is likely a genuinely relevant document, so we dropped the entire training sample.

4. Training

Stage 1: In-batch contrastive pre-training

To extract enough discriminative signal from weak supervision, we used a global batch size of 16,384. Gradient Cache resolved the memory problem, and we trained for 1 epoch with an InfoNCE loss that treats the other documents in the batch as negatives.

The important choice here is that each batch contains only one language. With more than ten thousand in-batch negatives, mixing Korean and English would let the model lean on the easy shortcut of “this document is in a different language from the query” instead of semantic relevance. To prevent this kind of hacking, we trained with single-language batches.

The checkpoint at the end of this stage is released as KURE-v2-unsupervised. Without any human retrieval relevance labels, it records an average nDCG@10 of 0.7283.

Stage 2: contrastive learning + KL divergence

The second stage combines a contrastive loss with a KL divergence loss.

• Contrastive loss: uses every document in the batch as a negative candidate, including other queries’ positives and mined negatives (batch size = 256).
• KL divergence loss: for the 1 positive and 7 negatives attached to each query, the student’s MaxSim score distribution is trained to follow the teacher reranker’s score distribution.

Both stages were trained with the PyLate library.

5. Evaluation setup and results

Evaluation was conducted on the 9 Korean retrieval tasks of MTEB(kor, v2). The tasks are listed below.

• Ko-StrategyQA: Korean ODQA multi-hop retrieval (translated from StrategyQA)
• AutoRAGRetrieval: retrieval over PDF-parsed documents in finance, public sector, healthcare, legal, and commerce domains
• MIRACLRetrieval: Wikipedia-based retrieval, about 1.5M documents
• PublicHealthQA: retrieval in the medical and public health domain
• BelebeleRetrieval: retrieval based on FLORES-200
• MrTidyRetrieval: Wikipedia-based retrieval, about 1.5M documents
• MultiLongDocRetrieval: long-document retrieval across diverse domains
• LawIRKo: Korean statute and case-law retrieval
• SQuADKorV1Retrieval: KorQuAD v1 paragraph retrieval

The results for the main models are as follows.

Detailed scores on all 9 tasks are available on the Korean MTEB leaderboard.

Three things stand out in these results.

1. For late-interaction models, parameter is not the big factor.

At 154M parameters, KURE-v2 beats the strongest single-vector model, a 7.6B model, by 2.33 points of average nDCG@10. Against a 27B model the gap is 4.93.

2. The advantage of late interaction is largest on long-document tasks.

On MLDR (MultiLongDocRetrieval), KURE-v2 scores 71.59 while the best single-vector model scores 51.57, a gap of 20.02. Another well-trained model, colbert-ko-en-v2, scores 69.92 and shows the same trend.

This shows that the benefit of the multi-vector architecture is not spread evenly across tasks: it is most pronounced on long documents, where the compression loss of a single vector grows largest.

3. Weak supervision alone already gets you a long way.

KURE-v2-unsupervised, trained with PFT only, records an average of 72.83 without any human relevance labels. SFT with hard negatives and reranker distillation then adds another 8.77 points.

6. Ablation Study

6–1. Results of the multi-layer linear head

Comparing right after PFT, the single head slightly leads the multi head, 73.46 versus 72.83. But once SFT completes, the ranking flips: single head 80.81, multi head 81.62.

These results alone cannot pin down the exact cause, but one interpretation is that the optimization benefit of the deeply factorized linear parameterization stays hidden under weak relevance signals, and only becomes effective when combined with reranker distillation over hard negatives.

6–2. Contrastive Learning and Distillation play different roles

With the multi-layer head, using only the contrastive loss in SFT gives 78.63, only KL distillation gives 80.66, and combining both gives 81.62.

The two losses are not redundant. The contrastive loss provides global discrimination using the many documents in the batch, while the distillation loss transfers the fine-grained relative ordering among hard negatives that the cross-encoder teacher can tell apart. The best results come when the ability to separate a wide candidate space is combined with the ability to precisely order the difficult candidates.

7. Serving: Setup

A late interaction model stores a set of token vectors per document. So when considering real adoption, two questions are unavoidable.

1. How large does the index get?

2. How fast can queries be served?

To find out, we ran experiments serving KURE-v2 with several ANN backends and compression techniques. First, a brief description of the four core techniques used in the experiments.

1. PLAID

PLAID is the representative search engine for ColBERT-family models, designed to cut down MaxSim computation as much as possible. At indexing time it clusters the document token vectors and stores each token as a centroid ID plus a low-bit residual. At search time it first narrows down candidate documents using query-centroid similarity, then reconstructs the residuals of the selected documents to compute full MaxSim scores.

In other words, PLAID compresses the token vectors and performs candidate generation at the same time, so it never has to sweep every token of the entire corpus with MaxSim on each query.

2. Hierarchical Token Pooling

In the late-interaction family, a single document can contain many token vectors that are similar to one another. Hierarchical Token Pooling groups each document’s token vectors with Ward-linkage agglomerative clustering and replaces each cluster with its mean vector.

For example, 3x pooling cuts the number of document token vectors roughly to a third. The MaxSim computation against the query vectors stays exactly the same. Because it reduces the number of vectors to store, it composes with PLAID or with binary quantization.

3. Asymmetric Binary Quantization

Asymmetric Binary Quantization maps each dimension of a document token vector to -1 or +1 by its sign and stores it as 1 bit per dimension. Compared to fp32 vectors, per-dimension storage shrinks by a factor of 32.

It is called “asymmetric” because only the documents are quantized to 1 bit, while the query vectors used in the final scoring stay in bf16 precision. The final MaxSim is computed between bf16 query vectors and 1-bit document vectors. This slashes storage while avoiding the loss that comes from quantizing the query as well.

However, if you exhaustively scan all the compressed document tokens, search cost still grows linearly with corpus size.

4. Binary IVF

BinaryIVF adds a Candidate Generation Stage so the 1-bit document vectors above are not scanned exhaustively. It uses faiss’s IndexBinaryIVF to partition the binary token vectors into inverted lists and probes only some of the lists at search time.

During candidate generation the query is also temporarily binarized, and nearby document tokens are found by Hamming distance, computed with XOR and popcount. Token hits are then aggregated per document to form candidates, and only the top 1,000 documents are reranked with the asymmetric MaxSim between bf16 query vectors and 1-bit document vectors. The quantized query is therefore used only for candidate generation, never in the final reranking.

To sum up the relationship between the four methods: PLAID is a complete engine that provides both compression and candidate generation. Hierarchical token pooling reduces the number of token vectors, and asymmetric binary quantization reduces the bits stored per vector. BinaryIVF turns the exhaustive scan over binary vectors into candidate search.

We served KURE-v2 with these methods and compared against a baseline of 5 single-vector models served with Faiss HNSW. All experiments ran serially on a single A100 80GB, and latency and QPS include both batch-size-1 query encoding and index search.

8. Serving: Results

The experimental results are as follows. Total Index Storage on the left represents the sum of index storage capacity required when storing 9 tasks. QPS is the average value for the 9 tasks when measured with a batch size of 1.

1. PLAID preserves the highest quality.

The results for fast-plaid with 4-bit residuals are:

• Average nDCG@10: 81.36
• Full index over 9 corpora: 24.2GB
• End-to-end throughput: 44 QPS

PLAID records the highest quality among all the KURE-v2 serving configurations we evaluated. If you need to pick one default engine with quality as the priority, PLAID is the safest choice.

2. Token pooling cuts the index by more than half at almost no cost.

The trade-off introduced by hierarchical token pooling was smaller than expected.

• 2x pooling: 24.2GB → 12.5GB // nDCG@10 81.36 → 81.45 // 44 → 49 QPS
• 3x pooling: 24.2GB → 8.4GB // nDCG@10 81.36 → 81.24 // 44 → 50 QPS

2x pooling shrank the index almost in half while lowering quality by only 0.04 points. With fewer document tokens, the MaxSim computation also shrank, so throughput actually went up.

3. Asymmetric quantization makes the multi-vector index smaller than HNSW.

Asymmetric binary quantization shrinks the index 4.8x versus PLAID, from 24.2GB to 5.0GB, while lowering average nDCG@10 by only 1.05 points.

Adding 3x token pooling on top brings the full index for all 9 corpora to 1.7GB with an average nDCG@10 of 79.57. That is smaller than every single-vector HNSW index compared under the same setup (fp16 vector, 7.0~25.4 GB), while scoring higher than the strongest single-vector model’s 79.10.

In other words, a multi-vector index is not necessarily larger than a single-vector one. Reduce the number of vectors and the bits per vector together, and the relationship can invert. That said, this configuration still pays the cost of an exhaustive scan, so large corpora need a candidate generator such as BinaryIVF (see below).

4. Experiments on MIRACL with 1.5M documents

To see how the picture changes at scale, we separately analyzed MIRACL (Korean), a corpus of about 1.5M documents.

Exhaustively scanning the 1-bit vectors (KURE-v2 Asym. Binary; pink) gives a p95 latency of 156ms, and even with 3x token pooling it takes 74ms. Compression reduces the amount of computation, but the structure of sweeping every document remains.

Applying BinaryIVF on top of the same binary vectors keeps the MIRACL index at 2.2GB while dropping p95 latency from 156ms to 38ms. Quality is essentially unchanged, and 38ms is even shorter than the 42ms p95 of the 4B-8B single-vector baselines.

So on large corpora you should use an index that prunes candidates first, like PLAID or BinaryIVF. Regardless of compression ratio, an exhaustive scan becomes the bottleneck as the corpus grows.

5. The end-to-end bottleneck is query encoding, not ANN.

Real service queries arrive as text. So you have to look at query encoding time, not just index search time.

The 4B-8B single-vector models spent 38–40ms on query encoding alone. Their HNSW search took only 1–2ms, yet their overall throughput was stuck at about 25 QPS.

KURE-v2’s 154M encoder, in contrast, encodes a query in 13.8ms. As a result, every KURE-v2 serving configuration except MUVERA recorded 44–57 QPS. In other words, it maintains higher retrieval quality while achieving roughly twice the end-to-end throughput of the 4B-8B single-vector models.

(Note) Measurement environment

• Hardware: 1x NVIDIA A100 80GB, AMD EPYC 7513 (64 cores), 1.2TB RAM
• Software: faiss-cpu 1.15.0, fast-plaid 1.6.0, sentence-transformers 6.0.0, PyTorch 2.8.0, flash-maxsim 0.3.0
• Protocol: batch size 1, serial execution
• Search latency: every query of each task measured once after 10 warmup queries
• Query encoding: 50 measurements after 5 warmup runs
• Encoding precision: bf16
• HNSW: HNSW fp16 IndexHNSWFlat, M=32, efConstruction=200, efSearch=64
• PLAID: nbits=4, n_ivf_probe=8, n_full_scores=4,096
• MUVERA: 6-bit SimHash with 10 repetitions, 8,192-dimensional fixed-dimensional encoding, exact reranking of the top 1,000 documents
• BinaryIVF: nlist=min(⌊√T⌋, 65,536), nprobe=32, top 128 token candidates per query token, exact reranking of the top 1,000 documents
• Token pooling: Ward-linkage agglomerative clustering applied to document vectors only
• QPS is computed as the inverse of mean end-to-end latency, and index size is measured as the full serialized files on disk.

Closing

The results of KURE-v2 show two things.

First, in retrieval, how you represent documents and let them interact matters as much as model size. A 154M-parameter late-interaction model beat 8B and 27B single-vector models, and the advantage was most pronounced in long-document retrieval, where single-vector compression loses the most.

Second, the deployment cost of multi-vector retrieval can be brought down. Using PLAID, hierarchical token pooling, asymmetric 1-bit quantization, and BinaryIVF for what each does best widens the menu of choices across quality, index size, and latency.

I wish KURE-v2 can help a lot of Information Retrieval and RAG systems.

Citation

@misc{kure-v2,
  title  = {KURE-v2: a Korean-English bilingual late-interaction retriever},
  author = {Jang, Youngjoon and Son, Junyoung and Lee, Taemin and Hong, Seongtae and Lim, Heuiseok},
  year   = {2026},
  url    = {https://huggingface.co/nlpai-lab/KURE-v2},
}

하지만 짧은 문장과 수천 토큰의 문서에 동일한 크기의 표현 공간이 주어진다는 한계가 있습니다. 정보량이 많은 문서 전체를 하나의 고정 차원 벡터로 압축하면, 검색에 필요한 세부 정보가 풀링 과정에서 약해질 수 있습니다.

ColBERT와 같은 Late Interaction 모델은 이러한 정보 병목을 줄이기 위해 질의와 문서의 각 토큰에 대한 벡터를 유지합니다. 검색 시에는 질의의 각 토큰이 문서에서 가장 잘 맞는 토큰을 찾는 MaxSim 연산을 사용합니다. 덕분에 세밀한 토큰 수준의 매칭이 가능하지만, 문서마다 많은 벡터를 저장해야 하므로 색인이 커지고 검색 비용이 증가합니다.

KURE-v2는 이 정확도와 배포 효율성 사이의 문제를 함께 다루기 위해 개발한 한국어–영어 이중언어 Late Interaction 검색 모델입니다. 154M 파라미터의 경량 인코더를 사용하며, 각 토큰을 128차원 벡터로 표현합니다.

핵심 결과를 먼저 요약하면 다음과 같습니다.

• MTEB(kor, v2) 검색 9개 태스크에서 평균 nDCG@10 0.8162을 기록합니다.
• 실험 시점 기준, 평가한 모든 Single-Vector 모델의 성능을 능가합니다.
• 토큰 풀링과 1-bit 양자화를 결합하면 9개 코퍼스의 전체 색인을 1.7GB까지 줄일 수 있습니다. 이는 동일한 실험에서 비교한 모든 Single-Vector HNSW 색인보다 작습니다.
• 배치 크기 1의 End-to-End 측정에서 KURE-v2는 43–55 QPS를 기록했습니다. 4B–8B 단일 벡터 모델의 약 25 QPS보다 대략 두 배 빠른 결과입니다.

모델과 학습, 배포 코드를 모두 공개합니다.

• 🐱 Github: https://github.com/nlpai-lab/KURE
• 🤗 Models: https://huggingface.co/nlpai-lab/KURE-v2

1. Late Interaction 모델이란?

일반적인 단일 벡터 검색 모델은 질의와 문서를 각각 하나의 벡터로 압축한 뒤, 두 벡터의 내적이나 코사인 유사도를 계산합니다. 검색은 빠르지만, 긴 텍스트의 모든 정보를 하나의 벡터에 담아야 한다는 병목이 생깁니다.

Late Interaction 모델은 질의와 문서의 상호작용을 인코딩 이후로 미룹니다. 두 입력을 서로 독립적으로 인코딩한다는 점은 단일 벡터 모델과 같지만, 마지막에 하나로 풀링하지 않고 토큰마다 작은 벡터를 유지합니다. 실제 검색 시점에 질의 토큰과 문서 토큰이 처음 상호작용하기 때문에 ‘Late Interaction’이라고 부릅니다.

대표적인 점수 함수가 ColBERT의 MaxSim입니다.

s(q, d) = ∑ᵢ maxⱼ(qᵢᵀdⱼ)

질의의 각 토큰 qᵢ가 문서에서 자신과 가장 잘 맞는 토큰 dⱼ를 하나씩 찾고, 그 최댓값을 모두 더합니다. 예를 들어 “근로 계약 해지 요건”이라는 질의가 들어오면, 각 질의 토큰이 긴 법령 문서 안에서 가장 관련 있는 표현을 개별적으로 찾습니다. 문서 전체의 의미를 벡터 하나로 요약하는 대신, 검색에 필요한 세부 표현을 토큰 단위로 직접 확인하는 셈입니다.

세 가지 검색 구조의 차이를 간단히 정리하면 다음과 같습니다.

• Single-Vector 모델: 질의 1개 벡터와 문서 1개 벡터를 한 번 비교. 가장 빠르고 색인이 작지만, 하나의 벡터가 정보 병목이 될 수 있습니다.
• Late Interaction (Multi-Vector) 모델: 질의와 문서를 미리 따로 인코딩하면서 토큰 벡터를 유지하고, 검색 시 MaxSim을 계산. 세밀한 매칭이 가능하지만 저장할 벡터가 많아집니다.
• Cross-Encoder (Reranker) 모델: 질의와 문서를 하나의 입력으로 함께 인코딩. 질의와 문서 간 상호작용이 가능하지만 문서마다 모델 추론을 다시 해야 하므로 전체 코퍼스의 1차 검색기로 쓰기에는 굉장히 비용이 많이 듭니다.

즉, Late Interaction은 Single-Vector 모델의 효율성과 Cross-Encoder의 세밀한 상호작용 사이에 있는 구조입니다.

이러한 장점에도 불구하고, 한국어를 명시적으로 지원하는 공개 Late Interaction 검색 모델은 여전히 드뭅니다. KURE-v2는 이 간극을 메우는 동시에, Late Interaction의 성능 이점을 실제적인 배포 비용 안에서 유지할 수 있는지를 검증합니다.

2. 모델 구조

KURE-v2는 16K 토큰 문맥을 지원하는 한국어–영어 이중언어 인코더 skt/A.X-Encoder-base를 기반으로 합니다.

각 토큰 표현은 투영 헤드를 거쳐 128차원 벡터로 변환되고, 정규화된 질의·문서 토큰 벡터 사이의 MaxSim으로 유사도를 계산합니다. 질의는 ColBERT 방식의 [MASK] 확장을 사용해 64토큰까지 패딩하며, 문서는 추론 시 최대 8,192토큰까지 지원합니다.

KURE-v2에서는 일반적인 단층 투영 hidden → 128 대신 hidden → 2 × hidden → hidden → 128과 같은 다층 투영 헤드를 사용했습니다.

앞의 두 층에는 residual connection을 적용했고, 모든 층에서 bias 혹은 non-linear function은 사용하지 않았습니다. 따라서 깊게 인수분해된 파라미터화와 residual connection은 최적화 경로와 암묵적 규제에 영향을 줄 수 있습니다. 뒤에서 살펴볼 Ablation Study에서는 이 차이가 SFT에서 실제 성능 향상으로 이어졌음을 보여줍니다. 다층 헤드는 약 4.8M개의 파라미터를 추가하며, KURE-v2의 전체 크기는 154M 파라미터입니다.

3. 데이터셋

학습은 두 단계로 진행했습니다.

1. Stage1 에서는 대규모의 약한 관련성 쌍으로 검색 표현을 학습하고 (weakly-supervised fine-tuning; Pre-FineTuning; PFT)
2. Stage2 에서는 hard negatives와 reranker 점수를 이용해 순위를 정교하게 다듬었습니다 (Supervised Fine-Tuning; SFT)

두 단계 모두 한 언어가 학습을 지배하지 않도록 한국어와 영어 데이터의 규모를 균형 있게 맞췄습니다.

Stage1 PFT Dataset: weakly-supervised pairs

Pre-FineTuning(PFT) 단계에서는 title-body, title-abstract, instruction-output처럼 관련성은 있지만, 정교한 검색 라벨은 없는 쌍을 사용합니다.

• 한국어 10M 쌍: AIHUB, 모두의 말뭉치, KULLM, mC4, 위키백과, 합성 데이터
• 영어 10M 쌍: MS MARCO, NQ, S2ORC 등을 포함한 공개 코퍼스 12종

한국어 데이터는 AIHUB와 모두의 말뭉치의 title-body 쌍, KULLM의 Intruction-output 쌍, 한국어 위키백과와 mC4에서 추출한 쌍, Magpie 방식의 합성 쌍으로 구성했습니다. 중복 제거 후 약 1,035만 쌍이 남았습니다.

영어 데이터는 https://huggingface.co/datasets/lightonai/embeddings-pre-training-curated 에서 실제 질의와 학술 텍스트 중심의 12개 코퍼스를 선별한 뒤, 한국어와 거의 동일한 규모로 표집하여 데이터 크기를 맞추었습니다.

Stage2 SFT Dataset: supervised pairs

Supervised Fine-Tuning(SFT) 단계에서는 각 query마다 positive1개와 hard negative 7개가 포함된 학습 샘플을 사용했습니다.

• 한국어 1.4M 개: KURE-v1 학습 데이터에서 출발해 LLM으로 정제
• 영어 1.5M 개: 하드 네거티브가 포함된 공개 검색 데이터셋 9종 (https://huggingface.co/datasets/lightonai/embeddings-fine-tuning#scores)

한국어 데이터는 KURE-v1 학습 시 구축한 query-positive 쌍에서 출발했습니다. 구체적으로는, 데이터를 직접 눈으로 보고, positive으로 매핑된 쌍이 실제로는 positive이지 않은 것 (위양성, false-positive)을 관찰했습니다. 이에, Qwen3–32B로 관련성을 다시 판정해 무관한 쌍을 제거했고, 약 146만 개의 더 깨끗한 샘플을 남겼습니다.

이후 NV-Retriever의 percpos방법을 차용하여 hard negative mining을 진행했으며, 양성과 네거티브를 포함한 모든 질의–문서 쌍에 Qwen3-Reranker-4B 점수를 부여했습니다. 이 점수는 두 가지 용도로 사용했습니다.

1. Student Model (KURE-v2)가 Teacher (Qwen3-Reranker-4B)의 문서별 상대 점수 분포를 학습하도록 하는 distillation 용도.
2. 위음성을 걸러내는 필터입니다. Teacher가 데이터에서 매핑된 positive보다 높게 평가한 negative가 있는 경우, 해당 negative는 실제로는 관련 문서일 가능성이 있다고 보고 그 학습 샘플을 아예 제거했습니다.

4. 모델 학습

Stage1: In-batch contrastive pre-training

weakly-supervised 신호에서 충분한 변별력을 얻기 위해 전역 배치 크기 16,384를 사용했습니다. Gradient Cache로 메모리 문제를 해결했고, 배치 내 다른 문서를 네거티브로 사용하는 InfoNCE Loss로 1 epoch 학습했습니다.

여기서 중요한 선택은 각 배치를 하나의 언어로만 구성한 것입니다. 1만 개가 넘는 배치 내 네거티브를 한국어와 영어로 섞으면, 모델은 의미 관련성보다 “이 문서는 질의와 언어가 다르다”는 쉬운 기준에 의존할 수 있습니다. 이러한 hacking을 방지하기 위해, 단일 언어 배치로 학습하였습니다.

이 단계가 끝난 체크포인트는 KURE-v2-unsupervised로 공개했습니다. 사람의 검색 관련성 라벨을 사용하지 않고도 평균 nDCG@10 0.7283을 기록했습니다.

Stage2: contrastive learning + KL Divergence

두 번째 단계에서는 contrastive loss과 KL divergence loss를 함께 사용했습니다.

• Contrastive Loss: 다른 query의 positive와 mined negative를 포함해, 배치 안의 모든 문서를 네거티브 후보로 활용합니다. (batch size = 256)
• KL Divergence Loss: 각 query에 연결된 positive 1개와 negative 7개에 대해, 학생의 MaxSim 점수 분포가 교사 리랭커의 점수 분포를 따르도록 학습합니다.

두 단계 모두 PyLate 라이브러리를 사용하여 학습했습니다.

5. 평가 세팅 및 평가 결과

평가는 MTEB(kor, v2)의 한국어 9개 검색 태스크에서 진행했습니다. 태스크들의 목록은 아래와 같습니다.

• Ko-StrategyQA: 한국어 ODQA multi-hop 검색 (StrategyQA 번역)
• AutoRAGRetrieval: 금융·공공·의료·법률·커머스 분야 PDF 파싱 문서 검색
• MIRACLRetrieval: Wikipedia 기반 검색, 약 150만 문서
• PublicHealthQA: 의료·공중보건 도메인 검색
• BelebeleRetrieval: FLORES-200 기반 검색
• MrTidyRetrieval: Wikipedia 기반 검색, 약 150만 문서
• MultiLongDocRetrieval: 다양한 도메인의 장문 검색
• LawIRKo: 한국어 법령·판례 검색
• SQuADKorV1Retrieval: KorQuAD v1 문단 검색

주요 모델의 결과는 다음과 같습니다.

전체 9개 태스크의 세부 점수는 MTEB 리더보드에서 확인할 수 있습니다.

이 결과에서 세 가지가 눈에 띕니다.

1. Late-Interaction 모델에서는 파라미터가 큰 요소가 아니다.

154M의 KURE-v2는 가장 강한 단일 벡터 모델인 7.6B 모델보다 평균 nDCG@10이 2.33 높았습니다. 27B 모델과 비교하면 격차는 4.93입니다.

2. Late Interaction의 이점은 Long doc 태스크에서 가장 크다.

MLDR (MultiLongDocRetrieval)에서 KURE-v2는 71.59, 최고 단일 벡터 모델은 51.57을 기록합니다. 차이는 20.02입니다. 잘 학습된 다른 colbert-ko-en-v2 모델 또한 69.92를 기록해 같은 경향을 보였습니다.

이는 다중 벡터 구조의 이점이 모든 태스크에 균일하게 나타나는 것이 아니라, 단일 벡터의 정보 압축 손실이 커지는 장문에서 특히 두드러진다는 점을 보여줍니다.

3. 약한 지도만으로도 일정한 성능에 도달한다

PFT만 거친 KURE-v2-unsupervised는 사람의 관련성 라벨 없이 평균 72.83을 기록했습니다. 이후 하드 네거티브와 리랭커 증류를 사용하는 SFT가 8.77을 추가로 끌어올렸습니다.

6. Ablation Study

6–1. Multi Linear head 의 결과

PFT까지만 비교하면 single head가 multi head보다 근소하게 앞섰습니다. 점수는 각각 73.46과 72.83이었습니다. 그러나 SFT를 마치면 순위가 뒤집혀 single head 80.81, multi head 81.62가 됩니다.

이 결과만으로 정확한 원인을 확정할 수는 없지만, 깊게 인수분해된 선형 파라미터화의 최적화상 이점이 약한 관련성 신호만으로는 드러나지 않다가, 하드 네거티브에 대한 리랭커 증류와 결합되면서 유효해졌다는 해석을 해볼 수 있습니다.

6–2. 대조 학습과 증류는 서로 다른 역할을 한다

다층 헤드를 기준으로 SFT에서 대조 손실만 사용하면 78.63, KL 증류만 사용하면 80.66, 두 손실을 결합하면 81.62였습니다.

두 손실은 중복되지 않습니다. 대조 손실은 수많은 배치 내 문서를 이용해 전역적인 변별력을 제공하고, 증류 손실은 크로스 인코더 교사가 구분한 하드 네거티브 사이의 세밀한 상대 순서를 전달합니다. 넓은 후보 공간을 구분하는 능력과 어려운 후보를 정교하게 정렬하는 능력이 결합될 때 가장 좋은 결과가 나왔습니다.

7. Serving: Setup

Late Interaction 모델은 문서마다 토큰 벡터 집합을 저장합니다. 따라서 실제 도입을 검토할 때는 두 가지 질문을 피할 수 없습니다.

1. 색인은 얼마나 커지는가?

2. 검색 상황에서 얼마나 빠르게 처리할 수 있는가?

이를 확인하기 위해 KURE-v2를 여러 ANN 백엔드와 압축 기법으로 서빙하는 실험을 진행했습니다. 먼저 실험에 사용한 네 가지 핵심 기법을 간단히 설명하겠습니다.

1. PLAID

PLAID는 ColBERT 계열 모델을 위한 대표적인 검색 엔진으로, MaxSim 연산을 최대한 줄이기 위한 방법입니다.. 색인 단계에서는 문서 토큰 벡터를 clustering하고, 각 토큰을 centroid ID와 low-bit residual로 저장합니다. 검색 단계에서는 query와 centroid의 유사도로 후보 문서를 먼저 줄인 뒤, 선택된 문서의 residual을 복원해 전체 MaxSim 점수를 계산합니다.

즉, PLAID는 토큰 벡터를 압축하는 동시에 후보 생성까지 수행하므로 전체 코퍼스의 모든 토큰을 매번 MaxSim 연산으로 탐색하지 않아도 됩니다.

2. Hierarchical Token Pooling

Late-Interaction 계열에서는 한 문서 안에 서로 비슷한 토큰 벡터가 많이 존재할 수 있습니다. Hierarchical Token Pooling은 문서별 토큰 벡터를 Ward 연결 병합 Clustering으로 묶고, 각 cluster를 평균 벡터 하나로 대체합니다.

예를 들어 3배 풀링은 문서 토큰 벡터의 개수를 대략 3분의 1로 줄입니다. 이때 query vector와의 MaxSim 계산 방식은 그대로 유지됩니다. 이 방법은 저장해야 할 벡터의 개수를 줄이므로 PLAID나 이진 양자화와 함께 사용할 수 있습니다.

3. Asymmetric Binary Quantization

Asymmetric Binary Quantization (비대칭 이진 양자화)는 각 문서 토큰 벡터의 차원을 부호에 따라 -1 또는 +1로 바꾸고, 이를 차원당 1bit로 저장합니다. fp32 벡터와 비교하면 차원별 저장량이 32분의 1로 줄어듭니다.

‘비대칭’이라고 부르는 이유는 문서만 1-bit로 양자화하고, 최종 점수 계산에 사용하는 질의 벡터는 bf16 precision 로 유지하기 때문입니다. 최종 MaxSim은 bf16 query vector와 1-bit document vector 사이에서 계산합니다. 저장 공간을 크게 줄이면서 질의까지 함께 양자화할 때 발생하는 loss을 피하는 방식입니다.

다만 이때 압축된 모든 문서 토큰을 전수 탐색하면 검색 비용은 여전히 코퍼스 크기에 선형으로 증가합니다.

4. Binary IVF

BinaryIVF는 위의 1-bit 문서 벡터를 전수 탐색하지 않도록 Candidate Generation Stage를 추가합니다. faiss의 IndexBinaryIVF를 사용해 이진 토큰 벡터를 여러 리스트로 나누고, 검색 시 일부 리스트만 탐색합니다.

후보 생성 단계에서는 질의도 일시적으로 이진화하고, 해밍 거리인 XOR와 popcount 연산으로 가까운 문서 토큰을 찾습니다. 이후 토큰 적중 수를 문서별로 집계해 후보를 만든 뒤, 상위 1,000개 문서만 bf16 query vector 와 1-bit document vector의 Asymmetric MaxSim으로 재순위화합니다. 따라서 양자화된 질의는 후보 생성에만 사용되고, 최종 재순위화에는 사용되지 않습니다.

네 방법의 관계를 정리하면, PLAID는 압축과 후보 생성을 모두 제공하는 완성형 엔진입니다. Hierarchical Token Pooling은 토큰 벡터의 개수를 줄이고, Asymmetric Binary Quantization은 벡터당 저장 비트를 줄입니다. BinaryIVF는 이진 벡터의 전수 탐색을 후보 검색으로 바꾸는 역할을 합니다.

위 방법들로 KURE-v2를 서빙하며, baseline으로는 Faiss HNSW를 사용하는 Single-Vector 모델 5종과 비교했습니다. 모든 실험은 A100 80GB 한 장에서 직렬로 수행했으며, 지연 시간과 QPS에는 배치 크기 1의 질의 인코딩과 색인 검색이 모두 포함됩니다.

8. Serving: Results

실험 결과는 다음과 같습니다. 좌측의 Total Index Storage는 9개 태스크 저장 시 차지하는 Index Storage 용량의 총합이며, QPS는 batch 1로 쟀을 때에 대한, 9개 태스크의 평균값입니다.

1. PLAID가 가장 높은 품질을 유지한다.

4-bit residual를 사용하는 fast-plaid의 결과는 다음과 같습니다.

• 평균 nDCG@10: 81.36
• 9개 코퍼스 전체 색인: 24.2GB
• e2e 처리량: 44 QPS

PLAID는 평가한 KURE-v2 서빙 구성 중 가장 높은 품질을 기록합니다. 품질을 우선하면서 하나의 기본 엔진을 선택해야 한다면 PLAID가 가장 안전한 선택입니다.

2. Token Pooling은 거의 손실 없이 색인을 절반 이하로 줄인다.

Hierarchical Token Pooling으로 인해 발생하는 trade-off는 예상보다 작았습니다.

• 2배 풀링: 24.2GB → 12.5GB // nDCG@10 81.36 → 81.45 // 44 → 49 QPS
• 3배 풀링: 24.2GB → 8.4GB // nDCG@10 81.36 → 81.24 // 44 → 50 QPS

2배 풀링은 품질을 0.04점만 낮추면서 색인을 거의 절반으로 줄였습니다. 문서 토큰 수가 감소하자 MaxSim 계산량도 줄어 처리량은 오히려 높아졌습니다.

3. Asym. Quantization은 multi-vector index를 HNSW보다 작게 만든다.

Asymmetric Binary Quantization은 PLAID 대비 색인을 24.2GB에서 5.0GB로 4.8배 축소하면서, 평균 nDCG@10은 0.98점만 낮춥니다.

3배 토큰 풀링까지 결합하면 9개 코퍼스의 전체 색인은 1.7GB, 평균 nDCG@10은 79.57이 됩니다. 이는 같은 설정에서 비교한 모든 단일 벡터 HNSW 색인 (fp16 벡터, 7.0~25.4 GB)보다 훨씬 작으면서, 가장 강한 Single-Vector 모델의 79.10보다도 높은 결과입니다.

즉, Multi-Vector 모델의 색인이 반드시 Single-Vector모델보다 크지 않을 수 있습니다. 벡터 개수와 벡터당 비트 수를 함께 줄이면 관계가 역전될 수 있습니다. 다만 이 구성만으로는 전수 탐색 비용이 남아 있으므로, 큰 corpus에서는 BinaryIVF와 같은 후보 생성기가 필요합니다. (아래 참고)

4. MIRACL 150만 문서에서의 실험

대규모 코퍼스에서의 차이를 확인하기 위해 약 150만 문서로 구성된 MIRACL (Korean)을 별도로 분석했습니다.

1-bit 벡터를 전수 스캔 (KURE-v2 Asym. Binary; 분홍색)하면 PLAID (초록색)에 비해 Index size는 줄지만, p95 지연이 156ms로 매우 컸고, 토큰을 3배 풀링해도 74ms가 걸렸습니다. 압축으로 연산량을 줄일 수는 있지만, 모든 문서를 훑는 구조 자체는 그대로이기 때문입니다.

같은 이진 벡터 위에 BinaryIVF (진한 붉은색)를 적용하자, MIRACL 색인을 2.2GB로 유지하면서 p95 지연이 156ms에서 38ms로 감소했습니다. 품질은 사실상 유지됐으며, 38ms는 4B–8B Single-Vector 베이스라인의 42ms보다도 짧은 p95 Latency입니다.

따라서 대규모 코퍼스에서는 PLAID나 BinaryIVF처럼 후보를 먼저 줄이는 색인을 사용해야 합니다. 전수 탐색은 압축률과 무관하게 코퍼스가 커질수록 병목이 됩니다.

5. e2e 검색에서의 병목은 ANN보다 Query Encoding이다.

실제 서비스의 query는 텍스트로 들어옵니다. 따라서 색인 검색 시간뿐 아니라 end-to-end에서의 query encoding 시간도 함께 봐야 합니다.

4B–8B 단일 벡터 모델은 query encoding에만 38–40ms를 사용했습니다. HNSW 검색은 1–2ms에 불과했지만, 전체 처리량은 약 25 QPS에 머물렀습니다.

반면 KURE-v2의 154M 인코더는 query를 13.8ms에 encoding 했습니다. 그 결과 MUVERA를 제외한 KURE-v2의 서빙 구성은 44–57 QPS를 기록했습니다. 즉, 더 높은 검색 품질을 유지하면서 4B–8B Single-Vector 모델보다 약 두 배 높은 e2e 처리량을 달성했습니다.

(참고) 측정 환경

• 하드웨어: NVIDIA A100 80GB 1장, AMD EPYC 7513 (64코어), RAM 1.2TB
• 소프트웨어: faiss-cpu 1.15.0, fast-plaid 1.6.0, sentence-transformers 6.0.0, PyTorch 2.8.0, flash-maxsim 0.3.0
• 프로토콜: 배치 크기 1, 직렬 처리
• 검색 지연: 워밍업 10회 후 태스크 전체 질의를 한 번씩 측정
• 질의 인코딩: 워밍업 5회 후 50회 측정
• 인코딩 정밀도: bf16
• HNSW: HNSW fp16 IndexHNSWFlat, M=32, efConstruction=200, efSearch=64
• PLAID: nbits=4, n_ivf_probe=8, n_full_scores=4,096
• MUVERA: 6-bit SimHash 10회 반복, 8,192차원 고정 차원 인코딩, 상위 1,000개 문서 정확 재순위화
• BinaryIVF: nlist=min(⌊√T⌋, 65,536), nprobe=32, 질의 토큰당 상위 128개 토큰 후보, 상위 1,000개 문서 정확 재순위화
• 토큰 풀링: 문서 벡터에만 Ward 연결 병합 적용
• QPS는 평균 e2e 지연의 역수로 계산, 색인 크기는 디스크에 직렬화된 전체 파일을 기준으로 측정

마치며

KURE-v2의 결과는 두 가지를 보여줍니다.

첫째, 검색에서는 모델 크기만큼이나 문서를 어떻게 표현하고 상호작용시키는지가 중요합니다. 154M 파라미터의 Late-Interaction 모델이 8B와 27B 단일 벡터 모델을 앞섰고, 그 이점은 단일 벡터의 압축 손실이 가장 큰 장문 검색에서 특히 뚜렷했습니다.

둘째, Multi-Vector 검색의 배포 비용은 개선할 수 있습니다. PLAID, Hierarchical Token Pooling, Asymmetric 1-bit Quantization, BinaryIVF를 목적에 맞게 사용하면 품질·색인 크기·지연 사이의 선택지를 넓힐 수 있습니다.

앞으로 KURE-v2가 정보 검색 시스템 및 RAG 시스템에 도움이 되었으면 좋겠습니다.

Citation

@misc{kure-v2,
  title  = {KURE-v2: a Korean-English bilingual late-interaction retriever},
  author = {Jang, Youngjoon and Son, Junyoung and Lee, Taemin and Hong, Seongtae and Lim, Heuiseok},
  year   = {2026},
  url    = {https://huggingface.co/nlpai-lab/KURE-v2},
}

The reviews are out, and you have a few days. Reviewer 1 wants a significance test you never ran. Reviewer 2 says your baselines are too weak. Reviewer 3 politely notes that your tables imply state-of-the-art results you never actually claimed.

Two things are true at that moment. First, someone has faced a related criticism before, in a paper you would probably never find, and their response may show you how to address it. Second, the experiment that would settle Reviewer 1 is one you should actually run right now: the code, the data, and the GPUs are already on your machine.

So I built two things.

ReviewSearch is a semantic search engine over 200,000 peer reviews and their author rebuttals. It finds who raised your reviewer’s concern before, and how they answered it.

RebuttalDraft hands that to a coding agent that runs the experiments your reviewers asked for, and writes the draft with the real numbers.

• ReviewSearch Demo site: yjoonjang-reviewsearch.hf.space
• Skills for Agents: rebuttal-skills

Part 1. ReviewSearch

1 · Motivation

For every criticism a reviewer can raise, some author has already answered carefully, under the same pressure you are under now, and often successfully. That knowledge exists. However, it is hard to search: it is spread across hundreds of thousands of reviews and threaded author responses that nobody indexes and nobody reads twice.

I wanted a search engine where the query is a concern: “The reviewer doubts the significance of a small improvement”, and the results are the reviews that raised that concern, each carrying the author’s rebuttal as its payload. That is a semantic retrieval problem, and a genuinely hard one, because no two reviewers phrase the same concern the same way and almost none of them use the words you would search for. Keyword search is nearly useless here. The whole project is an argument that this is worth doing well, not with an off-the-shelf embedder and a prayer.

2 · Building the corpus

I used the OpenReview dataset: for ICLR, NeurIPS, ICML and COLM, the reviews and the author responses are public, all under CC-BY-4.0. I pulled 245,394 review records, of which 205,988 carry an author rebuttal (84%).

Dealing with the data, there were two things to consider:

Every venue-year has a different review form. ICLR gives you summary / strengths / weaknesses / questions. NeurIPS 2025 merges strengths and weaknesses into one field. COLM calls weaknesses “reasons to reject.”
There is no clean schema; there is a normalizer with alias-fallback chains, and it is boring, and you have to write it.

What actually goes into the vector. A review is not one thing: it has a paper summary, praise, criticism, and questions. People search for the criticism. So the embedded text is title + summary + weaknesses + questions, and strengths are deliberately excluded. The rebuttal is not embedded at all. One vector per review, no chunking.

3 · Training & Evaluation data

In order to build a search engine, there should be an evaluation dataset to measure the retrieval performance. Building the eval set took longer than building either model, and it is the part I would defend hardest.

Queries. I generated them with an LLM from reviews, in five flavors:

1. Topic + Concern:
   a generic topic + specific complaint (e.g. “dense retriever p-test”)
2. Broad Concern:
   the complaint alone, no topic
   (e.g. “reviews demanding significance testing”)
3. Rebuttal-Oriented:
   framed by how it was answered
   (e.g. “paper that rebutted a novelty concern”)
4. Natural Question:
   a short discovery question
   (e.g. “which rebuttals addressed computational cost?”)
5. Reviewer-Voice Critique:
   a terse line in a reviewer’s voice
   (e.g. “the gains aren’t statistically significant”)

Papers were split at the paper level, so no paper leaks between train and eval. That gave 20k eval queries and 18k train queries.

Labels. A synthetic query comes with one “correct” review for free: the one it was generated from. That label is nearly useless: in a corpus of 206k reviews, dozens of others raise the same concern and are equally relevant. (Since there are broad concerns also.) Scoring against a single gold document punishes a retriever for being right.
So I built a judged pool: I encoded the whole corpus with Qwen3-Embedding-8B, retrieved the top 32 reviews per query, and made Qwen3.6–35B-A3B judge every (query, candidate) pair independently.

Then the judge turned out to be the hard part. My first judge was lenient and marked 9.7 of every 16 candidates relevant. Reading its output, the failure was specific: it accepted reviews that matched the topic but not the concern. A query about “reviewers asking for a significance test on a retrieval paper” was satisfied by any retrieval review at all. I rewrote the rubric to demand that topic and concern both match, with few-shot examples of that exact failure. After that prompting, 34% of the labels disappeared.

Judge system prompt:

You are a STRICT relevance judge for a peer-review search engine.

A search QUERY has two parts: a TOPIC (the research area or method) and a CONCERN (the specific
criticism the searcher wants — e.g. missing comparisons, weak baselines, unclear novelty, high
computational cost, unrealistic assumptions, missing ablations, no significance testing).

You are given the QUERY and ONE candidate review. Mark it RELEVANT only if BOTH hold:
  1. TOPIC matches — the review is about the same research area/method as the query; AND
  2. CONCERN matches — the review's MAIN weakness is the SAME KIND of concern the query asks for.

Decisive rule: topical overlap is NOT enough. If the review is on the same topic but its main
weakness is a DIFFERENT kind of concern than the query asks for, answer irrelevant.

Examples (query -> candidate's main concern -> verdict):
- "dynamic routing missing comparisons" -> "insufficient comparison to prior baselines" -> relevant
- "dynamic routing missing comparisons" -> "the theoretical assumption is unrealistic" -> irrelevant
- "dynamic routing missing comparisons" -> "the evaluation dataset has selection bias" -> irrelevant
- "diffusion model high compute cost" -> "sampling is too computationally expensive" -> relevant
- "diffusion model high compute cost" -> "the writing in the intro is unclear" -> irrelevant

Answer with ONLY one word: relevant OR irrelevant.

Judge user prompt:

QUERY: {query}

CANDIDATE REVIEW:
{candidate}

Answer (relevant/irrelevant):

Final gold:
19,833 judged queries, with an average of 14 relevant reviews per query.

4 · Retrieval models

[4.1 · Base model benchmark results]

Before training anything, I benchmarked six off-the-shelf dense embedders, BM25, and a learned-sparse retrieval model on the evaluation set. These are the results:

Dense. I did not pick the best zero-shot model, and that was deliberate. Although Qwen3-Embedding-0.6B led the benchmark, and embeddinggemma-300m came in third among the dense models, I picked embeddinggemma-300m for two serving reasons:

1. At 300M parameters, it is roughly half the size of Qwen3-Embedding-0.6B, making it substantially cheaper to run on a free CPU instance.
2. embeddinggemma is published with Quantization-Aware Training (QAT), so it survives int8 quantization intact.

This is what I needed to serve it on CPU. I bet that fine-tuning would more than close the 0.025 gap to Qwen3-Embedding, and it did.

Sparse. For hybrid search, I needed a sparse model which could also catch keywords. I chose opensearch-neural-sparse-encoding-doc-v3-gte because it is an inference-free SPLADE model: documents undergo transformer-based term expansion at indexing time, while queries require only tokenization and a lightweight lookup at serving time.

[4.2 · Fine-tuning the dense retriever]

As a hard contributor to Sentence-Transformers, I already had substantial experience training embedding models. (See my Korean blog about training embedding models 🤗)

Contrastive Learning vs. Distillation? Both were viable, but this task has a highly specific query–document relationship: concerns are queries, and reviews are documents. I therefore expected direct contrastive learning to fit the domain better.

Training Settings. I experimented with fine-tuning embeddinggemma with the following configs:

Detailed Setup.

• per_device_train_batch_size: 1024, 2-GPU DDP
• max_seq_len: 2048
• epoch: 1

Experimental Results (InfoNCE-style losses).

The first real result was blunt: training on positive pairs alone collapsed. In-batch negatives with no mining fell to 0.160 (below the zero-shot model). Hard negatives are essential here.

Which raises the question of how they should be mined, so I ran two experiments: mining them with the base model, embeddinggemma (exp2), or with a much stronger Qwen3–Embedding-8B (exp3). I could see that mining hard negatives with the base model (embeddinggemma) helped training a bit better than mining with Qwen3–8B (0.223 on exp2 vs. 0.208 on exp3).

Training on more negatives. A common assumption in training with contrastive learning is that the more negatives you have, the better the models should become. So I trained with the following setting:

I got a result that annoyed me: more hard negatives make the model monotonically worse. The standard explanation is false negatives: the deeper you mine, the more of your false negatives you get, and you are training the model to push away documents it should retrieve. If that is right, it is testable.

GISTEmbedLoss. To test if the false negatives are the real problem, I used GISTEmbedLoss. This is similar to InfoNCE, except it uses a guide model to first filter false negatives and use the filtered ones for contrastive learning. See the picture below:

Training Settings. So I experimented with fine-tuning embeddinggemma with GISTEmbedLoss, with the following configs:

Experimental Results (GIST-style losses).

For clearer comparison, the results are like this:

We could clearly see that

1. False negatives were what caused performance to degrade when increasing the number of negatives.
2. Guiding the student model with GISTEmbedLoss really helps the model improve its retrieval performance.

And I finally chose the best model exp9 (embeddinggemma, 7 hn, GIST, GOR, QAT, MRL) which recorded 0.237 on nDCG@20.

[4.3 · Fine-tuning the sparse retriever]

In the previous evaluation, I found that BM25 does not perform very well. So I fine-tuned a Learned-Sparse Retrieval (LSR) model, opensearch-neural-sparse-encoding-doc-v3-gte. “Inference-free” means the asymmetry is on purpose: documents get a full transformer pass that expands them into weighted vocabulary terms at index time, while the query side is just a trainable IDF lookup at serve time. You get learned term expansion at essentially zero query-time cost.

Training Settings. I experimented with fine-tuning the opensearch model with the following configs:

Experimental Results.

Two things transfer straight from the dense side. Positive pairs alone degrade the model, again (0.1548, below zero-shot). And hard negatives are again decisive: the best config, mined with the Qwen & embeddinggemma retriever, reaches 0.206~0.207, +20% over zero-shot.

The negative result. The natural next idea is ANCE-style self-mining: use the backbone sparse model to mine its own hard negatives, which can lead to fast convergence of loss. This setting lands at 0.175, barely above zero-shot. The reason is the false-negative from a different angle: the sparse model is a worse retriever than the dense one (0.207 vs 0.237). When a worse retriever mines negatives, a larger fraction of what it surfaces near the top is actually relevant-but-unlabeled, so its “hard negatives” are dirtier, and no margin rescues them. So in conclusion, mine with the best retriever you have, not with the one you are training.

Scaling the number of negatives. Scaling to n=5 and n=7 flattened out, exactly like plain MNRL on the dense side, since there is no guide model masking false negatives.

[4.4 · Hybrid Retrieval]

I used Reciprocal Rank Fusion (RRF), sweeping the fusion constant and the dense:sparse weight, starting from the best dense model.

With BM25, equal weighting actively hurts : its noise floods the top of the fused list, and you have to hold it at arm’s length with a 3:1 dense-dominant weight to extract the few exact-term hits it uniquely finds. With the fine-tuned sparse model, the optimum moves to 1:1. The weight your fusion wants is a direct readout of how much you trust the second retriever, and watching it move from 3:1 to 1:1 as the sparse model improved was the clearest sign the fine-tuning had actually done something.

5 · Serving

The search runs in a free Hugging Face Docker Space: 2 vCPUs, 16 GB of RAM, no GPU, no persistent disk. Serverless was the obvious first choice and the wrong one, since a cold-starting function cannot hold torch and a vector index in memory. A single always-on container is simpler. It runs python:3.12-slim, CPU torch, and a static Qdrant binary; at boot it downloads a 3.7 GB Qdrant snapshot of all 205,988 points from a Hugging Face dataset, recovers it, and serves both the API and the built web app from one origin.

[5.1 · The dense encoder is the only bottleneck]

The sparse retriever is inference-free at query time. It runs a tokenizer and an IDF lookup with no transformer forward pass (the point of the inference-free SPLADE in §4.3), so on CPU it costs a couple of milliseconds. In the measured pipeline the sparse encode and the fusion together come to about 2 ms and the index search to about 9 ms, so the dense-encoder forward pass is effectively the entire serving budget.

The dense query encoder is exported to ONNX and quantized to int8: 2.85× faster to encode, for a 0.4% nDCG drop the hybrid absorbs almost entirely. The ONNX runtime pins the container to transformers 4.57, and the SPLADE model loads cleanly on that version, so one environment serves both encoders.

Making the encode fast came down to threading, where the same mistake appears in two libraries. On a 2-vCPU box, the default thread counts are wrong.

Torch with set_num_threads(8) oversubscribes two cores and takes 3.4 seconds, so pinning it to two threads gives 225 ms. ONNX is slower than torch until you notice that onnxruntime (ort) ignores torch's thread settings and starts one thread per visible core. Setting intra_op_num_threads=2 brings it to 46 ms. That is a 74× range from a single configuration value, in a place a profiler does not point to.

[5.2 · Indexing]

I used an ANN index for the dense model. I benchmarked the full retriever against a naive (exhaustive) baseline that scores every query over the whole corpus (both performance & latency)

• Performance. The approximation is nearly free: the ANN index loses 0.8% nDCG against exhaustive search.

• Latency. And it is far faster: exhaustive search runs a full-corpus matrix multiply on every query, while the ANN does not.

That 0.8% is smaller than the variance between two fine-tuning runs, so full-fidelity storage and a free tier stay compatible.

Tuning the index hyperparameters barely changed the results:

The four sit on top of each other, so m and efConstruction do not matter. efSearch alone nudges toward the exact-search ceiling (dashed) and saturates by about 256, the deployed value.

Fusion happens on the client. Qdrant can fuse the dense and sparse results, but its built-in RRF fixes the rank constant at 60. Since that constant has no effect on latency, the retriever fetches both result lists in one round trip and fuses them itself. Metadata filters (accepted-only, a year range) are applied inside the ANN query.

The hybrid also makes quantization cheap. Reducing the dense vectors to int8 costs 7.6% nDCG on their own, but only 2.9% after fusion, because RRF combines ranks rather than scores. That is what allows the index to be compressed for a machine with 16 GB and no disk.

[5.3 · The index ships as a file]

The index does not run on the server because the free Space has no persistent disk. The index is a file instead. It is built once, offline, on a GPU: the entire Qdrant collection, 205,988 vectors and their payloads, is snapshotted into a 3.7 GB blob and pushed to a HuggingFace dataset. The serving container boots empty, downloads the snapshot, recovers it in a few seconds, and begins serving. Building the index and deploying the app become independent, and the container stays stateless. The cost of that design is the cold start: a Space that has gone idle must re-download the full 3.7 GB before answering, so the first query after a sleep waits roughly five minutes. Warm, a query takes about 58 ms.

Qdrant fits this setup for practical reasons. At 206k vectors the index is not the bottleneck, so the priority was operational simplicity rather than raw search speed. A hand-built index (FAISS with a learned-sparse engine) would sit in the same latency band, since it too is approximate, and add moving parts for no real gain. Qdrant is a single static binary with no JVM. It stores the dense and sparse vectors in one collection and keeps each review and its rebuttal as payload, so a result returns complete with no second lookup. It stores only the vectors the encoders produce; its built-in embedding supports a fixed set of public models, which does not include the fine-tuned ones.

Part 2. RebuttalDraft

1 · Why it had to be an agent

ReviewSearch answers who has raised your reviewer’s concern before, and how they got out of it. But searching for references is only half a rebuttal. The other half is running the experiment the reviewer asked for, and that needs an agent sitting in your repository.

I first built the rebuttal generator as a web app: paste your paper and your reviews, and it extracts each reviewer’s concerns, searches the corpus for how authors answered that same criticism, and writes a draft grounded in real precedent. It was good, and it had a ceiling I could not raise. The best it could ever produce for “Reviewer 1 wants a significance test” was:

We have run paired significance tests and report them below.
<<paste your p-values here>>

The web app cannot run your experiment. It has never seen your code. It has no GPU, no checkpoints, and no idea that the script which would answer Reviewer 1 has been sitting in your repo, written and never executed, since March.

A coding agent has all of those. It is already running in your workspace, with your code, your data, and your GPUs. That is not a small upgrade to the web app; it is a different product. So I rebuilt it as a skill.

2 · The Workflow

We have six stages.

1. Stage 0. Intake find the paper and the reviews
2. Stage 1. Concerns extract every distinct concern, per reviewer, with severity
3. Stage 2. Grounding search 200k reviews; keep the single best precedent per concern
4. Stage 3. Experiments — inventory existing resources, design missing experiments, estimate their cost, and present one consent table before running anything.
5. Stage 4. Draft write the reply, with the numbers it just measured
6. Stage 5. Critic a separate agent re-reads as a skeptic, checks every number, and revises

Two design constraints are load-bearing.

Privacy: ReviewSearch never receives your paper or review text. The skill sends only short, abstract concern queries to the search API. The paper and reviews are handled within your coding-agent workspace, subject to the data-handling policy of the agent provider.

Traceability: every number in the draft traces to a real artifact or to your paper, and experiments that cannot be run are reported as such rather than replaced with plausible-looking results.

3 · What RebuttalDraft produces

Every run writes a rebuttal/ directory containing one paste-ready response per reviewer, together with the evidence used to produce it.

rebuttal/
├─ R1.md            per-reviewer response (Overview → Draft → Guide)
├─ R2.md
├─ R3.md
└─ experiments/
   ├─ RESULTS.md    every measurement, tagged with where the number came from
   ├─ grounding.md  the precedents ReviewSearch retrieved, kept whole
   └─ run logs and artifacts

The reviewer files are the deliverable. RESULTS.md is the audit trail behind them: every number a draft cites traces back to a section here, marked with its provenance. grounding.md holds the precedent rebuttals the search returned, stored whole rather than compressed, so the agent can inspect how those authors structured and supported their responses. The run logs and generated artifacts stay on disk, so any figure in the reply can be checked by hand.

4 · Human approval stays in the loop

An agent in your repository can read your code, launch jobs, and spend your compute. That authority is bounded at both ends of the run.

Before anything executes, Stage 3 presents a single consent table: each experiment it proposes, the concern it answers, the exact command, the device, and an estimated cost. This consent table is the execution gate: nothing runs until you approve the corresponding row, and any proposed experiment can be rejected. Experiments that are already scripted but never run are flagged as recommended, with the reason stated plainly, since those are the cheapest path to a real answer.

After the experiments finish, the Critic stage re-reads the draft as a skeptic. It verifies, number by number, that every figure traces to a real result in RESULTS.md or to the paper, and that the draft never claims an experiment that was not run. The final files are still meant to be read by the authors before anything is submitted.

5 · Limitations

Retrieved rebuttals are just references, not templates. A response that resolved one reviewer’s concern may not fit a paper with different evidence or experimental constraints, so the draft is a starting point rather than a finished answer. The agent can only run experiments that the available repository, data, checkpoints, and compute actually support. When an experiment cannot be run, the draft states that limitation explicitly and responds using the evidence that is available.

It also expects some manual setup: you copy the paper and the review text into the project folder yourself, since it reads them from your workspace rather than fetching them anywhere. It does not replace the authors’ scientific judgment.

Try it!

Search the corpus in your browser: yjoonjang-reviewsearch.hf.space

Claude Code:

/plugin marketplace add yjoonjang/rebuttal-skills
/plugin install rebuttal

Codex:

git clone https://github.com/yjoonjang/rebuttal-skills.git ~/.codex/rebuttal-skills
mkdir -p ~/.agents/skills
ln -s ~/.codex/rebuttal-skills/skills ~/.agents/skills/rebuttal

Then open your paper’s own project, with the reviews in it, and ask for a rebuttal. It will show you the plan before it runs anything.

흔히 사용되는 Dense Embedding 모델의 경우, 임베딩은 일반적으로 대부분의 값이 0이 아닌 벡터 (e.g. 1024 차원)이다. 반면 Sparse Embedding 모델은 대부분의 값이 0인 고차원 벡터 (e.g. 30,000차원 이상)를 생성한다. 우리가 수십 년간 검색의 표준처럼 사용해 온 BM25 또한 하나의 Sparse Embedding을 생성하는 방식이다.

하지만 BM25에는 치명적인 단점이 있다. 바로 “형태는 다르지만 의미가 같은 단어”를 구별하지 못한다는 점이다. 예를 들어, BM25는 ‘스파게티’와 ‘파스타’를, ‘식당’과 ‘레스토랑’을 완전히 다른 토큰으로 취급한다. 이를 해결하기 위해 “Term Expansion(단어 확장)”에 관한 연구가 활발히 진행되어 왔다. (참고: 이전 글)

오늘은 그중에서도 가장 주목받는 SPLADE 모델과, 이를 더욱 효율적으로 개선한 Inference-Free SPLADE 모델에 대해 설명하고, 직접 한국어 모델을 개발한 경험을 나누고자 한다.

개발한 모델들은 아래 링크에서 직접 사용해 볼수 있다.

• 🤗 splade-ko-v1
• 🤗 inference-free-splade-ko-v1

시작하기에 앞서, 개발한 모델들과 기존 BM25, 그리고 Dense Embedding 모델인 KURE-v1, 그리고 Dense와 Sparse를 결합한 Hybrid Retrieval의 성능(NDCG@10)을 비교하면 다음과 같다.

위 표를 보면

1. BM25 대비 inference-free-splade-ko-v1과 splade-ko-v1은 단독 사용 시에도 엄청난 성능 향상을 보여준다. 특히 두 모델 모두 0.1B 모델임에도 불구하고, 0.6B Dense 모델인 KURE-v1과 비교하여 크게 차이 나지 않는 성능을 보여준다.
2. KURE-v1과 결합한 Hybrid RRF 상황에서도 BM25 + KURE-v1 (0.721) 조합보다 splade-ko-v1 + KURE (0.762) 조합이 약 4점 더 높은 점수를 기록한다.

이렇게 강력한 성능을 보여주는 SPLADE 모델들에 대해 알아보자.

1. SPLADE와 Inference-Free SPLADE

1–1. SPLADE

SPLADE 모델은 “BM25의 Term mismatch” 문제를 해결하기 위해 NAVER LABS EUROPE (한국 네이버 아님)에서 개발한 모델이다. 말 그대로 query와 document로부터 term을 확장하기 위한 노력을 많이 하는데, 논문 속 핵심 요소들을 살펴보자.

SPLADE 모델은 BERT의 Masked Language Modeling (MLM) 능력에 전적으로 의존한다. 구체적으로는 먼저 BERT에서 출력되는 토큰들의 hidden state를 MLM Head를 통과시킨다.

이 wij logit은 곧 “문맥에서 i token이 j token에게 주는 importance score”로 해석될 수 있다. 예시를 들어 보자. “The bank is near the river.” 라는 문장에서, “bank”는 문맥을 반영한 “강둑”이라는 단어이기 때문에, 이 자리에 “water”, “flow”와 같은 token들에 대해서는 높은 logit을, “money”, “deposit”와 같은 token들에 대해서는 상대적으로 낮은 logit을 매길 것이다.

이후, 이 importance score w에 ReLU를 적용하여 음수 값을 제거하고, log(1+x) 변환을 통해 값이 지나치게 발산하는 것을 막는다 (log-saturation). 마지막으로 모든 입력 토큰 i에 대해 Pooling을 수행하여, 문장 전체를 아우르는 최종 Vocab Size 크기의 벡터를 만들게 된다.

흥미로운 점은 SPLADE-v1에서는 Pooling 시, 이 logit들을 다 더하는 Sum 방식을 사용했지만, SPLADE-v2에서는 “가장 강력한 신호 하나를 택하는 Max Pooling 방식이 성능면에서 훨씬 우월함”을 실험적으로 입증했다는 것이다.

결론적으로 우리가 문장을 입력하면, 이러한 과정을 거쳐 입력 문맥이 가장 강력하게 지지하는 토큰들의 점수로 구성된 vocab 크기의 벡터를 얻게 된다. 다음 그림이 가장 직관적인 이해를 돕는다.

이제 이런 Lexical하게 확장된 벡터를 얻었다면, 학습은 어떻게 할까 ? 여기서는 2가지 Loss가 사용되는데, 하나는 Ranking Loss (Contrastive Loss), 다른 하나는 벡터를 sparse하게 만드는 Flops Loss이다.

[Ranking Loss] 먼저 데이터는 query (q)와 질의에 대한 답변을 제공하는 positive document (d+), 그리고 내용은 관련 있어 보이지만 정답은 아닌 hard negative document (d-)로 구성된 triplet 쌍을 준비한다. 이후 q와 d+는 가깝게, q와 d- 및 배치 내 documents를 멀어지도록 하는 In-Batch Negatives (IBN) 방식의 InfoNCE Loss로 학습한다. 자세한 설명은 이전 블로그를 참고하면 도움이 될 것이다.

• 참고: 유사도 함수 s로는 Cosine Similarity가 아닌 Dot Product를 사용한다.

[Flops Loss] Ranking Loss만으로 학습하면 벡터의 값이 0이 아닌 차원들이 많아질 수 있다 (물론 ReLU 함수에 의해 벡터가 어느 정도 sparse 해지긴 한다). 이는 곧 검색 엔진이 색인(Index)해야 할 양이 늘어나, 검색 속도가 느려짐을 의미한다.

따라서 SPLADE는 FLOPS Regularizer를 도입한다. 원리는 간단하면서도 강력하다. 배치 내에서 “각 토큰의 평균 빈도 수를 계산하고, 이 값을 제곱하여 loss로 사용 (최소화)하자”하는 것이다. 이렇게 하면 모델은 불필요한 단어의 가중치를 0으로 보내버리게 되고, 꼭 필요한 단어만 남기는 Sparsity을 학습하게 된다.

[Final Loss] 최종적으로는 이 두 Loss를 가중합하여 사용한다. 이때 Query와 Document에 적용하는 Regularizer 가중치를 따로 두는데, 보통 검색 속도에 더 큰 영향을 미치는 Query를 더 희소하게 만들기 위해 Query 쪽에 더 큰 가중치를 부여한다.

이렇게 학습된 splade-v3 모델로 “The weather is lovely today”라는 질의를 입력했을 때 나온 vocab size 크기의 벡터를 decoding 해보았을 때 결과가 다음과 같다.

• “weather”(빨간색) -> “summer” 활성화
• “lovely”(노란색) -> “beautiful”, “cool”, “pretty”, “nice” 활성화
• “today”(파란색) -> “currently”, “yesterday” 활성화

1–2. Inference-Free SPLADE

앞선 설명을 보면서 이런 의문이 들 수 있다. “어차피 Vocab Mismatch 해결이 목적이라면, Document만 확장해서 색인해두면 되는 거 아닌가 ?? Query는 굳이 무거운 모델에 통과시킬 필요 없이 들어온 단어 그대로 tokenize만 해서 써도 충분할 것 같은데…”

SPLADE-v2 저자들이 이러한 생각을 바탕으로 SPLADE-doc 모델을 제안했었다. 이름 그대로 Document만 SPLADE 모델로 확장하고, Query는 별도의 인코더 없이 단어의 유무(Bag-of-Words)만 따지는 방식이다.

• Query: 단순히 토크나이징된 단어들에 해당하는 위치만 1인 벡터 사용
• Document: SPLADE 모델을 통과하여 확장된 Sparse Vector 사용

이렇게 하면 검색 시 Query Encoding 시간이 0에 수렴하기 때문에, 기존 방식보다 훨씬 더 빠른 검색 속도를 확보할 수 있다 (query 입력 후 tokenize만 하면 됨). 논문 실험 결과, 일반 SPLADE-doc은 일반 SPLADE와 적은 성능 차이를 보이면서 큰 속도 향상을 기대할 수 있다.

위 표의 맨 아래에 보면, SPLADE-max와 SPLADE-doc의 성능이 크게 차이나지는 않지만, 그래도 꽤나 유의미하다는 것을 볼 수 있다 (약 2점) 이러한 성능 갭을 해결하기 위해 나온 것이 바로 Amazon Opensearch의 Inference-Free SPLADE이다. (논문)

논문은 기존의 FLOPS Regularization(희소하게 만드는 Loss)이 모든 토큰을 평등하게 처벌한다는 점을 지적한다. 특정 데이터로 penalize되는 비율을 분석해본 결과,

• “this”, “that” 같은 무의미한 토큰은 자주 등장함에도 페널티를 적게 받고,
• “algorithms”, “graph” 같은 중요한 토큰은 드물게 등장한다는 이유로 오히려 더 강하게 penalize 당하고 있었다.

이에 저자들은 현재의 FLOPS penalty 체계가 semantic을 반영하고 있지 않다고 말하며, IDF (Inverse Document Frequency)를 반영한 penalty를 제안한다. 그럼 어떤 효과를 같게 될까? 식으로 살펴보면 다음과 같다.

1. 먼저 vocab에 있는 모든 토큰들로 train data에 대한 IDF 값을 계산한다. 이 계산으로 우리는 vocab size |V| 크기의, idf 값들을 원소로 가지는 벡터 idf(t)를 얻는다. 이 idf(t)를 이용하여 유사도 계산 함수를 다음과 같이 변경한다.

2. 최종 Loss는 L_{rank-idf} + lambda * L_{FLOPS} 인데, 이를 미분하게 되면 다음과 같다.

1번에서 유사도 식에 idf 항이 들어가 있기 때문에, 우리는 비로소 L_{rank-idf}가 idf에 영향을 받게 됨을 알 수 있다.

즉, IDF가 높은 중요한 토큰은 Ranking Loss로부터 오는 Gradient가 커서 살아남고, IDF가 낮은 무의미한 토큰은 Ranking Loss의 영향력이 작아 FLOPS Penalty에 의해 쉽게 0으로 사라지게 된다.

이렇게 학습된 모델의 성능은 다음과 같다.

Inference-free (query는 고정)임에도 불구하고, 일반 sparse retriever와 거의 같은 성능을 보이며, 심지어는 ColBERTv2 보다도 높은 성능을 보인다. (물론 데이터에서의 차이와, 정교한 Distillation의 개입이 있었다.)

정리하면, 이 논문에서는 모든 token들에 대한 IDF weight을 미리 계산해둔 뒤, query에 해당하는 token들의 IDF weight를 이용해서 query embedding을 구성하면 학습 및 추론하면 성능 향상을 기대해볼 수 있다는 점을 시사한다.

2. SPLADE, Inference-Free SPLADE 학습기

이 섹션에서는 앞서 설명한 SPLADE와 Inference-Free SPLADE를 학습한 경험에 대해 이야기해보고자 한다.

2–1. splade-ko-v1

🤗 Model: splade-ko-v1

[Dataset] 먼저 데이터는 1M의 query-positive-negatives(6개)를 준비하였다. Hard negative mining 시에는 NV-Retriever의 방법을 따라, KURE-v1 모델로 relative margin 0.5를 주고 mining하였다. Margin과 관련해서는 이 PR를 확인하면 이해가 편할 것이다. 최종 데이터의 예시는 다음과 같다.

{
  "anchor": "오토마타를 사용해 검증할 필요가 있는 것은 시뮬레이션 제어 기능이니?", 
  "positive": "오토마타 이론을 적용한 Simulink 실행 모듈 구현\n이 모듈이 제공하려는 모든 기능을 오토마타를 이용하여 표현할 경우 전체 시스템이 매우 복잡하고 해석하기 부담스러운 모델이 된다. \n또한, 시뮬레이션 제어를 제외한 기능들은 굳이 오토마타를 이용하여 검증할 필요가 없다. \n따라서 3가지 시스템의 상태가 복합적으로 고려되어야 하는 시뮬레이션 제어 기능을 위주로 오토마타를 이용한 설계를 할 것이다. ", 
  "negative_1": "인크리멘탈 이벤트－구동 HDL 시뮬레이션에의 실제적 접근법\n이를 보완하기 위하여 별도의 하드웨어 검증플랫폼을 활용하는 에뮬레이션 또는 FPGA(Field Programmable Gate Array)를 활용하는 프로토타이핑을 통한 하드웨어기반 방식의 검증이 활용되고는 있지만, 합성가능하지 않은 설계객체(Verilog에서는 module, VHDL 에서는 entity) 들에 대한 처리 불가능, 유연성 결여, 디버깅의 어려움, 높은 가격 등의 문제들로 인하여 시뮬레이션과 같은 정도로 보편적으로 사용되고 있지는 못한 것이 현실이다. \n따라서, 이벤트-구동 HDL 시뮬레이터를 활용하는 시뮬레이션의 성능 향상의 요구가 구준히 지속적으로 제기되고 있다. ", 
  "negative_2": "3D 자동차 시뮬레이터 기반 상호작용형 ADAS 개발 및 검증 프레임워크\n다양한 내부 모듈들이 서로 상호작용 한다는 것은 복잡도가 증가한다는 의미를 갖고 있다. \n실제 환경에서 복합적인 센서들의 상호작용을 검증하는 것은 상당한 노력과 비용이 발생한다. \n따라서 차량 시뮬레이터를 사용하는 것은 시스템을 개발함에 있어서 매우 효율적인 방법이다.  그러나, 차량 시뮬레이터는 대부분 시나리오를 구성하기 위한 유연성을 제공하지만 모듈, 시스템을 수정하지 못하기 때문에 항상 같은 검증 시나리오만을 거치게 된다. ", 
  "negative_3": "최적의 Bang-Bang 입력을 이용한 볼-빔 시스템의 강인한 추적 제어\n따라서, 파라미터의 불확실성에 대해서는 강인성을 그리고 외란 및 입력정합조건에 부합하지 못하는 불확실성에 대해서도 성능을 향상시키는 새로운 강인한 제어 기법을 제안하였고, 시뮬레이션 및 실험 결과를 통해 입증하였다. ", 
  "negative_4": "오토마타 이론을 적용한 Simulink 실행 모듈 구현\n 요 약  본 연구에서는 오토마타 이론을 적용하여 Simulink 모델의 각종 정보를 보여주고 시뮬레이션 결과를 출력하는 모듈 구조를 제안하고, 이를 구현하였다. ", 
  "negative_5": "최적 위성추적 안테나 <tex>\\(H_{\\infty}\\)<tex> 제어 시스템의 설계에 관한 연구\n또한 컴퓨터 시뮬레이션을 통하여 설계된 \\( \\mathrm{H}_{\\infty} \\) 제어 시스템의 유용성을 확인하였다.  설계된 최적 위성추적 안테나 \\( \\mathrm{H}_{\\infty} \\) 제어 시스템은 기준모델과 마찬가지로 설계사양에 따른 만족스러운 \\(\\%\\) 오버슈트와 상승시간 및 정정시간을 가진다. \n반면에 가중치함수의 게인만을 고려한 경우 및 시행착오법의 경우에는 상대적으로 큰 기준모델추종오차가 존재한다. \n따라서 설계사양에 따른 만족스러운 응답특성을 얻기 위해서는 제시된 설계 알고리즘과 같이 게인 뿐만 아니라 가중치 함수의 동특성을 설계 파라메타 \\( \\gamma \\) 와 함께 동시에 최적화하는 것이 필요하다. ", 
  "negative_6": "인크리멘탈 이벤트－구동 HDL 시뮬레이션에의 실제적 접근법\n이와 같은 DUT에 존재하는 시그널들에 대한 가시도 확보를 위해서는 오류 발견을 위한 시뮬레이션 단계에서 요구되는 시뮬레이션 실행 횟수가 실제로는 2회 이상이며, 이 가운데 1회 이상이 DUT에 존재하는 시그널에 대한 덤프를 동반하는 시뮬레이션이 필요하다. "
}

[Train_Backbone] 백본 모델로는 한국어 토큰이 많이 학습되었으며, context length가 긴 skt의 A.X-Encoder-base 모델을 사용했다. SPLADE 모델 자체가 Tokenizer 및 모델의 MLM 성능에 굉장히 큰 영향을 받기 때문에, 잘 학습된 인코더를 사용하지 않으면 학습이 붕괴될 수 있다. (관련 이슈)

[Train_Loss] Loss는 InfoNCE Loss, sentence-transformers의 MultipleNegativesLoss로 진행한다. 이전 글에서 설명했듯이, contrastive learning은 기본적으로 배치 사이즈에 굉장히 큰 영향을 받는다. Dense Embedding model을 학습할 때는 Gradient Cache를 지원하기 때문에 쉽게 배치 사이즈를 키울 수 있었지만, Sparse model에서는 loss를 계산할 때 sparsity를 조절하는 Regularizer와 Gradient Cache의 개념이 충돌이 발생하여 사용하지 못한다 (참고). 다행인 점은, sentence-transformers v5.1.0 부터 gather_across_devices를 지원한다. 이는 device별로 document를 공유할 수 있게 하는 인자이기 때문에, (per_device_train_batch_size) * (device 개수)가 total batch size이다.

최종 학습에 사용한 하이퍼파라미터는 다음과 같다.

{
  "eval_strategy": "steps",
  "per_device_train_batch_size": 4,
  "per_device_eval_batch_size": 2,
  "learning_rate": 2e-05,
  "query_regularizer_weight": 5e-5,
  "document_regularizer_weight": 3e-5,
  "num_train_epochs": 2,
  "warmup_ratio": 0.1,
  "bf16": True,
  "negs_per_query": 6,
  "gather_device": True
}

[Evaluation] SPLADE에서 중요한 구성요소 중 하나는 “활성화되는 차원”의 개수이다. 활성화되는 차원이 너무 많으면 index 시에 많은 공간이 들 뿐더러 값이 여러 차원으로 분산되어 성능이 하락할 수 있다. 반면 활성화되는 차원이 너무 적으면 토큰이 제대로 확장되지 못해, 관련 문서 검색이 아예 불가해질 수 있다. 이렇기 때문에 차원의 균형을 맞추는 것은 굉장히 중요하다. (Indexing과 관련해서는 추후 글에서 다룰 예정이다.)

이러한 구성요소 (성능, activate dimension 수)들을 평가하기 위해, sentence-transformers의 SparseInformationRetrievalEvaluator를 사용한다. 상세한 평가 코드는 다음 huggingface에서 볼 수 있다.

평가 결과는 다음과 같다.

• splade-ko-v1 모델이 타 sparse 모델들에 비해 매우 높은 nDCG@10 점수를 보이는 것을 볼 수 있다.
• opensearch의 모델은 activate되는 document의 dimension이 너무 적어서, 유사한 단어를 효과적으로 확장하지 못하는 것으로 확인되고, PIXIE 모델의 경우 또한 query, document active dimension이 splade-ko-v1 모델보다 작은 것을 관찰할 수 있었다.

타 Dense 모델과의 nDCG@10 점수를 평가한 경우, 다음과 같은 양상을 볼 수 있다.

정말 놀랍게도, 0.1B라는 굉장히 작은 파라미터에도 불구하고, KURE-v1 급의 성능을 보여준다. 여기서 우리는 실제로 SPLADE 기반의 모델이 굉장히 강력한 검색 모델임을 알 수 있다.

2–2. inference-free-splade-ko-v1

🤗 Model: inference-free-splade-ko-v1

Inference-Free SPLADE 모델 학습 시에는 InfoNCE와 더불어 MaginMSE, KLDiv 등 여러 가지 Distillation 기반의 Loss들을 실험해 보았다.

[Dataset] 데이터는 앞서 사용한 데이터들을 사용하되, distillation을 위해 해당 데이터들에 대한 Reranker의 스코어를 뽑았다. 이때 성능이 매우매우 좋은 Qwen3-Reranker-8B를 사용했다. 시간을 단축하고자, DDP (Distributed Data Parallel) 기술처럼 각 gpu에 reranker 모델을 하나씩 서빙해두고, data를 모든 gpu에 분산시켜, gpu local rank마다 균등하게 slice 하도록 하여 데이터를 병렬로 처리하였다. 그 결과 약 2일 안에 스코어를 뽑을 수 있었다 (A6000 8장, vllm/FA2 사용, 1M query-positive-negatives(7개) 기준). 최종 데이터의 예시는 다음과 같다.

{
  "anchor": "오토마타를 사용해 검증할 필요가 있는 것은 시뮬레이션 제어 기능이니?", 
  "positive": "오토마타 이론을 적용한 Simulink 실행 모듈 구현\n이 모듈이 제공하려는 모든 기능을 오토마타를 이용하여 표현할 경우 전체 시스템이 매우 복잡하고 해석하기 부담스러운 모델이 된다. \n또한, 시뮬레이션 제어를 제외한 기능들은 굳이 오토마타를 이용하여 검증할 필요가 없다. \n따라서 3가지 시스템의 상태가 복합적으로 고려되어야 하는 시뮬레이션 제어 기능을 위주로 오토마타를 이용한 설계를 할 것이다. ", 
  "negative_1": "인크리멘탈 이벤트－구동 HDL 시뮬레이션에의 실제적 접근법\n이를 보완하기 위하여 별도의 하드웨어 검증플랫폼을 활용하는 에뮬레이션 또는 FPGA(Field Programmable Gate Array)를 활용하는 프로토타이핑을 통한 하드웨어기반 방식의 검증이 활용되고는 있지만, 합성가능하지 않은 설계객체(Verilog에서는 module, VHDL 에서는 entity) 들에 대한 처리 불가능, 유연성 결여, 디버깅의 어려움, 높은 가격 등의 문제들로 인하여 시뮬레이션과 같은 정도로 보편적으로 사용되고 있지는 못한 것이 현실이다. \n따라서, 이벤트-구동 HDL 시뮬레이터를 활용하는 시뮬레이션의 성능 향상의 요구가 구준히 지속적으로 제기되고 있다. ", 
  "negative_2": "3D 자동차 시뮬레이터 기반 상호작용형 ADAS 개발 및 검증 프레임워크\n다양한 내부 모듈들이 서로 상호작용 한다는 것은 복잡도가 증가한다는 의미를 갖고 있다. \n실제 환경에서 복합적인 센서들의 상호작용을 검증하는 것은 상당한 노력과 비용이 발생한다. \n따라서 차량 시뮬레이터를 사용하는 것은 시스템을 개발함에 있어서 매우 효율적인 방법이다.  그러나, 차량 시뮬레이터는 대부분 시나리오를 구성하기 위한 유연성을 제공하지만 모듈, 시스템을 수정하지 못하기 때문에 항상 같은 검증 시나리오만을 거치게 된다. ", 
  "negative_3": "최적의 Bang-Bang 입력을 이용한 볼-빔 시스템의 강인한 추적 제어\n따라서, 파라미터의 불확실성에 대해서는 강인성을 그리고 외란 및 입력정합조건에 부합하지 못하는 불확실성에 대해서도 성능을 향상시키는 새로운 강인한 제어 기법을 제안하였고, 시뮬레이션 및 실험 결과를 통해 입증하였다. ", 
  "negative_4": "오토마타 이론을 적용한 Simulink 실행 모듈 구현\n 요 약  본 연구에서는 오토마타 이론을 적용하여 Simulink 모델의 각종 정보를 보여주고 시뮬레이션 결과를 출력하는 모듈 구조를 제안하고, 이를 구현하였다. ", 
  "negative_5": "최적 위성추적 안테나 <tex>\\(H_{\\infty}\\)<tex> 제어 시스템의 설계에 관한 연구\n또한 컴퓨터 시뮬레이션을 통하여 설계된 \\( \\mathrm{H}_{\\infty} \\) 제어 시스템의 유용성을 확인하였다.  설계된 최적 위성추적 안테나 \\( \\mathrm{H}_{\\infty} \\) 제어 시스템은 기준모델과 마찬가지로 설계사양에 따른 만족스러운 \\(\\%\\) 오버슈트와 상승시간 및 정정시간을 가진다. \n반면에 가중치함수의 게인만을 고려한 경우 및 시행착오법의 경우에는 상대적으로 큰 기준모델추종오차가 존재한다. \n따라서 설계사양에 따른 만족스러운 응답특성을 얻기 위해서는 제시된 설계 알고리즘과 같이 게인 뿐만 아니라 가중치 함수의 동특성을 설계 파라메타 \\( \\gamma \\) 와 함께 동시에 최적화하는 것이 필요하다. ", 
  "negative_6": "인크리멘탈 이벤트－구동 HDL 시뮬레이션에의 실제적 접근법\n이와 같은 DUT에 존재하는 시그널들에 대한 가시도 확보를 위해서는 오류 발견을 위한 시뮬레이션 단계에서 요구되는 시뮬레이션 실행 횟수가 실제로는 2회 이상이며, 이 가운데 1회 이상이 DUT에 존재하는 시그널에 대한 덤프를 동반하는 시뮬레이션이 필요하다. ",
  "label": [7.649763107299805, -3.9823877811431885, -3.4091150760650635, -1.8980331420898438, -0.9259235858917236, -3.6210155487060547, -2.0109493732452393, -0.9311941862106323]
}

※ 본 모델의 스코어를 추출하기 위해, sentence-transformers 기반으로 사용하기 조금 더 용이한 tomaarsen/Qwen3-Reranker-8B-seq-cls 모델을 사용했는데, sentence-transformers 기반의 모델은 기본적으로 맨 끝에 sigmoid 함수를 달아 normalized score를 추출하도록 설계되어 있다. (관련 sentence-transformers PR)

하지만 Distillation을 위해 reranker score를 사용하는 경우, unnormalized score를 많이들 사용한다는 것을 알게 되었다 (관련 SPLADE 논문, 관련 PyLate PR). 이러한 관례?를 따라, 나는 inference 시 normalize=False로 설정했다. 구체적인 reranker score를 추출하는 코드는 다음과 같다.

# get_crossencoder_score.py
from vllm import LLM
from vllm import PoolingParams

from datasets import load_dataset, load_from_disk
import torch
import os
from transformers import AutoTokenizer
import fire

model_name = "tomaarsen/Qwen3-Reranker-8B-seq-cls"

prefix = '<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n'
suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"

query_template = "{prefix}<Instruct>: {instruction}\n<Query>: {query}\n"
document_template = "<Document>: {doc}{suffix}"

# TODO 하드코딩
max_model_len = 8192
truncation_len = 7680
num_negatives = 7


def main(
 data_path: str = "/data_x/EMBEDDING/DATA/SFT/v2-filtered-new/train_ko_250625_KURE-v1_hn50",
 output_dir: str = "/data_x/EMBEDDING/DATA/SFT/v2-filtered-new/train_ko_250625_KURE-v1_hn50_subsets"
) -> None:
 # RANK 읽기
 rank = int(os.environ.get("RANK", 0))
 world_size = int(os.environ.get("WORLD_SIZE", 1))
 
 print(f"[Rank {rank}/{world_size}] Starting process...")
 
 instruction = (
  "Given a web search query, retrieve relevant passages that answer the query"
 )

 if data_path.endswith(".jsonl"):
  data = load_dataset("jsonl", data_files=data_path, split="train")
 elif os.path.isdir(data_path):
  data = load_from_disk(data_path)
 else:
  raise ValueError(f"Unsupported data format: {data_path}")
 
 if data[0].get("task") is not None:
  dataset = data.remove_columns(["task", "lang"])
 else:
  dataset = data
 
 if data[0].get("negative_50") is not None:
  dataset = dataset.remove_columns([f"negative_{i}" for i in range(33, 51)])
 
 # 데이터를 world_size 개로 분할하여 이 rank가 담당할 부분만 선택
 total_samples = len(dataset)
 samples_per_rank = (total_samples + world_size - 1) // world_size 
 start_idx = rank * samples_per_rank
 end_idx = min(start_idx + samples_per_rank, total_samples)
 
 dataset = dataset.select(range(start_idx, end_idx))
 
 print(f"[Rank {rank}/{world_size}] Processing samples {start_idx} to {end_idx} ({len(dataset)} samples)")
 print(f"[Rank {rank}/{world_size}] Dataset example: {dataset[0]}")
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 model = LLM(
  model=model_name, 
  max_model_len=max_model_len, 
  task="score",
  dtype=torch.bfloat16,
  override_pooler_config={"activation": False} # For unnormalized scores!!!
 )

 
 batch_size = 1024
 all_scores = []
 
 for batch_idx in range(0, len(dataset), batch_size):
  batch_end = min(batch_idx + batch_size, len(dataset))
  batch_data = dataset[batch_idx:batch_end]
  
  queries = []
  documents = []
  
  for i in range(len(batch_data["anchor"])):
   anchor = batch_data["anchor"][i]
   query_text = query_template.format(
    prefix=prefix, 
    instruction=instruction, 
    query=anchor
   )
   
   docs = [batch_data["positive"][i]]
   docs.extend([batch_data[f"negative_{j}"][i] for j in range(1, num_negatives + 1)])
   
   for doc in docs:
    suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
    suffix_len = len(suffix_tokens)
    
    doc_text = document_template.format(doc=doc, suffix=suffix)
    combined = query_text + doc_text
    tokens = tokenizer.encode(combined, add_special_tokens=False)
    
    if len(tokens) > truncation_len:
     query_tokens = tokenizer.encode(query_text, add_special_tokens=False)
     available_for_doc = truncation_len - len(query_tokens) - suffix_len
     
     if available_for_doc > 0:
      doc_prefix = "<Document>: "
      doc_prefix_tokens = tokenizer.encode(doc_prefix, add_special_tokens=False)
      available_for_doc_content = available_for_doc - len(doc_prefix_tokens)
      
      if available_for_doc_content > 0:
       doc_tokens = tokenizer.encode(doc, add_special_tokens=False)[:available_for_doc_content]
       truncated_doc = tokenizer.decode(doc_tokens, skip_special_tokens=False)
       doc_text = document_template.format(doc=truncated_doc, suffix=suffix)
      else:
       doc_text = document_template.format(doc="", suffix=suffix)
     else:
      query_text = tokenizer.decode(query_tokens[:truncation_len // 2], skip_special_tokens=False)
      available_for_doc_after_query = truncation_len // 2 - suffix_len
      
      if available_for_doc_after_query > 0:
       doc_prefix = "<Document>: "
       doc_prefix_tokens = tokenizer.encode(doc_prefix, add_special_tokens=False)
       available_for_doc_content = available_for_doc_after_query - len(doc_prefix_tokens)
       
       if available_for_doc_content > 0:
        doc_tokens = tokenizer.encode(doc, add_special_tokens=False)[:available_for_doc_content]
        truncated_doc = tokenizer.decode(doc_tokens, skip_special_tokens=False)
        doc_text = document_template.format(doc=truncated_doc, suffix=suffix)
       else:
        doc_text = document_template.format(doc="", suffix=suffix)
      else:
       doc_text = document_template.format(doc="", suffix=suffix)
    
    queries.append(query_text)
    documents.append(doc_text)
  
  print(f"[Rank {rank}] Processed {batch_end}/{len(dataset)} samples ({batch_end/len(dataset)*100:.1f}%)")
  outputs = model.score(
   queries, 
   documents,
   pooling_params=PoolingParams(task="score", activation=False)
  )
  scores = [output.outputs.score for output in outputs]
  
  num_samples = len(batch_data["anchor"])
  total_docs_per_sample = num_negatives + 1  
  for i in range(num_samples):
   sample_scores = scores[i * total_docs_per_sample:(i + 1) * total_docs_per_sample]
   all_scores.append(sample_scores)
 
 dataset = dataset.add_column("label", all_scores)
 
 print(f"[Rank {rank}] Scoring complete! Dataset shape: {dataset}")
 print(f"[Rank {rank}] Sample: {dataset[0]}")
 
 data_name = os.path.basename(data_path.rstrip('/'))
 output_path = os.path.join(output_dir, f"{data_name}_rank{rank}")
 os.makedirs(output_path, exist_ok=True)
 dataset.save_to_disk(output_path)
 print(f"[Rank {rank}] Saved to {output_path}")

if __name__ == "__main__":
 fire.Fire(main)

export WORLD_SIZE=8

DATA_PATH="/data_x/EMBEDDING/DATA/SFT/v2-filtered-new/train_ko_250625_KURE-v1_hn50"
OUTPUT_DIR="/data_x/EMBEDDING/DATA/SFT/v2-filtered-new/train_ko_250625_KURE-v1_hn50_subsets"

SHORT_DATA_PATH=$(basename "$DATA_PATH")

echo "Starting parallel processing on 8 GPUs..."
echo "  Input data path: $DATA_PATH"
echo "  Output directory: $OUTPUT_DIR"

RANK=0 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=0 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank0 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank0.txt 2>&1 &
RANK=1 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=1 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank1 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank1.txt 2>&1 &
RANK=2 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=2 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank2 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank2.txt 2>&1 &
RANK=3 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=3 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank3 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank3.txt 2>&1 &
RANK=4 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=4 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank4 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank4.txt 2>&1 &
RANK=5 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=5 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank5 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank5.txt 2>&1 &
RANK=6 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=6 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank6 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank6.txt 2>&1 &
RANK=7 VLLM_ATTENTION_BACKEND=FLASH_ATTN CUDA_VISIBLE_DEVICES=7 VLLM_CACHE_ROOT=/tmp/vllm_cache_rank7 uv run get_crossencoder_score.py --data_path="$DATA_PATH" --output_dir="$OUTPUT_DIR" > logs/logs_rank7.txt 2>&1 &

wait

echo "All processes completed!"
echo "Results saved in:"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank0"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank1"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank2"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank3"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank4"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank5"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank6"
echo "  - ${OUTPUT_DIR}/${SHORT_DATA_PATH}_rank7"

[Train_Backbone] 백본 모델로는 이전 splade-ko-v1과 동일하게 skt의 A.X-Encoder-base 모델 (0.1B)을 사용했다.

[Train_Setting]

<1. IDF Weight의 사용 유무>

앞선 inference-free splade에 관한 설명을 보면, NAVER의 SPLADE-doc 모델은 idf를 사용하지 않는 반면, opensearch 저자들은 idf weight를 사용한다. 이를 고려하여, idf weight를 사용하는 경우/사용하지 않는 경우 모두에 대해 실험을 진행한다.

<2. IDF Weight 학습 여부>

opensearch 저자들은 idf weight를 고정하고 학습을 하지만, 생각해보면 idf weight마저 학습을 해볼 수 있으며, 이에 대한 실효성이 충분히 있을 수 있다. 이렇듯 idf weight 사용 시 해당 weight를 학습하는 경우/학습하지 않는 경우 모두에 대해 실험을 진행한다.

<3. Losses & Details>

Distillation의 효용성을 실험해 보기 위해, InfoNCE, MarginMSE, KLDiv Loss를 모두 비교한다.

• InfoNCE 학습 시에는 A100 GPU 8장, batch size 4, gather_across_devices를 사용하여 총 batch size가 32인 환경에서 학습한다.
• MarginMSE는 positive와 negative 차이의 유사도 점수의 차이 (margin)가 teacher score의 분포와 비슷하도록 학습하는 방식이다.
  일례로, query, positive, negative_1, negative_2, label가 있을 때,
  [
  s(query, positive)-s(query, negative_1),
  s(query, positive)-s(query, negative_2)
  ]
  를 label의 분포와 일치하도록 MSE Loss를 사용하여 학습한다.
• KLDiv는 student가 매기는 positive, negative의 score와 teacher score의 분포가 일치하도록 학습하는 방식이다.

본 실험에서는 동일한 hyperparameter 내에서, 이 3가지 loss들로 모델을 학습하며, 학습된 모델들의 평가 결과를 모두 비교해 본다.

[Results]

<1. IDF Weight를 사용하지 않는 경우>

첫번째로 idf weight를 사용하지 않지만 (이 경우, query lookup weight는 1로 초기화된다), 이 weight를 학습하는 경우에 대한 실험을 진행한다. 실험에 사용된 hyperparameter들은 다음과 같다.

• 공통: lr=5e-5, staticlr=5e-4 (query weight 학습 lr)
• Document regularizer (Dreg): 1e-5, 1e-6, 1e-7 모두에 대해 실험

모든 MTEB Korean Retrieval Task에 대한 실험 결과는 다음과 같다.

실험 결과,

• InfoNCE Loss로, lr 1e-6으로 학습한 모델과 KLDiv Loss로 lr 1e-7으로 학습한 모델이 거의 동일하게 가장 우수한 성능을 보인다. 그러나 일반 splade 모델 (splade-ko-v1)이 0.738의 성능을 보이는 것에 비하면 턱없이 부족한 성능을 보인다.
• Distillation 관점에서는 KLDiv이 MarginMSE보다 훨씬 더 좋은 결과를 보인다. 모든 hyperparm에서 KLDiv으로 학습된 모델이 유의미한 성능 향상을 보이며, 이는 다음 논문에도 리포팅된 사실이다.
• (당연하게도) Avg. Query active dimension은 일정하며, Avg. Document active dimension은 regularizer가 작아짐에 따라 커진다 (규제항이 작아지므로 더 많이 activate됨).

<2. IDF Weight를 사용하는 경우>

두번째로 idf weight를 사용하는 경우, 먼저 우리는 학습 데이터에 대해 tokenizer의 각 token의 idf weight를 구하고, json 등의 형태로 저장해 놓아야 한다. idf weight를 구하는 코드는 다음과 같다.

# get_idf.py

import argparse
import json
import math
import os
from collections import Counter

import datasets
from transformers import AutoTokenizer


def parse_args():
 parser = argparse.ArgumentParser(
  description="Compute IDF from BeIR/msmarco corpus using a Hugging Face tokenizer"
 )
 parser.add_argument(
  "--tokenizer",
  "-t",
  type=str,
  required=True,
  help="Hugging Face tokenizer id or local path (e.g., 'bert-base-uncased')",
 )
 parser.add_argument(
  "--corpus_path",
  type=str,
  required=True,
  help="Path to the corpus",
 )
 parser.add_argument(
  "--model_max_length",
  type=int,
  default=8192,
  help="Maximum length of the tokenizer",
 )
 parser.add_argument(
  "--output",
  "-o",
  type=str,
  required=True,
  help="Output JSON path for IDF mapping (token_id -> idf)",
 )
 parser.add_argument(
  "--num_proc",
  type=int,
  default=30,
  help="Number of processes for datasets.map",
 )
 parser.add_argument(
  "--batch_size",
  type=int,
  default=10000,
  help="Batch size for batched tokenization",
 )
 return parser.parse_args()


def main():
 args = parse_args()

 tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, use_fast=True)
 tokenizer.model_max_length = args.model_max_length
 special_ids = set(tokenizer.all_special_ids or [])

 if os.path.isdir(args.corpus_path):
  corpus = datasets.load_from_disk(args.corpus_path)
 elif args.corpus_path.endswith(".jsonl"):
  corpus = datasets.load_dataset("json", data_files=args.corpus_path, split="train")
 else:
  corpus = datasets.load_dataset(args.corpus_path, split="train")

 corpus = corpus.map(
        lambda x: {"text": (x["anchor"] + "\n\n" + x["positive"])},
        num_proc=args.num_proc,
        desc="Formatting corpus...",
    )

 def _tokenize_batch(batch):
  encodings = tokenizer(
   batch["text"],
   add_special_tokens=False,
   truncation=False,
   return_attention_mask=False,
   return_token_type_ids=False,
  )
  batch_df = Counter()
  valid_docs = 0
  for ids in encodings["input_ids"]:
   unique_ids = {tid for tid in ids if tid not in special_ids}
   if unique_ids:
    batch_df.update(unique_ids)
    valid_docs += 1

  batch_len = len(batch["text"])
  if len(batch_df) == 0:
   df_keys = [[] for _ in range(batch_len)]
   df_vals = [[] for _ in range(batch_len)]
   valid_docs_list = [0 for _ in range(batch_len)]
  else:
   keys_list = list(batch_df.keys())
   vals_list = [batch_df[k] for k in keys_list]
   df_keys = [keys_list] + [[] for _ in range(batch_len - 1)]
   df_vals = [vals_list] + [[] for _ in range(batch_len - 1)]
   valid_docs_list = [valid_docs] + [0 for _ in range(batch_len - 1)]

  return {"df_keys": df_keys, "df_vals": df_vals, "valid_docs": valid_docs_list}

 tokenized = corpus.map(
  _tokenize_batch,
  batched=True,
  batch_size=args.batch_size,
  remove_columns=[col for col in corpus.column_names if col != "text"],
  num_proc=args.num_proc,
  desc="Tokenizing corpus",
 )

 df_counter: Counter[int] = Counter()
 df_keys_col = tokenized["df_keys"]
 df_vals_col = tokenized["df_vals"]
 valid_docs_col = tokenized["valid_docs"]
 for keys, vals in zip(df_keys_col, df_vals_col):
  if keys:
   df_counter.update(dict(zip(keys, vals)))

 num_docs = sum(valid_docs_col)
 print(num_docs)

 idf_by_token_id = {
  tokenizer._convert_id_to_token(tid): math.log(
   1 + (num_docs - df + 0.5) / (df + 0.5)
  )
  for tid, df in df_counter.items()
  if df > 0
 }

 output_dir = os.path.dirname(os.path.abspath(args.output))
 if output_dir and not os.path.exists(output_dir):
  os.makedirs(output_dir, exist_ok=True)

 with open(args.output, "w", encoding="utf-8") as f:
  json.dump(idf_by_token_id, f, ensure_ascii=False)

 print(
  f"Wrote IDF for {len(idf_by_token_id)} tokens from {num_docs} documents -> {args.output}"
 )


if __name__ == "__main__":
 main()

# get_idf.sh
uv run get_idf.py \
 --tokenizer skt/A.X-Encoder-base \
 --model_max_length 8192 \
 --corpus_path /data_x/EMBEDDING/DATA/SFT/v2-merged/train_ko_formatted_30_2048_250613.jsonl \
 --output ./idfs/skt_A.X-Encoder-base_train_ko_formatted_30_2048_250613_idf.json \
 --num_proc 32 \
 --batch_size 8192

이렇게 구한 idf weight는 다음과 같다.

{“##가”: 0.12593505727433768, “모듈”: 5.345584361396796, “굳이”: 5.793233434985806, …….., “##裟”: 13.315633666372932, “袈”: 13.3156336}

이후 위 idf weight를 사용하여 학습을 진행한다. 학습 시 실험 변인은 다음과 같다.

• Loss (InfoNCE / KLDiv / MarginMSE)
• idf weight 학습 여부 (Freeze / Unfreeze)
• Query/Document regularizer (1e-6~1e-8)

각각의 변인을 고려하여 모델 학습 후 검색 성능 결과를 비교해보면 다음과 같다.

실험 결과,

• InfoNCE Loss를 사용하여, idf weight를 고정하고, lr 1e-8으로 학습한 모델이 압도적으로 좋은 성능 (0.721)을 보인다. 이 모델은 splade-ko-v1 (0.738)과 견주어도 크게 밀리지 않으며, PIXIE-Splade-Preview (0.718)보다 높은 성능을 보인다.
• 전체적으로 idf를 사용하는 모델들이, 사용하지 않는 모델들에 비해 훨씬 검색 성능이 높으며, regularizer에 따라 idf를 학습하는 편이 좋은지, 고정하는 편이 좋은지 다르다. (본 실험에서는 고정하는 편이 좋았다.)
• 그러나 activate되는 document dimension의 크기가 너무 커서, 실제로 얼마나 효용성이 있을지는 의문이다. (한 document 당 19,207 개의 숫자값을 색인해 놓아야 한다는 뜻. Indexing에 관해서는 향후 블로그에서 더 자세하게 다뤄보겠습니다.)
• 본 실험 또한 Distillation 관점에서는 KLDiv이 MarginMSE보다 더 유의미한 결과를 보인다.

마무리하며..

오늘은 SPLADE 모델과 Inference-Free SPLADE 모델을 한국어 데이터로 학습하는 과정에 있어 경험적인 실험들을 공유해 보았다. 구체적으로는

1. 전체적인 학습 과정
2. 대규모 데이터 처리를 위한 vLLM 기반의 병렬 스코어링 파이프라인
3. Inference-Free SPLADE 모델 학습 시 3가지 Loss의 비교
4. Inference-Free 환경에서 IDF Weight 고정 전략이 성능에 미치는 영향

등을 중점적으로 살펴보았다.

이번 실험을 통해 얻은 가장 고무적인 성과는 0.1B 수준의 경량화된 모델로도 0.6B급 이상의 Dense Embedding 모델과 대등하거나 더 나은 성능을 낼 수 있다는 가능성을 확인한 점이다. 이는 GPU 리소스가 제한적인 On-premise 환경이나, 빠른 응답 속도가 필수적인 검색 시스템에서 SPLADE가 강력한 대안이 될 수 있다.
또 글의 초입부에서 첨부한 그래프처럼, Hybrid Retrieval 시스템 구축 시 BM25 사용에 비해 더 향상된 검색 성능을 확인할 수 있다.

물론 Inference-Free SPLADE 모델에서 성능을 높일수록 Document Active Dimension이 증가하여 인덱싱 비용이 늘어나는 trade-off가 관찰되기도 했다. 따라서 실제 서비스에 적용할 때는 가용 가능한 메모리와 허용 가능한 Latency 사이에서 적절한 균형점을 찾는 것이 중요할 것 같다.

여전히 검색 기술은 Dense Retrieval이 주류를 이루고 있지만, 정확한 키워드 매칭이 중요한 도메인이나 검색 결과의 설명 가능성(Explainability)이 요구되는 분야에서 Sparse Retrieval은 대체 불가능한 장점을 가진다. 이번에 공개한 splade-ko-v1과 inference-free-splade-ko-v1 모델, 그리고 이 학습기가 한국어 검색 모델을 연구하는 분들에게 작게나마 도움이 되기를 바란다.

Embedding model (Bi-encoder)은 input으로 text를 받아 output으로 고정된 차원의 실수들의 집합인 Embedding을 출력하는 반면, Reranker (Cross-encoder)는 input으로 두 쌍의 text (e.g. query[SEP]document)를 받아 해당 쌍의 유사도 score를 출력한다.

더 정확하게 설명하면, Embedding model (e.g. BERT)은 텍스트를 입력으로 받으면, 각 입력 토큰별로 하나의 임베딩을 출력하는데, 최종 embedding 출력 시에는 이 토큰들에 다양한 pooling method (e.g. mean pooling, [cls] token pooling)을 적용하여 하나의 임베딩으로 만든다.
반면 Reranker는 Embedding model과 동일한 구조를 갖지만, 마지막에 임베딩을 하나의 점수로 변환하는 선형 변환을 통해 스코어를 출력한다.

이렇듯, reranker는 항상 두 쌍의 pair를 모델의 input으로 사용하기 때문에 Time complexity 문제로 첫 번째 검색 (First-stage Retrieval)에 사용하기에는 상당히 부적합하다. 따라서 최근 Retrieval(검색) 시나리오는 다음과 같다.

1. [Embedding] 다량의 문서들을 각각 임베딩으로 만들어 VectorDB에 색인 (Indexing) 해둔다.
2. 검색 시 query가 입력되면 해당 query를 임베딩하여 기존에 색인된 모든 임베딩들과의 유사도 (cosine similarity)를 계산한다. (이 경우, 숫자들의 dot-product만 계산하기 때문에 빠르게 계산할 수 있다.)
3. 유사도가 높은 top-k 문서들을 검색해온다. (reranker에 들어갈 scope을 줄이는 과정으로 볼 수 있다.)
4. [Reranking] query와 k개의 문서(q-d1, q-d2, …, q-dk)들을 하나씩 pair로 만들어 Reranker에 통과시키면, k개의 유사도 점수가 출력된다.
5. 해당 k개의 유사도를 내림차순으로 정렬하여 top n개를 뽑으면 그것이 모델로 판단한 document pool에서 query와 가장 유사한 n개라고 할 수 있다.

First stage retrieval에 사용되는 Embedding model을 학습하는 방법과 현대 embedding model들이 사용하는 loss function에 대해서는 이전 포스트에 작성했었다. 오늘은 Reranker model을 학습하는 방법론에 대해 작성해보도록 하겠다.

Reranker 학습에 대해 공부하게 된 이유는 다음 PR을 보면서 부터이다.

[`feat`] CrossEncoder Training refactor - MultiGPU, loss logging, bf16, etc. by tomaarsen · Pull Request #3222 · UKPLab/sentence-transformers

Sentence-transformers 측에서 reranker (cross-encoder) 학습 라이브러리를 만들고자 위와 같은 PR을 만들었다. 저 PR에 있는 내용들이 너무 새롭고 재미있어서 매일 무엇이 업데이트 되는지 들여다 보았는데, 가장 흥미로웠던 것은 처음 보는 2개의 loss인 ListnetLoss, LambdaLoss였다.

도대체 저 2개의 loss가 무엇인가 하여, 이에 대해 공부하게 되었다. 오늘은 이렇게 reranker 학습에 대해 공부하며 알게 된 Learning To Rank 라는 ranking과 관련한 오래된 분야에 대해 이야기하고자 한다. 참고한 논문들은 다음과 같다.

• Learning to Rank using Gradient Descent
• Learning to Rank: From Pairwise Approach to Listwise Approach
• Listwise Approach to Learning to Rank — Theory and Algorithm
• Position-Aware ListMLE: A Sequential Learning Process for Ranking
• The LambdaLoss Framework for Ranking Metric Optimization

먼저 Learning To Rank, LTR에 대해 간단하게 얘기하며 시작해보겠다. LTR은 검색 결과 혹은 추천 목록의 순위를 최적화하기 위해 ML을 사용하는 기법이다. “순위(Rank)”에 대한 최적화를 진행하기 때문에, 순위와 관련된 label이 필요하며, 이 때문에 LTR은 추천 시스템과도 깊게 연관되어 있다. (e.g. 사용자가 가장 좋아할만한 아이템을 상위에 배치하는 등)

이러한 LTR은 학습 방법론에 따라 크게 3가지로 나뉜다.

1. Pointwise: 개별 아이템의 점수를 예측해 순위를 정하는 방식
2. Pairwise: 두 아이템 간의 상대적 순서를 비교해 학습하는 방식
3. Listwise: 전체 리스트의 순위를 최적화하는 방식

각각에 대해 알아보도록 하자.

1. Pointwise Ranking

첫번째는 Pointwise Ranking이다. 이는 모든 데이터를 개별적으로 보자는 접근법이다. 예시를 들어, 다음과 같이 2개의 query가 있는데, 각각의 query에 대해 2개, 3개 문서가 각각 매칭되었다고 가정하자.

Pointwise Ranking에서는 각 query에 매칭된 문서들을 풀어서
(query-document) pair 형태로 구성한다.

이후, 각각의 문서들은 모두 label 1을 가지고 학습된다. 즉, cross-encoder 모델이 예측하는 q-d의 유사도가 1에 가깝게 되도록 Cross-Entropy Loss로 학습된다.

이 방식의 단점은 query와 매칭되는 문서 목록들의 Ranking은 직접적으로 학습하지 못한다는 점이다. 즉 특정 문서들과의 유사도만 학습이 가능하고, ranking에 대해서는 학습할 수 없다.

2. Pairwise Ranking — RankNet

이러한 문제를 해결하기 위해 RankNet이 등장한다. RankNet에서는 N개의 문서 중 2개 쌍을 선택하여 (총 경우의 수: NC2), 각 쌍에 대해 어떤 문서가 상위에 위치해야 하는지를 학습한다.

예시로, 다음과 같이 5개의 문서가 있다고 가정하자.

RankNet에서는 이 5개의 문서로 각각 pair를 만든 후 - [(doc-1, doc-2), (doc-1, doc-3), (doc-1, doc-4), …, (doc-4, doc-5)], 각 pair에 속한 doc과 query와의 유사도 스코어로 비교한다.

이는 Bradley-Terry 모델을 활용하여 확률을 표현한 방식이다. 참고로 Bradley-Terry 모델에 대한 설명은 다음 예시로 쉽게 끝낼 수 있다.

이렇듯 한쪽이 다른 한쪽보다 클 확률을 softmax 형식으로 나타낸 것이 Bradley-Terry 모델이다.

이와 같이, 위 식의 P(i>j)는 i번째 문서가 j번째 문서보다 높게 랭킹될 확률이다. 실제로 식에서 s_i가 크면, 분모가 점점 작아지고 확률이 1에 가까워진다.

우리는 이를 이용하여 다음과 같은 Cross-entropy를 활용한 Loss function으로 모델을 학습시킬 수 있다.

여기서 Pˉij는 실제로 문서 i가 j보다 높은 랭킹을 가져야 하는지를 나타내는 gold 확률 값이다. 이는 일반적인 CE (Cross Entropy) Loss에 Ranking을 녹여낸 방식이다.

그러나 이 방법의 단점은, 근본적으로 모든 document의 순위를 학습하지 않고, 두 문서간의 상대적인 순위만 학습한다는 것이다.

3. Listwise Ranking — ListNet

이러한 문제를 해결하기 위해 Listwise Approach가 등장한다. Listwise Ranking, ListNet은 전체 목록의 순위를 최적화하는 알고리즘이다.

앞서 N개의 문서 중 2개를 골라 총 NC2번의 학습을 하는 것과는 달리, ListNet에서는 N개의 문서를 모두 활용하여 1번의 학습만을 진행한다. ListNet에서는 다음의 2가지 개념을 제시한다.

3-1. Permutation Probability

먼저, 문서 n개의 상대적인 ranking (e.g. 검색 엔진에서의 문서들의 클릭 수)이 정해져 있다고 가정하자. (Gold label)

Permutation probability의 경우, 모든 문서들의 permutation (총 n! 가지)을 만들어 보며, 각각의 경우에 대해 model이 예측한 relevance 점수로 gold label과의 CE Loss를 계산하여 평균을 낸다.

예시로, 5개의 문서 [doc-1, doc-2, doc-3, doc-4, doc-5]가 있고, 이들의 ranking이 softmax 거친 값을 [0.6, 0.2, 0.08, 0.07, 0.05] 라고 하자. 그럼 이 경우, 먼저 query와 5개 문서 각각의 유사도를 model이 예측하게 하고, softmax를 구한다. 예측된 값을
[{doc-1}: 0.4, {doc-2}: 0.3, {doc-3}: 0.1, {doc-4}: 0.08, {doc-5}: 0.12] 라고 하자.

그럼 이들의 모든 permutation (5!; 줄세우기 라고 생각하면 편하다)를 구하며 gold label과의 CE Loss를 구하고, 그들의 평균을 최종 loss로 사용한다.

• [0.6, 0.2, 0.08, 0.07, 0.05]와 [0.4, 0.3, 0.1, 0.08, 0.12]의 CE
• [0.6, 0.2, 0.08, 0.07, 0.05]와 [0.4, 0.3, 0.1, 0.12, 0.08]의 CE
• …
• [0.6, 0.2, 0.08, 0.07, 0.05]와 [0.12, 0.08, 0.1, 0.3, 0.4]의 CE

3-2. Top-One Probability

그런데 사실 이렇게 하는 것은 cost 측면에서 정말 비효율적이다. 방금 설명한 예시에서 문서가 5개만 되어도 5!=120가지 경우에 대해 CE Loss를 구해야하는데, 그것이 너무 비용이 많이 들기 때문이다. 따라서 논문의 저자들은 Top-One Probability를 제안한다. 이는 전체 순서를 고려하지 않고, 각 document가 첫 번째 rank에 올 확률만 계산하겠다는 것이다.

예를 들어, 10개 document의 gold label에 softmax를 적용한 label이 [0.7, 0.1, 0.02, 0.05, …, 0.03] 이라고 가정하자. ListNet에서는 모델로 query와 각각의 document의 유사도 스코어를 뽑고 (아래 그림에서 [0.6, 0.3, …, 0.02]), 해당 스코어들의 softmax를 통해 각 document가 첫 번재 rank에 ranking될 확률을 표현한다.

이렇게 한 뒤, softmax를 거친 gold label [0.7, 0.1, 0.02, 0.05, …, 0.03]과 모델이 예측한 score들의 softmax를 거친 score [s1, s2, …, s10]와의 CE Loss만을 계산하면 끝이다. 이때는 loss 계산이 단 1번이라 효율적이며, 각각의 점수는 gold label에 가까워지게끔 학습되어 ranking까지 반영될 수 있다.

4. ListMLELoss

그러나 위와 같은 ListNet도 명확한 한계가 있다. 각 문서가 Top-1이 될 확률만 고려하기 때문에, 전체 순위(정확한 순열 정보)를 완벽히 반영하지 못한다는 것이다. 예를 들어, 정답 순위가 [1, 2, 3]이고, 모델이 [3, 1, 2]라고 예측했을 때,ListNet은 단순히 각 문서가 1등이 될 확률만 학습하므로, 전체 리스트의 정렬 오류를 제대로 반영하지 못한다.

따라서 ListMLE는 이를 보완하기 위해 단순히 “Top-1”을 예측하는 것이 아니라, 각 순위에서 올바른 문서를 선택하는 확률을 학습한다.

Plackett-Luce Model

Plackett-Luce Model은 주어진 리스트에서 각 항목이 특정 위치에 올 확률을 곱셈 형태의 조건부 확률로 정의하는 방식이다. 리스트에서 1등은 누구인지를 먼저 결정하고, 그다음 2등을, 그다음 3등을 결정하는 방식인 것이다.

예를 들어, [A, B, C] 문서가 있고, 각 문서의 유사도 점수는 [sA=3.0, sB=1.0, sC=2.0]라고 하자. 그리고 [A, C, B]로 정렬될 확률을 구해보자.
먼저, A가 1등이 될 확률은 다음과 같다.

남은 두 문서 B, C 중, C가 2등이 될 확률은 다음과 같다.

마지막으로 남은 문서 B 중 B가 3등이 될 확률은 1이다.

이 확률을 곱하면, 그것이 최종적으로 [A, C, B]로 정렬될 확률이 된다.

ListMLE에서는 이 확률을 최대화시키는 방법으로 학습한다. (Negative Log Lieklihood 적용)

이렇듯 ListMLE는 결국 “순서”를 최적화하도록 하는 loss이다.

5. Position Aware ListMLELoss (PListMLELoss)

PListMLELoss는 ListMLELoss만 알고 있으면 간단하다. PListMLE는 ListMLE가 position 간의 weight가 없어, top에 위치한 문서들의 중요성을 학습하지 못한다는 생각으로부터 나온 loss이다.

따라서, 기존 ListMLELoss에 가중치 alpha만 곱해주면 된다.

그리고 실제로 이로 인해 더 중요한, 더 높이 위치해야 하는 문서에 큰 weight를 주어 학습하여, model의 ranking 능력을 향상시켜주었다.

PListMLE가 강력한 힘을 발휘하는 때는 gold가 2개인데, 그 gold간에 순서가 있는 경우이다. 예를 들어, Reranker로 tool retrieval task를 학습시키려고 한다고 하자.

Tool retrieval task는 주어진 query에 알맞은 tool을 검색하는 태스크이다. 실제 세계에서의 tool 사용에 대해 생각해 보면, 하나의 query에 대해 2개의 tool을 사용해야 하며, 해당 tool 간의 순서가 있는 경우가 있다. 이럴 때, 2개의 tool 모두 query와의 유사도가 높아야 하는 tool이지만, 순서가 분명 존재한다. 바로 이러한 시나리오에 PListMLE를 사용하면 안성맞춤이다.

여담으로 ListMLE와 PListMLELoss는 필자가 앞서 언급한 PR에 구현하였다:

Add Position-Aware ListMLELoss by yjoonjang · Pull Request #6 · tomaarsen/sentence-transformers

해당 loss는 순서를 최적화시켜야 하는 loss이기 때문에, 사용자의 입력 순서를 보존하도록 하는 항을 추가하였다. 따라서 (P)ListMLE Loss로 학습하려면, 학습 데이터 정의 시, rank를 정의하여 입력으로 넣으면 좋을 것 같다.

6. LambdaLoss

사실 ListNet 혹은 (P)ListMLE가 가장 이상적이고, 우리가 reranker 학습 시에 적용하기 쉽고 효율적인 방식이라고 생각할 수 있다. 하지만 LambdaLoss가 나오게 된 배경은 우리가 측정하는 metric에 최적화된 loss를 사용할 필요가 있다는 것이다. 그래서 이 논문에서는 metric별로 다른 weight를 주는데, 오늘은 NDCG metric과 관련된 weight만 설명하도록 하겠다.

6–1. NDCG metric

먼저, NDCG (Normalized Discounted Cumulative Gain) metric에 대해 알아보자. 용어 정리부터 하자면

• DCG(Discounted Cumulative Gain): 순위별로 관련성 점수를 가중치와 함께 합산한 값. 상위 순위일수록 더 높은 가중치를 부여

• IDCG(Ideal DCG): 이상적인 순위를 기준으로 계산된 DCG
• NDCG: DCG / IDCG (0~1; 1은 완벽한 순위임)

이렇게 NDCG는 DCG 값을 이상적인 IDCG로 나눈 값이다. 예시를 들어 보자.

이상적인 검색 결과(주황색)과 실제 검색 결과 (파랑색)이 다음과 같다고 가정하자.

DCG와 IDCG를 위 식을 활용해 계산하면다음과 같고,

NDCG = DCG / IDCG이므로 결과는 다음과 같이 0.9761이 나오게 된다.

실제로 C, B의 위치만 바뀌었고 그 둘의 실제 관련성 점수 (label)이 크게 차이나지 않아 높은 값이 나오게 된 것이다.

6–2. LambdaLoss

이러한 NDCG metric에 어떻게 최적화하여 학습할 수 있을까? 이 논문에서는 이를 수학적으로 증명하는데, 함께 따라가 보자.

먼저 변수에 대한 정의가 필요하다.

• π: 특정 리스트의 순위 (조합)
• y: 실제 관련성 (relevance) label
• s: 모델이 예측한 점수

앞선 방법들 (Pairwise, RankNet, ListNet)에서는 ranking model에서 점수 s가 주어지면 단 하나의 순위 리스트가 결정되는 것으로 가정했다. 그러나 같은 점수 s를 가진 문서들이 여러 가지 가능한 순위 리스트 π를 가질 수 있다. 즉, 우리는 π를 예측한 점수 s에 대한 확률적인 형태 P(π|s) 로 표현해야 한다.

이를 활용하여, 우리의 주 목적인 “s가 주어진 경우 실제 관련성 y가 나올 확률 P(y|s)”를 다음과 같이 재정의할 수 있다.

그리고 Negative Log Likelihood (NLL)을 따라 loss를 다음과 같이 표현할 수 있다.

이제 우리의 목표는 저 loss를 최소화시키는 것이다. 이를 위해 구성요소인 P(y|s, π)와 P(π|s)를 살펴보자.

P(y|s, π)

π를 고려하지 않고, 모델이 document i,j 에 대해 예측한 점수 s_i, s_j에 대해 s_i, s_j에 의거하여 y_i > y_j가 나올 확률 P(y|s) = P(y_i>y_j | s_i, s_j )는 다음과 같다.

이를 활용하면, Loss 식을 다음과 같이 재정의할 수 있다.

그러나 우리는 P(y|s, π)라는 식에 맞게, 순위 정보 π가 y에 주는 영향도 함께 고려해야 한다. 따라서, document i, j의 순서를 바꿨을 때 생기는 NDCG 점수의 차이값 ΔNDCG(i, j)를 곱해준다. 그럼 Loss 식이 다음과 같이 재정의된다.

이러한 ΔNDCG(i, j)는 G (Gain)와 D (Discount)에 대한 항으로 표현이 가능하다. G, D는 다음과 같다.

이를 기반으로 ΔNDCG(i, j)를 예쁘게 정리한 식은 다음과 같다.

이를 위 Loss 식에 넣어주면, P(y|s, π)를 고려한 Loss는 다음과 같다.

P(π|s)

이제 뒤의 항을 정의해야한다. P(π|s)는 스코어 s를 기반으로 rank π가 결정될 확률을 뜻한다. 본 논문에서는 π에 가우시안 분포를 적용하여 확률 분포를 뽑는 SoftRank 방식이나, ListMLE에서 사용한 Plackett-Luce (PL) 모델을 먼저 소개하지만, 실제로는 Hard-Assignment 방식을 채택한다.

Hard-Assignment 방식은 굉장히 간단하다. 그냥 점수 s를 기준으로 내림차순으로 정렬된 순위 π∗ 만을 선택하자는 것이다.

그래서 실제로 s에 대한 π의 확률은 한 가지 경우에만 1인 함수가 된다. 그럼 이제 π가 π*인 경우, 최종 loss 식은 다음과 같다.

결국에는 그냥 RankNetLoss에 ΔNDCG 라는 weight만 준 loss가 되었다.

5–3. 예시

이해를 위해 예시를 들어보자. 다음과 같이 이상적인 검색 결과를 파랑색, 실제 검색 결과를 주황색이라고 하자.

이상적인 검색 결과(파랑색)의 NDCG 스코어는 1이고, 실제 검색 결과 (주황색)의 NDCG 스코어는 다음과 같다.

그럼 먼저 A, B를 보자. ΔNDCG(A, B) = 1 - 0.92 = 0.08 이다. 즉, A, B가 서로 거꾸로 바뀌어서 정렬되면 NDCG 점수가 0.08점 감소한다는 뜻이다.

이제 이를 바탕으로 loss를 계산하자.

각각의 항을 식에 넣어, loss가 0.08이 나왔다. 이제 update를 위해 gradient를 계산하자. S_A에 대한 gradient를 계산하면 다음과 같다.

반대로, S_B에 대한 gradient를 계산하면, 부호만 반대로 나오기 때문에 +0.04가 된다.

이를 바탕으로 score를 업데이트하면, 다음과 같다. (lr=0.1라고 가정)

이제 문서의 순위는 바뀌었다 !!

이러한 방식으로, LambdaLoss는 RankNet에 실제 metric의 delta 값을 weight로 주어 실제 metric에 맞게끔 학습시키는 전략이다.

7. RankLLM

이렇게 Pointwise Ranking부터 Pairwise, Listwise를 거쳐 LambdaLoss까지 알아보게 되었다. 그런데 이 학습 방식 (loss)들의 공통점은 Rank가 label된 데이터가 있어야 한다는 것이다. 앞서 소개한 논문들은 MS 또는 Google 이 냈기 때문에, 이 회사들은 자체 랭킹 데이터 (e.g. 클릭 수)를 가지고 있다. 그렇다면 일반적인 자연어를 다루는 우리는 어떻게 이런 효과적인 Ranking label을 만들 수 있을까 ? 이에 대해 알아보자. (참고 논문: Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents)

RankGPT 혹은 RankLLM으로도 잘 알려진 이 방법의 아이디어는 ChatGPT로부터 주어진 query에 대한 문서들의 relevance 순위를 추출하고, 해당 순위를 encoder, decoder 모델에 distillation 시키자는 것이다.

그럼 먼저 ChatGPT로부터 문서들의 순위를 어떻게 추출하는지 알아보자.

방법은 정말 간단하게도 위와 같다. query와 passage를 구분자 [1], [2],.. 등으로 주고, 그 구분자에 부등호 표시를 통해 대소비교를 진행한 string을 출력하도록 한다 (그림의 노란색 부분을 출력하도록 함.)

이후, 논문에서는 이 순위를 RankNet을 사용해서 Reranker, LLM에 distillation 시킨다.

이렇게 학습시킨 모델들의 ranking 결과는 다음과 같다.

놀랍게도 distillation을 받은 monoT5 (요즘은 잘 안쓰는) 정도의 모델도 gpt-4 ranking 급의 성능을 보이는 것을 볼 수 있다. 따라서 만약 자연어 형태의 랭킹이 labeling된 데이터가 필요하다면 이런 식으로 랭킹을 생성해볼 수 있을 것 같다.

이렇게 오늘은 Reranker들이 학습하는 다양한 loss 함수와, 이를 학습시키기 위해 자연어 데이터로부터 rank를 생성하는 방법에 대해 알아보았다. 다행히 최근 반영된 sentence-tranformers의 reranker training에서는 rank가 없는 데이터 (즉, 단순히 [0,1]의 label로 유관, 무관이 구분되어 있는 데이터)로 충분히 학습 가능하다. 작은 모델로 실험해본 결과, LambdaLoss > PListMLELoss > ListNetLoss > RankNetLoss > ListMLELoss 의 순서로 성능이 좋게 학습된다고 하니, 이를 활용해보면 좋을 것 같다.

그런데 최근 들어 Hybrid Retrieval을 많이 사용하면서, Sparse Retrieval의 필요성에 대해서 알게 되었다. AutoRAG 블로그에서 보듯이, Dense Retrieval와 Sparse Retrieval이 할 수 있는 역할이 다르다는 것을 많이 느꼈다.

그러던 중 Sigrid Jin님의 글을 접했고, 현대의 Sparse Retrieval에는 어떤 방법론들이 있는지 공부해 보고 싶어졌다. 그리하여 오늘은 그간 공부한 Sparse Retrieval의 역사에 대해서 정리하는 글을 작성해본다. (각 방법들의 key idea들 위주로 작성했으니, 자세히 보고 싶은 부분들은 페이퍼를 꼭 보시길 추천드립니다.)

참고한 논문들은 다음과 같다.

1. TF-IDF, BM25: https://dl.acm.org/doi/abs/10.1561/1500000019
2. Doc2query: https://arxiv.org/abs/1904.08375
3. DeepCT: https://arxiv.org/abs/1910.10687
4. SparTerm: https://arxiv.org/abs/2010.00768
5. SPLADE: https://arxiv.org/abs/2107.05720
6. LENS: https://arxiv.org/abs/2501.09749

1. TF-IDF, BM25, Indexing

사실 Sparse Retrieval에 대해 공부하기 전까지는 항상 BM25만 사용하는 줄 알았다. 그만큼 몇십년동안 강한 베이스라인을 유지해온 IR 방법론이다. BM25에 대해서 알기 전에 먼저 TF-IDF에 대한 이해가 선행되어야 한다.

1–1. TF-IDF

TF-IDF (Term Frequency-Inverse Document Frequency)는 단어의 빈도 (TF)와 역 문서 빈도 (IDF)를 사용하여 단어들의 중요도를 부여하는 방법이다. TF-IDF의 요소들은 다음과 같이 요약할 수 있다.

• TF(t, d) = (문서 d에서 단어 t가 등장한 횟수) / (문서 d에 등장한 모든 단어의 수)
• DF(t) = 특정 단어 t가 등장한 문서의 수
• IDF(t, D) = ln(총 문서의 개수 / 1 + DF(t)) = ln(D/(1+DF(t))
• TF-IDF(t, d, D) = TF(t, d) * IDF(t, d)

즉, 특정 문서에는 많이 등장하지만, 전체 문서에는 많이 등장하지 않는 단어가 높은 TF-IDF 값을 갖는다. 예로 들면 “오스트랄로피테쿠스”는 높은 값, “이, 그, 저, 어” 등의 지시 대명사는 낮은 값을 갖게 된다고 예측해볼 수 있다. (전체 문서에 많이 등장할 확률이 높으므로)

1–2. BM25

BM25는 이러한 TF-IDF를 조금 개선한 방법이다. 위의 TF-IDF를 보면, 문서가 길어질수록 단어 빈도가 높아질 가능성이 있기 때문에, 긴 문서에서 TF-IDF 점수가 무한히 크게 늘어나는 경향이 있다. 따라서 BM25는 문서 길이와 평균 문서 길이를 반영하여 점수를 조정하고, 이를 통해 긴 문서가 과도하게 높은 점수를 받지 않도록 한다.

또, TF-IDF 사용 시, 단어가 문서 내에서 반복적으로 등장하면 그 점수가 선형적으로 증가하는 특징이 있다. (e.g. 아파트 아파트 아파트 아파트..) 따라서 BM25는 단어가 반복적으로 등장하더라도 일정 수준 이상에서는 그 중요도가 더 이상 크게 증가하지 않도록 조정한다. (일종의 smoothing 효과)

이러한 두 특징을 반영한 BM25의 공식은 다음과 같다:

위의 식에서, 좌변이 BM25에서 정의하는 IDF, 우변이 BM25에서 정의하는 TF이다.

• 먼저 우변을 보면, f(q_i, D)가 기존에 정의하던 TF(= t / d)인데, 그 뒤에 k1이라는 하이퍼파라미터가 생긴 것을 볼 수 있다. 이 k1은 “단어가 한 번 더 등장했을 때 얼마나 추가적인 점수를 부여할지"를 결정한다.
  TF-IDF에서의 TF = f(q_i,D)로, 여기서 f(q_i, D)가 높아질수록 점수는 선형적으로 증가한다. 그러나 BM25 우변에서는 f(q_i, D) 값에 크게 상관없이, k1이 낮으면, 단어가 문서 내에서 반복적으로 등장하더라도 그 값을 조정하는 역할을 한다.
• 실제로 f(q_i, D)=10 일때와 1000일때를 보면, TF-IDF에서는 100배 차이가 난다. 반면, BM25에서 k1이 0.5라고 가정했을 때, 우변 식의 값은 (b 관련 항 무시)
  -> 10*(1.5) / 10+0.5 = 1.43
  -> 1000*(1.5) / 1000 + 0.5 = 1.5 정도로, 1.X배 차이이다.
  따라서, 우변의 식은 언어적 반복을 막는 역할을 한다.
• 또 우변의 분모에서 (1-b+b*(|D|/avgdl)) 을 볼 수 있는데, 이 부분이 문서 길이와 평균 문서 길이(avgdl)을 반영해서 점수를 조정해주는 부분이다. D가 avgdl보다 크면, (-b)보다 b 앞의 계수가 커지므로 분모가 더 커져, 전체적인 BM25 점수는 작아지는 구조이다. 단순히 문서가 길다는 이유로 높은 점수를 받는 것을 방지하는 것이다.
• IDF (좌변)은 다음과 같이 정의된다.

• N은 전체 문서 수, n(q_i)는 단어 q_i가 나타나는 문서 수, q_i는 query의 i번째 단어를 뜻한다.
• BM25에서는 TF-IDF에서 정의했던 IDF와는 다르게, 자주 등장하는 단어는 덜 중요하게, 드물게 등장하는 단어는 더 중요하게 가중치를 부여한다.
• n(q_i)가 작을수록 (특정 단어가 적은 문서에서 등장할수록), 분모는 작아지고 분자는 커지므로 IDF 값은 커진다. 반대로 n(q_i)가 클수록 분모는 커지고 분자가 작아지므로 IDF 값이 낮아진다.
• (분모와 분자에 0.5를 더해주는 이유는 n(q_i)가 0일수도 있기 때문에 !)

1–3. Indexing

BM25는 “문서의 관련성을 평가하는 랭킹 함수, 스코어링 방식"이다. 보통은 검색 전에 “색인, Indexing” 작업이 이루어지며,이 과정에서 모든 문서를 토큰화하고 통계를 수집하여 역색인 (Inverted Index)을 생성한다.

역색인 (Inverted Index)이란, 검색을 빠르게, 효율적으로 하기 위해 사용되는 데이터 구조로, 특정 keyword가 어떤 문서에 포함되어 있는지를 빠르게 찾아낼 수 있도록 설계된 방식이다. 일반 색인에서는 doc_id를 key로, 그에 매핑된 단어들과 TF를 value로 저장해 둔다면, inverted index에서는 단어를 key로, 그 단어가 포함된 doc_id와 TF를 value로 저장해둔다.

예로 들어, 다음과 같은 문서들이 있는 경우,

1. Forward Index (일반 색인)에서는 다음과 같이 doc_id를 key로,

• doc1: {“apple”: 1, “banana”:1, “cherry”:1}
• doc2: {“banana”:1, “cherry”:1, “date”: 1}
• doc3: {“cherry”:1, “date”: 1, “apple”:1}

2. Inverted Index (역색인)에서는

• apple: {“doc1”: 1, “doc3”: 1}
• banana: {“doc1”:1 , “doc2”: 1}
• cherry: {“doc1”:1 , “doc2”: 1, “doc3”:1}
• date: {“doc2”: 1, “doc3”:1}

으로 색인하여, “banana”를 검색했을 때 어느 문서에 해당 단어가 몇 번 포함되어 있는지 바로 알 수 있도록 한다.

이렇게 Inverted Index가 만들어진 후, 검색 시점에 query가 들어와 tokenizer로 tokenizing되고, query에 있는 token과 Inverted Index에 있는 각 document와의 BM25 점수를 계산하여, 해당 점수를 기반으로 top scored document를 검색해 온다.

2. Doc2query

그러나 BM25에는 치명적인 단점이 있다. 바로 “Vocab Mismatch problem”이다. 말 그대로, BM25는 “스파게티”로 검색했을 때 “파스타”를 가져오지 못한다.

이러한 문제를 해결하기 위해, Sparse Retrieval에서는 많은 Expansion 방법론들이 제기되었다. 그 중 Doc2query는 가장 naive하게 document를 확장한 방법이다. (2019년도 페이퍼)

2–1. Doc2query Framework

여기서 제기된 방법은, 제목 그대로 생성형 모델을 사용해서 document를 기반으로 만들어질 수 있는 여러 query들을 생성하는 것이다.

이 논문에서는 passage-query pair가 주어져, Transformer (from scratch) 모델에 passage를 입력하여 query를 생성하도록 한다. 2019년 당시에는 괜찮은 생성형 모델이 없었기에, Transformer 모델을 사용했다고 한다.

모델이 학습된 후, document를 입력하여 top-k random sampling으로 10개의 query를 생성하도록 한 뒤, 기존 document corpus에 더한다. Document와 query를 구분하도록 하는 토큰도 하나 없이, 정말 무식하게 string끼리 더하는 방법이다. 이후 확장된 새로운 document가 indexing되고, BM25로 retrieve 하여 평가한다.

2–2. Evaluation

당시에는 First-stage retrieval로 BM25, Reranking으로 BERT를 사용했다. Baseline에서 위쪽 부분은 Reranker, 중간 부분은 First-stage retrieval, 아래 부분은 First-stage retrieval + Reranking이다. (맨 위 2개는 당시 최고의 reranker들, RM3는 query expansion 기법이다 — 말 그대로 query 기반으로 후보 query들을 만들어서 검색하는 방식.)

MS MARCO와 TREC-CAR 데이터셋으로 실험했을 때, doc2query가 First-stage retrieval에서 당시 SOTA를 달성하게 된다.

2–3. Qualitative Analysis

질적 분석을 보면, 생성형 모델이 input document에 있는 단어들을 많이 차용해서 query를 생성하기 때문에 term이 re-weighting (해당 단어가 더 강조)되고, document에 나오지 않는 단어들을 생성하기도 하여, 정말 이전에 문제되었던 vocab mismatch problem을 해결하기도 한다.
(실제로 새로 생성된 단어:새로운 단어 의 비율은 31:69 라고 한다.)

3. DeepCT

이러한 방법에도 불구하고 BM25와 Doc2query는 문맥적 의미에 대한 반영이 부족하다는 문제점을 지적받았다. 즉, Term importance에 대해 좀 더 신경쓸 필요가 있다는 것이다.

예로 들어, 아래 예시에서

“stomach”라는 단어가 위아래 passage에 똑같이 2번씩 나오는데, 위 passage는 정말 “stomach”에 관한 글이고, 아래 passage는 몸의 구성에 대해 설명하며 stomach를 부가적으로 설명한다.

DeepCT에서는 많은 query들이 document의 key idea에 기반하여 나온다고 가정하며, 어떤 term이 문장의 central한 의미를 반영하는지 아는 것이 중요하다고 강조한다. 중심적인 의미를 반영하는 DeepCT의 목적은 2가지이다.

1. Passage retrieval을 위한 중요한 passage term 찾기 (DeepCT)
2. query에서 중요한 의미를 갖는 term 찾기 (DeepCT-Index)

3–1. DeepCT Framework

DeepCT에서는 단어의 문맥적 의미를 추출해서 term weight를 계산하기 위해 BERT를 사용한다. 간단히 말해서 문맥으로부터 각 단어의 importance score를 뽑는건데, 이는 다음 식을 통해 이루어진다.

T_{t, c}는 BERT로 구한 text c 안에서의 term t의 embedding, w&b는 linear combination weights&biases 이다. 따라서 y^_{t, c}는 predicted word importance score가 된다. 이 예측된 값을 각 token의 GT term weight와의 MSE Loss를 학습한다.

그렇다면 passage에서 특정 token의 GT importance score (GT term weight)는 어떻게 정의할까 ? 저자들은 Query Term Recall (QTR)이라는 방법을 이전 연구로부터 차용한다. 식은 다음과 같다.

이 식은 “passage d와 유관한 queries (여러 개) 중 term t를 포함하고 있는 query의 비율”을 뜻한다. 이 QTR은 우리가 주로 어떤 문서의 핵심 개념에 대한 질의를 하듯이, queries가 document의 key idea를 잘 반영할 수 있다는 가정에서 시작한다. Document와 관련된 query에서 나타나는 단어들은 다른 document에 있는 단어들보다 더 중요하다는 것이다.

이제 GT를 정의했으니, predicted term weights와 GT term weights의 MSE Loss로 학습을 시작한다.

학습이 완료되면, 단일 passage만으로 model inference가 가능해지고, passage를 입력했을 때 passage 내 각 term들의 weight가 출력된다.

검색 전에는 모든 passages에 대해 DeepCT 모델로 inference 한 후, predicted weights를 scaling해서 integer range로 만들어둔다.

이렇게 Integer range로 변환된 TF 값을 TF_DeepCT(t, d)라고 하고, 검색 전에 이 값을 inverted index로 만들어 둔다. (역)색인될 때, 각 term의 특정 document에서의 importance score가 색인되는 것이다.

E.g.)

• 기존의 inverted index: “DNA”: {“doc_1”: 3, “doc_5” : 2}
• DeepCT inverted index: “DNA”: {“doc_1”: 9, “doc_5”: 7}
  (이렇게 importance score가 색인됨)

3–2. DeepCT-Query

여기에서 더 나아가, 저자들은 query로부터 중요한 term을 포착할 수 있어야 한다고 주장한다. 이를 위해 이전 방법과 비슷하게 query로부터 term weight를 추출하고자 한다.

Query의 GT importance score는 Term Recall (TR)로 정한다. 식은 다음과 같다:

이전과 비슷하게, D_q는 query와 유관한 documents (여러개). D_{q,t}는 유관한 documents 중 term t를 포함하는 documents이다. Query term이 더 많은 유관한 문서에서 등장한다면 더 중요하다는 가정에서 시작하였다.

이전과 비슷하게 query-documents pair를 준비하여 predicted 값과 GT 의 MLE Loss를 학습시키고, inference 시에 query가 입력되면 query의 각 token의 weights를 계산하도록 한다.

검색 시 query가 입력되면, 모델을 이용해 query의 각 term 들의 importance score (weights)를 뽑고, 해당 weights를 사용해 BOW queries와 SDM queries (bigrams와 window 내의 co-occurrences를 반영한)를 생성하도록 한다.

예로 들어, query가 “apple pie”라고 했을 때,

• term weights: (0.8 apple) (0.7 pie)
• bigram 기반 확장 = “apple”, “pie”가 연속해서 등장하는 경우
  -> (0.56 apple pie)
• window 내에서 공존 = “apple”, “pie”가 특정 window 거리에서 함께 등장하는 경우

따라서 최종적으로 “apple pie”는 DeepCT model, SDM을 거쳐
# weight (0.8 apple 0.7 pie 0.5 #1(apple pie) 0.3 #uw8(apple pie)) 과 같은 weighted query로 변환된다. (윈도우 크기가 8인 경우의 예시) 그리고 term weight가 음수면 삭제된다.

따라서 우리는 그냥 term의 weight만 뽑은 BoW-DeepCT-Query와 SDM까지 거친 SDM-DeepCT-Query를 가지게 된다.

3–3. Experiment Setup

학습 및 평가 dataset은 다음과 같다.

• MS-MARCO: 8.8M
  - query 하나당 relevant passage 1개
  - test set이 MS에 의해 비공개라 dev set이 main eval set
• TREC-CAR: 29.7M
  - query: concatenation of Wiki title + section 1개
  - 평균 2.5개의 relevant passages

비교 Baseline은 다음과 같다.

• TF index (그냥 기본 TF-based index)
• TextRank: unsupervised graph-based term weighting approach
• Doc2query (앞에서 설명한 방법)
• DeepCT-Index
• DeepCT_W-Index: DeepCT에서 BERT를 context-independent word embeddings로 바꾼 버전
  이 당시에는 word2vec 모델로 각 word embedding에서 passage embedding을 빼서 “문맥에 독립적인 단어 임베딩”을 구하는 것이 입증된 방안이었음.
• DeepCT_E-Index: DeepCT에서 BERT를 ELMo로 바꾼 버전

Indexing과 Retrieval 방법론은 다음과 같다.

• BM25, QL로 first-stage ranking
• Conv-KNRM, BERT reranker로 reranking..도 하는데 너무 길기에 first-stage retrieval만 다루도록 하겠다.

평가 방식은 다음과 같다:

• MRR@10 — MSMARCO
• MAP@1000 — TREC-CAR

3–4. Results

위의 setup으로 결과를 확인해 보면,

1. DeepCT-Index가 가장 좋았다. — Doc2query와 같은 데이터로 학습되었지만, SOTA였던 Doc2query를 넘었다.
2. DeepCT_W-Index, DeepCT_E-Index 모두 BERT 기반보다 성능이 낮다 — BERT model의 중요성을 강조한다.

또, DeepCT-Query의 유효성에 대한 결과를 확인해 보면,

평균적으로 DeepCT-Query를 사용하고, SDM으로 단어들을 확장시켰을 때 가장 좋은 성능을 기록하는 것을 볼 수 있다.

3–5. Qualitative Analysis

본 논문에서는 Query와 유관한 document, 무관한 document에 대해서 정성적 분석을 진행한다.

사진에서 보듯이, “susan boyle”에 대해서 물어보는 query에 대해서, 유관한 document (실제로 Susan Boyle에 대해서 소개하는 document)에는 해당 term들의 weight가 높게, 무관한 document (예시로 Susan Boyle이 등장하는 document)에서는 해당 term들의 weight가 덜 높게 측정된다.

또, 각 평가 데이터셋에서 weight of each passage’s highest-weighted term을 가지고 평균내서 weight distribution을 보여준다. (top 10개 뽑아서 해당 weight로 그래프 그림)

위 그림에서는 분포가 Flat한 original tf index에 반해, DeepCT-Index는 정말 단어들의 중요성이 잘 강조되고 있는 것을 볼 수 있다. 따라서 DeepCT-Index가 적은 중심 단어들은 강조하고, 다른 단어들은 억제하는 것을 볼 수 있다.

저자들도 얘기했지만, 이 논문의 단점은 “Query가 항상 document의 central idea로부터 온다는 것을 가정하고 시작한다"는 점이다. 만약 수능의 킬러 문항처럼, 핵심 개념이 아닌 지엽적인 곳으로부터 찾아야 하는 query가 들어온다면, DeepCT의 방법은 오히려 역효과를 내게 될지도 모른다.

4. SparTerm

DeepCT의 방법은 기존의 BM25의 한계였던 “Vocab mismatch” 문제를 해소하지 못한다. 따라서, SparTerm은 semantic level matching을 높이기 위한 방안을 고안한다.

4–1. SparTerm Architecture

SparTerm에서는 term 확장을 위해 다음과 같은 아키텍쳐를 사용한다.

SparTerm architecture는 Importance Predictor와 Gating Controller라는 2가지 모듈로 구성되어 있으며, 최종 representation은 두 모듈의 곱으로 산출된다.

• F는 dense vector를 생성하는 importance predictor,
• G는 sparse vector를 생성하는 binary gating vector이다.

각각에 대해 알아보도록 하자.

4–1–1. Importance Predictor

Importance predictor는 input passage p가 주어진 경우, p의 각 term에 대해 vocab들의 semantic importance를 출력한다. 위 그림에서 p의 token별로 PLM으로부터 hidden state를 뽑고, 해당 hidden state를 token-wise importance predictor에 입력하여 token-wise importance distribution을 얻는다. i번째 token의, 모든 vocab에 대해 얻는 dense importance distribution은 다음과 같이 정의할 수 있다.

• h_i는 embedding vector
• Transform 함수는 GELU activation과 layer normalization을 거친 선형 변환
• E는 shared word embedding matrix, b는 bias term이다.

우리는 이 식을 BERT의 MLM pretraining에서 본 적이 있다. 위의 식에서 h_i는 i번째 token이 BERT에 입력되어 마지막 layer를 통과해서 나오는 hidden state으로 볼 수 있다.

이후 통과하는 Token-wise Importance Predictor (그림 참고)는 BERT의 MLM head와 같은 역할을 한다. BERT의 MLM head를 통과하면 해당 token의 모든 vocab에 대한 distribution을 얻게 된다.

따라서, 이 부분은 BERT로 initialize할 수 있다. 이렇게 각 토큰에 대해, BERT vocab의 importance distribution을 계산한 후, ReLU를 통과시켜 합계를 구한다.

그럼 이 부분을 요약하면, 모든 문맥을 반영한 vocab distribution으로 설명할 수 있다.

4–1–2. Gating Controller

Gating Controller는 passage를 표현하기 위해서 어떤 term을 activate 시켜야하는지 결정하는 binary gating signal을 출력한다. 당연히, 원본 passage에 있는 단어들은 모두 activate 시켜야 한다. 그리고 우리가 앞서 제기한 “Lexical mismatch” 문제를 해결하기 위해 passage topic과 관련한 term들도 activate하면 좋다.

따라서, SparTerm에서는 다음과 같은 2가지 gating controller 제안한다.

1. literal-only gating
2. expansion-enhanced gating

• Literal-only gating은 앞서 말했듯이, 원본 문장에 나타나는 것만 activate 시키는 것이기 때문에, BoW activation은 다음과 같다.

• 키포인트는 Expansion-enhanced gating이다. 여기서는 다른 term들을 포함함으로써 lexical mismatch gap을 줄일 것이라고 기대한다. Expansion-enhanced gating이 이루어지는 과정은 다음과 같다.

1. 앞의 importance predictor와 같이, passage-wise dense term gating distribution G를 얻는다. (G는 passage 기반 전체 vocab에서 각 token이 나타날 logit 값을 뜻한다.)
2. G’ = Binarizer (G)​
   G를 G의 모든 vocab 중에, 확률이 k 이상일 경우 1 아니면 0 (여기서 k = 0.7)인 Binarizer에 통과시킨다. 이렇게 되면 문맥을 고려한, 상대적으로 높은 확률을 가지는 token들만 남는다.
3. G_e = G’ ⊙ ¬BoW(p)​
   - BoW(p): vocab 중에 passage에 있는 어휘에 1, 없는 어휘에 0으로 되어 있는 vector
   - ¬BoW(p): BoW(p)의 여집합
   - 따라서, G_e는 Binarizer를 통과한 term이 BoW(p)의 여집합에 포함된 경우(p 안에 없는 어휘인데 Binarizer를 통과한 경우)를 뜻한다.
   <그냥 높은 logit 값을 가지는, 원본 passage에 없는 token만 남기는 과정이다>
4. Final expansion-enhanced gating vector G_{le} = G_e + BoW(p)​
   - 확률 높은 원본 passage에 없는 token과 원본 passage들에만 있는 token들의 sparse vector를 결합하여 하나의 sparse vector로 만드는 과정이다.

4–2. Training

SparTerm 프레임워크에서는 Importance Predictor와 Gating Controller를 각각 따로 학습시켜야 한다.

먼저 Importance Predictor는 다음 loss function을 기반으로 학습한다.

Dense retrieval에서 많이 보던, Triplet loss 형식이다. Term distribution을 올바르게 하는데, 왜 이렇게 학습을 하나 싶을 수 있겠지만, 본 논문에서는 passage ranking에 사용되는 distant supervisory signal로 importance를 학습해도 효과가 좋다고 한다.

Gating Controller는 우선 term expansion이 가능한 4가지 시나리오를 제시한다.

직관적으로 생각해 보았을 때, MLM pretrainig으로 term expansion이 가능한 시나리오는 Synonym (유의어)과 Co-occurred words (자주 등장하는 단어)이다. 따라서, 본 논문에서는 이 시나리오는 BERT에서 충분히 학습이 되었을 것이라고 가정하고 passage2query와 summarization expansion에 집중한다. passage2query는 passage에서 나올 수 있는 query들을 예측, summarization은 passage를 summarize할 때 나올 수 있는 단어들을 예측하는 시나리오이다.

그리하여, 논문에서는 passage-query, passage-summary 데이터만 가지고 학습한다. Loss는 2개의 cross-entropy loss를 사용한다.

• p: passage, t: target text (query / summary)
• T: binary BoW vectors of t
• G: dense gating probability distribution for p

1. G_j와 관련된 첫 번째 항을 먼저 보자.
   G_j는 target text t 에 <등장하지 않은> 단어들 중 (m|T_m=0), passage로 구한 probability distribution을 뜻한다. 즉 “Target text에 등장하지 않는 단어들의 logit 분포" 인 것이다.
   Loss가 작아지기 위해 (1-G_j)가 커져야 하고, G_j는 작아져야 하므로, Target text에 등장하지 않는 단어들의 logit은 작아지도록 학습된다. 데이터 자체가 Expansion을 위한 데이터이므로, 준비된 데이터 외의 logit은 모두 줄이는 것이다.
2. G_k와 관련된 두 번째 항을 보자.
   G_k는 target text t에 <등장하는> 단어들 중, passage로 구한 probability distribution이다. 이 logit은 앞 상황과 반대로 커지도록 학습한다. (마찬가지로 준비된 데이터가 Expansion을 위한 데이터이므로)

이렇게 Importance Predictor와 Gating Collector를 위한 loss 식을 더해주면 최종 loss가 된다.

4–3. Experimental Setup, Training Details

• Dataset: MSMARCO Passage retrieval dataset, MSMARCO Document retrieval dataset
• Metric: MRR@10, Recall@10, Recall@100
• Training Details
  - Importance Predictor, Gating Controller 모두 BERT-base로 initialize, weight는 공유하지 않는다.
  - 단계별로 학습을 진행하는데, 먼저 Gating Controller를 Fine-tuning한다. 이때 MSMARCO passage retrieval train dataset으로 50k step 동안 학습한다.
  - 이후, Gating Controller의 parameter를 고정해두고 Triplet Loss 형태의 loss로 SparTerm을 fine-tune 한다.
• 검색 전, 각 document를 SparTerm에 통과시킨 sparse representation으로 inverted index를 만들고, 검색 시에는 들어오는 query를 SparTerm에 통과시켜 확장한 후, sparse representation으로 topK 검색하도록 한다.
  즉 vocab 크기의 query의 sparse representation과 vocab 크기의 document의 sparse representation의 matmul을 구하는 것이다.
• Baselines
  - BM25
  - DeepCT
  - Doc2query
  - Doc2query-T5: T5 기반으로 Doc2query 방법 그대로 한 것. SOTA였음
  - SparTerm (literal-only): passage에 제공된 단어만으로 gating controller 만든 것
  - SparTerm (expansion-only): passage에 제공되지 않은 단어만으로 gating controller를 만든 것
  - SparTerm (expansion-enhanced): passage에 제공된 단어와 제공되지 않은 단어를 모두 활용하여 gating controller 만든 것

4–4. Results

• SparTerm (expansion-enhanced)이 MRR에서는 가장 좋았고, Recall 에서 Doc2query-T5를 제외한 베이스라인들 중에서 가장 좋았다.
• T5가 Doc2query(Transformers)보다 훨씬 높은 것으로 보아, 저자들은 SparTerm도 BERT보다 괜찮은 PLM으로 바꾸면 더 좋을 것이라는 예측을 한다.
• Expansion 없이, SparTerm(literal-only)가 DeepCT outperform한다.

(더 많은 실험이 있지만 생략..)

4–5. Qualitative Analysis

DeepCT와 SparTerm모두 query가 입력되면 passage에 있는 단어들 중 가장 중요한 단어에 weight를 높게 매기는 것을 볼 수 있다. 다만 DeepCT는 적은 단어에 weight를 크게 주는 반면, SparTerm은 조금 더 많은 단어를 포착할 수 있다. 예를 들어, “hives”에 대해 물어보는 질문에서, SparTerm은 “allergic reactions”에 weight를 주지만, DeepCT는 그러지 못하는 것을 알 수 있다.

또 SparTerm에서는 “sign”으로 “symptom”이라는 단어를 예측하고, “pregnancy”로 “women”이라는 단어를 예측하는 것을 보아, 확실히 term exapnsion을 거친게 효과가 있음을 보여준다.

이 분포는 그래프 아래에 있는 각 단어들을 expansion했을 때 top5로 꼽히는 단어들을 보여준다. 모두 연관있거나, 원본 단어를 기반으로 자주 나타나는 단어들이 나와, SparTerm의 우수성을 입증한다.

5. SPLADE

이후 나온 SPLADE는 정말 간단하다. 단순히 SparTerm 모델에 log activation과 sparsity regularization를 더함으로써 성능을 훨씬 높일 수 있었다.

SPLADE에서 지적하는 SparTerm의 문제는 모델이 쓸데없이 복잡하고, end-to-end로 학습될 수 없다는 것이었다. 앞서 말했던 것처럼 SparTerm에서는 Gating Controller가 1차로 학습된 후, 전체 SparTerm 프레임워크를 2차 학습해야 하는 구조이다. 게다가 SparTerm의 결과를 보면, 충분히 literal-only 만으로도 잘하는데, expansion 하는 것이 충분히 효과가 있는 것인지에 대한 의문을 가진다.

5–1. SPLADE 방법론

그래서 SPLADE에서는 SparTerm을 간단하게 개선한다. 기존의 Gating Controller를 없애고, Importance Predictor의 식에서 log만 추가한다.

이렇게 함으로써 직관적으로는 log-saturation effect가 weight가 너무 큰 단어들은 smoothing 해줄 수 있을 것이라는 효과를 생각해 볼 수 있게 된다.

Loss 식도 조금 변경한다. 기존에 query-pos-neg 3가지만 가지고 Triplet Loss로 학습시켰던 것을 query-pos에 in-batch negative로 InfoNCE Loss를 학습시킨다. (아마 SparTerm 당시에는 in-batch negative이라는 밝혀진 전략이 없었던 것 같다.)

또 SPLADE에서는 Gating Controller를 제거하였기 때문에, 벡터가 sparse해지지 못한다. 이렇게 되면 vocab size 크기의 벡터들을 계속 indexing 해 두어야 할 텐데, 이건 정말 무겁다. 이를 위해 FLOPS regularizer loss를 사용하는데, 이게 좀 간단하면서 신기하다.

식에서 보듯이, 모든 embedding 차원 (=vocab size)의 평균을 구한 뒤, 각 차원들의 제곱합을 구해서 그걸 loss로 사용한다. 이렇게 했을 때 기대하는 바는, 많은 차원들이 0이 되는 (sparse 해지는) 것이다. SPLADE에서는 이를 query 따로, document 따로 가중치를 두고 적용한다.

그리하여 최종 loss는 다음과 같이 설계된다:

이렇게 함으로써, L_{rank-IBN}을 통해 유사한 토큰들이 강조되도록, L_reg을 통해 무의미한 토큰들은 0이 되도록 하여, sparse한 벡터가 되도록 유도한다.

5–2. Training details & Experimental setup

• 학습과 평가를 위해 각각 MSMARCO passage ranking dataset의 train, dev set을 사용하고, TREC 2019 eval set 또한 평가에 사용한다.
• metric은 Recall@1000, MRR@10 and NDCG@10
• 모델은 BERT-base로 initialize되어, batch size 124로 150k step 동안 학습된다.
• 검색 전, 각 document는 SPLADE 모델을 통과하여 vocab 크기의 logit 값이 적힌 sparse vector를 가진다. 검색 시, query 또한 같은 방식으로 확장되어, query와 document의 sparse vector를 matmul하여 최종 output을 구한다.

• 베이스라인은 다음과 같다.
  - BM25
  - DeepCT
  - Doc2query-T5
  - SparTerm (lexical-only, expansion 둘다)
  - ANCE (Dense)
  - TCT-ColBERT (Dense)
  - SparTerm-lexical-only (w/ InfoNCE)
  - SparTerm-exp (SparTerm에 log saturation 주고 InfoNCE로 학습시킨거)
  - SPLADE

5–3. Results

1. TREC DL Recall@1000을 제외하면 SPLADE가 모든 데이터, 모든 metric에서 퍼포먼스가 좋다.
2. (당시) SOTA dense retriever과도 비교해도 견줄만하다.
3. ST lexical-only를 InfoNCE 방법으로 학습한게 일반 ST expansion보다 높다. (InfoNCE vs. Triplet loss의 효과 비교)

5–4. Qualitative Analysis

위의 예시를 보면, model을 통과했을 때 중요한 term에 대해서는 re-weighting을 수행하고, 불필요한 term은 없애는 (아마도 term weight ≤ 0) 것을 볼 수 있다.

또한 legs -> leg / bow legs -> treatment 등으로 보아 expansion이 실제로도 효과적으로 이루어짐을 알 수 있다.

5–5. SPLADE-v2

그런데 SPLADE에서 token importance score를 구할 때, 굳이 sequence에 있는 모든 token에서의 logit을 더할 필요가 있을까 ? 그냥 sequence에 있는 token 중 가장 특정 vocab의 logit이 크게 나오는 부분만 사용하면 되지 않을까? 라는 의문을 가지고 Pooling 전략을 바꾸어 SPLADE-v2를 만든다. 식으로 표현하면, 원래는 이랬던 식이

아래처럼 변하는 것이다.

따라서 SPLADE-v1의 vocab weight를 구하면 sequence에 있는 모든 단어들을 모델에 통과시켰을 때 내뱉는 첫 번째 vocab의 logit을 다 합치면 되지만, SPLADE-v2는 그 중 가장 큰 logit만 사용하겠다는 것이다. 이를 SPLADE-max 라고 한다.

실험 결과, SPLADE-max(SPLADE-v2)가 더 뛰어난 성능을 보인다.

6. LENS

SPLADE의 첫 등장 (2022) 이후 몇년 간 SPLADE가 발전을 거듭하며 sparse retrieval에서 부동의 SOTA 모델로 자리잡았다. 또 ANN 기반의 dense retrieval이 활성화되면서, 다양한 임베딩 모델 (E5, BGE, GTE, Jina, LLM2vec, NV-Embed)들이 retrieval task를 매우 잘 수행하는 모델이 되었다.

여기서 주목할 만한 점은 기존에 encoder-only 모델만을 검색에 사용하던 것과는 달리, decoder 기반의 모델도 검색에 사용하기 시작했다는 것이다.

LLM2vec에서는 unidirectional attention을 수행하던 LLM을, bidirectional attention을 수행하도록 변경한 뒤, Masked Next Token Prediction과 Pre-finetuning, Fine-tuning을 수행하여 dense retrieval에 최적화된 검색 모델이 되도록 학습한다. 평가에서도 우수한 성능을 입증하며 “큰 모델이 최고다" 라는 것을 보여 준다.

그럼 LLM으로 dense retrieval도 가능하니, sparse retrieval도 해볼만하지 않을까? 라며 2025년 1월에 LENS가 등장하게 된다.

6–1. LENS가 제기하는 문제점

LENS에서는 먼저 LLM을 sparse retrieval에 사용할 때 나타날 수 있는 잠재적인 문제점을 제기한다. 저자들이 제기하는 문제는 tokenizer에 있었는데, 다음과 같다.

1. 같은 어휘소 (lexeme)으로부터 발생하는 중복성과 관련된 문제와, 해당 중복성이 token matching에도 영향을 줄 수 있다는 문제
   (e.g. “What”, “what”, “ what”이 모두 다른 token으로 처리되는 문제)
2. 한 단어들이 쪼개져서 추가적인 매칭이 이루어져야 한다는 문제
   (e.g. “education” -> “edu”, “cation”)
3. 큰 문서 단위에서 학습된 tokenizer로 인해 rare한 term이 포함되어 vocab의 사이즈를 늘리는 문제

확실히 내가 생각해 보아도, LLM이 쓰는 큰 tokenizer로 sparse vector를 만들면, 차원이 너어무 커질 것 같았다. 따라서, 저자들은 tokenizer를 그대로 사용하는 것이 아닌, tokenizer에 있는 token을 기준으로 clustering 시켜, cluster centroid embeddings를 원본 token embeddings에 대체시킨다. 이렇게 비슷한 token들을 merging 하면, 중복을 줄이고 sparse embedding size를 줄일 수 있다.

6-2. LENS Framework

LM head에서 비슷한 token들끼리 cluster 하기 위해 KMeans clustering을 진행한다. (여기서의 k는 우리가 원하는 embedding size이다.) 그럼 LM head의 원래 token embeddings는 cluster centroid로 대체된다. 즉, input text가 tokenizer를 통해 변환된 후 생성되는 초기 input embedding은 수정되지 않고, 모델 내부에서 LM head의 output embedding만 clustering을 통해 간소화한다는 것이다. 따라서 이제 input은 큰 vocab size 크기의 embedding 대신에, 작은 clustered size 크기의 lexicon-based embedding이 된다.

또 LENS 프레임워크에서는 기존 unidirectional attention이었던 LLM attention을 bidirectional attention으로 바꿔보는 등, attention mechanism에 변화를 주며 실험해보기도 한다.

6–3. Representation Generation

그럼 실제로 sparse embedding은 어떻게 만들까 ? 먼저 q-p 쌍을 준비해, instruction을 수용하는 요즘 LLM-based embedding model처럼 q_ins를 만든다.

예로 들어 instruction은 다음과 같다. (NV-Embed 참조.)
"Given a question, retrieve passages that answer the question"

이후 SPLADE-v2에서처럼 log-saturation과 max-pooling을 모두 고려한 logit vector를 추출한다.

• 이때 query의 경우, 실제 query에 대한 logit만을 사용해야 한다. 그러므로 q_ins 중 <query> 부분에 해당하는 token의 logit만 뽑아서 사용한다.
• 또, next token prediction으로 학습하는 LLM이기 때문에, {보고자 하는 토큰의 index-1}에서의 logit을 추출해야 보고자 하는 토큰의 logit을 뽑을 수 있다.
• passage는 따로 instruction을 주지 않으니, 바로 BOS 토큰부터 생성되는 logit을 본다.

6–4. Training

학습 또한 SPLADE에서 한 것처럼 InfoNCE Loss로 학습한다.

q_ins와 p의 embedding (cluster 크기)를 뽑아 위 Loss로 학습시킨다.

6–5. Experiment & Training config

• backbone model: Mistral-7B-v0.1
• data: BGE-en-ICL에서 썼던 공개된 데이터- retrieval, reranking, clustering, classification, STS (1 positive, 7 hard negative)
• KMeans의 cluster 수: 4000, 8000 (LENS-4000, LENS-8000)
• Training config
  - LoRA 사용
  - 1 positive, 7 hard negative 사용
  - in-batch negative 사용 (retrieval task에 대한 batch size 512, 나머지 task에 대한 batch size 256)
  - BGE-reranker로부터 ranking score distillation. (KL-divergence. 이건 BGE-en-ICL 학습 방법 그대로 따라하려고 해서 그런 것 같음)
  - max_seq_len=512
• Evaluation: MTEB, AIR-Bench (Jina에서 미공개 데이터로 1달에 한번씩 갱신하는 평가 벤치마크)
• Baselines
  - E5-mistral-7b-instruct
  - NV-Embed-v1/v2
  - gte-Qwen2–7B-instruct
  - stella
  - SFR-Embedding-2_R
  - GritLM-7B
  - BGE-en-ICL
  - SPLADE 같은 모델들은 retrieval task 용도로만 만들어졌고, 너무 못해서 안넣음
• zero-shot scenario만 가정

6–6. Results

1. fully-public training data로 학습한 모델들 중 가장 높은 Avg 성능
2. Dense embedding과 비교해도 더 높은 성능을 보여준다
3. 4000 차원의 LENS-4000과 4096 차원의 BGE-en-ICL이 비슷한 성능 (물론 zero-shot에서만)
4. Non-fully Public Training Data를 포함해도 2등. 1등을 차지한 NV-Embed-v2보다 training pipeline이 단순하고, synthetic data를 사용하지 않기 때문에 효율성 관점에서는 더 좋다고 함.

심지어 Jina의 비공개 벤치마크 AIR-Bench에서는 LENS-8000이 NV-Embed-v2를 이기는 모습이었다.

6–7. Qualitative Analysis

같은 cluster에 속하는 term들을 보면, 정말 비슷한 단어들끼리 모여 있음을 볼 수 있다. “생성” 관점에서는 diverse한 표현을 사용하는 것이 좋을 수는 있어도, 임베딩과 같은 “표현” 관점에서는 이렇게 비슷한 표현끼리는 묶는 것이 확실히 좋은 것 같긴 하다는 생각이 들었다.

또 LENS-8000이 MSMARCO 데이터셋에 대해 뽑은 top-5 cluster를 보면, input text와의 유사도가 상당히 높은 것을 볼 수 있고, 심지어는 3번째에 “causes of hypoxia in adults”에서 “oxygen”으로 term expansion을 하는 것도 볼 수 있다.

6–8. Analysis

마지막으로 여러 가지 분석을 보여 주는데, 먼저 cluster 개수에 따른 성능 차이를 보면

cluster가 너무 크면 비슷한 단어들이어도 서로 다른 곳에 위치하기 때문에 성능도 안좋고, embedding size도 커지니까 효율성 관점에서도 좋지 않은 것을 알 수 있다.

unidirectional와 bidirectional attention, pooling 방식을 last-token, sum, max의 3가지로 나눠 놓고 보았을 때, bidirectional attention이 훨씬 좋음을 모든 task, 모든 pooling method에서 입증한다. (근데 task당 데이터가 왜 하나밖에 없는지는 의문)

또, max-pooling이 가장 좋았음 (bidirectional attention에서는 근소 우위)을 보여 주며, 기존의 SPLADE 방법론을 따르는게 여전히 좋다는 것을 입증한다.

마치며..

이렇게, 오늘은 sparse retrieval의 과거부터 현재까지의 모든 방법론들을 알아보았다. 특히 마지막, LENS가 sparse embedding도 차원을 줄일 수 있고, LLM을 사용하여 sparse retrieval을 수행할 수 있다는 것을 보여주며, 주는 시사점이 굉장히 많은 것 같다.

이렇듯 sparse retrieval도 BM25 이외의 방법이 굉장히 많은데, 앞으로 한국어 모델들도 많이 나와 sparse retrieval을 활용할 수 있으면 한국어에서도 검색이 더 잘 되지 않을까 싶다.

하루가 멀다하고 매일 새로운 모델들이 나오는 요즈음, 사용할 임베딩 모델을 고르는 것은 굉장히 피곤하다. 그래서 Open-LLM-leaderboard처럼 임베딩 모델들을 평가하여 공개하면 좋을 것 같다는 생각을 꾸준히 해왔다.

MTEB는 텍스트 임베딩 모델을 평가하기 위해 사용되는 라이브러리로, 다양한 태스크에 대해 다양한 데이터셋을 활용하여 굉장히 편리하게 평가할 수 있다. (딸깍) 이러한 측면에서 나의 니즈와 굉장히 잘 맞았고, MTEB를 유심히 살펴 보았다.

MTEB는 오픈 프로젝트이며, 누구나 새로운 데이터셋을 등록할 수 있고, 자유롭게 issue를 남겨 error fix 등을 할 수 있다. 오늘은 이와 관련된 나의 경험, 그리고 MTEB를 활용하여 한국어 검색 태스크를 수행한 코드 및 결과를 공유하고자 한다.

MTEB-ko-retrieval Leaderboard는 다음 Github repository에 공개되어 있다.
+ issue에 평가하고 싶은 모델 남겨주시면 평가를 진행해 드리니, 필요하시면 남겨주세요 !! (⭐ 스타도 많이 눌러주세요 ⭐️)

GitHub - nlpai-lab/KURE: KURE: 고려대학교에서 개발한, 한국어 검색에 특화된 임베딩 모델

1. Evaluation

먼저 MTEB를 활용하여 embedding model들을 평가하는 코드부터 살펴보자.

1–1. Evaluation code

사실 평가는 굉장히 간단하다. 다음과 같이 MTEB 라이브러리를 활용한 코드를 실행하면 자동으로 평가가 진행된다. (코드는 Github에도 공개돼 있다.)

"""Benchmarking all datasets constituting the MTEB Korean leaderboard & average scores"""
from __future__ import annotations

import os
import logging
import argparse

import torch

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

import mteb
from mteb import MTEB, get_tasks
from dotenv import load_dotenv

from mteb.encoder_interface import PromptType
from mteb.models.sentence_transformer_wrapper import SentenceTransformerWrapper
from mteb.models.instruct_wrapper import instruct_wrapper

load_dotenv() # for OPENAI

parser = argparse.ArgumentParser(description="Extract contexts")
parser.add_argument('--quantize', default=False, type=bool, help='quantize embeddings')
args = parser.parse_args()

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger("main")

TASK_LIST_CLASSIFICATION = []

TASK_LIST_CLUSTERING = []

TASK_LIST_PAIR_CLASSIFICATION = []

TASK_LIST_RERANKING = []

TASK_LIST_RETRIEVAL = [
    "Ko-StrategyQA",
    "AutoRAGRetrieval",
    "MIRACLRetrieval", # 시간이 오래 걸림 주의
    "PublicHealthQA",
    "BelebeleRetrieval",
    "MrTidyRetrieval", # 시간이 오래 걸림 주의
    "MultiLongDocRetrieval",
    "XPQARetrieval"
]

TASK_LIST_STS = []

TASK_LIST = (
    TASK_LIST_CLASSIFICATION
    + TASK_LIST_CLUSTERING
    + TASK_LIST_PAIR_CLASSIFICATION
    + TASK_LIST_RERANKING
    + TASK_LIST_RETRIEVAL
    + TASK_LIST_STS
)

model_names = [
    # my_model_directory
]
model_names = [
    "Salesforce/SFR-Embedding-2_R", # 4096
    "Alibaba-NLP/gte-Qwen2-7B-instruct", # 8192
    "BAAI/bge-multilingual-gemma2", # 8192
    "intfloat/e5-mistral-7b-instruct", # 32768
    "intfloat/multilingual-e5-large-instruct", # 512
    "openai/text-embedding-3-large", # 8191
    "Alibaba-NLP/gte-multilingual-base", 
    "intfloat/multilingual-e5-base", # 512
    "intfloat/multilingual-e5-large", # 512
    "jinaai/jina-embeddings-v3", # 8192
    "jhgan/ko-sroberta-multitask", # 128
    "BAAI/bge-m3", # 8192
    "nlpai-lab/KoE5", # 512
    "dragonkue/BGE-m3-ko", # 8192
    "Snowflake/snowflake-arctic-embed-l-v2.0" # 8192,
    "nlpai-lab/KURE-v1" # 8192
] + model_names

def evaluate_model(model_name):
    try:
        model = None
        if not os.path.exists(model_name): # hf에 등록된 모델의 경우
            if "m2v" in model_name: # model2vec 모델의 경우: 모델명에 m2v를 포함시켜주어야 model2vec 모델로 인식합니다.
                static_embedding = StaticEmbedding.from_model2vec(model_name)
                model = SentenceTransformer(modules=[static_embedding])
            else:
                if model_name == "nlpai-lab/KoE5":
                    # mE5 기반의 모델이므로, 해당 프롬프트를 추가시킵니다.
                    model_prompts = {
                        PromptType.query.value: "query: ",
                        PromptType.passage.value: "passage: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts)
                elif model_name == "BAAI/bge-multilingual-gemma2":
                    # mbge-gemma2의 경우, mteb에서 지원하지 않습니다. 따라서, instruct_wrapper를 사용합니다.
                    instruction_template = '<instruct>{instruction}\n<query>'
                    model = instruct_wrapper(
                            model_name_or_path=model_name,
                            instruction_template=instruction_template,
                            attn="cccc",
                            pooling_method="lasttoken",
                            mode="embedding",
                            torch_dtype=torch.float16,
                            normalized=True,
                    )
                elif model_name == "Snowflake/snowflake-arctic-embed-l-v2.0":
                    # mteb에서 Snowflake 모델을 지원하지 않으므로, Snowflake에서 사용하는 "query: " prefix를 임의로 추가합니다.
                    model_prompts = {
                        PromptType.query.value: "query: ",
                    }
                    model = SentenceTransformerWrapper(model=model_name, model_prompts=model_prompts)
                else:
                    # mteb에 등록된 모델의 경우, 프롬프트/prefix 등을 포함하여 평가할 수 있습니다. 등록되지 않은 경우, sentence-transformer를 사용하여 불러옵니다.
                    model = mteb.get_model(model_name)
        else: # 직접 학습한 모델의 경우
            file_name = os.path.join(model_name, "model.safetensors")
            if os.path.exists(file_name):
                if "m2v" in model_name: # model2vec의 경우: 모델명에 m2v를 포함시켜주어야 model2vec 모델로 인식합니다.
                    static_embedding = StaticEmbedding.from_model2vec(model_name)
                    model = SentenceTransformer(modules=[static_embedding])
                else:
                    model = mteb.get_model(model_name)

        if model:
            print(f"Running task: {TASK_LIST} / {model_name}")
            evaluation = MTEB(
                tasks=get_tasks(tasks=TASK_LIST, languages=["kor-Kore", "kor-Hang", "kor_Hang"])
            )

            # 48GB VRAM 기준 적합한 batch sizes
            if "multilingual-e5" in model_name or "KoE5" in model_name:
                batch_size = 512
            elif "jina" in model_name:
                batch_size = 8
            elif "bge-m3" in model_name or "KURE" in model_name:
                batch_size = 32
            elif "gemma2" in model_name:
                batch_size = 256 
            elif "Salesforce" in model_name or "Snowflake" in model_name:
                batch_size = 128
            else:
                batch_size = 64
              
            evaluation.run(
                model,
                output_folder=f"results/{model_name}",
                encode_kwargs={"batch_size": batch_size},
            )
    except Exception as ex:
        print(ex)

if __name__ == "__main__":
    for model_name in model_names:
        evaluate_model(model_name)

• 만일 새로운 모델을 평가하고 싶다면 model_names 변수에 해당 모델의 이름을, 새로운 데이터셋으로 평가하고 싶다면 TASK_LIST 변수에 해당 데이터셋의 이름을 추가하면 된다. (Retrieval task가 아니어도 된다.) + 만일 위 코드에서 model2vec 변환을 거친 모델을 평가하고 싶다면 model 이름에 “m2v”를 넣어주면 된다.
• 또, 새로운 모델을 평가하는데 prompt/prefix를 넣어줘야 한다면, evaluate_model 함수에 작성한 것처럼 customize해서 사용할 수 있다.
• languages 변수에는 어떤 언어로 평가할지 설정할 수 있다. (MTEB의 각 task의 metadata-language를 보면 어떤 언어를 지원하는지 나와 있다.)
• 또, batch size를 높여 평가 시간을 단축할 수 있다. (모델마다 다름 주의)

1–2. Why MTEB?

그렇다면 왜 굳이 MTEB를 활용할까 ? 아무래도 가장 큰 장점은 prompt/ prefix의 반영과 편리함인 것 같다.

decoder-based retriever의 경우 prompt가 추가되어 학습되었으며, encoder-based retriever의 경우 prefix가 붙어 학습된 모델들이 있다. 예를 들어, Microsoft의 multilingual-e5 시리즈는 각 query 앞에 “query: ”, document 앞에 “passage: ” 라는 prefix가 붙어서 학습되었다.

이러한 모델들 중, MTEB에서 지원하는 오픈 모델들의 경우 (대부분 대기업의 multilingual model 들), 모델을 불러서 평가할 때 자동으로 prompt, prefix를 달아준다. 그러나 일부 모델 혹은 fine-tuning된 모델들은 지원하지 않는데, 그마저도 쉽게 달 수 있게 만들어두었다. (코드의 evaluate_model 함수 참고)

또, 평가하고자 하는 task, languages도 이름만 바꿔주면 평가할 수 있으며, 모든 지표에 대한 결과를 json 파일로 만들어주기 때문에, 굉장히 편리하다.

2. Dataset Contribution

MTEB는 dataset contribution 또한 굉장히 오픈되어 있다. 6월 정도까지는 데이터셋을 일정 이상 기여하고, MTEB repository에 많이 contribute 하면 MMTEB paper의 공동 저자로 넣어줬었는데, paper가 제출되어서 그런지 지금은 해당 활동을 진행하지 않는 것 같다 ㅠㅠ. 아무튼, 이번에는 데이터셋을 contribute 해본 경험을 작성해보도록 하겠다.

최근에 가장 잘 사용하고 있었던 임베딩 모델 평가 데이터셋 - Markers 회사의 AutoRAG에서 공개한 AutoRAG-example-korean-embedding-benchmark 를 MTEB repository에 올려보고자 하였다.

이 데이터는 원래 Allganize에서 한국어 RAG evaluation을 평가하기 위해 공개했던 데이터셋인데, AutoRAG 측에서 정제를 거쳐서 embedding model benchmark로 활용하였다. 공공기업의 문서를 페이지별로 PDF parsing 했다는 점에서 real-world scenario를 굉장히 잘 반영한 데이터라고 생각했다. (다만, retrieval pool이 작다는 단점이 있긴 하다.)

2–1. BEIR format

MTEB에 dataset을 추가하기 위해서는, 먼저 데이터셋을 BEIR format 으로 바꾸어야 한다. 즉 default, corpus, query의 3가지 subset이 있어야 한다.
(데이터셋 예시를 보면 더 이해가 잘된다.)

• default에는 검색을 위한 query-id, corpus-id, score (score는 다 1)가,

• corpus에는 retrieval pool (corpus)에 해당하는 corpus-id와 text가,

• query에는 query에 해당하는 query-id와 text들이 있어야 한다.

2–2. Making an issue

이렇게 변환을 거친 후, MTEB Github repository에 issue를 남기면 된다. 내가 남긴 issue는 다음과 같다.

fix: Add Korean AutoRAGRetrieval by yjoonjang · Pull Request #1388 · embeddings-benchmark/mteb

가장 귀찮지만, MTEB에서 가장 중요하게 체크하는 것이 TaskMetadata이다. 이 항목은 해당 task의 metadata를 작성하는 란인데, 나는 다음과 같이 작성했다. Github code

# mteb/tasks/Retrieval/kor/AutoRAGRetrieval.py
from __future__ import annotations

from mteb.abstasks.TaskMetadata import TaskMetadata

from ....abstasks.AbsTaskRetrieval import AbsTaskRetrieval


class AutoRAGRetrieval(AbsTaskRetrieval):
    metadata = TaskMetadata(
        name="AutoRAGRetrieval",
        description="This dataset enables the evaluation of Korean RAG performance across various domains—finance, public sector, healthcare, legal, and commerce—by providing publicly accessible documents, questions, and answers.",
        reference="https://arxiv.org/abs/2410.20878",
        dataset={
            "path": "yjoonjang/markers_bm",
            "revision": "fd7df84ac089bbec763b1c6bb1b56e985df5cc5c",
        },
        type="Retrieval",
        prompt="Retrieve text based on user query.",
        category="s2p",
        modalities=["text"],
        eval_splits=["test"],
        eval_langs=["kor-Hang"],
        main_score="ndcg_at_10",
        date=("2024-08-03", "2024-08-03"),
        domains=["Government", "Medical", "Legal", "Social"],
        task_subtypes=["Article retrieval"],
        license="mit",
        annotations_creators="human-annotated",
        dialect=[],
        sample_creation="created",
        bibtex_citation="""@misc{kim2024autoragautomatedframeworkoptimization,
      title={AutoRAG: Automated Framework for optimization of Retrieval Augmented Generation Pipeline}, 
      author={Dongkyu Kim and Byoungwook Kim and Donggeon Han and Matouš Eibich},
      year={2024},
      eprint={2410.20878},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2410.20878}, 
}""",
    )

이렇게 작성하고 나면, 300M 크기의 작은 모델들 (intfloat/multilingual-e5-small, sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2)을 가지고 평가를 돌려보도록 시킨다. (TaskMetadata를 제대로 작성하지 않으면 평가도 안돌아간다 ㅠ)

이후, 결과가 괜찮게 나왔다면 리포트한 후, 리뷰를 요청하여 승인되면 dataset contribute가 마무리된다 !!

3. Error Resolving

마지막으로, MTEB로 다양한 모델들을 평가하면서 Error를 다룬 경험이 인상적이었어서, 이와 관련된 경험을 남기고 마무리한다.

먼저 내가 마주한 에러는 openai의 text-embedding-3-large 모델을 평가할 때, 임베딩하려는 text가 openai-text-embedding이 지원하는 최대 길이인 8191을 넘어서 발생하는 문제였다. 대부분의 오픈 모델들은 알아서 잘라 주는데, api로 제공되어서 그런지 알아서 잘라주지 못하고 다음 에러를 반환했다:

ERROR:mteb.evaluation.MTEB:Error while evaluating MIRACLRetrieval: Error code: 400 - {'error': {'message': "This model's maximum context length is 8192 tokens, however you requested 8461 tokens (8461 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.", 'type': 'invalid_request_error', 'param': None, 'code': None}}
Error code: 400 - {'error': {'message': "This model's maximum context length is 8192 tokens, however you requested 8461 tokens (8461 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.", 'type': 'invalid_request_error', 'param': None, 'code': None}}

곧바로 이를 issue로 남겼고, 작업을 시작했다.

fix: add sentence trimming to OpenAIWrapper by yjoonjang · Pull Request #1526 · embeddings-benchmark/mteb

사실 뭐 큰 문제는 아니고, 그냥 text-embedding-3-large가 사용하는 tokenizer를 불러와 sentence를 trim 하고, 다시 api로 넘겨주면 되는 문제였다. 작업을 하고 리뷰를 요청했는데, 생각보다 굉장히 빠르고 열심히 리뷰해줘서 고마웠다.

위 issue를 열고 진행하면서, 같은 문제를 겪고 있던 MTEB maintainer (Muennighoff)의 코드를 일부 반영하여 수정하기도 하고, tiktoken 때문에 발생하는 에러로 인해 openai를 위한 라이브러리 버전 관리도 해보았다.

회사를 다니지 않으면 이런 경험은 쉽게 해보지 못할 것 같은데, 그러한 점에서 작지만 굉장히 인상적인 경험이었다.

이렇게 오늘은 MTEB에 대한 간단한 소개와 MTEB에 contribution을 남긴 나의 경험들에 대해 정리해 보았다. 어렵지 않은 작업이라, 관심 있는 분들은 한번 해보신다면 좋을 것 같다.

아 그리고 MTEB를 활용한 MTEB-ko-retrieval Leaderboard도 출시했으니 많은 관심(⭐ 스타 ⭐) 부탁드립니다 !!

리뷰하는 모든 loss 함수들은 sentence-transformers에 구현돼 있다.

1. InfoNCE (MultipleNegativesRanking) Loss

첫번째는 가장 근본인 InfoNCE loss이다. (sentence-transformers에는 MultipleNegativesRankingLoss로 구현돼 있다.)

이 함수는 “query와 유사한 positive는 가깝게, query와 관련없는 negative는 멀게" 학습시키는 loss이다. 식으로 표현하면 다음과 같다.

분모에 들어가는 N은 배치 사이즈이다. 그러므로 InfoNCE loss로 학습시킬 때는 항상 배치 내의 다른 negative들 (이들을 in-batch negatives 라고 한다.)과 멀어지도록 학습된다.

e.g.1) batch size가 1024, {query-positive} 쌍이 입력으로 들어오는 경우, 0번째 query는 0번째 positive를 제외한 나머지 1023 개의 positives과 멀어지도록 학습된다.

e.g.2) batch size가 1024, {query-positive-hard_negative} 쌍이 입력으로 들어오는 경우, 0번째 query는 0번째 positive를 제외한 나머지 positives와 hard negatives (총 1023 positives + 1024 hard_negatives = 2047 passages. ) 과 멀어지도록 학습된다. 또, hard_negative 수가 늘어남에 따라 batch_size 배수만큼 in-batch negatives가 늘어난다. (실제로 mE5와 bge-m3에서는 각 query별로 7개의 hard_negative를 사용한다.)

그러나 사실 InfoNCE loss에서 in-batch negative는 논문마다 다르다.

• E5와 bge-m3에서는 각 query의 hard_negative 만을 in-batch negative로,

• gte에서는 positives와 hard_negatives를 모두 in-batch negative로,

• gecko에서는 여러 positives와 현재 계산중인 query의 hard_negative만을 in-batch negative로 가져간다.

• sentence-transformers에서는 gte와 같이 모든 positives, hard_negatives가 포함되도록 구현돼 있다. (https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/losses/CachedMultipleNegativesRankingLoss.py)
  그리고 그냥 MultipleNegativesRankingLoss를 쓰는 것보다 CachedMultipleNeagtivesRankingLoss를 쓰는 것이 batch size 늘리는 것에 유용하다. (gradient cache 관련 글)

Q. 근데 in-batch negative로 포함된 다른 query들의 positive이나 다른 query들의 hard_negative가 현재 계산중인 query의 positive sample일 수 있지 않나요 ??

나는 이 생각을 굉~~장히 많이하고 걱정도 많이 했다. 만약 negative sampling 했을 때 query를 대답할 수 있는 문서이면 어떡하지 ?? 이런 생각이 굉장히 많이 들었던 것 같다.

이를 보완하기 위해 나온 것이 오늘 소개할 GISTEmbedLoss이다.

2. GISTEmbedLoss

먼저 이 loss 함수는 “GISTEmbed: Guided In-sample Selection of Training Negatives for Text Embedding Fine-tuning” 라는 논문에서 소개되었다.

GISTEmbed: Guided In-sample Selection of Training Negatives for Text Embedding Fine-tuning

먼저, 소개하기 전에 다음 그림을 보자.

앞서 우리가 많은 종류의 InfoNCE loss를 보았지만, 보통의 InfoNCE loss를 C (현재 query의 positive를 제외한 모든 queries, positives, hard_negatives를 in-batch negative로 여기는) 라고 하자.

그럼 분명 “Capital cities”라는 anchor에 대해서 “Where is Manila?”, “Philippines”, “What’s the capital of Canada” 등은 negative가 아닌 positive가 될 수 있다. GISTEmbedLoss는 이러한 document를 학습 중에 필터링한다.

어떻게 필터링할까?

답은 생각보다 간단하다. Guide (Teacher) model을 두고, anchor와 in-batch negative의 유사도를 구해, anchor와 positive의 유사도보다 높은 경우, -inf로 보내, 계산에 포함되지 않도록 한다.

이런 식으로, 배치 내의 false negatives (위양성)을 찾아 제거하면, 더 정교한 학습이 가능하게 된다.

이 GISTEmbedLoss 또한 sentence-transformers에 구현돼 있고, 이 loss 역시 gradient cache가 결합된 CacheGISTEmbedLoss를 지원한다.

sentence-transformers/sentence_transformers/losses/CachedGISTEmbedLoss.py at master · UKPLab/sentence-transformers

간단하게 코드를 살펴보면, guided_sim, 즉 anchor와 positive의 유사도보다 큰 유사도를 가지는 document는 모두 -inf로 보내버린다.

다만 단점은, 모델을 2개 쓰고, 모든 similarity를 구하는 과정이 2번이 되므로, 시간이 굉장히 오래 걸리고, InfoNCE만큼 큰 mini_batch를 설정하지 못한다는 단점이 있다. (mini_batch와 관련해서는 gradient cache 글을 참고하길 바란다.)

※(번외) GISTEmbedLoss를 쓰면 중복방지를 할 수 있지 않을까?

어느날 이런 의문을 가지게 되었다. ‘중복된 negative를 보는 것이 사실 모델 학습 관점에서 크게 좋지 않을텐데, GISTEmbedLoss를 사용할 때 guide model을 활용해서 구한 anchor-positive와 유사도가 높거나 같은 데이터를 필터링하면 중복 방지를 할 수 있지 않을까?’

그래서 바로 issue를 남겼고,

Deduplication integrated in CachedGISTEmbedLoss · Issue #3063 · UKPLab/sentence-transformers

위의 코드를 ‘>’ 대신 ‘>=’ 으로 바꾸어 보았다.

deduplication in GISTEmbedLoss and CachedGISTEmbedLoss by JINO-ROHIT · Pull Request #3074 · UKPLab/sentence-transformers

그리고 해당 코드를 maintainer가 실험해 주었는데 (한낱 학생이 고친 코드를 실험해주다니, 감동이었다), 결과는 실패로 돌아갔다. 예상되는 이유는 anchor-positive와 유사도 같지만, positive와는 또 다른 hard_negative가 같은 유사도로 측정이 되고 필터링되어 오히려 악영향을 준 것이 아닐까 생각했다. 왜냐하면 query와의 유사도가 같다고 해서 반드시 같은 text는 아니기 때문이다.

만약 중복을 방지할 것이라면, sentence-transformers의 NoDuplicatesBatchSampler를 사용하는 것이 좋아보인다.

3. MultipleNegativesSymmetricRankingLoss

다음은 MultipleNegativesSymmetricRankingLoss이다. 이름도 참 길다. 이 loss는 InfoNCE를 다음과 같이 조금 변형한 loss 이다.

쉽게 생각해서, q-p pair로 계산된 loss와, p-q로 계산된 loss를 더하는 것이다. (다만 MultipleNegativesSymmetricRankingLoss 에서는 평균을 구한다.)

코드를 봐도 구현이 쉽게 되어있다.

이 또한 gradient cache가 결합된 버전이 있다.

sentence-transformers/sentence_transformers/losses/CachedMultipleNegativesSymmetricRankingLoss.py at master · UKPLab/sentence-transformers

다만 이 loss는 negative를 활용하지 않기 때문에, hard_negative과 같은 정교한 표현은 학습이 안될 수도 있다.

이 외에도 훨씬 더 많지만..오늘은 내가 자주 사용하는 loss들을 리뷰해 보았다. 이 모든게 sentence-transformers에 구현돼 있다니..임베딩 모델이라면 이제 sentence-transformers 환경을 피하기 매우 어려울 것 같다.

앞으로 또 흥미있는 loss들을 사용하게 된다면 이 글에 계속 추가해두도록 하겠다 !

1. Mixed Precision Training

Mixed Precision과 관련해서는 이전 글에서 자세하게 다뤄본 적이 있으므로, 지금은 간단하게 어떤 방식의 학습법이었는지 복기하고 가도록 하자.

Mixed Precision Training은 fp16과 fp32를 섞어 학습시키는 방법으로, 학습 속도를 빠르게 할 수 있었다. 학습 과정은 다음과 같다.

1. fp32 model weight들을 fp16으로 하나 더 저장한다.
2. fp16 input과 fp16 model의 forward를 통해 fp16 activation (출력값) 생성
3. fp16 activation으로 loss를 계산 -> fp16 loss 생성
4. fp16 loss를 backward -> fp16 gradient 생성
5. fp16 gradient를 fp32로 확장
6. fp32 model weight update에 사용
7. update된 model weight를 다시 fp16 형식으로 복사해서 덮어씌움
8. 1–7 반복

이렇게 fp32 model은 최종 parameter update만, fp16 model은 그 외 나머지 계산을 모두 함으로써 학습을 시키는 것이 mixed precision training이다.

2. ZeRO-DP (ZeRO Data Parallelism)

이렇게 mixed precision을 통해 학습을 빠르게 할 수 있게 되었지만, fp16 형식의 model weight를 하나 더 저장해야 하므로, 메모리에 부담이 커지는 상황이 발생했다.

Mixed precision 학습 방식에 말했듯이 parameter update는 fp32 model에서 일어나기 때문에, AdamW 등의 adaptive optimizer를 사용하는 경우 momentum, bias 등의 optimizer state (여기서는 os로 표기)를 fp32로 저장해야한다.

또, forward 이후에 발생하는 activation들을 담아 놓을 메모리, 분산 처리에서 통신에 발생하는 bucket들을 담아 놓을 메모리도 필요하다.

이렇듯, 메모리를 효율적으로 관리하는 것이 학습에 있어서 핵심 기술로 작용하는데, ZeRO (Zero Redundancy Optimization)에서는 ZeRO-DP와 ZeRO-R을 소개한다. 오늘은 그 중 ZeRO-DP에 대해서 알아보도록 하겠다.

위 사진과 같이, ZeRO-DP는 4개의 stage로 나뉘어 적용된다. 각각의 stage의 특징은 다음과 같다.

1. Stage 0

• ZeRO-DP를 적용 X

2. Stage 1

• Optimizer States (os)만 partitioning하여 여러 GPU에 할당한다.
• 메모리 소비량 4배 감소
• 기존과 비슷한 양의 Communication Cost

3. Stage 2

• Stage 1에 더하여 gradient까지 partitioning하여 여러 GPU에 할당한다.
• 메모리 소비량 2배 더 감소 (총 8배 감소)
• 기존과 비슷한 양의 Communication Cost

4. Stage 3

• Stage 2에 더하여 parameter까지 partioning하여 여러 GPU에 할당한다.
• 메모리 소비량 분할 수준에 따라 선형적 감소
• 기존보다 1.5배 많은 Communication Cost

이들 중 ZeRO-DP Stage 3 방식에 대해 조금 더 자세하게 알아보도록 하자. 아래부터는 사진을 통해 설명을 제공할 예정인데, 사진은 모두 Microsoft에서 제공한 영상 으로부터 가져왔음을 밝힌다.

2–1. Setting

먼저 학습하려는 모델이 16-layer transformer 모델이라고 가정하고,

우리가 학습하려는 큰 데이터셋과 가용할 수 있는 GPU가 4개라고 가정한다.

첫 번째 단계에서는 데이터가 배치사이즈로 분할되어 순차적으로 각 GPU에 할당된다. (DataParallel과 같은 방식)

그리고 각 GPU는 다른 데이터들을 같은 모델에 대해 optimize 하는 것이 목적이다.

기존 DP에서는 모델을 각 device마다 복사하여 다음과 같은 설정에 놓이게 된다.

또, forward 이후 생겨는 activations들을 담기 위한 buffer도 필요하다.

그러나 ZeRO-DP는 DP와는 다르게 각 device마다 모델을 할당하지 않고, 모델을 각 device에 나누어 할당한다. 모델이 분할되어 할당된 모습은 다음과 같다.

2–2. Train - forward

이제 세팅이 끝났으니, 학습 과정을 살펴보자.

먼저 모델이 M0, M1, M2, M3의 4개로 분할되어 각 GPU에 할당되었으니,
- GPU0은 M1, M2, M3가,
- GPU1은 M0, M2, M3가,
- GPU2는 M0, M1, M3가,
- GPU3은 M0, M1, M2가 없는 상황이다.

2–2–1. M0 forward

따라서, 학습하고자 하는 모델 partition을 각각의 GPU에 옮겨주어야 한다. 먼저 M0를 모든 GPU에 복사하자. (이렇게 복사하는 과정을 broadcast라고 한다.)

이렇게 M0가 모든 device에 복사되었으니, 각각의 데이터 배치를 이용해서 forward를 진행한다.

이렇게 M0에 대한 forward가 끝났으면, GPU1, GPU2, GPU3에서 해당 모델 partition M0를 지운다.

이 과정을 모든 모델 partition에 그대로 적용한다. (M1까지만 나타내겠음)

2–2–2. M1,2,3 forward

이걸 M2, M3에 대해서도 적용하면 forward 이후의 최종 모습은 다음과 같다.

다만 마지막 M3 모델 partition 복사본은 제거하지 않는데, 이는 바로 backward를 하여 gradient를 계산하기 위함이다.

2–3. Train - caclulate_loss & backward

2–3–1. M3 backward

Forward pass가 종료되었으므로 각 dataset에 대해 Loss를 계산한다.

이후 계산된 loss를 가지고 backward를 진행한다. 이때는 backward를 진행하므로 fp16 gradient memory 를 담을 temporary buffer (주황색)를 갖게 된다.

먼저 M3에 대해 backward를 진행하여 gradient를 계산한다.

이후, GPU0, GPU1, GPU2가 각 device에서 각 data batch에 대해 계산한 M3 gradient를 GPU3에게 넘기고, GPU3는 해당 gradient들을 더하여 축적한다. (Gradient accumulation과 같은 방식)

이렇게 모든 gradient가 전달되어 축적이 완료되면, 각 device에서 M3 모델 partition과 관련된 gradient, parameter, activation tensor를 모두 지운다.

2–3–2. M2 backward

다음으로는 M2 부분에 대한 backward를 위해 GPU2에서 다시 각 device들에게 모델 partition M2를 전달한다.

계산되는 gradient를 담기 위한 temporary buffer를 준비하고 backward를 진행한다.

이렇게 각 data batch에 대해 계산한 M2 gradient를 GPU2에게 넘기고, GPU2는 해당 gradient들을 더하여 축적한다.

Gradient accumulation이 완료됐으면 각 device에 복사된 M2와 각 device가 가지고 있던 temporary memory를 제거한다.

2–3–3. M1, M0에 대해서도 같은 backward를 수행한다.

그럼 최종적으로 backward 후에 다음과 같이 gradient들이 device마다 accumulate되고, activation들은 지워진 상태가 된다.

2–4. Optimization

이제 계산된 gradient를 기반으로 parameter를 update하자. optimizer.step()을 통해 update를 진행한다.

(fp16 gradient -> fp32로 casting -> AdamW 등의 optimizer로 fp32 model parameter update)

이후, 이 updated fp32 model weights를 fp16으로 casting하고

이걸 다시 각각의 fp16 model parameters에 반영한다. (덮어씌운다.)

이렇게 완성된 1 iteration 후의 모델은 다음과 같다 !

2–5. 결과

이렇게 ZeRO-DP를 사용한 결과는 다음과 같다.

ZeRO 방식을 사용한 것이 TFLOps, Speed-Up 관점에서 훨씬 더 좋은 성능을 보여주었다.

2–6. 코드 실습

ZeRO stage1과 stage3(위에서 설명한) 방식 둘다 실습을 해보자. 코드는 GITHUB에서 확인이 가능하다.

ZeRO는 Deepspeed라는 라이브러리로 쉽게 구현이 가능하다.
(근데 deepspeed가 설치가 생각보다 쉽지 않다.)

Latest News

실습할 전체 코드는 다음과 같다.

from argparse import ArgumentParser
from datasets import load_dataset
from torch.optim import Adam
from torch.utils.data import DataLoader
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import deepspeed
import torch.distributed as dist


model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

parser = ArgumentParser()
parser.add_argument(
    "--deepspeed_config", default="/mnt/raid6/yjoonjang/projects/distributed-train/src/zero_dp_config.json", type=str
)
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()

optimizer = Adam(model.parameters(), lr=3e-5, weight_decay=3e-7)

engine, optimizer, _, scheduler = deepspeed.initialize(
    args=args,
    model=model,
    optimizer=optimizer,
)

datasets = load_dataset("squad").data["train"]["context"]
datasets = [str(sample) for sample in datasets]
data_loader = DataLoader(datasets, batch_size=8, num_workers=8)

engine.train()

for i, data in enumerate(data_loader):
    tokens = tokenizer(
        data,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=1024,
    )

    loss = engine(
        input_ids=tokens.input_ids.cuda(),
        attention_mask=tokens.attention_mask.cuda(),
        labels=tokens.input_ids.cuda(),
    ).loss

    engine.backward(loss)
    engine.step()

    if i % 10 == 0 and dist.get_rank() == 0:
        print(f"step:{i}, loss:{loss}")

    if i >= 300:
        break

args 인자로 deepspeed_config 파일을 받도록 해 놓았기 때문에, train 하는 명령어 뒤에 --deepspeed_config 라는 인자로 어떤 json 파일을 사용할 것인지만 지정해 놓으면 된다. 참고로 다양한 config.json 파일들은 아래 사이트에서 확인할 수 있다.

DeepSpeed Configuration JSON

이제 stage1, stage3에 대한 실습을 진행하자.

1. ZeRO-DP stage1

zero_dp_config_stage1.json은 다음과 같다.

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 1,
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "total_num_steps": 300,
      "warmup_min_lr": 0,
      "warmup_max_lr": 3e-5,
      "warmup_num_steps": 30
    }
  },
  "fp16": {
    "enabled": true,
    "initial_scale_power": 32,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 1
  },
  "zero_allow_untested_optimizer": true,
  "wall_clock_breakdown": false,
  "steps_per_print": 9999999999
}

그리고 아래와 같은 명령어를 통해 파일을 실행시킨다.

CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed
--num_gpus=4 4_zero_dp_tutorial.py
--deepspeed_config=/mnt/raid6/yjoonjang/projects/distributed-train/src/zero_dp_config_stage1.json

실행시켜보면 굉장히 많은 로그들이 뜨는데, 그 중 중요한 로그들만 첨부하자면,

[INFO] [[logging.py:96](<http://logging.py:96/>):log_dist] [Rank 0] Creating torch.float16 ZeRO stage 1 optimizer

[INFO] [[logging.py:96](<http://logging.py:96/>):log_dist] [Rank 0] DeepSpeed Final Optimizer = DeepSpeedZeroOptimizer

이정도가 있다. 즉, stage1으로 잘 실행되는 것을 볼 수 있고,

학습이 잘 되어 step마다 loss도 잘 찍히는 것을 볼 수 있다.

근데 자세히 보면

[INFO] [loss_scaler.py:183:update_scale] [deepspeed] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: {}, reducing to {}

이런 메시지가 엄청 뜨는 것을 볼 수 있다. 그 이유는 너무 작은 수에 대한 표현을 하지 못하는 fp16을 위해 loss scaling factor를 곱하는데, 이 loss scaling factor가 fp16이 표현할 수 있는 범위 이상으로 커지기 때문에 loss가 inf가 되어버리는 현상이 발생한 것이다. 그래서 이런 경우에, loss를 처음에는 엄청 큰 값을 곱했다가, inf가 되어 overflow가 발생하면 점차 1/2배 되는 값을 곱해가면서 조정한다.

2. ZeRO-DP stage3

이번에는 stage3에 대한 실습을 해보자. 코드는 그대로 같고, json 파일만 다음과 같이 작성한다.

zero_dp_config_stage3.json

{
  "fp16": {
    "enabled": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 3e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
  "type": "WarmupDecayLR",
  "params": {
      "total_num_steps": 1000,
      "warmup_min_lr": 0.0,
      "warmup_max_lr": 5e-5,
      "warmup_num_steps": 100
    }
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,

    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": 5e7,
    "stage3_param_persistence_threshold": 1e5,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "memory_efficient_linear": true,
    "gather_16bit_weights_on_model_save": true
  },
  "gradient_accumulation_steps": 1,
  "gradient_clipping": true,
  "steps_per_print": 5,
  "train_batch_size": 64,
  "train_micro_batch_size_per_gpu": 16,
  "wall_clock_breakdown": false
}

이후에는 아래 명령어를 실행한다.

CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed
--num_gpus=4 4_zero_dp_tutorial.py
--deepspeed_config=/mnt/raid6/yjoonjang/projects/distributed-train/src/zero_dp_config_stage3.json

이번에도 굉장히 많은 로그들이 뜨는데,

zero_config 3, zero stage 3 라고 뜨는 것을 볼 수 있다.

그리고 stage3는 신기하게

loss값과 더불어서 시간, 할당된 메모리 등에 대한 정보도 함께 출력된다.

2–7. 코드 까보기

이렇게 실습을 마쳤는데, 막상 나오는 로그를 이해한게 많이 없어서, 코드를 좀 까보고 싶어졌다.

먼저 전체적인 코드에서 deepspeed.initialize를 통해 deepspeed가 활성화되기 때문에, 이 deepspeed의 initialize라는 함수가 어떻게 되어 있는지를 한번 보고, 각 json 파일에 작성된 부분이 어떻게 코드와 연결되는지, 무엇을 담당하는지 등을 좀 알아보도록 하겠다.

먼저 코드 중 deepspeed.initialize를 활성화하는 부분을 보면,

engine, optimizer, _, scheduler = deepspeed.initialize(
    args=args,
    model=model,
    optimizer=optimizer,
)

이렇게 deepspeed.initialize의 결과로 engine, optimzer, scheduler를 받는다.

그래서 initialize 하는 부분의 코드를 좀 보면,

여기서 다른 부분은 DP 세팅과 비슷한데,

이렇게 deepspeed_config로 engine을 설정해 주는 부분이 있다. HybridEngine이 뭔가 했는데, Multi-dimensional Parallelism이라고, Tensor Parallelism과 Pipeline parallelism을 함께 사용하는 방식인 것 같았다.

아무튼 우리는 현재 PipelineParallelism을 사용하려고 하지는 않기 때문에, DeepSpeedEngine이라는 코드를 까 본다.

쭉 내리다 보면 _configure_zero_optimizer이라는 코드가 있다. 이 함수를 통해 zero 단계에 맞추어 os, g, p를 gpu에 분할한다.

먼저 ZeRO stage 1,2는 다음 함수를 따른다.

optimizer = DeepSpeedZeroOptimizer(
              optimizer,
              self.param_names,
              timers=timers,
              static_loss_scale=self.loss_scale(),
              dynamic_loss_scale=self.dynamic_loss_scale(),
              dynamic_loss_args=self.dynamic_loss_scale_args(),
              clip_grad=self.gradient_clipping(),
              contiguous_gradients=contiguous_gradients,
              reduce_bucket_size=self.zero_reduce_bucket_size(),
              use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(),
              allgather_bucket_size=self.zero_allgather_bucket_size(),
              dp_process_group=self.seq_data_parallel_group,
              expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None,
              expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None,
              reduce_scatter=self.zero_reduce_scatter(),
              overlap_comm=overlap_comm,
              offload_optimizer_config=self.zero_offload_optimizer(),
              mpu=self.mpu,
              postscale_gradients=self.postscale_gradients(),
              gradient_predivide_factor=self.gradient_predivide_factor(),
              gradient_accumulation_steps=self.gradient_accumulation_steps(),
              ignore_unused_parameters=self.zero_ignore_unused_parameters(),
              partition_grads=zero_stage == ZeroStageEnum.gradients,
              round_robin_gradients=round_robin_gradients,
              has_moe_layers=self.has_moe_layers,
              fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(),
              gradient_accumulation_dtype=gradient_accumulation_dtype,
              communication_data_type=self.communication_data_type,
              elastic_checkpoint=self.zero_elastic_checkpoint()
            )

(번외)

round_robin_gradients가 어떤 인자인가 너무 궁금해서 찾아보았는데 이 인자를 처리하는 코드는 다음과 같았다.

if self.round_robin_gradients:
    round_robin_tensors, round_robin_indices = self._round_robin_reorder(
        self.bit16_groups[i], dist.get_world_size(group=self.real_dp_process_group[i]))
else:
    round_robin_tensors = self.bit16_groups[i]
    round_robin_indices = list(range(len(self.bit16_groups[i])))

self.round_robin_bit16_groups.append(round_robin_tensors)
self.round_robin_bit16_indices.append(round_robin_indices)

그리고 _round_robin_reorder를 통해 gradient 부분을 나누어 할당하는데, 쉽게 예시로 설명해보면: gradient가 [1,2,3,4…8] 이 있고 4개의 gpu가 있을 때
- round_robin_gradients가 True면 [1,5], [2,6], [3,7], [4,8] 으로 나뉘고,
- round_robin_gradients가 False면 [1,2], [3,4], [5,6], [7,8] 으로 나뉘게 된다.
근데 솔직히 뭐가 좋은건지는 모르겠다.

암튼 DeepSpeedZeroOptimizer는 이런식으로 인자를 받아 텐서로 분할하고 device에 나누어 주는 역할을 한다.

그리고 놀랍게도 ZeRO stage3는 stage3만의 optimizer class가 있다.

그리고 DeepSpeedZeroOptimzer_Stage3라는 클래스에 들어가 보면, _create_fp16_partitions_with_defragmentation이라는 함수를 확인할 수 있다. 이 함수에서

해당 코드를 통해 parameter들을 partitioning하고 GPU에 할당한다는 것을 확인할 수 있다.

오늘은 이렇게 ZeRO-DP 방식이 어떠한 분산 학습 방식인지, 각 stage별로 어떤 항목을 어떻게 partitioning하는 것인지에 대해 알아보았다. 또한 코드 실습을 통해 deepspeed를 경험해 볼 수 있었으며, deepspeed 라이브러리를 까봄으로써 어떻게 모델이 partitioning되어 GPU에 할당되는지에 대해서도 알아보았다.

Deepspeed 라이브러리를 까보면서 느꼈지만, Trainer 급으로, 어쩌면 Trainer보다 더 많은 attribute들이 있는 것 같다. 이 attribute들도 언젠가 하나의 글로 정리해볼 수 있었으면 좋겠다.

Model Parallelism은 보통 Tensor Parallelism을 지칭하는 말 같은데, 사실은 Pipeline Parallelism과 Model Parallelism의 2가지로 나뉘어진다. 이들이 뭐가 다른지, 그리고 뭐가 더 효율적이고 사용하기 좋은지 알아보자.

1. Pipeline Paralleism (PP)

Pipeline Parallelism은 쉽게 말해, 모델을 layer 별로 쪼개어 서로 다른 GPU에 올린 것이라고 볼 수 있다.

그러나 decoder 모델 (또는 다른 생성 모델)들의 특성상, 앞 layer의 결과물이 다음 layer의 입력으로 들어와 연산을 통해 next token prediction이 가능하기 때문에, 병렬적으로 연산하지 못한다는 단점이 있다.

1–1. GPipe

이렇게 병렬적으로 연산하지 못하여 발생하는 학습 시간 문제를 해결하기 위해 입,출력 값을 micro batch 단위로 나누어 학습시키기도 한다. 이 방법을 구글에서 GPipe라는 이름으로 개발했다고 한다.

Micro-batch로 나누는 개수가 많아질수록, model이 연산을 쉬는 구간 (bubble time)이 적어진다.

그러나 이런 GPipe 병렬 처리도 문제가 존재한다.

위 사진처럼, GPU끼리 한번에 모델 forward, 한번에 backward를 수행하며, 모델 parameter를 update하는 시간 마저 따로 분리되어 있다 (오른쪽에 있는 검정색 선). 이렇게 parameter를 update하는 과정을 pipeline flush라고 하는데, GPipe 에서는 이 pipeline flush마다 어떠한 연산도 하지 않기 때문에 버리는 시간이 너무 많아진다.

1–2. 1F1B Pipelining (PipeDream)

그래서 Microsoft에서는 1번의 forward 이후에 바로 backward를 하자는 1F1B Pipelining, PipeDream을 제안한다.

PipeDream은 여러 개의 GPU들이 동 시간대에 forward, backward를 담당하는 것이 아니라, GPU 별로 상황에 맞게 forward, backward를 담당하는 방식이다. (위의 그림을 보면 한번에 이해가 될 것이다.) 이로 인해 GPipe처럼 프로세스가 parameter update로 인해 쉬게 되는 현상이 없어진다.

그러나 이런 PipeDream 방식에도 문제가 있다. 아래 그림을 보자.

오른쪽의 그림을 보면, 1번 layer에 대한 parameter update (빨간색 1번 줄 )가 일어나는 도중에 계속 forward, backward가 진행된다. 이는 각 layer를 지날 때마다 모델 파라미터가 변하는 것이기 때문에, 정상적인 학습이 진행되는 것이라고 얘기할 수 없다.

그래서 이를 위해서는 여러 버전의 paremter state를 저장하고 있어야 한다. 그래야 backward 시에 당시 forward에 어떤 paremeter를 썼는지 추적할 수 있기 때문이다.

통상적으로 정상적인 모델 fine-tuning 시에 알려진 VRAM 차지 용량은 모델 parameter 개수의 16배이다. 이는 모두 forward시에 저장되는 모델 parameter 및 activation 때문인데, 저장되는 parameter가 늘어난다면… 얼마나 큰 VRAM이 필요할지 상상조차 하기 싫다.. (물론 PP도 기본적으로 parallelism이 적용되는 것이기 때문에 여러 버전의 parameter state를 저장해도 parallelism을 사용하지 않을 때보다는 차지하는 VRAM이 적을 것이다.)

그래서 Pipeline Parallelism (PP)에 관해서는 다음과 같이 정리할 수 있다.

- GPipe: 메모리 효율적, 프로세싱 비효율적
- PipeDream: 메모리 비효율적, 프로세싱 효율적

1–3. 더 향상된 버전

그래서 이런 단점을 또 극복하고자 더 향상된 Pipeline Parallelism을 구현할 수 있다.

하나는 기본적인 PipeDream에 gradient accumulation을 섞는 방법이다. 이렇게 gradient accumulation을 사용함으로써 저장되는 weight 버전이 단 2개로 줄어들 수 있다.

두 번째는 PipeDream에 앞서 본 Gpipe의 pipeline flush을 결합한 방법이다. 이는 기본적으로 activation memory를 줄이기 위한 방법이다.
파이프라인 구성을 위 (b) 사진처럼 바꾸면, 1,2번 layer에 대한 activation을 저장해 놓다가, 3번 forward가 일어나는 시점에는 1번의 backward가 끝나기 때문에 2,3번 layer에 대한 activation을 저장하면 되고, 4번 forward가 일어나는 시점에는 2번 backward가 끝나기 때문에 3,4번 layer에 대한 activation을 저장하면 된다.

즉, 기존 방식이 activation을 4개 layer에 대해서 저장했다면, 이 방식은 2개 layer에 대해서만 저장하면 되므로, 굉장히 memory-efficient한 방식이다.

1–4. 더 더 향상된 버전..

이걸 더 향상시킨 버전이 있다.. 그림을 살펴보자.

사진만 보다시피 굉장히 굉장히 분할을 많이 해놨다. 기본 개념은 “각 stage에서 수행하는 layer의 연산마저 분할하자"는 것이다. 그림을 보면, 위의 그림에서는 Device1 한 칸에 1번 layer가 할당되어 있지만, 아래 그림에서는 그 칸이 2개로 나뉘어 1,2번 layer가 할당되어 있다. 이렇게 하면 bubble 시간을 극도로 줄일 수 있다고 하는데…
개인적으로 뭔가 좀 뇌절같은 느낌이 든다..ㅎ

1–5. PP 정리

이렇게 Pipeline Parallelism (PP)에 대해 알아보았다.

총 정리를 해보면, PP는 모델을 layer별로 나누어 서로 다른 GPU에 할당하는 방식이다. 이전 layer의 출력이 다음 layer의 입력으로 들어오는 모델 특성상, 이전 layer의 출력을 기다려야 하는 시간이 많은데, 이를 최적화 하기 위한 다양한 기법들을 알아보았다.

PP는 개인적으로 잘 사용하는 분산 처리 방식이 아니어서, 코드로 구현해보지 않았다.

2. Tensor Parallelism

두번째로, Tensor Parallelis (TP)에 관해 살펴보자.

TP는 쉽게 말해 모델 자체 (전체 파라미터)를 여러 개의 부분으로 쪼개어 서로 다른 GPU에 올리는 방식이라고 볼 수 있다.

Layer별로 나누어 GPU에 올리는 PP와는 완전히 다른 방식이다. PP는 세로로 자르는 방식이라면 TP는 가로로 자르는 방식..?!

이를 처음 구현한 것은 Nvidia의 Megatron-LM 프레임워크이다.

2–1. Column, Row Parallelism

Megatron-LM에서는 모델의 파라미터를 어떻게 자르는지를 기준으로 parallelism을 적용하는데, 자르는 방법은 column, row의 2가지이다.

1. Column Parallelism

위의 그림과 같이, 파라미터를 column으로 자르는 경우, 모델 파라미터에 해당하는 tensor를 column 기준으로(세로로) 잘라 각각 다른 GPU에 올리고, 같은 input 를 복사 (broadcast)하여 각각의 GPU마다 올라간 파라미터와 계산한다. 이후, cat (all-gather)을 통해 최종 계산 결과를 얻는다.

2. Row Parallelism

반대로 파라미터를 row로 잘라 계산하는 경우, 각각 분할된 파라미터를 서로 다른 GPU에 올려 두고, input을 scatter할 때 input을 column으로 잘라서 scatter한다. 이후 결과물들의 element-wise 덧셈을 통해 최종 산출물을 얻을 수 있다.

Column Parallelism을 쓰던지, Row Parallelism을 쓰던지, 최종 계산 결과물은 같다.

2–2. Transformer block Parallelism

그럼 이 parallelism을 transformer 아키텍쳐에 대해서는 어떻게 적용할 수 있을까.

Megatron-LM에서는 transformer 구조에서 Layer Norm 레이어를 제외한 다른 모든 레이어들(Attention, MLP)에 Column, Row Parallelism을 적용한다.
(Layer Norm 레이어에 적용하지 않는 이유는 파라미터 수가 적기 때문이다.)

그럼 이 transformer 아키텍쳐에 parallelism을 적용한 방법에 대해서 알아보자.

2–2–1. MLP Layer

위 그림의 오른쪽과 같이 Linear -> GeLU -> Linear -> Dropout 순으로 진행되는 것이 MLP Layer이다.

이때 2개의 Linear layer는 각각 column, row parallelism이 적용된다. 즉 앞의 linear layer는 column parallelism, 뒤의 linear layer는 row parallelism이 적용되는 것이다.

왜 굳이 이 순서대로 적용이 될까? 이유는 2가지가 있다.

a. All-gather과 scatter 통신 프로세스의 생략

위의 그림은 2개의 linear layer에 column-row 순서로 parallelism이 적용되었을 때의 시나리오를 보여 준다. 그림을 보면, column parallelism의 출력이 나온 후 (초록색), all-gather하여 하나의 tensor로 합치고 (빨간색), 다시 scatter하는 것 (초록색)을 볼 수 있다. 그럼 의문이 든다. 굳이 왜 합쳤다가 다시 나누지 ??

column parallelism의 출력이나 (XW), row parallelism의 입력이나 (Y1, Y2) 생긴게 아예 똑같다. 즉, gather 후 다시 scatter하는 것을 하지 않아도 되는 것이다.
그래서 결론적으로는 이렇게 column-row 순서로 parallelism을 적용한 경우에만 gather, scatter 연산을 생략할 수 있다.

b. GeLU 연산을 고려한 병렬화

앞서 말했듯이, MLP layer에서 벌어지는 상황은
input -> linear layer -> GeLU -> linear layer -> Dropout 이다.

지금 우리는 GeLU 좌우에 있는 linear layer에 대해 어떻게 parallelism을 적용할 지 살펴본 것이고, 실제로 재대로 된 시나리오는 아래와 같다.

이렇듯 막무가내로 scatter, gather 연산을 생략할 수는 없고, GeLU 연산도 병렬화되어 제대로 수행될 수 있는지 확인해보아야 한다.

즉 scatter, gather 연산을 제거했을 때도 아래와 같이 되어야 하는 것이다.

앞서 우리는 row, column parallelism에 대해서 알아보았다. 그때 row parallelism은 두 텐서를 element-wise하게 더함으로써, column parallelism은 두 텐서를 concatenate 함으로써 최종 출력물을 얻는다는 것을 알게 되었다. 즉 GeLU를 낀 채로 row parallelism이 가능하게끔 하기 위해서는 아래 식이 성립해야 하고,

GeLU(XW1 + XW2) = GeLU(XW1) + GeLU(XW2)

GeLU를 낀 채로 column parallelism이 가능하게끔 하기 위해서는 아래 식이 성립해야 한다는 것이다.

GeLU(XW1 ⊚ XW2) = GeLU(XW1) ⊚ GeLU(XW2)

그러나, GeLU는 비선형 함수이기 때문에 row parallelism에 해당하는 식은 성립하지 않고, column parallelism에 해당하는 식은 성립한다:

GeLU(XW1 + XW2) ≠ GeLU(XW1) + GeLU(XW2),
GeLU(XW1⊚XW2) = GeLU(XW1)⊚GeLU(XW2)

따라서 column parallelism이 수행된 후에 GeLU 연산이 수행되어야만 reduce, scatter 프로세스를 생략할 수 있고, 이것이 우리가 column-row 순서로 parallelism을 적용해야 하는 이유이다.

2–2-2. Multi-head Attention Layer (Self-Attention Layer)

이제 MLP layer에 대해 알아보았으니, multi-head attention layer (논문에서는 self-attention layer라고 한다.)에 parallelism을 어떻게 적용하는지에 대해 알아보자.

Multi-head attention layer에서는 Input -> Linear1 -> Split heads -> Scaled Dot-Product Attention -> Concat -> Linear2 로 진행된다.

Multi-head attention에서도 MLP layer와 마찬가지로 두 개의 linear layer에 대해서 parallelism을 적용한다. 이전과 마찬가지로 선순위에 있는 Q, K, V linear projection에는 column parallelism을, 마지막 output projection 부분에는 row parallelism을 적용한다. 이렇게 다시 column-row parallelism을 만들면서 gather, scatter 프로세스를 없앨 수 있다.

2–2–3. Vocab parallel embedding

마지막으로 vocab에 대해서도 분산 처리를 진행하는데, 신기하게도 vocab size를 기준으로 반으로 나눈다.

이렇게 50,000 크기의 vocab이 있다고 했을 때, seq_len이 6인 input이 들어온다고 가정하자. 그럼 보통의 경우 위의 그림처럼 각 단어의 임베딩을 따와 [6, embed_dim] 크기의 tensor로 변환된다.

그러나 vocab parallel embedding을 적용하면 먼저 전체 vocab을 각 GPU에 나눠서 할당한다. 즉, vocab을 GPU의 개수만큼 분할하여, 각 GPU는 전체 vocab 중 일부분만 담당하게 된다.

이후 입력 단어에 대해, 각 GPU는 자신이 담당하는 단어에만 임베딩을 계산하고, 다른 GPU에서 담당하는 단어들은 처리하지 않기 때문에 masking 처리가 된다. 이때 scatter 프로세스를 통해 각 GPU가 맡은 단어들만 적절하게 할당된다.

마지막으로, 각 GPU에서 계산된 임베딩을 다시 gather하여 최종 임베딩을 완성한다.

방금까지는 input이 들어와 parallelism으로 처리되는 임베딩에 대해 알아 보았으니, 이제는 output이 어떻게 처리되는지 알아보자.

2–2–4. Vocab parallel cross entropy

분산 처리를 하지 않는 일반적인 상황에서는 위와 같이 (seq_len, vocab_size) 크기의 output tensor가 나오고, 그것을 target과 비교하여 cross-entropy loss를 계산한다.

그러나 앞서 말했듯이, Megatron-LM은 vocab parallel embedding을 사용하기 때문에, vocab_size / GPU 개수의 크기의 vocab들이 GPU에 나누어져 들어가 있다. 즉, 위의 사진처럼 logits가 다 찬 것이 아니고, 서로 다른 GPU에 분포되어 있는 것이다. (아래 사진 참고)

그럼 이런 경우에는 어떻게 해야할까? 답은 생각보다 간단하다. Masking되어 나온 부분에 대해서는 target도 똑같이 masking 시켜 loss를 계산하지 않고, 병렬적으로 계산한 후 모아주는 것이다. 그림으로 보면 다음과 같다.

이를 vocab parallel cross entropy 라고 한다.

이렇게 vocab embedding, vocab에 대한 cross entropy loss까지 분산처리를 함으로써 메모리 최적화를 확실하게 한다.

2–3. 코드

이번에도 Trainer 내부 구현을 간단하게 살펴보자. Trainer에서는 accelerate 라이브러리의 device_map속성을 설정해 줌으로써 model parallelism을 활성화할 수 있다.

전체적인 코드는 GITHUB에서 확인 가능하다.

위 코드와 같이 device_map="auto"로 설정하여 모델을 로드하고, pipeline.model.hf_device_map을 출력해보면 다음과 같이 되는 것을 볼 수 있다.

그럼 device_map=”auto”는 어떤 속성이 있길래 이렇게 하나의 명령어만으로 모델을 GPU에 나누어 올릴 수 있을까? 코드를 한번 보자.

Accelerator의 코드를 까보던 도중, 아래와 같은 부분이 모델을 쪼개어 gpu에 올리는 부분임을 알게 되었다.

그렇다면 이 infer_auto_device_map함수를 한번 보자. GITHUB

저자들이 친절하게 남겨준 주석을 보면, 모델을 로드할 때 gpu 공간이 부족하면 cpu에 오프로딩하고, cpu마저 공간이 부족하면 디스크에 오프로딩하도록 해 놓았다고 한다. 과연 이 과정을 어떻게 코드로 구현했을까 ? 한번 살펴보자.

1. 세팅

# Get default / clean up max_memory
    max_memory = get_max_memory(max_memory)
    if no_split_module_classes is None:
        no_split_module_classes = []
    elif not isinstance(no_split_module_classes, (list, tuple)):
        no_split_module_classes = [no_split_module_classes]

    devices = list(max_memory.keys())
    if "disk" not in devices:
        devices.append("disk")
    gpus = [device for device in devices if device not in ["cpu", "disk"]]

    # Devices that need to keep space for a potential offloaded layer.
    if "mps" in gpus:
        main_devices = ["mps"]
    elif len(gpus) > 0:
        main_devices = [gpus[0], "cpu"]
    else:
        main_devices = ["cpu"]

    module_sizes = compute_module_sizes(model, dtype=dtype, special_dtypes=special_dtypes)
    tied_parameters = find_tied_parameters(model)

    if check_tied_parameters_in_config(model) and len(tied_parameters) == 0:
        logger.warn(
            "The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function."
        )

    device_map = OrderedDict()
    current_device = 0
    current_memory_used = 0
    device_memory_used = {}
    device_buffer_sizes = {}

    # Direct submodules and parameters
    modules_to_treat = (
        list(model.named_parameters(recurse=False))
        + list(model.named_children())
        + list(model.named_buffers(recurse=False))
    )
    # Initialize maximum largest layer, to know which space to keep in memory
    max_layer_size, max_layer_names = get_max_layer_size(modules_to_treat, module_sizes, no_split_module_classes)

먼저 사용할 수 있는 어떤 디바이스가 있는지, 각 디바이스당 가용할 수 있는 최대 메모리 등을 설정해두고, main_device를 찾고 각종 파라미터를 초기화해둔다. modules_to_treat는 처리할 모듈, 파라미터, 버퍼 등의 목록이다. 참고로 tied_parameters가 뭔가 하고 찾아보았는데, 모델 내 같은 파라미터를 가진 모듈을 뜻했다. 이렇게 tied_parameters를 찾는 이유는 어차피 같은 파라미터이니, 모델 학습에서 계산 시 메모리 용량을 줄이기 위함이라고 한다.

2. 메모리 관련 세팅

# Ready ? This is going to be a bit messy.
    while len(modules_to_treat) > 0:
        name, module = modules_to_treat.pop(0)
        if verbose:
            print(f"\nTreating module {name}.")
        # Max size in the remaining layers may have changed since we took one, so we maybe update it.
        max_layer_names = [n for n in max_layer_names if n != name and not n.startswith(name + ".")]
        if len(max_layer_names) == 0:
            max_layer_size, max_layer_names = get_max_layer_size(
                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
                module_sizes,
                no_split_module_classes,
            )
        # Assess size needed
        module_size = module_sizes[name]

        # We keep relevant tied parameters only: one of the tied parameters in the group is inside the current module
        # and the other is not.
        # Note: If we are currently processing the name `compute.weight`, an other parameter named e.g. `compute.weight_submodule.parameter`
        # needs to be considered outside the current module, hence the check with additional dots.
        tied_param_goups = [
            tied_group
            for tied_group in tied_parameters
            if any(name + "." in k + "." for k in tied_group) and not all(name + "." in k + "." for k in tied_group)
        ]

        if verbose and len(tied_param_goups) > 0:
            print(f"  Found the relevant tied param groups {tied_param_goups}")

        # Then we keep track of all the parameters that are tied to the current module, but not in the current module
        tied_params = sum(
            [[p for p in tied_group if name + "." not in p + "."] for tied_group in tied_param_goups], []
        )

        if verbose and len(tied_params) > 0:
            print(f"  So those parameters need to be taken into account {tied_params}")

        device = devices[current_device]
        current_max_size = max_memory[device] if device != "disk" else None
        current_memory_reserved = 0
        # Reduce max size available by the largest layer.
        if devices[current_device] in main_devices:
            current_max_size = current_max_size - max_layer_size
            current_memory_reserved = max_layer_size

(코드에서 “Ready ? This is going to be a bit messy.”라고 한다.)

여기서는 처리할 모듈을 순회하며 각 모듈의 메모리 사용량과 tied parameters를 고려하여 디바이스에 효율적으로 배분하는 데 필요한 메모리 관리와 관련된 로직을 다룬다.

3. 모델 크기와 GPU의 VRAM을 고려한 분산 처리 방식

# Case 1 -> We're too big!
if current_max_size is not None and current_memory_used + module_size > current_max_size:
    # Split or not split?
    modules_children = (
        []
        if isinstance(module, nn.Parameter) or isinstance(module, torch.Tensor)
        else list(module.named_children())
    )
    if verbose:
        print(
            f"Not enough space on {devices[current_device]} to put {name} (space available "
            f"{current_max_size - current_memory_used}, module size {module_size})."
        )
    if len(modules_children) == 0 or module.__class__.__name__ in no_split_module_classes:
        # -> no split, we go to the next device
        if verbose:
            print("This module cannot be split, going to the next device.")

        device_memory_used[device] = current_memory_used + current_memory_reserved
        current_device += 1
        modules_to_treat = [(name, module)] + modules_to_treat
        current_memory_used = 0
    else:
        # -> split, we replace the module studied by its children + parameters
        if verbose:
            print(f"Splitting {name}.")
        modules_children = list(module.named_parameters(recurse=False)) + modules_children
        modules_to_treat = [(f"{name}.{n}", v) for n, v in modules_children] + modules_to_treat
        # Update the max layer size.
        max_layer_size, max_layer_names = get_max_layer_size(
            [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
            module_sizes,
            no_split_module_classes,
        )

Case 1: GPU 메모리가 부족한 경우

• current_max_size는 현재 디바이스에서 사용할 수 있는 최대 메모리 공간이다. 만약 current_memory_used + module_size가 current_max_size보다 크다면, 즉 현재 모듈을 GPU에 할당할 수 없다면 모듈을 나눌 수 있는지 여부를 확인한다.
• modules_children는 모듈의 하위 모듈이나 파라미터를 추적한다. 만약 모듈이 파라미터나 텐서라면 나눌 수 없지만, 모듈의 하위 모듈이 있거나 나눌 수 있는 경우에는 이를 하위 모듈로 분할할 수 있다. (no_split_module_classes에 속한 경우에는 분할이 불가능하다)
• 나눌 수 없는 경우
• 만약 모듈을 나눌 수 없으면, 현재 디바이스에서 처리가 불가능하다는 의미이므로 “다음 디바이스로 이동”한다. 즉, 다른 GPU나 CPU로 넘어가서 다시 할당을 시도한다.
• 이때 current_device를 증가시켜서 다음 디바이스를 사용하도록 하고, 현재 모듈은 다시 modules_to_treat 리스트에 추가한다.
• 나눌 수 있는 경우
• 만약 모듈을 나눌 수 있다면, 해당 모듈을 하위 모듈 또는 파라미터로 나누어서 처리한다.
• modules_children 리스트를 생성하여 모듈을 자식 단위로 나눈 후, 이를 modules_to_treat에 다시 추가한다.
• 그리고 다시 최대 레이어 크기를 업데이트하여, 새로운 레이어 크기와 메모리 사용량을 반영한다.

# Case 2, it fits! We're not entirely out of the wood though, because we may have some tied parameters.
elif len(tied_params) > 0:
    # First locate all tied modules
    tied_module_names = []
    tied_modules = []
    for tied_param in tied_params:
        tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n in tied_param][0]
        tied_module_names.append(modules_to_treat[tied_module_index][0])
        tied_modules.append(modules_to_treat[tied_module_index][1])
    if verbose:
        print(
            f"  It looks like {name} is going to fit on {devices[current_device]} but we have tied "
            f"parameters to account for.\n  - Names {tied_params}\n  - Module names {tied_module_names}"
        )

    # Let's see if it all fits first
    module_size_with_ties = module_size
    for tied_param, tied_module_name in zip(tied_params, tied_module_names):
        module_size_with_ties += module_sizes[tied_module_name] - module_sizes[tied_param]

    if current_max_size is None or current_memory_used + module_size_with_ties <= current_max_size:
        # We really really fit!
        if verbose:
            print(f"Putting {name} and {tied_module_names} on {devices[current_device]}.")
        current_memory_used += module_size_with_ties
        device_map[name] = devices[current_device]
        for tied_module_name in tied_module_names:
            if tied_module_name in [m[0] for m in modules_to_treat]:
                # The module may have been removed by a previous iteration of this loop.
                tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name][
                    0
                ]
                modules_to_treat.pop(tied_module_index)
            device_map[tied_module_name] = devices[current_device]

        if not offload_buffers and isinstance(module, nn.Module):
            current_buffer_size = compute_module_total_buffer_size(
                module, dtype=dtype, special_dtypes=special_dtypes
            )
            device_buffer_sizes[device] = device_buffer_sizes.get(device, 0) + current_buffer_size

    else:
        # We don't fit with the tied modules. Next question is: can we split one of the tied modules to make it
        # smaller or do we need to go on the next device?
        if verbose:
            print(
                f"Not enough space on {devices[current_device]} to put {name} and {tied_module_names} (space "
                f"available {current_max_size - current_memory_used}, needed size {module_size_with_ties})."
            )
        split_happened = False
        for tied_module_name, tied_module in zip(tied_module_names, tied_modules):
            tied_module_children = list(tied_module.named_children())
            if len(tied_module_children) == 0 or tied_module.__class__.__name__ in no_split_module_classes:
                # can't break this one.
                continue

            if verbose:
                print(f"Splitting {tied_module_name}.")
            tied_module_children = list(tied_module.named_parameters(recurse=False)) + tied_module_children
            tied_module_children = [(f"{tied_module_name}.{n}", v) for n, v in tied_module_children]
            tied_module_index = [i for i, (n, _) in enumerate(modules_to_treat) if n == tied_module_name][0]

            modules_to_treat = (
                [(name, module)]
                + modules_to_treat[:tied_module_index]
                + tied_module_children
                + modules_to_treat[tied_module_index + 1 :]
            )
            # Update the max layer size.
            max_layer_size, max_layer_names = get_max_layer_size(
                [(n, m) for n, m in modules_to_treat if isinstance(m, torch.nn.Module)],
                module_sizes,
                no_split_module_classes,
            )
            split_happened = True
            break

        if not split_happened:
            # If the tied module is not split, we go to the next device
            if verbose:
                print("None of the tied module can be split, going to the next device.")

            device_memory_used[device] = current_memory_used + current_memory_reserved
            current_device += 1
            modules_to_treat = [(name, module)] + modules_to_treat
            current_memory_used = 0

Case 2: 할당할 수 있는 공간이 충분한 경우

이 코드를 눈여겨 보자. 이 코드는 모델의 일부 모듈과 tied parameters를 고려하여 메모리 사용량을 최적화하면서 여러 디바이스에 효율적으로 할당하는 과정을 처리한다.

• 연관된 tied_modules 검색
• tied_params에서 공유된 파라미터가 속한 다른 모듈들을 찾아서, 이 모듈들도 함께 디바이스에 할당할 수 있는지 확인한다. tied_modules는 이러한 파라미터를 공유하는 모듈들의 이름과 해당 모듈 객체를 저장하는 리스트다.
• 모듈과 tied parameters의 메모리 크기 계산
• module_size_with_ties는 현재 모듈과 tied parameters를 고려한 전체 크기이다. 공유된 파라미터의 크기를 중복으로 계산하지 않기 위해, 공유된 파라미터를 제외한 모듈 크기를 계산한다. (집합에서 중복 제거하는 느낌)
• 디바이스에 할당할 수 있는지 여부 확인
• 현재 디바이스의 가용 메모리(current_max_size) 내에 module_size_with_ties가 들어갈 수 있다면, 모듈과 tied parameters를 해당 디바이스에 할당한다.
• 할당이 성공하면, device_map에 모듈과 해당 디바이스의 매핑을 기록합니다.
• 예를 들어, device_map[“module1.weight”] = “cuda:0”과 같은 식으로 기록될 수 있다.
• 메모리가 부족한 경우
• 만약 module_size_with_ties가 현재 디바이스의 가용 메모리보다 크다면, tied 모듈 중 일부를 분할해서 처리할 수 있는지 확인한다.
• tied_module_children은 tied 모듈을 나눌 수 있는 하위 모듈이나 파라미터를 추적한다. 모듈을 나눌 수 있다면, 이를 분할해서 다시 modules_to_treat 리스트에 추가하고, 메모리 크기를 재계산한다.
• 만약 나눌 수 없는 경우, 다음 디바이스로 넘어가서 할당을 시도한다.

    else:
        if verbose:
            if current_max_size is None:
                print(f"Putting {name} (size={module_size}) on {devices[current_device]}.")
            else:
                print(
                    f"Putting {name} (size={module_size}) on {devices[current_device]} "
                    f"(available={current_max_size - current_memory_used})."
                )
        current_memory_used += module_size
        device_memory_used[device] = current_memory_used + current_memory_reserved
        device_map[name] = devices[current_device]

        if not offload_buffers and isinstance(module, nn.Module):
            current_buffer_size = compute_module_total_buffer_size(
                module, dtype=dtype, special_dtypes=special_dtypes
            )
            device_buffer_sizes[device] = device_buffer_sizes.get(device, 0) + current_buffer_size

if clean_result:
    device_map = clean_device_map(device_map)

non_gpu_buffer_size = device_buffer_sizes.get("cpu", 0) + device_buffer_sizes.get("disk", 0)
if non_gpu_buffer_size > 0 and not offload_buffers:
    is_buffer_fit_any_gpu = False
    for gpu_device, gpu_max_memory in max_memory.items():
        if gpu_device == "cpu" or gpu_device == "disk":
            continue

        if not is_buffer_fit_any_gpu:
            gpu_memory_used = device_memory_used.get(gpu_device, 0)

            if gpu_max_memory >= non_gpu_buffer_size + gpu_memory_used:
                is_buffer_fit_any_gpu = True

    if len(gpus) > 0 and not is_buffer_fit_any_gpu:
        warnings.warn(
            f"Current model requires {non_gpu_buffer_size} bytes of buffer for offloaded layers, which seems does "
            f"not fit any GPU's remaining memory. If you are experiencing a OOM later, please consider using "
            f"offload_buffers=True."
        )

return device_map

마지막으로 else 블록에서

• module_size와 현재 디바이스의 가용 메모리(current_max_size - current_memory_used)를 비교하여, 해당 모듈을 특정 디바이스에 할당한다.
• device_map[name] = devices[current_device]는 해당 모듈이 어느 디바이스에 할당되었는지를 기록하는 딕셔너리이다.
  (ex. device_map["model.layers.1"] = “cuda:0”)
• current_memory_used는 현재 디바이스에 할당된 메모리의 양을 추적하고, 새로운 모듈이 할당될 때마다 메모리 사용량을 업데이트한다.

이후, device_map을 다듬고 최종적으로 device_map을 반환한다.

이러한 내부 과정을 통해 우리가 이런 device map을 볼 수 있었던 것이다.

2–4. TP 정리

이렇게 Tensor Parallelism에 대해서 알아보았다.

정리를 해보면, TP는 모델 파라미터를 분할하여 서로 다른 GPU에 할당하는 방식이다. 이를 위해 column-row 순서의 parallelism을 사용하며, vocab parallel embedding 등의 기법을 통해 다양한 최적화 기법을 보였다.

또한 accelerate 라이브러리에서 device_map을 통해 tensor parallelism이 어떻게 적용되는지에 관한 코드를 분석해보았다.

3. 리뷰

최종적으로 PP와 TP를 리뷰해보았다. 주로 Model Parallelism이라고 하면, Pipeline Parallelism에 관한 얘기는 없고, Tensor Parallelism을 지칭하는 것 같다. 내가 보기에도 그리고 사용하기에도 Tensor Parallelism이 편한 것 같지만, multi-dimension 분산 처리를 하면 PP도 사용된다고 한다.

다음에는 Deepspeed 리뷰로 분산 처리를 마무리 짓겠다 ! 많관부 !