Skip to content

[RNE Rewrite] Add LFM2.5 ColBERT (late-interaction multi-vector) embedding pipeline #1306

Description

@msluszniak

Summary

Add a LFM2.5-ColBERT-350M late-interaction (multi-vector) embedding pipeline under the new flow. Follow-up to the single-vector embeddings work in #1292 / #1247.

Why it isn't covered by the text embedder

ColBERT is a late-interaction model: forward returns a per-token matrix [numTokens, 128] (Linear 1024→128 + per-token L2-norm), not a single pooled vector. Relevance is scored with MaxSim (for each query token, the max cosine against all document tokens, summed) rather than a single cosine. So createTextEmbedder / useTextEmbedder cannot represent it — it needs its own pipeline.

The export scripts already exist (exporter_models/lfm2_5_colbert_350m, xnnpack + mlx) and an HF repo (react-native-executorch-lfm2.5-colbert-350m, currently only v0.9.0 / old flow).

Scope

Export (export-scripts, lfm2_5_colbert_350m)

  • Port both backends to the new flow (dynamic seq + get_dynamic_dims_forward companion), matching the embedding export
  • MLX must be built with the software-mansion-labs/executorch fork + MLX-native quantization — a pip/torchao-built MLX delegate fails at execute with Error::InvalidProgram (same root cause as the [RNE Rewrite] Add text and image embeddings pipelines #1292 MLX fix; the MLX backend is version-locked to the runtime MLXBackend.xcframework)
  • Publish xnnpack + mlx to HF react-native-executorch-lfm2.5-colbert-350m v0.10.0

Library

  • New task createColbertEmbedder + useColbertEmbedder hook returning per-token embeddings [numTokens, 128]
  • MaxSim late-interaction scoring helper (query vs. document token matrices)
  • models.ts entries (xnnpack 8da4w + mlx int4) + index exports; download-libs already enables the MLX backend for textEmbeddings

Demo (nlp)

  • Late-interaction retrieval screen (or extend Text Embeddings) using MaxSim; MLX chip disabled on Android (MLX is iOS-only)

Notes

  • The exporter NOTES document that XNNPACK lowering bakes the dynamic seq dim to static [1,512] / [1,512,128] on the current ET build; decide between static + runtime padding vs. a build that preserves the dynamic seq (the embedding MLX variant shipped dynamic; both static+pad and dynamic were verified on device).
  • Related: standalone on-device vector store [RNE Rewrite] Add a small standalone on-device vector store #1303.

Metadata

Metadata

Assignees

No one assigned

    Labels

    improvementPRs or issues focused on improvements in the current codebaserefactoring

    Type

    Fields

    No fields configured for Task.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions