RAG embedding quality: replace USE-Lite (100-dim) with a device-tiered embedding + reranking stack

[sagar-develop]

Summary

RAG retrieval sometimes grounds on the wrong document (observed: a query about one source is answered from an unrelated one). Root cause is primarily the embedding model, not the LLM: USE-Lite produces only 100-dimensional vectors, which are too coarse to separate semantically close documents (e.g. car vs health insurance, two similar PDFs). There is also no reranker and the chunker is character-based, so even when the right chunk is retrieved it can be split/diluted.

This issue proposes upgrading the embedding pipeline and making embedding scale with device capability, mirroring the LLM tiering the app already has.

Current state (from code analysis)

• Embedder: Universal Sentence Encoder (USE-Lite), universal_sentence_encoder.tflite (~6 MB) via MediaPipe Tasks TextEmbedder — lib/src/androidMain/kotlin/com/sagar/aicore/MediaPipeEmbeddingEngine.kt.
• Dimension: 100-dim, hard-coded; not exposed by the interface. EmbeddingEngine (lib/src/commonMain/kotlin/com/sagar/aicore/EmbeddingEngine.kt) is symmetric/task-agnostic (embed(text) only — no query vs document distinction).
• Vector index: ObjectBox @HnswIndex(dimensions = 100L, distanceType = COSINE, …) on DocumentChunkEntity (sample-app/.../data/db/Entities.kt). HNSW dimension is a compile-time literal, so changing dims requires a new entity.
• Retrieval: DefaultDocumentRetriever — hybrid vector + BM25 fused via Reciprocal Rank Fusion. RagConfig defaults: chunkSize=500, chunkOverlap=50, relevanceMaxDistance=0.75, vectorPoolSize=30, keywordPoolSize=120, maxContextChars=4000. No per-document cap; no reranker.
• Chunker: TextChunker is a character sliding window (not token-aware).
• Device tiers already exist for LLMs in NativeLmModelCatalog.kt (minDeviceRamMb, ~3.5 GB → 12 GB+), but embedding does not tier — every device uses USE-Lite.
• A planning doc already exists: docs/EMBEDDING_GEMMA_PLAN.md (EmbeddingGemma 256-dim default, USE-Lite low-end fallback, task-aware interface, parallel GemmaChunkEntity, re-index from stored text). This issue extends/formalizes that.

Proposed direction

1. Adopt EmbeddingGemma with Matryoshka (the key lever)

EmbeddingGemma (308M, Gemma-3-based) is purpose-built for on-device RAG: #1 on MTEB under 500M params, <200 MB RAM with QAT, 2048-token context, 100+ languages, and a LiteRT/litert-community build — so it can run on the same LiteRT runtime we already use for the LLM (no new ONNX dependency strictly required; evaluate LiteRT vs ONNX). Crucially it uses Matryoshka Representation Learning (MRL): one model whose output can be truncated to 768 / 512 / 256 / 128 dims with minimal quality loss — perfect for device tiering.

2. Tier the embedding by device capability

Pick the embedding model + truncation dim at first run from minDeviceRamMb, reusing the existing tier logic:

Device tier	Embedder	Dim	Reranker
Entry (~3–4 GB)	USE-Lite (fallback) or EmbeddingGemma @128	100 / 128	none
Mid (~6–7 GB)	EmbeddingGemma	256	none
High / Flagship (~10 GB+)	EmbeddingGemma	512–768	cross-encoder rerank

3. Add task-aware embedding (query vs document)

EmbeddingGemma is asymmetric — it expects task prefixes (verify exact strings against the model card, but roughly task: search result | query: {q} for queries and title: none | text: {chunk} for documents). Extend EmbeddingEngine with an EmbeddingTask { QUERY, DOCUMENT } parameter and apply the right prefix in the retriever vs ingestor. This alone typically lifts retrieval meaningfully.

4. Add a reranker on capable tiers

After fusion, run a lightweight cross-encoder reranker (e.g. ms-marco-MiniLM-L-6-v2 or bge-reranker-v2-m3) over the top ~30–100 candidates to produce a precision-tuned top-k (+5 to +15 NDCG@10 typical). Gate it to High/Flagship tiers to keep low-end latency acceptable.

5. Better, structure-aware chunking

Move from a pure character window toward token-aware / sentence/paragraph-boundary chunking so tables and numeric rows aren't split mid-record. Consider larger chunks now that EmbeddingGemma supports 2048 tokens.

Quick wins (no re-import, ship first)

These were prototyped and reverted; they help even before the embedder swap:

• Per-document cap in DefaultDocumentRetriever so one large source can't fill every top-k slot (maxChunksPerDocument in RagConfig).
• Bump retrieval k (5 → 8), widen vectorPoolSize/keywordPoolSize, raise maxContextChars.

Migration / infra notes

• HNSW dim is fixed per entity → add a new entity per dimension (e.g. GemmaChunkEntity @256) alongside the existing 100-dim one; migrate by re-embedding from stored chunk text (no re-extraction/re-OCR needed).
• Expose dimension and task on EmbeddingEngine so the index and embedder can't silently diverge.
• Add iOS embedding impl (currently none) — EmbeddingGemma LiteRT works cross-platform.

Open questions / risks

