MPE preprint version

[FarmersWrap]

Summary

This PR adds Multi-Prefix Embedding (MPE) support for long-document dense retrieval in Tevatron, together with reproduction scripts for the preprint experiments.

The main idea is to represent a long passage with multiple prefix embeddings extracted at EOS-separated chunk boundaries, then score query-document relevance with MaxSim over the passage prefixes.

Main changes

• Add chunking-related data arguments:

  • passage_chunk_size
  • passage_chunk_size_range
  • passage_chunk_size_variable
  • passage_chunk_independent
  • encode_use_pre_chunked

• Extend training and encoding collators to support:

  • fixed-size chunking,
  • deterministic random chunking,
  • pre-chunked input,
  • independent chunk encoding,
  • EOS position tracking for prefix embedding extraction.

• Extend DenseModel to support chunked passage representations:

  • extract passage embeddings at EOS positions,
  • return [num_passages, max_chunks, hidden_dim] plus chunk_mask,
  • support independent chunk encoding,
  • synchronize chunk dimensions across DDP ranks.

• Update trainer logic to:

  • pass EOS positions / chunk counts into the model,
  • gather chunked representations safely under DDP,
  • compute document-level MaxSim similarity for contrastive training.

• Update search logic to:

  • detect chunked passage embeddings,
  • retrieve over chunk embeddings,
  • aggregate chunk-level scores back to document-level scores using MaxSim.

• Add examples/mpe reproduction scripts for:

  • data preparation,
  • MLDR-EN training,
  • MLDR-EN evaluation,
  • BrowseComp-Plus evaluation,
  • LongEmbed evaluation,
  • result collection.

Methods covered by the reproduction scripts

The scripts compare:

• Single-Vector
• MaxP
• MaxP-Train
• MPE Fixed-64
• MPE-Rand-32to1024

Notes for reviewers

This PR touches both core retriever code and paper reproduction scripts. The core change is the chunked passage representation path; the scripts under examples/mpe are included to reproduce the preprint results.

Areas worth reviewing carefully:

• Whether MPE-specific code paths are sufficiently gated and do not affect standard dense retrieval.
• Whether debug prints/logging should be removed or downgraded.
• Whether the DDP shape synchronization and gather logic is robust.
• Whether random chunking behavior should be further documented.
• Whether reproduction scripts should remain in this PR or be split from the core library changes.

AI generated

[MXueguang]

Thanks — multi-prefix embeddings / MaxSim over chunk boundaries is a feature we'd like to land, and moving the scripts under examples/mpe/ is already an improvement over #211. But there are several blockers, one of them urgent, so I'm requesting changes.

Urgent: decrypted BrowseComp-Plus queries committed in plaintext
examples/BrowseComp-Plus/topics-qrels/queries.tsv contains 830 decrypted BC+ queries. BrowseComp-Plus deliberately distributes queries encrypted to prevent web-crawl/training contamination; committing the plaintext to a public repo defeats that. Since the file is in the branch history, deleting it in a follow-up commit isn't enough — the branch needs to be rewritten (or recreated) without it ever having been committed.

Silently breaks SPLADE training (clean-merge semantic conflict)
The PR moves loss computation out of EncoderModel.forward (which now returns loss=None during training) into TevatronTrainer.compute_loss. But SpladeTrainer.compute_loss on current main (from 84433fb) reads output.loss from the model — after this merge that's None, and git merges the file cleanly so nothing flags it. It also bypasses the compute_similarity()/compute_loss() override contract that EncoderModel subclasses rely on. Please keep the forward-loss contract intact — e.g. route MaxSim through a compute_similarity override in a chunk-aware subclass — so SPLADE/GradCache/subclass paths are untouched.

Not gated even when the feature is off
With all chunk flags at defaults, training still goes through the rewritten forward + relocated trainer loss, the restructured TrainCollator.__call__ (whose no-chunk fallback hardcodes add_special_tokens=True, reverting main's add_special_tokens=self.data_args.add_special_tokens), and the unconditional tokenizer.eos_token_id = tokenizer.pad_token_id mutation in both drivers (in train.py it runs before the pad-None fallback, so Llama/Mistral-style tokenizers end up with pad = eos = None). "MPE off" must mean "behavior unchanged".

Other issues

• trust_remote_code=True added to all three load_dataset calls in dataset.py — a security-relevant default change unrelated to the feature; please drop it.
• from rich import print at the top of collator.py and driver/encode.py (rich is not a Tevatron dependency → ImportError for all retriever users), ~10 debug print()s, a bare except:, a 40-line commented-out logging block, and module-level torch.set_printoptions.
• The tests from #211 didn't make it into this PR, so this version is untested — please bring them over (hermetic, no Hub downloads in fixtures).
• The branch predates the July Megatron/SPLADE merge; collator.py has two hard textual conflicts with main (lazy multimodal imports, tokenization block).

Suggested path forward: rebase/re-implement on current main with the feature strictly opt-in, restore the forward-loss contract, drop queries.tsv from history, port the tests, and close #211 as superseded so this PR is the single vehicle. Happy to review a reworked version — the underlying feature is a good fit for Tevatron.

[FarmersWrap]

Thanks for the thorough review. The urgent item is addressed: I've rewritten the branch history so examples/BrowseComp-Plus/topics-qrels/queries.tsv never existed on it — the commit that introduced the file was pruned entirely, and no commit on the branch contains it anymore (verified by scanning every tree in the branch history). A .gitignore entry now prevents it from being re-committed; the README keeps the instructions for regenerating it locally via decrypt_dataset.py --generate-tsv.

One caveat: the pre-rewrite commits may remain reachable by SHA on GitHub until the dangling objects are purged — I'll contact GitHub support about purging them from my fork.

I'm now working through the remaining items: restoring the forward-loss contract (routing MaxSim through a chunk-aware DenseModel subclass so SPLADE/GradCache paths are untouched), making the feature strictly opt-in with defaults byte-identical to main, dropping trust_remote_code, removing the debug/rich leftovers, porting the tests as hermetic (no Hub downloads), and rebasing onto current main. I'll push the reworked branch to this PR and close #211 as superseded.

[FarmersWrap]

Superseded by #229, which was opened from a branch whose history never contained the decrypted queries.tsv.