pocket-mind/pocket-embed-22m-128

Pocket Embed 22M — 128-dimensional English search

A frozen MiniLM encoder with a small fitted projection for compact search indexes. 128-dimensional vectors use 512 bytes each in float32: one third of the original 384-dimensional vector storage. Index metadata and overhead are additional.

On two confirmation tasks (ArguAna and SciDocs), it retained 95.9% of base384 mean nDCG@10 and scored 7.0% above raw base128 truncation. These are measured comparisons on this small evaluation set, not a general quality guarantee. Results for the original four tasks appear below.

Additional tests showed lower retention on FiQA: 88.6% for FP32 and 86.4% for INT8. Quora retention was 98.7% and 98.3%. BGE Small scored higher on both additional tasks; Pocket FP32 also scored slightly below raw128 on Quora. See the full additional audit for scores and confidence intervals. The compact index is a storage/quality tradeoff, not a claim of stronger retrieval than these full-dimensional encoders.

Try browser search: a free static demo that runs inference on your device. The first search downloads about 23 MB of INT8 model weights plus runtime files.

Quality retention varies by task: 97.8% on ArguAna, 91.4% on SciDocs. Both score lower than the full 384-dimensional baseline. The gain over raw 128-dimensional truncation on SciDocs is uncertain: its paired 95% interval includes zero. At 64 dimensions, SciDocs also scores slightly below raw baseline truncation. Use 128 dimensions for the evaluated storage/quality tradeoff.

The encoder architecture is unchanged. The fitted projection reduces embedding storage and the dimensions used in vector comparison. It does not reduce encoder parameter count or token encoding work. The optional INT8 ONNX export separately reduces the weight download; the default PyTorch weights remain float32. This uses established PCA and similarity distillation methods, with no claim of a new architecture.

Usage

pip install sentence-transformers==5.1.2 transformers==4.57.1

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("pocket-mind/pocket-embed-22m-128", device="cpu")
documents = ["Bike repairs are available on Saturdays.",
             "The library lends books and audiobooks."]
vectors = model.encode(documents, normalize_embeddings=True)
query = model.encode(["Where can I fix my bicycle?"], normalize_embeddings=True)
best = int((query @ vectors.T).argmax())
print(documents[best])
print(vectors.shape)  # (2, 128)

No prompt or prefix is required. Always use this same model for queries and documents; these vectors are incompatible with the original MiniLM vectors. Re-encode your corpus when migrating. Longer inputs are truncated at 256 word pieces. Chunk long documents before encoding.

Optional 64-dimensional output:

vectors64 = model.encode(documents, truncate_dim=64, normalize_embeddings=True)

64 dimensions trade more retrieval quality for storage; the release gate applies only to 128 dimensions. The projection is evaluated at both sizes.

Search a local corpus

The CPU indexing example builds an exact-search index from JSONL passages using the pinned INT8 graph, then searches it locally. It stores model revisions and checksums with the index and can search offline once the graph and tokenizer are cached. The chunk command splits long documents into bounded overlapping passages and preserves source character offsets. The original FP32 graph is also supported. Corpus text and queries are processed on your device.

python examples/pocket_search.py chunk documents.jsonl passages.jsonl
python examples/pocket_search.py index passages.jsonl my-index
python examples/pocket_search.py --offline query my-index "Where can I fix my bicycle?"

See the example README for dependency versions, JSONL input format, memory limits and 256-token truncation. Download the example files before running.

Browser and direct ONNX inference

The browser export contains a float32 ONNX opset 17 graph of the entire pipeline, including pooling, projection, and normalization. It returns normalized 128-dimensional sentence embeddings. It is validated against PyTorch across variable batches, empty inputs and 256-token truncation (maximum observed absolute difference 3.2e-7). This export is for direct ONNX Runtime inference; do not load it as SentenceTransformer backend="onnx", which expects a raw transformer graph and applies the remaining layers separately. Input names and reproduction instructions are in browser/README.md.

The optional browser/model_int8.onnx graph uses dynamic QInt8 weights and returns float32 normalized 128-dimensional vectors. It is 23.1 MB, versus 90.6 MB for FP32, a 74.5% smaller graph download. It retains 99.52% of the FP32 projection mean nDCG@10 on the existing ArguAna and SciDocs tasks, or 95.45% of the base384 mean. These are deployment regression checks on the same tasks, not additional confirmation evidence. No encoding speed claim.

