README.md

---
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: <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.

```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.