README.md

---
license: mit
tags:
- executorch
- xnnpack
- pte
- on-device
- feature-extraction
- sentence-similarity
base_model:
- BAAI/bge-m3
---
# 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](https://huggingface.co/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:

```python
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
```

```bash
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](https://github.com/john-rocky/executorch-models))