--- license: mit language: - en tags: - embeddings - sentence-similarity - retrieval - medical - biomedical - bitnet - 1.58-bit - ternary - gguf - llama.cpp - llm2vec library_name: llama.cpp pipeline_tag: feature-extraction base_model: microsoft/bitnet-b1.58-2B-4T-bf16 --- # 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](https://huggingface.co/microsoft/bitnet-b1.58-2B-4T-bf16) 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): ```bash 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:** ```bash 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`: ```bash ./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: ` - **Document:** `` ### 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. ```python 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.