Rabe3/1-bit-embedding-general

1-bit Medical Embedding Model (BitNet b1.58 · ternary · CPU)

A 1.58-bit ternary (weights in {−1, 0, +1}) medical/biomedical text-embedding model, adapted from Microsoft's BitNet b1.58 2B4T decoder into a bidirectional sentence encoder via the LLM2Vec recipe, then exported to a 1.1 GB ternary GGUF that runs on CPU with llama.cpp (no GPU required).

• Format: GGUF, TQ1_0 ternary quantization · 1.1 GB (from 4.8 GB bf16)
• Embedding dimension: 2560 (Matryoshka-trained: the first 768 / 512 / 256 / 128 dims are independently usable)
• Pooling: mean · Attention: bidirectional (non-causal)
• Tokenizer: LLaMA-3 128K byte-level BPE (bundled inside the GGUF)
• Domain: biomedical literature (PubMed) and clinical QA

Why 1-bit? BitNet stores weights as ternary values, so the model is ~4× smaller than an fp16 model of the same size and is designed for efficient CPU inference — useful for cheap, local, or large-scale vector search.

---

Quick start — serve with llama.cpp (CPU only)

1. Build llama.cpp (unmodified upstream — no patches needed):

git clone https://github.com/ggml-org/llama.cpp
cd llama.cpp
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j

2. Download the model:

huggingface-cli download Rabe3/1-bit-embedding-general medbit-2b-embed.gguf --local-dir .

3. Embed text. The model is bidirectional + mean-pooled, so you must pass --attention non-causal --pooling mean:

./build/bin/llama-embedding \
  -m medbit-2b-embed.gguf \
  --pooling mean \
  --attention non-causal \
  --embd-normalize 2 \
  --embd-output-format array \
  -p "Metformin is first-line therapy for type 2 diabetes."

Batch a file (one text per line) with -f texts.txt. Output is an L2-normalized 2560-dim vector; cosine similarity = dot product.

Asymmetric instructions (recommended for retrieval)

The model was trained (E5/LLM2Vec style) with an instruction prefix on queries and none on documents:

• Query: Represent this clinical question for retrieving relevant biomedical abstracts: <your query>
• Document: <the passage, no prefix>

Matryoshka (shorter vectors)

Truncate the 2560-dim output to the first 768 / 512 / 256 / 128 dims and re-normalize — all remain usable for cheaper storage / faster search.

---

Evaluation

Held-out biomedical retrieval (PubMed title → abstract, pairs never seen in training):

Setting	R@1	R@10
GGUF ternary, CPU (llama.cpp), 100-way	0.85	0.98
PyTorch bf16 (GPU reference), 100-way	0.90	0.98
PyTorch bf16 (GPU reference), 1000-way	0.85	0.97

The ternary CPU export retains retrieval quality: R@10 is identical to the GPU reference; R@1 is within ~5 points. Ternary-vs-bf16 embedding cosine fidelity ≈ 0.85 (ranking preserved).

---

How it was built (LLM2Vec on a ternary decoder)

Base model: microsoft/bitnet-b1.58-2B-4T-bf16 (2.4 B params, ternary weights, bf16 master weights for fine-tuning). All fine-tuning kept the base frozen and trained LoRA adapters (merged after each phase), so the model stayed ternary throughout (quantization-aware).

1. Bidirectional patch + MNTP (Phase 1). Replaced the causal attention mask with a full (bidirectional) mask and trained Masked Next-Token Prediction on 24 M PubMed titles+abstracts to adapt the decoder to bidirectional encoding and inject medical knowledge. (1.2 k steps; loss 5.05 → 2.15.)
2. Weakly-supervised contrastive (Phase 3). InfoNCE with large in-batch negatives (via GradCache), Matryoshka loss over {768, 512, 256, 128}, and asymmetric query/document instruction prefixes, on medical positive pairs (PubMed title ↔ abstract, PubMedQA question ↔ context). ~1.5 k steps, warmup + cosine LR; this is what breaks the raw-decoder anisotropy and produces usable embeddings.

Mean pooling over the final hidden states; embeddings are L2-normalized.

Training data (all public)

• Corpus (MNTP): MedRAG/pubmed (23.9 M title+abstract snippets), MedRAG/textbooks.
• Weak pairs (contrastive): PubMed title ↔ abstract; qiaojin/PubMedQA (pqa_artificial) question ↔ context.
• Cleaning: unicode-normalize, English-filter (fastText lid.176), length filter, exact dedup.

Export to ternary GGUF

Merged the LoRA into the bf16 master weights, bridged the transformers-native BitNetModel tensor layout to the one llama.cpp's BitNet converter expects, and converted to TQ1_0. The llama.cpp runtime is unmodified — its standard non-causal + mean-pooling path serves the bidirectional embeddings directly.

---

Intended use & limitations

Use for: biomedical/clinical retrieval, semantic search, clustering, similarity — strongest on literature-style (PubMed) and QA-style medical text.

Limitations (read before deploying):

• Lightly trained. Trained for a few thousand steps on a compute-limited setup (V100, no bf16 tensor cores, no torch.compile), not a full multi-day / full-corpus run. It is a strong, honest baseline, not a tuned SOTA system; expect headroom from more training.
• Not benchmarked on external suites yet (BEIR NFCorpus/SciFact/TREC-COVID, BIOSSES, MTEB). Numbers above are held-out internal pairs.
• Public data only — no MIMIC/clinical notes, so it skews toward literature and exam/QA phrasing and is relatively weaker on raw clinical-note text.
• Ternary export gap ≈ 5 R@1 points vs the bf16 model.
• Runtime: runs on stock llama.cpp (generic ternary kernels), not Microsoft's bitnet.cpp optimized I2_S/TL1 kernels — correct and CPU-native, but not BitNet's peak advertised throughput.
• Not a medical device; not for clinical decision-making.

License

MIT (following the base model). Built on microsoft/bitnet-b1.58-2B-4T-bf16.

Acknowledgements

LLM2Vec (BehnamGhader et al., 2024) · BitNet b1.58 (Microsoft, 2025) · E5 · MedCPT · llama.cpp (ggml-org).

---

Full-precision (bf16) weights

The bf16 master weights are also provided (in the bf16/ subfolder) for GPU inference with 🤗 Transformers or further fine-tuning. Note: BitNet keeps a ternary forward pass even from these bf16 weights (online quantization) — bf16 is the storage/master-weight format used for training.

import torch
from transformers import AutoModel, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Rabe3/1-bit-embedding-general", subfolder="bf16")
model = AutoModel.from_pretrained("Rabe3/1-bit-embedding-general", subfolder="bf16",
                                  torch_dtype=torch.bfloat16).cuda().eval()
# NOTE: this is a decoder patched to BIDIRECTIONAL attention for embeddings; use mean pooling
# over the last hidden state and L2-normalize. See the repo scripts for the exact embedder.

Files: bf16/model.safetensors (~4.8 GB, bf16), bf16/config.json (BitNet, ternary-online), tokenizer.