Fix encoder-mask reuse in uncached multi-vector losses by Afloat16 · Pull Request #4127 · huggingface/sentence-transformers · GitHub
Skip to content

Fix encoder-mask reuse in uncached multi-vector losses - #4127

Merged
tomaarsen merged 2 commits into
huggingface:mainfrom
Afloat16:fix/mve-loss-feature-isolation-20261005
Oct 14, 2026
Merged

tomaarsen merged 2 commits into
huggingface:mainfrom
Afloat16:fix/mve-loss-feature-isolation-20261005

Conversation

@Afloat16

@Afloat16 Afloat16 commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Preserve encoder input masks across calls to the three uncached multi-vector losses by shallow-copying each input column dictionary before encoding. The returned scoring masks and tensor storage remain unchanged.

Identified through code analysis and reproduced with a composed multi-vector training objective: with classic query expansion (attend=False), expansion positions are excluded from encoder attention but included in MaxSim. Writing that scoring mask back into the caller's features changes encoder attention for a later loss term. Related dense-mask work in #3944 does not isolate these loss inputs.

The regression now uses the concise three-loss case suggested by the maintainer in the existing test_misc.py module.

Validation

Reproducer:

python -m pytest tests/multi_vector_encoder/losses/test_misc.py -k test_losses_preserve_encoder_masks -q

Fresh Linux CPU checks on Python 3.12.14, Torch 2.14.1+cpu, Transformers 5.19.0, NumPy 2.5.4:

  • All three new cases fail against upstream production files at 4a3b5cd6ec718e421f57e824a41ed3fd99595df6.
  • python -m pytest tests/multi_vector_encoder/losses/test_misc.py -q --tb=short: 48 passed, no skips.
  • Repository-pinned Ruff 0.15.6 lint/format on changed Python files, compilation and git diff --check: passed.

The current production expressions are unchanged from the previously validated fix; this follow-up simplifies regression coverage and removes comments. The full repository suite, distributed/GPU execution, native pretrained VLM/Flash Attention and inference-provider/backend checks were not run in this follow-up. No performance improvement is claimed.

The previous-head Python 3.11/3.12 CI failures occur before test collection because pyarrow requires NumPy >=2 but those environments have 1.26.4. Updated-head hosted CI remains pending; no all-CI pass is claimed.

@tomaarsen

Copy link
Copy Markdown
Member

@Afloat16

Copy link
Copy Markdown
Contributor Author

Thanks, Tom — I pushed your concise three-loss regression into test_misc.py, removed the separate larger test module, and removed the three production comments. The shallow-copy fix is unchanged.

Fresh CPU validation (Python 3.12.14, Torch 2.14.1+cpu, Transformers 5.19.0, NumPy 2.5.4):

  • The three new cases fail against the upstream production files at 4a3b5cd.
  • The complete native test_misc.py module passes: 48 tests, no skips.
  • Pinned Ruff 0.15.6 lint/format, compilation and git diff --check pass.

At the previous head 09162bd, the four Python 3.11/3.12 CI failures stop before collection with pyarrow requires NumPy 2.0 or newer, found 1.26.4. The six other checks passed. The updated-head CI still needs to run; no full-suite or all-CI pass is claimed.

@Afloat16
Afloat16 marked this pull request as ready for review October 14, 2026 13:35
@tomaarsen
tomaarsen merged commit 489a4bd into huggingface:main Oct 14, 2026
12 checks passed
@tomaarsen

Copy link
Copy Markdown
Member

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.

2 participants