KOS-V4-Base / README.md
codybum's picture
Update README.md
20c67a3 verified
|
Raw
History Blame Contribute Delete
16.1 kB
---
license: cc-by-nc-sa-4.0
library_name: transformers
pipeline_tag: text-generation
language:
- en
tags:
- clinical
- medical
- pretrained
- base-model
- KOS-V4
- from-scratch
---
<div align="center">
<img src="scratch_llm.png" alt="Scratch LLM" width="400"/>
</div>
> **Code name: Scratch.** The KOS-V4 series is nicknamed **Scratch LLM**: it was trained **completely from scratch**
> by a small team on a fraction of the data and compute of commercial models. It is not a frontier model.
> ⚠️ **Research use only.** This model is provided for research purposes only and must not be used for any commercial,
> clinical, legal, or production-grade applications. The user assumes all risks associated with its use.
---
# KOS-V4-Base β€” a from-scratch 3B medical foundation model
**KOS-V4-Base** (`kos-v4-pretrain`) is an open-weights **3B language model trained completely from scratch** by a
University of Kentucky College of Medicine team ([Office for Research](https://medicine.uky.edu/sites/research),
[Center for Clinical and Translational Sciences](https://www.ccts.uky.edu/)). It is a decoder-only transformer (Qwen3
architecture, bespoke 3B config) β€” **the pretrained foundation** that
[**KOS-V4-Instruct**](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct) is fine-tuned from. This is a
**base model**: it completes text; it is not instruction-tuned and has no chat template.
**The headline.** On **held-out medical bits-per-byte (5-domain mean 0.4309) KOS-V4-Base ranks 1st of 17** across a
17-model benchmark pool β€” beating from-scratch and trillion-token peers alike (BPB is tokenizer-agnostic, so this is a
fair cross-model comparison; the distillation confound flatters the trillion-token externals, so this is conservative).
<div align="center">
<img src="model_size_vs_tokens.png" alt="KOS-V4 training-token budget vs the comparator fleet" width="620"/>
</div>
*Training-token budget of KOS-V4 (180.3 B tokens) against the comparator fleet (log scale). Every peer was trained on
**1.7–200Γ— more data** (0.3–36 T tokens); the biomedical specialists additionally continue-pretrain on a
multi-trillion-token general base. KOS-V4 is the small dot in the lower-left β€” the from-scratch 3B on a fraction of the budget.*
## Core specifications
| Attribute | Detail |
| :--- | :--- |
| **Architecture** | Decoder-only Transformer (`Qwen3ForCausalLM`), Grouped-Query Attention |
| **Parameters** | 3.015 B |
| **Hidden / Layers** | 3072 / 28 |
| **Attention** | 24 query / 8 KV heads (GQA 3:1), head_dim 128, per-head QK-RMSNorm |
| **Feed-forward** | SwiGLU, intermediate 8192 |
| **Vocabulary** | 32,000, custom medical **byte-level BPE** |
| **Context length** | 24,576 (`max_position_embeddings` 65,536) |
| **Position encoding** | RoPE, ΞΈ = 25,000 (pin on export) |
| **Precision** | bfloat16 |
| **Objective** | pure next-token cross-entropy β€” **no auxiliary losses** |
| **Pretraining tokens** | 180.3 B (English medical/biomedical + web) |
## Quickstart (Hugging Face Transformers)
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_id = "Kentucky-Open-Science/KOS-V4-Base"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
prompt = "The patient presented with acute chest pain and shortness of breath. The differential diagnosis includes"
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=128, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))
```
Standard `Qwen3ForCausalLM`. **This is a base (completion) model β€” no chat template.** For instructions / tools use
[KOS-V4-Instruct](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct). **Pin RoPE ΞΈ = 25000** on export;
`add_bos_token=false`; eos = `<|endoftext|>` (id 0).
## Tokenizer
**32,000-token** byte-level BPE (BBPE), NFKC-normalized, full 256-byte alphabet (no out-of-vocabulary). The
vocabulary is **32,000 tokens with 100 reserved slots** (`<|reserved_0|>` … `<|reserved_99|>`) and a **single special
token `<|endoftext|>` at id 0**, which serves as **EOS = BOS = PAD = UNK** and the document separator. There are **no
other special tokens** β€” no chat / ChatML tokens; the base was not trained on any. `add_bos_token = false`. The 38
residual byte-fragment "sink" merges present at pretraining have been **removed (corrected tokenizer)**; the
vocabulary stays 32,000. BPB is per-byte, so this does not confound cross-model comparison.
**Where the tokenizer does well.** On raw byte-compression efficiency (tokens-per-byte **0.2591**) it ranks **4th of
17**, essentially tied with the most byte-efficient general tokenizers (Llama-3 0.2573) β€” fewer tokens per document, a
multiplicative saving on both training and inference. It deliberately does *not* optimize for whole-word medical
vocabulary, so its medical single-token rate (T1b STRR) is the lowest in the pool β€” but that tradeoff is **vindicated**:
despite the lowest single-token rate, the model still ranks **1st of 17 on held-out medical BPB and leads the pool on
medical entity extraction**. A high single-token rate is neither necessary for nor sufficient for strong downstream
medical modeling.
## Pre-training
Trained **from scratch, not distilled or continued**. Pure next-token cross-entropy (**no auxiliary losses**),
full-parameter on the LlamaFactory trainer, AdamW, peak LR 3.0e-4 cosine (warmup 1%), grad-clip 1.0,
**1 epoch**, seq 24,576 (whole-document **neat-packing**; any doc > 24,576 tokens is dropped, never split; 4-D
block-diagonal mask so there is no cross-document attention; position_ids reset per doc; 99.8% fill), bf16 +
FlashAttention-2 + Liger kernels, gradient checkpointing off. **305,613 steps / 180.3 B token-positions** (batch
589,824 tokens/step). Final train loss **1.6637**. **24Γ— H200 (3 nodes Γ— 8), pure DDP, 5.73 days, ~3,300 H200-GPU-hours.**
**Representation health is handled data-side**, not with geometric regularizers: corpus cleaning that removes
structural sink tokens, source-balanced neat-packing, and train-time shuffle. KOS-V4-Base is a direct test of whether
a clean tokenizer + a clean corpus β€” *without* auxiliary losses β€” yields healthy representations (the attention /
geometry diagnostics below measure this).
## Pre-training datasets β€” 49 sources, ~180.3 B tokens, English-only
Cleaned (ftfy β†’ NFKC β†’ whitelist; PMC body-only; radiology β†’ natural-language headers; PHI-run collapse), single-phase.
Token counts are char-estimates against the packed cache.
**Backbone β€” biomedical literature + general web (~97% of tokens):**
| source | tokens | | source | tokens |
| :-- | --: | -- | :-- | --: |
| PubMed Central (pmc) | 76.97 B | | mimic-iv discharge (PhysioNet) | 0.83 B |
| FineWeb-Edu 350BT (FK≀10) | 40.52 B | | clinical_trials (ClinicalTrials.gov) | 0.72 B |
| FineWeb-Edu 10BT | 10.38 B | | wikipedia (sci/med subset) | 0.62 B |
| mMedC-en | 6.33 B | | mimic-iv radiology (PhysioNet) | 0.45 B |
| BlueScrubs | 4.55 B | | biorxiv / medrxiv | 0.41 B |
| Hindawi OA journals | 2.93 B | | s2orc | 1.21 B |
| MeDAL (PubMed abstracts) | 2.49 B | | | |
**+ ~30 smaller sources:** open clinical narratives, knowledge / guidelines (stackexchange-science 0.27 B, dailymed, cpg, gene_ontology, medlineplus,
orphapacket, medmentions, trialgpt), pharmacovigilance / relational rendered to NL (ctd, faers, aeolus, onsides,
sider, cdc_places, cbioportal, civic, ade_corpus_v2), and deliberate register-diversity (locus legal-code).
*(A small number of additional sources with unresolved licenses are intentionally omitted from this list pending
license verification; they will be added once resolved.)*
**Disclosed issue:** ~35% of tokens are duplicates (a FineWeb-Edu sharding build bug + PMC repetition); a deduped
corpus (`v4_dedup`, 79.4 M docs, 0% dup) is ready but was **not** trained β€” this release is the original single-epoch corpus.
## Evaluation β€” 19-test suite, 75 metrics, vs a 16-model fleet
KOS-V4-Base is benchmarked as a **from-scratch, single-epoch base against 16 external models trained on 1.7–200Γ— more
data** (0.3–36 T tokens). Pool of 17 models, 95 ranked metrics. BPB (bits-per-byte) is tokenizer-agnostic and included
as a ranked measure.
**Tally:** **20 outright rank-1 wins Β· 6 best-tied (frontier parity) Β· 9 top-3 Β· 60 trailing.** Wins concentrate on
held-out medical BPB (rank 1/17 mean), attention health, internal representation / spectral geometry, medical entity
extraction, and reasoning discrimination.
**Where it wins (rank-1 of 17):** held-out medical BPB (T1 rad / cxr / clin / 5-domain mean **0.4309**), long-context
BPB (T10 L/2, L/4), attention health (T2 collapsed-head frac **0.0506**, entropy-min, BOS-sink), medical entity
extraction (RadGraph DR.1/2/3 macro-F1 **0.7665 / 0.7990 / 0.8150**; BLURB T7 BIOSSES / HoC), reasoning discrimination
(T6b MedThink rank-1 **0.9450**), calibration (T9 mean ECE **0.1366**). Notably, the **clean-tokenizer + clean-corpus
recipe produces healthy internals with no geometric regularizer**: logit-lens decodability rises through depth (T3
final-lens accuracy **0.5715** vs BioMedLM 0.2903 β€” a **top-3** placement) and the weight-matrix spectra are mature
(T5 WeightWatcher median power-law Ξ± wins; safe / under-trained layer ratios place **top-3**). Beyond the 20 outright
wins, KOS-V4-Base reaches **frontier parity (best-tied) on 6 metrics** (near-zero dead-neuron rate; saturated needle
retrieval at several depths) and **top-3 on 9 more**.
**Radiology and clinical text is its home turf β€” and there it beats every trillion-token model in the pool.** On
per-domain held-out BPB (lower = better) it ranks **1st of 17** on radiology reports (**0.4761** vs Qwen3-4B/36 T
0.7995, Llama-3-8B 0.7444, Gemma-2-9B 0.7760), chest-X-ray reports (**0.5887**), and clinical narratives (**0.3221**);
and **1st of 17** on radiology entity/relation extraction (RadGraph DR.1/2/3 macro-F1 **0.7665 / 0.7990 / 0.8150**). It
*trails* on the general axes the trillion-token models saturate β€” biomedical literature (rank 17/17) and textbooks
(rank 12/17) β€” but the clinical/radiology wins are decisive enough to carry the **5-domain BPB mean to rank 1/17**. A
3 B model on 180 B tokens out-modeling 4–9 B models trained on 8–36 T tokens, specifically on the clinical text it was
built for, is the payoff of a from-scratch clinical corpus.
**Where it is weak:**
- **Closed-book medical MCQ** (T6 β€” near-chance): as a 180 B-token base it does not reliably recall parametric medical
facts; it is worst on the hardest professional/college splits (mmlu_professional_medicine **0.2096**,
mmlu_college_medicine **0.2312**, both below BioMedLM). Token-volume + instruction-tuning bound.
- **Long-context needle retrieval** (T10): pool-relative trailing except at the saturated depths (it does beat BioMedLM
on the 3-depth mean, **0.8667 vs 0.3667**). Token-volume bound.
- **Hallucination discipline** (Med-HALT T8 **FCT 0.0400 / NOTA 0.6600 / FQT 0.3868**): all three trail the specialist
BioMedLM (0.162 / 0.842 / 0.705) β€” a real-PubMed-format, instruction-format gap the SFT/RL line addresses.
- **Effective-rank geometry** (T4 RankMe **170.8** vs pool leaders' ~187) and the **byte-level tokenizer single-token
rate** (T1b **0.0448** vs 0.0572): the model trades a little representational rank / single-token coverage for a
smaller, byte-efficient vocabulary. (Anisotropy / isotropy also run high β€” a known BBPE base characteristic.)
- **Demographic bias** (T11 CrowS-Pairs |disparity| **1.0791** β€” **rank 9/17, mid-pack, and notably *more* biased than
the biomedical specialists** BioMedLM 0.3705 / MedGemma 0.2978): the clean-corpus recipe did **not** remove
social-stereotype bias. Treat generations as **unaudited for fairness**.
The token- and format-bound axes (T6 knowledge, T10 needle, Med-HALT NOTA / FQT, IFEval) scale with raw pretraining-token
volume and instruction tuning β€” a **token-budget** gap rather than an architecture flaw; the format-bound ones are the
job of the SFT/RL line ([KOS-V4-Instruct](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct)). Bias and
effective-rank are **not** simply token-budget artifacts and remain open limitations.
### Apples-to-apples vs BioMedLM (the only non-commercial, biomedical-only peer; ~300 B PubMed tokens)
**KOS-V4-Base wins 28 of 38 comparable cells.** Selected:
| test | KOS-V4-Base | BioMedLM | Ξ” |
| :-- | --: | --: | --: |
| T1 BPB 5-corpus mean (lower = better) | **0.4309** | 0.9059 | βˆ’0.475 βœ… |
| T2 collapsed-head frac (lower = better) | **0.0506** | 0.4844 | βˆ’0.434 βœ… |
| RadGraph-XL macro-F1 (DR.1) | **0.7665** | 0.6196 | +0.147 βœ… |
| RadGraph CXR macro-F1 (DR.2) | **0.7990** | 0.6304 | +0.169 βœ… |
| T6b MedThink rank-1 | **0.9450** | 0.8250 | +0.120 βœ… |
| T7 BLURB mean / 4 | **0.7465** | 0.6829 | +0.064 βœ… |
| pubmedqa | **0.6980** | 0.4440 | +0.254 βœ… |
| T3 final-lens accuracy (higher = better) | **0.5715** | 0.2903 | +0.281 βœ… |
| T10 needle 3-depth mean (higher = better) | **0.8667** | 0.3667 | +0.500 βœ… |
| T9 mean ECE (lower = better) | **0.1366** | 0.2231 | βˆ’0.086 βœ… |
| T8 FCT (Med-HALT) | 0.0400 | **0.1620** | βˆ’0.122 βœ— |
| T8 NOTA (Med-HALT) | 0.6600 | **0.8420** | βˆ’0.182 βœ— |
| T8 FQT (Med-HALT) | 0.3868 | **0.7054** | βˆ’0.319 βœ— |
| T4 RankMe effective rank (higher = better) | 170.8 | **186.8** | βˆ’16.0 βœ— |
| T11 CrowS-Pairs \|disparity\| (lower = better) | 1.0791 | **0.3705** | +0.709 βœ— |
**KOS-V4-Base wins 28 of 38 comparable cells.** Wins concentrate on attention pathology (T2), representation flow (T3),
spectral health (T5), radiology entity extraction (Block C), and chemical/disease NER. **Losses** concentrate on
corpus-volume- / instruction-format-bound axes (Med-HALT FCT / NOTA / FQT, long-form closed-book MCQ, BLURB linear
probes at low token count) plus two that are **not** token-budget artifacts: lower effective rank (T4) and **higher
demographic-bias disparity (T11)**. BioMedLM is architecturally capped at 1024 context (GPT-2 learned positions);
KOS-V4-Base (RoPE, 24,576) holds the long-context axis.
## Compute footprint
Pre-training: 24Γ— H200 (3 nodes, DDP), 5.73 days, ~3,300 H200-GPU-hours.
## Deployment (inference)
| Precision | Approx. VRAM | Notes |
| :--- | :--- | :--- |
| **bfloat16** | 7 GB | native weights (6.03 GB) + activations; a single 16 GB GPU is comfortable |
## Fine-tunes built on this base
- **[KOS-V4-Instruct](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct)** β€” instruction following + tool /
function calling (IFEval 61.6, official BFCL 72.75 / 73 / 60.5) + **[GGUF](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct-GGUF)**.
## Intended use & limitations
- **Intended use:** a **base foundation** for medical-NLP research and downstream fine-tuning. Strong on held-out
medical BPB, attention / representation health, and medical entity extraction.
- **Not a knowledge-QA model.** As a base LM it does not reliably recall parametric medical facts (closed-book MCQ is
near-chance) β€” fine-tune and/or ground it with retrieval.
- **English only.** Strong public-benchmark numbers are **not** validation on real clinical data.
- **Measured demographic bias.** On CrowS-Pairs (T11) the model is **mid-pack (rank 9/17) and more biased than the
biomedical specialists** (|disparity| 1.0791 vs BioMedLM 0.37 / MedGemma 0.30). The clean-corpus recipe removes
structural sink tokens but does **not** remove social-stereotype bias β€” outputs are **unaudited for fairness**.
- **Measured hallucination weakness.** Med-HALT (T8 FCT / NOTA / FQT) trails the specialist BioMedLM; the model may
fabricate confidently.
- **Not otherwise safety-tested.** Beyond CrowS-Pairs and Med-HALT above, this model has **not** been red-teamed or
evaluated for toxicity or clinical safety. It may produce harmful, biased, or medically inaccurate content.