Skip to content

Optimize BGE-M3 compute_score with pairwise scoring (#1624) - #1625

Open
Rajan-vish wants to merge 1 commit into
FlagOpen:masterfrom
Rajan-vish:optimize-bge-m3-pairwise-scoring-1624
Open

Rajan-vish wants to merge 1 commit into
FlagOpen:masterfrom
Rajan-vish:optimize-bge-m3-pairwise-scoring-1624

Conversation

@Rajan-vish

@Rajan-vish Rajan-vish commented Nov 11, 2026 •

Copy link
Copy Markdown

Fixes #1624

Summary of changes

When computing scores for requested query-passage pairs in M3Embedder.compute_score_single_device, the embedder previously calculated the full cross-matrix (B x B) for dense, sparse, and ColBERT representations before taking the diagonal. For ColBERT especially, constructing the full B x B token interaction tensor scales memory as O(B^2 * L_q * L_p * d), causing unnecessary OOM errors on larger batches.

This PR:

  • Adds compute_dense_score_pairwise, compute_sparse_score_pairwise, and compute_colbert_score_pairwise in EncoderOnlyEmbedderM3Model. The ColBERT pairwise calculation uses orch.einsum('bin,bjn->bij', q_reps, p_reps) directly on corresponding pairs, reducing intermediate tensor memory to O(B * L_q * L_p * d) (a B-fold reduction).
  • Updates M3Embedder.compute_score_single_device to use these pairwise scoring methods directly when available.
  • Adds unit tests in ests/test_m3_pairwise_scoring.py confirming exact numerical alignment between pairwise scoring and the diagonal of full all-pairs scoring.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

BGE-M3: Explore scoring only the requested query-passage pairs

1 participant