• P2P/mesh sync across tiers: two devices with different embedding dims can't share a vector index. Standardize on an interop dimension (likely 256) for shared/synced packs; let flagship optionally use higher dims locally only. Needs a decision.
• Cold-start cost: EmbeddingGemma is heavier than USE-Lite — measure ingest latency on entry devices; that's why the low tier keeps USE-Lite/128-dim.
• Reranker model size/latency budget on-device — benchmark before enabling by default.

References

• Introducing EmbeddingGemma — Google Developers Blog
• EmbeddingGemma model overview — Google AI for Developers
• litert-community/embeddinggemma-300m — Hugging Face
• Welcome EmbeddingGemma — Hugging Face blog
• MarkTechPost: EmbeddingGemma 308M, SOTA MTEB <500M
• Rerankers: BGE reranker guide · Best reranker models 2026
• Internal: docs/EMBEDDING_GEMMA_PLAN.md

[sagar-develop]

Shipped in v0.10.0 (PR #32, merged to main; release https://github.com/sagar-develop/litertlm-kmp/releases/tag/v0.10.0).

What landed vs the proposal

1. EmbeddingGemma with Matryoshka — done. USE-Lite (100-dim) replaced by EmbeddingGemma-300M as the primary embedder, with Matryoshka truncation (768 → 512/256/128) + re-normalization so one downloaded model serves every tier.

• Runtime decision: ONNX Runtime, not LiteRT. The issue left "LiteRT vs ONNX" open. I went with ONNX Runtime for two reasons: (a) the litert-community EmbeddingGemma bundles didn't load cleanly on the test device (same GPU/NPU executor issues the LLM hits), and (b) ONNX Runtime is telemetry-free (Microsoft, no Play/Google deps), which protects the app's zero-telemetry promise. New files: OnnxEmbeddingEngine, OnnxReranker.
• Tokenizers are pure-Kotlin. onnxruntime-extensions' in-graph tokenizer doesn't support GemmaTokenizer, so GemmaBpeTokenizer (BPE, byte-fallback) and BertWordPieceTokenizer (reranker) were reimplemented and validated bit-for-bit against HF transformers.

2. Device-tiered embedding — done. EmbedderRecommendation.forDevice(ramMb) mirrors the LLM tiering:

Tier	Embedder	Dim	Reranker
< 6 GB	USE-Lite	100	–
6–8 GB	EmbeddingGemma	256	–
8–10 GB	EmbeddingGemma	256	✓
≥ 10 GB	EmbeddingGemma	512	✓

Per-dim ObjectBox HNSW entities (GemmaChunk128/256/512) since the HNSW dim is a compile-time literal; repository routes by active dim; existing projects self-heal via a per-document re-index migration from stored chunk text. Model + companions (ONNX external-data blob + tokenizer) download on-device through the catalogue (gated → HF token flow), surfaced with a Recommended badge.

3. Task-aware embedding — done. EmbeddingEngine.embed(text, task: QUERY|DOCUMENT, title); the ingestor uses DOCUMENT prefixes, the retriever QUERY prefixes.

Reranker — done. Cross-encoder ms-marco MiniLM-L6 (ONNX, Apache-2.0, ungated, ~90 MB) on ≥8 GB tiers, re-scoring the top fused pool.

The wrong-document grounding fix (the issue's headline symptom)

The root cause was more than coarse vectors — it was the fusion layer. With several same-domain documents (a car, a life, and a health insurance policy in one project), BM25 lexical overlap on "insurance"/"premium" let a car-insurance question ground on the life policy (answering with the wrong premium). Fixes:

• Document-relevance (dominance) gate — ground on the document the vector arm clearly favours (best chunk + near-ties); restrict the keyword arm to those documents so shared terms can't reintroduce the wrong source.
• Title-match override — a distinctive query term that names a document by title ("car" → a CarPolicy source) grounds on that document over a higher-scoring but wrong one (generic words like "policy"/"insurance" are excluded).
• Stateful-KV grounding flush — a separate bug surfaced where grounded answers truncated to 1–2 tokens after a few turns: the chat session's KV cache accumulated each turn's grounding block until it overflowed the on-device context window. Grounded turns now re-prefill only the bounded visible transcript.

Verification (on-device, Realme CPH2723, release build, clean project)

• "what is my car insurance premium" → ₹8,504 (correct car premium), grounded on the car doc
• "who is the insurer of my car policy" → TATA AIG (previously wrongly "ICICI Lombard" from the health doc)
• Routing correct across car / life / health docs; 9 grounded turns in one session with no truncation
• Reranker fixed a health-insurer recall miss ("ICICI Lombard … Health Shield 360")
• 11/11 retriever unit tests green (dominance gate, title-match, fusion, reranker reorder)

Known residual

One honest miss: a "sum assured" query on a 206-page life policy returns "specified in the Schedule" rather than the figure. This is a first-stage recall / PDF-extraction limitation — the value lives in a garbled extracted table, so it never enters the candidate pool (the reranker can't reorder what isn't retrieved). Best addressed by better PDF table extraction; tracking separately if it proves common.

Full engineering writeup: _session/material/blog-embedding-enhancements.md. Architecture diagrams: ARCHITECTURE.md.

Closing as implemented and shipped in v0.10.0.