AFM-4.5B-OpenMed-RL-CoT

AFM-4.5B-OpenMed-RL-CoT

Lightweight medical finetune on top of Arcee’s AFM-4.5B for education and research use. Trained using a straightforward 3-step process (SFT → DPO → GRPO-CoT) for optimal CoT enrichment.

More information about our methodology will be available in a forthcoming blog post.

All experiments were performed on AMD MI300x GPUs, with computing credits generously provided by Hot AISLE.

⚠️ Medical safety
This model is not a clinician. It can hallucinate and should not be used for diagnosis or treatment. Always involve qualified medical professionals.

---

TL;DR

• Base: arcee-ai/AFM-4.5B – Arcee’s 4.5B instruction model intended for cloud-to-edge deployment.
• Training (high level):
  1. SFT proprietary synthetic medical datasets + tool-calling (search) traces
  2. DPO using MedMCQA-derived preferences (multiple-choice signal)
  3. GRPO for chain-of-thought enrichment, using MedReason verifiable rewards; short rationales encouraged, final answer checked.
• Eval (EleutherAI harness; author’s settings, bs=64)
  • MMLU: 61.40 (vs 55.53 base)
  • MMLU-Pro: 33.16 (vs 32.61 base) – harder 10-choice variant.
  • IFEVAL: 59.59 (vs 63.67 base) – verifiable instruction following.

Note: Arcee’s internal evals may use different harnesses; avoid cross-harness comparisons.

---

What’s inside

Specialization steps

1. Domain SFT (medical + tools)
   Instruction-style synthetic medical Q&A + conversions; supervised search/tool-use traces to teach function-calling patterns compatible with chat templates.

2. Preference alignment — DPO
   Uses MedMCQA correctness as a proxy preference signal to bias toward concise, clinically reasonable options.

3. Reasoning enrichment — GRPO (CoT)
   Group Relative Policy Optimization without a critic; groups of sampled solutions are scored by verifiable rewards (answer correctness + light format checks). Trained with MedReason QA signal.

---

Intended use & limitations

Intended: Medical SLM's research, tool-augmented retrieval demos.

Out of scope: Unsupervised patient care, generating prescriptions, and time-critical guideline decisions.

---

Evaluation

Author-run with the EleutherAI lm-evaluation-harness; seeds, prompts, and templates affect absolute scores.

• MMLU-Pro increases difficulty (10 options; more reasoning-heavy); small deltas are still meaningful.
• IFEVAL checks verifiable constraints (length, keyword counts, format, etc.).

Reproduce (example commands)

# MMLU classic
lm_eval --model hf \
  --model_args pretrained=openmed-community/AFM-4.5B-OpenMed-RL-CoT,parallelize=True,dtype=bfloat16,trust_remote_code=True \
  --task mmlu \
  --batch_size=64 \
  --apply_chat_template \
  --output_path=results \
  --fewshot_as_multiturn 


# MMLU-Pro (10-choice)
lm_eval --model hf \
  --model_args pretrained=openmed-community/AFM-4.5B-OpenMed-RL-CoT,parallelize=True,dtype=bfloat16,trust_remote_code=True \
  --tasks leaderboard_mmlu_pro  \
  --batch_size=64 \
  --apply_chat_template \
  --output_path=results \
  --fewshot_as_multiturn 

# IFEVAL (verifiable instruction following)
lm_eval --model hf \
  --model_args pretrained=openmed-community/AFM-4.5B-OpenMed-RL-CoT,parallelize=True,dtype=bfloat16,trust_remote_code=True \
  --tasks leaderboard_ifeval \
  --batch_size=64 \
  --apply_chat_template \
  --output_path=results \
  --fewshot_as_multiturn

---

Quickstart (Transformers)

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "openmed-community/AFM-4.5B-OpenMed-RL-CoT"
tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")

messages = [
  {"role": "system", "content": "You are a careful medical assistant. Cite sources and warn this is not medical advice. Think step-by-step."},
  {"role": "user", "content": "Briefly: cellulitis vs erysipelas differences?"}
]
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
print(tok.decode(out[0], skip_special_tokens=True))

Data & training notes

• SFT data: Proprietary synthetic medical data + search traces.
• DPO signal: Preferences derived from MedMCQA multiple-choice correctness.
• GRPO reward: Answer-checking + format verifiers; MedReason used to shape faithful, short CoT.
• No known PHI; please open an issue if you spot any.

---

Compatibility & licenses

• Base model: AFM-4.5B (Arcee). Refer to the base card/blog for architecture and usage details. License for AFM releases is Apache 2.0;

---

Additional note

We also provide a merged openmed-community/AFM-4.5B-OpenMed version after step 3 (GRPO). In our harness, it shows worse CoT behavior but a significant gain on IFEVAL.