Task	FP32 Pocket 128 nDCG@10	INT8 Pocket 128 nDCG@10	INT8 / base384
ArguAna	0.4932	0.4911	97.44%
SciDocs	0.1979	0.1966	90.83%
SciFact	0.6104	0.5977	92.66%
NFCorpus	0.2938	0.2963	93.57%

INT8 passed the predefined regression gate: >=99.5% of FP32 confirmation mean,

=95% of base384 confirmation mean, >=90% of base384 on every measured task, and >=2% gain over raw base128 confirmation mean. It is close to the FP32 retention threshold; the table shows task-specific changes. See browser/quantized_decision.json and browser/quantized_results.json for exact values, recipe, size and checksums. The browser demo uses INT8; the Python SentenceTransformer example and retrieval table below use FP32. Use the same export for query and document encoding when building an index.

Retrieval measurements

nDCG@10 uses linear relevance gains, matching the trec_eval convention. Full official BEIR test queries and passage corpora, cosine search, title plus text for documents, max sequence length 256. No reranking.

Task	Role	Base 384	Base 128 truncation	Pocket 128	Pocket 64
arguana	confirmation	0.5041	0.4547	0.4932	0.4515
scidocs	confirmation	0.2164	0.1909	0.1979	0.1552
scifact	diagnostic	0.6451	0.6016	0.6104	0.4937
nfcorpus	diagnostic	0.3167	0.2741	0.2938	0.2408

Paired bootstrap 95% intervals for Pocket128 minus base128 nDCG@10:

• arguana: +0.0385, CI [+0.0275, +0.0503].
• scidocs: +0.0069, CI [-0.0002, +0.0142].
• scifact: +0.0088, CI [-0.0187, +0.0365].
• nfcorpus: +0.0196, CI [+0.0064, +0.0340].

The full metrics, 2,000-resample intervals, corpus/query counts and archive hashes are in results.json. release_decision.json records the thresholds fixed before this run and the observed outcome.

SciFact and NFCorpus are diagnostic tasks: their earlier results informed the decision to freeze the encoder. They are not fresh confirmation evidence.

How it was fitted

Base revision: 1110a243fdf4706b3f48f1d95db1a4f5529b4d41. SQuAD revision: 7b6d24c440a36b6815f21b70d25016731768db1f. Fit projection candidates on normalized original-model embeddings of unique SQuAD train questions and passages. Candidates: centered PCA, uncentered PCA, and uncentered-PCA initialized cosine-similarity distillation at 128 and 64 dimensions. The original transformer weights were frozen. Selected distill_500, using 512 deterministic SQuAD validation queries against the complete validation passage corpus. The candidate was selected and saved before evaluating confirmation labels. A standard Dense projection and Normalize module produce native 128-dimensional outputs. Model export and reload were checked numerically. No custom remote code.

Reproduce with the files in reproduce. In an environment with PyTorch and a GPU:

pip install -r reproduce/requirements.txt
python reproduce/project.py --config reproduce/config.json --output ./results --cache ./data

See selection.json, environment.txt and experiment_README.md for the protocol.

Efficiency and limits

Total parameters: 22,762,496. Native embedding dimension: 128. One-thread CPU encoding measured 24.4 passages/s versus 25.0 for the baseline, on 128 SQuAD validation passages, batch 32, median of three repetitions. Timing is specific to this cluster CPU; no encoding-speed improvement is claimed.

English short-text retrieval only. No multilingual, instruction-following, long-document, adversarial robustness or safety-critical evaluation. Compression can change ranking and degrade quality on your domain. MiniLM base pretraining includes SQuAD and many other datasets; base-model benchmark overlap is unknown. Exact fine-tuning overlap checks do not rule out paraphrases or upstream pretraining overlap. This is a limited evaluation, not an MTEB leaderboard submission. Test on your own queries before migrating.

Licenses and attribution

Original MiniLM encoder weights remain Apache-2.0; see LICENSE_BASE and BASE_MODEL_CARD.md. Our fitted projection is released under CC BY-SA 4.0, the license used in this repository’s metadata. Research code is Apache-2.0. See NOTICE for attribution.

Projection fitting data: SQuAD by Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang (CC BY-SA 4.0). Adaptation: dimensionality reduction fitted to encoder representations of question/context text; no raw dataset text is included in this release.

Evaluation: BEIR, ArguAna (Wachsmuth et al.), SciDocs (Cohan et al.), SciFact (Wadden et al.), NFCorpus (Boteva et al.). The BEIR dataset cards list CC BY-SA 4.0. They were used for evaluation only.

Built by Pocket Mind. Report reproducible retrieval failures in this model’s Community tab; include task, dimensions, library version and an example you have permission to share.