Interpretability

sentence_transformers.multi_vector_encoder.interpretability provides a per-query-token MaxSim heatmap utility for ColPali-style image documents. Useful for spot-checking which patch positions in an image contribute most to a given query.

Heatmaps

maxsim_heatmap is the one-shot entry point. get_n_patches supplies its n_patches argument, and real_query_token_slice selects the query tokens worth visualizing.

sentence_transformers.multi_vector_encoder.interpretability.maxsim_heatmap(image: PILImage | str, query_embedding: Tensor, image_embedding: Tensor, n_patches: tuple[int, int], image_mask: Tensor | None = None, aggregate: Literal['sum', 'amax', 'none'] = 'sum', alpha: float = 0.5) → PILImage | list[PILImage][source]

One-shot MaxSim heatmap for a (query, image-document) pair.

Parameters:

  • image – PIL image, URL, or local file path.

  • query_embedding – (Qt, D) per-token query embeddings.

  • image_embedding – (Dt, D) per-token image-document embeddings.

  • n_patches – (n_cols, n_rows) image grid shape, as returned by get_n_patches().

  • image_mask – optional mask filtering image_embedding to image patches only.

  • aggregate – "sum" adds every query token’s similarity at each patch, giving an aggregate similarity heatmap rather than an attribution of the MaxSim score, which counts only each query token’s single best patch. "amax" shows the strongest single-token match per patch. "none" returns one heatmap per query token as a list.

  • alpha – constant overlay opacity in [0, 1].

Returns:

A single PIL image for aggregate in {"sum", "amax"}, or a list for "none".

sentence_transformers.multi_vector_encoder.interpretability.get_n_patches(model: MultiVectorEncoder, image_size: tuple[int, int]) → tuple[int, int][source]

Infer the (n_patches_x, n_patches_y) = (n_cols, n_rows) image-patch grid that model’s processor produces for an image of image_size ((width, height), i.e. PIL’s image.size). Plugs directly into the n_patches parameter of maxsim_similarity_map() and maxsim_heatmap().

Dynamic grids (Qwen-VL family) are read from the image processor’s own image_grid_thw output for a blank image of that size, fixed square grids (PaliGemma / Gemma3) from processor.image_seq_length. Idefics3 / SmolVLM split-image processors raise NotImplementedError: their sub-patch token order is not a plain row-major grid.

sentence_transformers.multi_vector_encoder.interpretability.real_query_token_slice(model: MultiVectorEncoder, query: str) → slice[source]

Return the slice into encode_query’s output that selects the real content tokens.

Chat-template prefixes (e.g. <bos>), ColBERT query markers (e.g. “[Q] “), suffixes (e.g. <|im_end|>), and MultiVectorEncoder’s query-expansion tokens (<mask> / <pad>) wrap the actual query and carry attention-sink signals that distort heatmap visualisations. Slicing them out keeps only the real tokens:

s = real_query_token_slice(model, query)
query_embedding = model.encode_query([query])[0][s]

Works for both right-padded (PaliGemma) and left-padded (ColQwen2 / ColModernVBert) backbones by comparing the encoded sequences of an empty and the actual query.

Building blocks

maxsim_heatmap composes these two. Use them directly to reach the raw similarity tensor, e.g. for a custom colormap, a matplotlib figure, or an aggregation other than the built-in ones.

sentence_transformers.multi_vector_encoder.interpretability.maxsim_similarity_map(query_embedding: Tensor, image_embedding: Tensor, n_patches: tuple[int, int], image_mask: Tensor | None = None, normalize: bool = False) → Tensor[source]

Per-query-token similarity over a 2D image-patch grid.

Parameters:

  • query_embedding – (Qt, D) per-token query embeddings.

  • image_embedding – (Dt, D) per-token image-document embeddings. Pass image_mask if Dt includes non-image tokens, otherwise Dt == n_patches[0] * n_patches[1].

  • n_patches – (n_cols, n_rows) = (width, height), as returned by get_n_patches(). Patches are row-major.

  • image_mask – boolean mask over the Dt axis (True for image-patch tokens).

  • normalize – rescale each per-query-token map to [0, 1].

Returns:

Tensor of shape (Qt, n_rows, n_cols), float32 regardless of the input dtype.

sentence_transformers.multi_vector_encoder.interpretability.render_similarity_map_on_image(image: PILImage | str, similarity_map: Tensor, alpha: float = 0.5, normalization_range: tuple[float, float] | None = None) → PILImage[source]

Overlay a 2D similarity map onto an image with the mako colormap.

Parameters:

  • image – PIL image, URL, or local file path (loaded via transformers.image_utils.load_image()).

  • similarity_map – (n_rows, n_cols) similarity tensor.

  • alpha – constant overlay opacity in [0, 1].

  • normalization_range – (min, max) for the colour scale. Defaults to the map’s own range. Pass a shared range across multiple per-token maps so they render proportionally.

Returns:

A new RGBA PIL image.