Skip to content

Exclude passage padding from M3 training MaxSim - #1610

Open
Afloat16 wants to merge 1 commit into
FlagOpen:masterfrom
Afloat16:fix/m3-colbert-passage-padding
Open

Afloat16 wants to merge 1 commit into
FlagOpen:masterfrom
Afloat16:fix/m3-colbert-passage-padding

Conversation

@Afloat16

@Afloat16 Afloat16 commented Oct 2, 2026

Copy link
Copy Markdown

Passage padding currently participates in the M3 training ColBERT MaxSim reduction. _colbert_embedding sets padding vectors to zero, but zero is still a candidate: when a query token has negative similarities to every valid passage token, padding wins the maximum and suppresses that token's gradient. The score then depends on the padded passage length.

For example, a native tiny BERT encoder with normalized ColBERT vectors produces a score of 0.547694 for a passage and 0.652896 after adding three padding tokens. The valid passage vectors are identical. The first query token's gradient norm changes from 0.333333 to 0.

The fix passes the passage attention mask through the unified training objective and excludes padding similarities before the maximum. The mask includes CLS, matching the existing query-mask convention; scoring removes its first position because ColBERT vectors already exclude CLS. This preserves valid zero vectors and leaves EOS handling, query-length normalization and temperature unchanged. The optional p_mask argument preserves existing unpadded calls.

Validation uses normal imports of the complete package with actual PyTorch and Transformers:

  • 39 native CPU tests pass with PyTorch 2.14.1+cpu and Transformers 4.57.1.
  • An independent valid-token MaxSim definition matches scores and gradients in BF16, FP32 and FP64, including zero/nonzero padded vectors and valid zero vectors.
  • Real tiny BERT unified training gives identical losses, encoder/sparse/ColBERT gradients and SGD updates for padded and unpadded passages, covering supervised, KL/M3 distillation, self-distillation and sub-batch encoding. List features also propagate their masks correctly.
  • Eight unpadded positional-API controls pass both versions. On the original source, 13 existing-forward padding comparisons fail; 18 additional tests reject the newly introduced p_mask keyword. The total baseline is 31 failures and 8 passes.

These are synthetic-vocabulary, random-weight tiny-model checks; they do not establish pretrained BGE-M3 retrieval quality or GPU/distributed training performance.

Related to #515: multiplying passage padding vectors by zero preserves MaxSim only when the valid maximum is nonnegative. This change addresses training padding, independently of the inference EOS proposal in #1558.

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.

1 participant