Modules

sentence_transformers.multi_vector_encoder.modules defines the building blocks specific to multi-vector models. Combined with the shared backbone in sentence_transformers.base.modules (see Base > Modules), they make up the standard ColBERT-style stack: Transformer -> Dense -> MultiVectorMask -> Normalize.

See also Training Overview.

MultiVectorMask

class sentence_transformers.multi_vector_encoder.modules.MultiVectorMask(skiplist_words: list[str] | None = None, *, skiplist_tasks: str | list[str] | None = None, keep_only_token_ids: list[int] | None = None)[source]

Module that overwrites features["attention_mask"] with the per-row scoring mask for late-interaction (ColBERT-style) models.

Place this at the end of the module sequence in a MultiVectorEncoder. Reads task from forward kwargs and features["query_expansion_positions"] (a (B, T) mask set by the Transformer during preprocess when query expansion is on) to decide:

  • Real tokens (the tokenizer’s attention_mask) count. For task="query" with query expansion active during preprocess, the expansion positions count too. This is the ColBERT trick: expansion positions contribute to MaxSim even if the Transformer’s attention didn’t see them.

  • For tasks listed in skiplist_tasks (documents only by default, where “no task” counts as document): tokens whose IDs are in the skiplist are dropped.

  • For any non-query task, when keep_only_token_ids is set the mask additionally restricts to those IDs (typically the image-patch token id for ColPali-style image documents, roughly halving index storage by zeroing out text-prefix token embeddings).

  • When the batch has no input_ids (e.g. raw image tensors with no text), the attention_mask is used unchanged.

Reusing the attention_mask key means downstream consumers (encode(), losses, and any future module that respects attention_mask, e.g. Pooling in a hybrid setup) just work.

Parameters:

  • skiplist_words – Tokens to drop from scoring. Defaults to [] (no skiplist). Pass list(string.punctuation) to match the original PyLate / Stanford-NLP ColBERT behaviour of skipping punctuation, or any other custom list. The listed tokens stop matching at all on the skiplist_tasks side (documents by default): a domain choice that suits English prose but not symbol-heavy retrieval such as code. Legacy PyLate / Stanford-NLP loaders apply string.punctuation automatically so existing saved checkpoints keep their historical behaviour. Set at construction: changing it on a built model additionally requires calling resolve_with_tokenizer() with the model’s tokenizer.

  • skiplist_tasks – Task names the skiplist applies to. A single string counts as a one-element list. Defaults to ["document"] (ColBERT-style: only document tokens are dropped), with inputs encoded without a task counting as "document". Use ["query", "document"] for variants that also drop query-side skiplist tokens from scoring.

  • keep_only_token_ids – Allowlist of token IDs to keep in document scoring. Defaults to None (no allowlist: every non-skiplisted real token is scored). Set this to the model’s image-patch token id (processor.tokenizer.convert_tokens_to_ids(processor.image_token), which works for both the ColPali and ColQwen2 families, unlike image_token_id) to reproduce colpali-engine’s mask_non_image_embeddings=True behaviour: only image patch embeddings contribute to MaxSim, roughly halving the document index size. The allowlist is applied in addition to the skiplist. input_ids must be present for it to take effect.

BaseTokenPooling

class sentence_transformers.multi_vector_encoder.modules.BaseTokenPooling(*, tasks: str | list[str] | None = None)[source]

Abstract base for token pooling strategies. Subclasses implement pool_one() (per-sample pooling) and set config_keys for save/load. See the module docstring for the three ways to apply a pooling (pipeline module, per-call token_pooling= kwarg, standalone pool()).

Parameters:

tasks – Task names this pooling applies to. A single string counts as a one-element list. Inputs encoded for any other task pass through unchanged. Defaults to ["document"] (ColBERT-style: only compress the document index). Use ["query", "document"] for strategies that compress queries too, e.g. CRISP-style fixed-k clustering. Inputs encoded without a task count as "document".

forward(features: dict[str, Tensor], task: str | None = None) → dict[str, Tensor][source]

Forward pass of the module. This method should be overridden by subclasses to implement the specific behavior of the module.

