mlboydaisuke/bge-m3-ExecuTorch

bge-m3 — ExecuTorch (dense + sparse + multi-vector, one pass)

Three retrieval signals out of one forward pass. Every other embedding model on this shelf returns a vector; this one returns a vector, a set of per-token lexical weights, and a per-token matrix — and they are meant to be combined.

input_ids, attention_mask  [1, 512] int64
    -> dense    [1, 1024]        the CLS row, L2-normalised
    -> sparse   [1, 512]         one weight per token, masked
    -> colbert  [1, 511, 1024]   one vector per token, L2-normalised, CLS excluded

• Source: BAAI/bge-m3 — 568M parameters, XLM-RoBERTa large, 24 layers, 100+ languages
• License: MIT
• No prefix. Unlike E5 and Qwen3-Embedding on this shelf, bge-m3 wants the text as it is, on both sides.

Using the three heads

Dense is ordinary vector search: cosine against other dense vectors.

Sparse is lexical matching, BM25-shaped. The graph gives one weight per token position; the vocabulary-space vector is one line of indexing in the caller:

weights = {}
for w, t in zip(sparse[0], input_ids[0]):
    if t in (0, 1, 2, 3):            # <s>, <pad>, </s>, <unk> — never scored
        continue
    weights[int(t)] = max(weights.get(int(t), 0.0), float(w))   # max over repeats
score = sum(w * other[t] for t, w in weights.items() if t in other)

That scatter stays outside the graph on purpose: emitting [1, 250002] on every call would be a megabyte of almost entirely zeros to save the caller those six lines.

ColBERT is late interaction: for each query token take its best match among the document tokens, and sum. Row i of the output is token i+1 of the input — the CLS row is dropped before projection, which is what the reference implementation does and is easy to get wrong by one.

Verification

build	file	size (MB)	Mac ms*	dense	colbert	sparse weight shift
XNNPACK fp32	bge_m3_xnnpack_fp32.pte	2271.5	233.3	1.000000	1.000000	0.0000
Core ML	bge_m3_coreml_all.pte	1137.2	64.8	0.999990	0.999976	0.0008
XNNPACK fp16	bge_m3_xnnpack_fp16.pte	1136.3	484.3	0.999999	0.999998	0.0004

*Mac arm64, median of 10, one 512-token sequence — a reference point for relative cost, not a device number. Torch eager fp32 on the same machine is 182.8 ms, so the Core ML build is 2.8x eager, 100% delegated in a single subgraph. XNNPACK fp32 is 63.7% delegated across 100 subgraphs; its fp16 build is slower than fp32 because XNNPACK has no fp16 kernels for this graph and inserts casts.

Dense and colbert are worst-case cosine against the eager model over six sentences; sparse is the largest change to any single token's weight.

The recipe was checked against the authors' implementation before anything was exported. All three heads have a detail that does not throw when wrong — dense is CLS and not mean, colbert drops the CLS row, and sparse_linear is a Linear(1024, 1) giving a scalar per token rather than a projection into vocabulary space. Against FlagEmbedding's BGEM3FlagModel on six sentences:

dense    max_abs_diff  2.645e-07
sparse   max_abs_diff  3.427e-07
colbert  max_abs_diff  4.061e-07

And the published number reproduces. The model card computes a lexical matching score of 0.19554901123046875 between its two example sentences. Running those same sentences through the fp32 .pte and the scatter above gives 0.1955 — which is the only independent check there is on a step that happens outside the graph.

Both retrieval heads separate an answer from an unrelated sentence:

dense    0.6259 answer  vs  0.3625 unrelated
sparse   0.1955 answer  vs  0.0115 unrelated

python convert/check_bge_m3.py fp32 --reference   # against FlagEmbedding
python convert/check_bge_m3.py fp32               # or fp16, int8, coreml

Two decisions worth knowing about

The window is 512, not 8192. bge-m3 accepts 8192 tokens, and the colbert head returns one 1024-vector per token — so an 8192 window would be a 32 MB output on every call for a passage that is almost always shorter. 512 covers an ordinary passage; longer input is the caller's chunking problem.

The sparse head is masked in the graph, which upstream does not do. Upstream returns the raw relu and relies on the caller dropping special tokens at scatter time. Measured on one 31-token sentence padded to 512, the fp32 model puts weights of up to 0.196 on padding positions. A caller who forgets to drop them scatters that onto the pad token's vocabulary slot. Zeroing them here changes no score — the scatter discards them either way — and removes a silent trap.

That masking also fixed the measurement. Before it, this build's sparse head read correlation -0.162 against fp32 eager, which looks like a broken head; on the 31 real token positions it was +0.998, and the other 481 were padding neither arm's caller ever reads.

Not shipped

int8 converts and holds — worst head 0.985 — but it comes out at 1363.3 MB against fp16's 1136.3 MB. Dynamic int8 quantises the linear weights and leaves the token embedding table in fp32, and with a 250k vocabulary at 1024 dimensions that table is 1024 MB of the 2271 MB model, 45%. This shelf's rule of thumb: int8 beats fp16 only when the embedding table is under about a third of the weights.

torch.export -> to_edge_transform_and_lower(partitioner) -> .pte (conversion scripts: executorch-models)