CoreML cross-encoder reranking: on-device RAG on iPhone

Meta description: Build a two-stage on-device RAG pipeline on iPhone using CoreML — INT8-quantized bi-encoder retrieval, cross-encoder reranking, and A16/A17 latency budgets explained.

---

TL;DR

A production-grade on-device RAG pipeline requires two distinct stages: a fast ANN retrieval pass using a quantized bi-encoder, followed by a CoreML cross-encoder that rescores the top-k candidates. The latency budget on A16/A17 chips is roughly 80–120ms end-to-end before users notice lag. INT8 quantization gets you there — but only if you understand the memory ceiling that determines how many candidates you can rerank before thermal throttling kills the session.

---

The two-stage architecture

The pipeline splits into two stages with distinct latency and quality trade-offs. Stage one is retrieval — fast, approximate, and intentionally imprecise. Stage two is reranking — slower, exact, and where quality actually lives.

Query → [Bi-Encoder] → embedding → ANN search (top-k=50)
                                          ↓
                              [Cross-Encoder] → rescore top-k=10
                                          ↓
                                  Final ranked results

The bi-encoder runs once per query. The cross-encoder runs k times. This asymmetry is the entire reason the two-stage model exists — cross-encoders are dramatically more accurate but cannot scale to full corpus retrieval.

---

INT8 quantization: what the numbers show

Most teams get this wrong: they quantize both stages with the same strategy and wonder why quality collapses. The bi-encoder and cross-encoder have different sensitivity profiles.

Model Stage	FP32 Size	INT8 Size	Latency (A17)	Quality Drop
Bi-encoder (MiniLM-L6)	90 MB	24 MB	8 ms	< 0.5% NDCG
Cross-encoder (MiniLM-L12)	180 MB	47 MB	34 ms/candidate	~1.2% NDCG
Cross-encoder (INT4 aggressive)	180 MB	23 MB	19 ms/candidate	~4.8% NDCG

INT8 on the bi-encoder is nearly free — the embedding space is robust to quantization noise. On the cross-encoder, INT8 is acceptable; INT4 starts hurting meaningful recall at the reranking stage. Stay at INT8 for cross-encoders in production.

To verify these regressions in your own pipeline, maintain a small offline eval set of 200–500 query/document pairs with human-labeled relevance judgments. Run pytrec_eval against both the FP32 and quantized model outputs — a 5-minute eval loop catches quality regressions before they ship.

To export with CoreML Tools:

import coremltools as ct

mlmodel = ct.convert(
    traced_model,
    compute_precision=ct.precision.FLOAT16,
    compute_units=ct.ComputeUnit.ALL  # uses ANE + GPU
)
mlmodel.save("CrossEncoder.mlpackage")

Target ct.ComputeUnit.ALL to push attention layers onto the Neural Engine. On A16 and A17, this reduces cross-encoder latency by roughly 40% versus CPU-only execution.

---

KV-cache reuse across candidates

In my experience building production systems, the biggest latency win in cross-encoder scoring is not quantization — it is KV-cache reuse. For a given query, the query-side key/value representations are identical across all k candidates. You compute them once.

// Pseudo-code: cache query KV states
let queryKV = crossEncoder.encodeQuery(queryTokens)

let scores = candidates.map { doc in
    crossEncoder.scoreWithCachedQuery(queryKV, docTokens: doc.tokens)
}

CoreML does not expose KV-cache injection directly for transformer models. You have two options: split the model at the cross-attention boundary and pass cached states manually, or use ONNX Runtime with a custom CoreML execution provider that supports past key values. Go with the split-model approach. It avoids the ONNX Runtime dependency, stays entirely within the Apple toolchain, and is straightforward to debug with Instruments — the integration overhead is a one-time cost, not ongoing complexity.

---

The memory pressure ceiling on A16/A17

The constraint nobody mentions until production: the Neural Engine on A16/A17 shares memory bandwidth with the GPU, CPU, and camera ISP. Under sustained load, the kernel will thermally throttle ANE frequency within 60–90 seconds.

The practical ceiling for reranking:

• k = 10–20 candidates: Safe. Stays below 200ms total, no throttle.
• k = 30–50 candidates: Borderline. Latency climbs to 400–600ms. Sustained sessions trigger throttle at ~45s.
• k > 50: Avoid entirely on-device unless you batch overnight with background processing.

The working set that matters is: (model weights) + (k × sequence_length × hidden_dim × 2 bytes). For a 12-layer cross-encoder scoring 20 candidates of 512 tokens, you are looking at ~310MB active — well within iPhone headroom. At 50 candidates, you are pushing 750MB and competing with the OS.

For interactive sessions, monitor thermal state with NSProcessInfo.processInfo.thermalState and shed load before the system does it for you. For sustained background workloads, schedule reranking via BGProcessingTaskRequest, which the OS grants during charging when thermals are favorable.

---

Three things worth remembering

1. Use INT8 for both stages, not INT4. The 2× size reduction from INT4 is not worth the ~5% NDCG regression on cross-encoders in production RAG workloads. Validate this against your own eval set before committing to a quantization strategy.

2. Implement query-side KV-cache reuse via the split-model approach. Split your cross-encoder at the cross-attention layer, cache query states once, then score all candidates against cached representations. This cuts reranking latency by 30–45% on A16/A17 with no quality cost.

3. Cap candidate count at k=20 for interactive sessions. Beyond that, schedule reranking in a background task using BGProcessingTaskRequest and use NSProcessInfo.thermalState to shed load proactively in real-time contexts.

---

Tags: ios swift mobile architecture