File size: 5,283 Bytes
2e1f688
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
448c45a
 
49f49ab
 
 
2e1f688
49f49ab
 
 
 
 
2e1f688
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c0599b9
 
 
 
2e1f688
448c45a
 
 
49f49ab
 
 
 
 
 
 
 
 
 
 
 
 
 
448c45a
 
 
2e1f688
 
448c45a
2e1f688
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
---
license: mit
tags:
- executorch
- xnnpack
- pte
- on-device
- text-ranking
- text-classification
base_model:
- BAAI/bge-reranker-base
---
# bge-reranker-base — ExecuTorch

A cross-encoder reranker: a query and one document in, one relevance score out. The second
stage of on-device retrieval — an embedding model fetches candidates cheaply, this reads
each candidate together with the query and scores it properly.

- **Source**: BAAI/bge-reranker-base — 278M parameters, 12 XLM-RoBERTa layers, hidden 768, 250k vocabulary
- **License**: MIT
- **Input**: `input_ids` and `attention_mask`, both `[1, 512]` int64
- **Output**: `[1, 1]` fp32 — the raw logit. `sigmoid(x)` maps it to 0..1 and does not
  change the ordering.

## Variants

| build | file | size (MB) | worst score error vs eager | Mac median (ms)* | backend takes |
|---|---|---|---|---|---|
| fp32 | `rerank_bge_reranker_base_xnnpack_fp32.pte` | 1112.3 | 0.0000 logits | 54.9 | 78.3% |
| fp16 | `rerank_bge_reranker_base_xnnpack_fp16.pte` | 556.4 | 0.0067 logits | 95.6 | 67.8% |
| Core ML (fp16, iOS) | `rerank_bge_reranker_base_coreml_all.pte` | 557.7 | 0.0419 logits | 19.9 | 100.0% |

\*Mac arm64, one query-document pair at 512 tokens, **fastest of five medians of ten** — a
reference point for relative cost, not a device number. The host shares its cores with
other work, and a single median does not survive that; contention only ever adds time, so
the fastest repetition is the one that means something. PyTorch eager fp32, measured the
same way: 62.3 ms.

Correlation is not reported because it cannot be: the output is a single number, and the
correlation of a one-element vector is undefined. The column above is the error in the
units the model is used in — logits — over 6 real query-document pairs, and every build
listed reproduces eager's ranking order exactly.

## What it does, on the shipped fp32 build

Query: *"How many people live in Berlin?"*

| rank | score | document |
|---|---|---|
| 1 | +10.308 | ベルリンの人口はおよそ350万人です。 |
| 2 | +10.302 | In 2019 the city recorded 3.7 million residents within its metropolitan area. |
| 3 | +9.940 | Berlin has a population of 3,520,031 registered inhabitants in an area of 891.82 km². |
| 4 | -2.708 | The capital of France is Paris, a city of about 2.1 million people. |
| 5 | -6.198 | Berlin is well known for its museums, its nightlife and its history. |
| 6 | -10.194 | Water boils at 100 degrees Celsius at sea level. |

The narrowest gap between adjacent ranks here is 0.0057 logits — the top two both answer the question, so their order is a coin toss and a build that swapped them would not be wrong.

## It ranks across languages

The candidate list above includes a Japanese passage that answers the English query. This
model puts it **first** at +10.31; ms-marco-MiniLM-L6, the English-only reranker on
this shelf, scores the same passage -10.96 and puts it fifth of
6. That is what the 250k-token vocabulary buys, and it is also why this file is
12 times larger.

## The attention is eager, and that is the faster export

`F.scaled_dot_product_attention` does not survive export as one operation. The edge
dialect lowers it through `_safe_softmax`, whose guard against a row with no unmasked key
at all leaves **11 operations XNNPACK cannot take, in every attention
block**`scalar_tensor`, `where`, `mul.Scalar`, `logical_not`, `eq`, `full_like`, `any.dim`. Each one cuts the subgraph in two.

The switch is `attn_implementation="eager"`: transformers then builds the mask
itself, as `torch.finfo(dtype).min`, instead of handing `F.sdpa` a **boolean** mask
for PyTorch to fill with `-inf`.

The guard is emitted whether or not it can ever fire, and here it cannot: it triggers only
on `-inf`, and this arm never produces one. So the two differ only about rows that have no
unmasked key at all — sdpa zeroes them, this one gives them a uniform row — and those are
padding rows, which the pooling discards and which every real query row masks out anyway.
Measured with all but eight positions masked, as adversarial as this shape gets, the two
graphs agree to 1.1e-05.

XNNPACK fp32 goes from **62.5% to 78.3%** delegated.

## Not shipped

- **int8 (dynamic) is not shipped**: at 856.1 MB it is larger than the fp16 build's 556.4 MB, and its score error is 0.7316 logits. Dynamic int8 quantizes the linear weights and leaves the token embedding table in fp32, while fp16 halves that table too. The table here is 768 MB of a 1112 MB model, and the arithmetic says int8 only comes out smaller when the table is under a third of the weights (371 MB) — measured on ten models on this shelf, the rule called all ten correctly.

## Verification

```bash
python convert/export_rerank.py bge_reranker_base
python convert/check_rerank.py bge_reranker_base fp32
```

The check has two halves. One is agreement with the model run in eager, in logits and in
ranking order. The other is that the ranking is useful at all: the passage that answers the
question has to outscore a passage about the same subject that does not — agreement alone
would pass a build that ranked by document length in both arms.

(conversion scripts: [executorch-models](https://github.com/john-rocky/executorch-models))