Rabe3/1-bit-embedding-general

1---2license: mit3language:4- en5tags:6- embeddings7- sentence-similarity8- retrieval9- medical10- biomedical11- bitnet12- 1.58-bit13- ternary14- gguf15- llama.cpp16- llm2vec17library_name: llama.cpp18pipeline_tag: feature-extraction19base_model: microsoft/bitnet-b1.58-2B-4T-bf1620---21 22# 1-bit Medical Embedding Model (BitNet b1.58 · ternary · CPU)23 24A **1.58-bit ternary** (weights in {−1, 0, +1}) **medical/biomedical text-embedding model**, adapted from25Microsoft's [BitNet b1.58 2B4T](https://huggingface.co/microsoft/bitnet-b1.58-2B-4T-bf16) decoder into a26bidirectional sentence encoder via the **LLM2Vec** recipe, then exported to a **1.1 GB ternary GGUF** that runs27on **CPU with `llama.cpp`** (no GPU required).28 29- **Format:** GGUF, `TQ1_0` ternary quantization · **1.1 GB** (from 4.8 GB bf16)30- **Embedding dimension:** 2560 (Matryoshka-trained: the first 768 / 512 / 256 / 128 dims are independently usable)31- **Pooling:** mean · **Attention:** bidirectional (non-causal)32- **Tokenizer:** LLaMA-3 128K byte-level BPE (bundled inside the GGUF)33- **Domain:** biomedical literature (PubMed) and clinical QA34 35> **Why 1-bit?** BitNet stores weights as ternary values, so the model is ~4× smaller than an fp16 model of the36> same size and is designed for efficient **CPU** inference — useful for cheap, local, or large-scale vector search.37 38---39 40## Quick start — serve with `llama.cpp` (CPU only)41 42**1. Build llama.cpp** (unmodified upstream — no patches needed):43 44```bash45git clone https://github.com/ggml-org/llama.cpp46cd llama.cpp47cmake -B build -DCMAKE_BUILD_TYPE=Release48cmake --build build -j49```50 51**2. Download the model:**52 53```bash54huggingface-cli download Rabe3/1-bit-embedding-general medbit-2b-embed.gguf --local-dir .55```56 57**3. Embed text.** The model is **bidirectional + mean-pooled**, so you must pass `--attention non-causal --pooling mean`:58 59```bash60./build/bin/llama-embedding \61  -m medbit-2b-embed.gguf \62  --pooling mean \63  --attention non-causal \64  --embd-normalize 2 \65  --embd-output-format array \66  -p "Metformin is first-line therapy for type 2 diabetes."67```68 69Batch a file (one text per line) with `-f texts.txt`. Output is an L2-normalized 2560-dim vector; cosine70similarity = dot product.71 72### Asymmetric instructions (recommended for retrieval)73The model was trained (E5/LLM2Vec style) with an instruction prefix on **queries** and none on **documents**:74 75- **Query:** `Represent this clinical question for retrieving relevant biomedical abstracts: <your query>`76- **Document:** `<the passage, no prefix>`77 78### Matryoshka (shorter vectors)79Truncate the 2560-dim output to the first **768 / 512 / 256 / 128** dims and re-normalize — all remain usable for80cheaper storage / faster search.81 82---83 84## Evaluation85 86Held-out biomedical retrieval (PubMed **title → abstract**, pairs never seen in training):87 88| Setting | R@1 | R@10 |89|---|---|---|90| GGUF **ternary, CPU** (llama.cpp), 100-way | **0.85** | **0.98** |91| PyTorch bf16 (GPU reference), 100-way | 0.90 | 0.98 |92| PyTorch bf16 (GPU reference), 1000-way | 0.85 | 0.97 |93 94The ternary CPU export retains retrieval quality: R@10 is identical to the GPU reference; R@1 is within ~5 points.95Ternary-vs-bf16 embedding cosine fidelity ≈ 0.85 (ranking preserved).96 97---98 99## How it was built (LLM2Vec on a ternary decoder)100 101Base model: `microsoft/bitnet-b1.58-2B-4T-bf16` (2.4 B params, ternary weights, bf16 master weights for102fine-tuning). All fine-tuning kept the base **frozen** and trained **LoRA adapters** (merged after each phase),103so the model stayed ternary throughout (quantization-aware).104 1051. **Bidirectional patch + MNTP (Phase 1).** Replaced the causal attention mask with a full (bidirectional) mask106   and trained **Masked Next-Token Prediction** on ~24 M PubMed titles+abstracts to adapt the decoder to107   bidirectional encoding and inject medical knowledge. (~1.2 k steps; loss 5.05 → 2.15.)1082. **Weakly-supervised contrastive (Phase 3).** **InfoNCE** with large in-batch negatives (via **GradCache**),109   **Matryoshka** loss over {768, 512, 256, 128}, and asymmetric query/document instruction prefixes, on medical110   positive pairs (PubMed title ↔ abstract, PubMedQA question ↔ context). ~1.5 k steps, warmup + cosine LR;111   this is what breaks the raw-decoder anisotropy and produces usable embeddings.112 113Mean pooling over the final hidden states; embeddings are L2-normalized.114 115### Training data (all public)116- **Corpus (MNTP):** `MedRAG/pubmed` (23.9 M title+abstract snippets), `MedRAG/textbooks`.117- **Weak pairs (contrastive):** PubMed title ↔ abstract; `qiaojin/PubMedQA` (`pqa_artificial`) question ↔ context.118- Cleaning: unicode-normalize, English-filter (fastText lid.176), length filter, exact dedup.119 120### Export to ternary GGUF121Merged the LoRA into the bf16 master weights, bridged the transformers-native `BitNetModel` tensor layout to the122one `llama.cpp`'s BitNet converter expects, and converted to `TQ1_0`. The `llama.cpp` runtime is **unmodified** —123its standard non-causal + mean-pooling path serves the bidirectional embeddings directly.124 125---126 127## Intended use & limitations128 129**Use for:** biomedical/clinical **retrieval, semantic search, clustering, similarity** — strongest on130literature-style (PubMed) and QA-style medical text.131 132**Limitations (read before deploying):**133- **Lightly trained.** Trained for a few thousand steps on a compute-limited setup (V100, no bf16 tensor cores,134  no `torch.compile`), not a full multi-day / full-corpus run. It is a **strong, honest baseline**, not a tuned135  SOTA system; expect headroom from more training.136- **Not benchmarked on external suites yet** (BEIR NFCorpus/SciFact/TREC-COVID, BIOSSES, MTEB). Numbers above are137  held-out internal pairs.138- **Public data only** — no MIMIC/clinical notes, so it skews toward literature and exam/QA phrasing and is139  relatively weaker on raw clinical-note text.140- **Ternary export gap** ≈ 5 R@1 points vs the bf16 model.141- **Runtime:** runs on stock `llama.cpp` (generic ternary kernels), **not** Microsoft's `bitnet.cpp` optimized142  I2_S/TL1 kernels — correct and CPU-native, but not BitNet's peak advertised throughput.143- Not a medical device; **not for clinical decision-making.**144 145## License146MIT (following the base model). Built on `microsoft/bitnet-b1.58-2B-4T-bf16`.147 148## Acknowledgements149LLM2Vec (BehnamGhader et al., 2024) · BitNet b1.58 (Microsoft, 2025) · E5 · MedCPT · `llama.cpp` (ggml-org).150 151---152 153## Full-precision (bf16) weights154 155The bf16 **master weights** are also provided (in the `bf16/` subfolder) for GPU inference with 🤗 Transformers156or further fine-tuning. Note: BitNet keeps a *ternary* forward pass even from these bf16 weights (online157quantization) — bf16 is the storage/master-weight format used for training.158 159```python160import torch161from transformers import AutoModel, AutoTokenizer162 163tok = AutoTokenizer.from_pretrained("Rabe3/1-bit-embedding-general", subfolder="bf16")164model = AutoModel.from_pretrained("Rabe3/1-bit-embedding-general", subfolder="bf16",165                                  torch_dtype=torch.bfloat16).cuda().eval()166# NOTE: this is a decoder patched to BIDIRECTIONAL attention for embeddings; use mean pooling167# over the last hidden state and L2-normalize. See the repo scripts for the exact embedder.168```169 170**Files:** `bf16/model.safetensors` (~4.8 GB, bf16), `bf16/config.json` (BitNet, ternary-online), tokenizer.171