The forward method takes a dictionary of features as input and returns a dictionary of features as output. The keys in the features dictionary depend on the position of the module in the model pipeline, as the features dictionary is passed from one module to the next. Common keys in the features dictionary are:

  • input_ids: The input IDs of the tokens in the input text.

  • attention_mask: The attention mask for the input tokens.

  • token_type_ids: The token type IDs for the input tokens.

  • token_embeddings: The token embeddings for the input tokens.

  • sentence_embedding: The sentence embedding for the input text, i.e. pooled token embeddings.

Optionally, the forward method can accept additional keyword arguments (**kwargs) that can be used to pass additional information from model.encode to this module.

Parameters:

  • features (dict[str, torch.Tensor | Any]) – A dictionary of features to be processed by the module.

  • **kwargs – Additional keyword arguments that can be used to pass additional information from model.encode.

Returns:

A dictionary of features after processing by the module.

Return type:

dict[str, torch.Tensor | Any]

pool(embeddings: list[Tensor], *, task: str | None = None, attention_mask: Tensor | None = None, padding_side: str = 'right') → list[Tensor][source]
pool(embeddings: Tensor, *, task: str | None = None, attention_mask: Tensor | None = None, padding_side: str = 'right') → Tensor

Apply the pool strategy to a list of 2D or a 3D padded tensor.

Parameters:

  • embeddings – A list of (t_i, D) tensors, or a 3D (B, T, D) padded tensor.

  • task – Task the embeddings were encoded for. If it is not in tasks, the input is returned unchanged. None (default) counts as "document", so standalone document compression needs no extra argument. encode() forwards its task here for per-call pooling.

  • attention_mask – Optional (B, T) boolean mask for the 3D case. Required unless the padding is zero-valued (in which case the input boundary is detected by padding_side, but a real token whose embedding is exactly zero at the boundary of the content region will be clipped).

  • padding_side – "left" or "right". Used to detect the input boundary when no mask is passed, and always used to re-pad the output when the input was 3D.

Returns:

List of (num_out, D) tensors when the input was a list. A padded 3D tensor when the input was 3D (padded on padding_side).

HierarchicalTokenPooling

class sentence_transformers.multi_vector_encoder.modules.HierarchicalTokenPooling(pool_factor: int = 1, *, num_protected_tokens: int = 1, tasks: str | list[str] | None = None)[source]

Ward-linkage hierarchical clustering on cosine similarity. Keeps the first num_protected_tokens untouched (typically the [CLS]), clusters the rest into num_tokens // pool_factor groups, and replaces each cluster with its mean.

Assumes L2-normalized embeddings (place after a Normalize in the pipeline).

Reference compatibility: this implementation matches PyLate after its condensed-Ward fix (PyLate > 1.3.4). It intentionally does NOT reproduce released PyLate <= 1.3.4 (square-matrix linkage quirk, protected tokens re-appended at the end) nor colpali-engine’s HierarchicalTokenPooler (different linkage input, cluster means re-normalized to unit norm, no protected-token concept), so indexes pooled with those tools will not be byte-reproducible. For the closest colpali-engine setup, pass num_protected_tokens=0 and re-normalize afterwards.

Parameters:

  • pool_factor – Keep roughly 1 / pool_factor of each document’s tokens. 1 (default) disables pooling (the module becomes a no-op).

  • num_protected_tokens – Leading tokens excluded from pooling (typically [CLS]). Default 1. colpali-engine has no protected-token concept: use 0 when matching its setup.

  • tasks – Task names this pooling applies to. Defaults to ["document"] (only compress documents).

LambdaTokenPooling

class sentence_transformers.multi_vector_encoder.modules.LambdaTokenPooling(pool_fn: Callable[[Tensor], Tensor], *, tasks: str | list[str] | None = None)[source]

User-supplied pool function applied per-sample.

Cannot be baked into a saved checkpoint (a Python callable isn’t serializable), and will not round-trip through MultiVectorEncoder.encode()’s multi-process path (pool= arg) if pool_fn is a lambda or nested function. Use in the pipeline for experimentation, or standalone / per-call for ad-hoc compression.

Parameters:

pool_fn – A callable that takes a (num_tokens, dim) tensor and returns a (num_out, dim) tensor. Applied once per document in a batch.

Example:

def halve(emb: Tensor) -> Tensor:
    # Average consecutive pairs of tokens, dropping the tail if odd.
    n = emb.size(0)
    return emb[: n - n % 2].view(n // 2, 2, -1).mean(dim=1)

pooling = LambdaTokenPooling(pool_fn=halve)
pooled = pooling.pool(document_embeddings)