mxbai-edge-colbert-v0-32m

mxbai-edge-colbert-v0-32m — LiteRT

mixedbread-ai/mxbai-edge-colbert-v0-32m converted to LiteRT (.tflite) for on-device inference. Mixedbread's edge-sized late-interaction (ColBERT) retriever: one 64-dimensional unit vector per token from a 32M-parameter ModernBERT backbone, scored with MaxSim — fully offline, on CPU.

The projection head (384→768→768→64, folded into a single matrix) and per-token L2 normalization are inside the graph; MaxSim scoring is a few lines host-side (below).

* load + invoke all four signatures, XNNPACK, 8 threads, M4 Max Mac. The fp16 file is bit-exact against PyTorch on every gate below but XNNPACK expands fp16 weights to fp32 per signature subgraph, so it peaks at ~10× its file size; int8 is the on-device recommendation, and its measured retrieval cost is small (below).

Host contract (read this before using the output)

Reimplemented from the PyLate reference stack and verified against the base card's own published example — scores reproduced to all four printed decimals and both printed shapes:

• Lowercase the text first (do_lower_case is set for this model). Skipping this changes the tokenization.
• Pad with token id 50284 ([MASK]) — that is what PyLate pads with. Note config.json's pad_token_id: 50283 is not what the reference stack uses; the vendor's own onnx_config.json says 50284. (With padding masked out, the graph output on real tokens is provably independent of the pad id — but use 50284 anyway.)
• Queries: lowercase, tokenize with max_length=47, then insert [Q] (id 50368) at position 1 (after [CLS]) and a matching 1 in the attention mask. No query expansion — every valid position is scored, punctuation included.
• Documents: same with max_length=511 and [D] (id 50369). Keep vectors at valid (mask=1) positions, then drop vectors whose token is one of the 32 ASCII punctuation characters (the ColBERT skiplist; [CLS]/[D]/[SEP] are kept).
• Score = MaxSim: for each query vector, take the max dot product over the document's kept vectors, then sum over query vectors. Vectors are unit length, so dot product = cosine.

Signatures

Batch-1, right-padded static shapes: input_ids int32 [1, S], attention_mask int32 [1, S] (1 = real token, then 0s), for S in 48 / 128 / 256 / 512 (48 matches the vendor's query length).

Route each text to the smallest signature that fits and keep only the valid positions. Padding is masked inside the graph, so the same text returns bitwise identical vectors through every signature, and pad-region token ids cannot influence the valid positions at all.

Usage (Python)

import string
import numpy as np
from ai_edge_litert.interpreter import Interpreter
from transformers import AutoTokenizer

PAD_ID, Q_ID, D_ID = 50284, 50368, 50369
tok = AutoTokenizer.from_pretrained("mixedbread-ai/mxbai-edge-colbert-v0-32m")
SKIP = {tok.convert_tokens_to_ids(c) for c in string.punctuation}
it = Interpreter(model_path="mxbai-edge-colbert-v0-32m_wi8fc.tflite", num_threads=8)

LENS = sorted(int(n.split("_")[1]) for n in it.get_signature_list())
runners = {s: it.get_signature_runner(f"encode_{s}") for s in LENS}

def encode(text, prefix_id, max_len):
    e = tok(text.lower(), truncation=True, max_length=max_len - 1)["input_ids"]
    ids = [e[0], prefix_id] + e[1:]          # insert [Q]/[D] after [CLS]
    S = next(s for s in LENS if len(ids) <= s)
    x = np.full((1, S), PAD_ID, np.int32)
    m = np.zeros((1, S), np.int32)
    x[0, :len(ids)] = ids
    m[0, :len(ids)] = 1
    v = list(runners[S](input_ids=x, attention_mask=m).values())[0][0][:len(ids)]
    if prefix_id == D_ID:                     # skiplist: drop punctuation vectors
        v = v[[i not in SKIP for i in ids]]
    return v

def maxsim(q, d):
    return float((q @ d.T).max(axis=1).sum())

q = encode("Which planet is known as the Red Planet?", Q_ID, 48)
docs = ["Mars, known for its reddish appearance, is often referred to as the Red Planet.",
        "Venus is often called Earth's twin because of its similar size and proximity."]
print([maxsim(q, encode(d, D_ID, 512)) for d in docs])

Documents longer than 512 tokens must be chunked.

Quality

Four independent checks against the PyTorch fp32 reference.

1. The base card's own usage example (Red Planet query × Venus/Mars/Jupiter/Saturn): the converted fp32 reference reproduces the card's published MaxSim scores [11.2081, 11.5308, 11.4104, 11.4756] to all four printed decimals, and the published embedding shapes (12 query vectors, 18 kept vectors for the first document). Every variant keeps Mars on top — at 32M the margins are ~0.1, so this is a genuinely tight rank-order check.

2. NanoSciFact retrieval (English, 600-doc corpus, 50 claims, MaxSim): PyTorch nDCG@10 0.8644 / recall@5 0.920 / hit@1 0.780; fp16 identical to PyTorch; int8 0.8564 / 0.920 / 0.780 — recall@5 and hit@1 unchanged, −0.008 nDCG@10.

3. Cross-variant retrieval — the deployment shape. Document banks encoded with the PyTorch model, queries with the int8 artifact (index built on a server, queried on device): nDCG@10 0.8578 vs 0.8644 control, recall@5 and hit@1 unchanged.

4. Mechanics: cross-signature outputs bitwise identical (all four signatures, every variant); pad-content invariance exactly 0; all positions finite in fp32 at the longest signature (the sliding-window edge case is guarded inside the graph); per-token norms within 2e-7 of 1.

Corpora are subsampled, so absolute numbers are not comparable to published benchmarks.

Speed

CPU/XNNPACK, median of 10 runs, 75%-full signatures:

A static signature computes all S positions regardless of how many are real, so route each text to the smallest signature that fits.

Conversion

Encoder lane — a direct multi-signature litert_torch trace of the HF ModernBERT backbone (not an LLM export), with the full-attention and sliding-window (±64) masks built by hand inside the traced wrapper. The three PyLate Dense modules (384→768→768→64, all-linear) fold exactly into a single 64×384 projection, verified against the module stack. Gated on: bitwise agreement with the vendor's own mask path, pad-content invariance, all-position finiteness (sliding-window edge rows), cross-signature bitwise agreement, and the base card's published scores and shapes.

Script and full notes: hf-to-litertlm.

License

Apache 2.0, inherited from the base model by Mixedbread.

Modification notice: these files are converted, not original. The weights were exported to LiteRT and quantized (int8 dynamic-range / fp16); the projection head and per-token L2 normalization were folded into the graph. No fine-tuning or weight modification beyond quantization was performed.