BrainHealthAI

BrainMed-8B

README

License: apache-2.0

Results

Accuracy (%) on a six-benchmark medical suite. The BrainMed-8B row was measured here; every other row is transcribed from published tables and was not re-run.

Table with columns: Model, MedBullets op4, MedBullets op5, MedXpertQA, MedQA, MedMCQA, PubMedQA, Avg, Source
Model	MedBullets op4	MedBullets op5	MedXpertQA	MedQA	MedMCQA	PubMedQA	Avg	Source
Medical-Llama3-8B	33.4	25.3	9.0	40.3	46.8	48.0	33.8	paper
Mistral-Instruct-7B	43.5	33.4	11.4	48.2	44.9	50.1	38.6	paper
BioMistral-7B	46.4	33.1	12.4	45.0	40.2	66.9	40.7	paper
Medical-CoT-8B	39.3	34.1	12.6	49.0	42.6	68.0	40.9	paper
DeepSeek-Distill-8B	41.9	35.1	13.5	55.4	49.0	73.9	44.8	paper
OpenBioLLM-8B	39.2	35.7	10.7	57.7	54.1	74.1	45.3	paper
Llama3.1-Instruct-8B	43.2	40.9	14.3	58.7	56.0	75.2	48.0	paper
Qwen2.5-Instruct-7B	50.0	41.6	12.6	57.0	55.6	72.7	48.2	paper
Huatuo-o1-SFT-8B	53.3	49.7	17.3	70.2	58.2	76.1	54.1	paper
Huatuo-o1-RL-8B	55.2	51.3	16.7	72.6	60.4	79.2	55.9	paper
MedReason-8B	57.5	55.5	19.0	71.8	60.7	79.4	57.3	paper
BrainMed-8B-SFT	58.77	56.82	18.08	72.11	61.20	77.10	57.35	measured
BrainMed-8B	62.01	54.87	18.63	76.67	64.07	79.10	59.23	measured

All 10 benchmarks

Table with columns: Benchmark, n, Raw, Decontaminated
Benchmark	n	Raw	Decontaminated
MedQA (4-opt)	1273	76.67	76.73
MedMCQA	4183	64.07	64.07
PubMedQA	1000	79.10	79.10
MMLU-Pro (Med)	1535	65.73	65.44
MedBullets op4

Results table

BrainMed-8B against published 7B-8B medical LLMs

Full report, per-run summaries and paired tests: evaluation.

Sample answers

Greedy decoding, with the system prompt below.

Multiple choice, strict format

Multiple choice

Differential diagnosis

Differential diagnosis

The model is trained and evaluated on medical question answering. It is weaker on open-ended acute management — the corpus contains QA pairs, not treatment protocols — and should not be relied on there.

Usage

python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "BrainHealthAI/BrainMed-8B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

SYSTEM = ("You are a medical reasoning assistant. Work through the clinical problem step by "
          "step inside <think>...</think>, grounding every step in established medical "
          "knowledge, then give the final, complete answer inside <answer>...</answer>.")

messages = [{"role": "system", "content": SYSTEM},
            {"role": "user", "content": "How is an eclamptic seizure managed?"}]
inputs = tok(tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True),
             return_tensors="pt").to(model.device)
print(tok.decode(model.generate(**inputs, max_new_tokens=1024)[0], skip_special_tokens=True))

The system prompt is part of the contract: the model was trained under it and answers as <think>…</think><answer>…</answer>. Dropping it measurably lowers accuracy.

Training

Full-parameter supervised fine-tuning. No adapter, no LoRA, no frozen layers.

Table with columns: Parameter, Value
Parameter	Value
Backbone	`FreedomIntelligence/HuatuoGPT-o1-8B` (Llama-3.1 architecture)
Trainable parameters	8,030,261,248 across 291 tensors — 100% of the model
Vocabulary size	128,256
Hidden size / layers	4,096 / 32
Attention heads	32 query, 8 key-value (grouped-query attention)
FFN intermediate size	14,336
Training rows	44,351
Training tokens	35,096,633 per epoch → over 3 epochs

Validation loss was evaluated every 50 optimizer steps on a held-out split of 923 rows and tracked separately per data source.

Training and validation loss

One result worth recording: the checkpoint with the lowest validation loss (epoch 2) scored 1.44 points below epoch 3 on the benchmarks. On this corpus, validation loss tracks conformity to the training style rather than medical accuracy, and is not a reliable selection criterion.

Post-training: weight interpolation (WiSE-FT)

Fine-tuning raised the challenging clinical benchmarks and cost accuracy on knowledge recall. Rather than retrain, the shipped weights recover the trade-off by linear interpolation in weight space between the backbone and the fine-tuned checkpoint — the WiSE-FT procedure (Wortsman et al., Robust fine-tuning of zero-shot models, CVPR 2022):

markdown
θ(α) = (1 − α) · θ_backbone + α · θ_fine-tuned

How it is implemented

The operation is element-wise over every parameter tensor, with no data, no gradients and no forward pass involved:

python
from safetensors.torch import load_file, save_file

for shard in shards:                       # 4 shards, 291 tensors, 8.03 B parameters
    wb, wf = load_file(base/shard), load_file(finetuned/shard)
    merged = {k: ((1 - a) * wb[k].float() + a * wf[k].float()).to(wb[k].dtype)
              for k in wb}
    save_file(merged, out/shard, metadata={"format": "pt"})

Implementation notes that matter in practice:

Shard layouts must match. Both checkpoints are verified to carry identical shard files and identical tensor names before anything is merged; a mismatch aborts.
Accumulate in float32. The tensors are bf16; averaging directly in bf16 loses precision on the mantissa. Each pair is upcast, combined, then cast back to the original dtype.
Non-floating and shape-mismatched tensors are copied, not averaged — integer buffers and any resized embedding cannot be meaningfully interpolated.
Cost. ~2 minutes of tensor arithmetic per α on CPU, then one evaluation pass. No training GPU time at all.

Reference implementation: GitHub.

Why it works

Both endpoints descend from the same pre-training initialisation, so they sit in the same loss basin and the straight path between them crosses no loss barrier — the linear mode-connectivity result of Frankle et al. (2020), and the transfer-learning analysis of Neyshabur, Sedghi & Zhang (2020), which shows fine-tuned models remain in the basin of their pre-trained initialisation. Averaging weights inside a basin lands in a flatter region, which is where the robustness comes from (Izmailov et al., 2018), and can exceed each of the averaged models (Wortsman et al., Model soups, 2022).

Sweep

α was swept and every point evaluated on the full benchmark suite:

Table with columns: α, 0.0 (backbone), 0.3, 0.5, 0.7, 1.0 (fine-tuned)
α	0.0 (backbone)	0.3	0.5	0.7	1.0 (fine-tuned)
Avg over 6 benchmarks	58.79	59.23	58.73	58.90	57.35

The interpolation exceeds both endpoints, most visibly on MedBullets-op4 (62.01, against 59.09 and 58.77) and MedXpertQA (18.63, against 17.05 and 18.08): gains and losses do not progress at the same rate along the path, so at α = 0.3 most of the gain is captured while the losses have barely begun. This model ships at α = 0.3; BrainMed-8B-SFT is the α = 1.0 endpoint, published separately as the plain supervised fine-tune.

Data

44,351 chain-of-thought traces from two verified sources, unified under one <think>/<answer> constitution. The build pipeline extracts rather than generates: question and reasoning are byte-identical to the source across all 45,274 rows, asserted at build time. The only text transformation is a canonical The answer is X. appended to 17,083 rows (38.5%) where the letter is unambiguously recoverable from options already present in the question — verified to be a pure suffix, and skipped whenever ambiguous.

Evaluation protocol

10 benchmarks, 11,822 questions, extracted from the reference repositories with row counts and SHA-256 digests verified — none rebuilt or resampled. Greedy decoding, strict answer prompt, max(head, tail) extraction, answer-format compliance reported per benchmark.

Citation

bibtex
@software{brainmed8b,
  title  = {BrainMed-8B: a reproducible medical reasoning fine-tuning and evaluation pipeline},
  author = {BrainResearch},
  year   = {2026},
  url    = {https://github.com/williamsassa/BrainMed-8B}
}

@misc{wu2025medreason,
  title={MedReason: Eliciting Factual Medical Reasoning Steps in LLMs via Knowledge Graphs},
  author={Wu, Juncheng and others}, year={2025}, eprint={2504.00993}, archivePrefix={arXiv}
}

@misc{chen2024huatuogpto1,
  title={HuatuoGPT-o1, Towards Medical Complex Reasoning with LLMs},
  author={Chen, Junying and others}, year={2024}, eprint={2412.18925}, archivePrefix={arXiv}
}

@inproceedings{neyshabur2020transfer,
  title={What is being transferred in transfer learning?},
  author={Neyshabur, Behnam and Sedghi, Hanie and Zhang, Chiyuan},
  booktitle={NeurIPS}, year={2020}, eprint={2008.11687}, archivePrefix={arXiv}
}

Intended use and limitations

Research and decision support for qualified users. Not a medical device, not clinically validated, not for autonomous clinical use, and not for decisions about any real person. Benchmark accuracy measures multiple-choice performance and does not establish clinical safety; independent review of open-ended answers found real gaps in acute-management advice.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.