MedGemma-GRPO

MedGemma-GRPO

Model Description

This model is a fine-tuned version of the MEDGemma model, adapted for tasks involving clinical case reasoning. The goal of this project was to explore the effects of combining Supervised Fine-Tuning (SFT) with Group Relative Policy Optimization (GRPO) to guide the model's ability to follow a logical diagnostic path. The model was trained on a synthetic dataset of medical case studies, which, due to the nature of the large models used for generation, has a notable bias towards complex cardiovascular, thrombotic, and autoimmune diseases (e.g., Lupus, Pericarditis).

Fine-tuning Procedure

The model was first adapted to the medical domain using Supervised Fine-Tuning (SFT) on a synthetic dataset. Subsequently, Group Relative Policy Optimization (GRPO) was applied, utilizing a custom rule-based reward function. This reward function was designed to guide the model towards accurate diagnoses and logically structured reasoning, while also penalizing overly verbose outputs.

Using with Unsloth

from unsloth import FastModel
import torch

max_seq_length=1024
model, tokenizer = FastModel.from_pretrained(
    model_name = "lastmass/MedGemma-GRPO",
    max_seq_length = max_seq_length, 
    load_in_4bit = False,  
    load_in_8bit = False, 
    full_finetuning = False, 
)
reasoning_start = "<start_working_out>" 
reasoning_end   = "<end_working_out>"   
solution_start  = "<SOLUTION>"
solution_end    = "</SOLUTION>"

system_prompt = \
f"""You are given a problem.
Think about the problem and provide your working out.
Place it between {reasoning_start} and {reasoning_end}.
Then, provide your solution between {solution_start}{solution_end}"""

user_input=”“

messages = [{
    "role": "user",
    "content": [{"type" : "text", "text" :system_prompt + user_input + reasoning_start,}]
}]
text = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt = True, 
)

from transformers import TextStreamer
_ = model.generate(
    **tokenizer([text], return_tensors = "pt").to("cuda"),
    max_new_tokens =max_seq_length, 
    # Recommended Gemma-3 settings!
    temperature = 0.0,
    streamer = TextStreamer(tokenizer, skip_prompt = True),
)

Intended Use & Limitations

This model is intended for research purposes only, as an exploration into enhancing the reasoning capabilities of smaller language models (4B) in a specialized domain.

⚠️ Disclaimer:

This model is NOT a medical device and should NEVER be used for actual medical diagnosis, treatment, or advice. The outputs are generated by an AI and may contain inaccuracies or hallucinations, and they reflect the significant biases present in the training data. Always consult a qualified healthcare professional for any medical concerns.

Example Prompts and Outputs

Here are a couple of examples demonstrating the model's reasoning capabilities on complex cases with diagnostic traps.

Question

A 58-year-old overweight male with a history of hypertension presents to the emergency department with a 12-hour history of excruciating pain, swelling, and redness in his right great toe. He states the pain started abruptly overnight and is so severe that even the bedsheet touching his toe is unbearable. He admits to drinking several beers and eating a large steak dinner the night before. His medications include hydrochlorothiazide for his blood pressure.
On examination, he has a low-grade fever of 38.1°C (100.6°F). The first metatarsophalangeal (MTP) joint of his right foot is markedly swollen, erythematous, warm, and exquisitely tender to palpation.
Initial laboratory results show:
CBC: White blood cell count is elevated at 13.5 x 10^9/L.
Serum Uric Acid: Markedly elevated.
CRP: Elevated.

What is the most likely diagnosis? What specific findings from his history and medications are significant risk factors for this condition? Finally, what is the gold standard for a definitive diagnosis, and what would it reveal?

Answer

<start_working_out>
Alright, let's think this through. We have a 58-year-old man who's overweight and has a history of hypertension.
He's come in with excruciating pain, swelling, and redness in his right great toe.
The pain started abruptly overnight and is so severe that even the bedsheet touching his toe is unbearable.
He also mentions drinking several beers and eating a large steak dinner the night before.On examination,
he has a low-grade fever of 38.1°C and the first metatarsophalangeal (MTP) joint of his right foot is markedly swollen, erythematous, warm, and exquisitely tender to palpation.
His lab results show an elevated white blood cell count, a markedly elevated serum uric acid, and an elevated CRP.
These findings are pointing towards an inflammatory process, likely related to uric acid.Given his symptoms and lab results, the most likely diagnosis is gout.
Gout is characterized by the deposition of monosodium urate crystals in joints, often affecting the first MTP joint.
The sudden onset of severe pain, the presence of swelling and redness, and the elevated uric acid levels are all classic signs of gout.
Now, let's consider the risk factors. His history of hypertension and being overweight are significant risk factors for gout.
Obesity and hypertension are both associated with increased uric acid levels, which can lead to the formation of uric acid crystals and subsequent gout attacks.
The gold standard for a definitive diagnosis of gout is a joint aspiration. This procedure involves taking a sample of synovial fluid from the affected joint and examining it under a microscope for the presence of monosodium urate crystals.
The presence of these crystals confirms the diagnosis of gout.So, in summary, the most likely diagnosis is gout, and the joint aspiration would reveal the presence of monosodium urate crystals, confirming the diagnosis.
<end_working_out>
<SOLUTION>The most likely diagnosis for the 58-year-old male patient is gout.
The presentation of excruciating pain, swelling, redness, and warmth in the first metatarsophalangeal (MTP) joint of the right foot,
 along with the abrupt onset of symptoms, elevated white blood cell count, markedly elevated serum uric acid, and elevated CRP,
are all characteristic of gout.Significant risk factors for this condition include his history of hypertension and being overweight.
 Obesity and hypertension are both associated with increased uric acid levels, which can lead to the formation of uric acid crystals and subsequent gout attacks.
The gold standard for a definitive diagnosis of gout is a joint aspiration. This procedure involves taking a sample of synovial fluid from the affected joint and examining it under a microscope for the presence of monosodium urate crystals.
The presence of these crystals confirms the diagnosis of gout.In a joint aspiration, the presence of monosodium urate crystals would be revealed,
confirming the diagnosis of gout.
</SOLUTION>

Question

A 48-year-old female presents to her primary care physician stating, "I just feel empty." She reports a 4-month period of pervasive low mood. She has lost interest in her hobbies, including gardening and book club, which she used to love.
She describes significant fatigue, stating she feels "exhausted all the time" despite sleeping more than usual.
She has difficulty concentrating at her job as an accountant and has made several simple errors recently. Her appetite has decreased, and she has had an unintentional weight loss of 10 lbs. She expresses feelings of worthlessness, saying "I feel like a failure as a wife and mother."
She denies any specific plans to harm herself but admits to sometimes "wishing she wouldn't wake up in the morning."
Her past medical history is unremarkable. A physical examination and basic metabolic panel are entirely within normal limits.
What is the most likely diagnosis for this patient? Based on the DSM-5 criteria, list at least five key symptoms from her history that support this diagnosis.
Finally, what is a crucial next step in assessing this patient's immediate safety?

Answer

<start_working_out>
Alright, let's think about this. A 48-year-old woman is feeling empty and low in mood for the past four months.
She's lost interest in things she used to enjoy, like gardening and book club. That's a big red flag for depression.
She's also feeling really tired, even though she's sleeping more.
That's another classic symptom of depression.
And she's having trouble concentrating at work, which is impacting her performance.
She's lost her appetite and has lost weight unintentionally.
That's concerning and could be a sign of depression.She's feeling worthless and says she feels like a failure as a wife and mother.
That's a sign of low self-esteem, which is often associated with depression.She mentions wishing she wouldn't wake up in the morning.
That's a sign of suicidal ideation, which is a serious concern.
Her past medical history is unremarkable, and her physical exam and basic metabolic panel are normal.
This suggests the symptoms are likely psychological rather than physical.Based on the DSM-5 criteria,
she has five key symptoms: persistent low mood, loss of interest in activities, fatigue, difficulty concentrating, and feelings of worthlessness.
Given these symptoms, the most likely diagnosis is Major Depressive Disorder.Now, it's crucial to assess her immediate safety.
She admits to suicidal thoughts, which is a significant risk factor.The next step is to conduct a thorough safety assessment.
This involves asking direct questions about her suicidal thoughts, plans, and intent.
It's important to have a safety plan in place and to ensure she has access to mental health support.
This patient needs immediate attention to ensure her safety and well-being.
<end_working_out>
<SOLUTION>The most likely diagnosis for this patient is Major Depressive Disorder, based on the DSM-5 criteria.
The patient presents with several key symptoms that align with this diagnosis:
1. **Persistent low mood:** The patient reports feeling "empty" and experiencing pervasive low mood for the past four months.
2. **Loss of interest in activities:** She has lost interest in her hobbies, including gardening and book club, which she used to enjoy.
3. **Fatigue:** She describes feeling "exhausted all the time," even though she is sleeping more than usual.
4. **Difficulty concentrating:** She has difficulty concentrating at her job as an accountant, leading to simple errors.
5. **Feelings of worthlessness:** She expresses feelings of worthlessness, stating "I feel like a failure as a wife and mother."
In addition to these symptoms, the patient's unintentional weight loss of 10 lbs and suicidal ideation (wishing she wouldn't wake up in the morning) further support the diagnosis of Major Depressive Disorder.
A crucial next step in assessing this patient's immediate safety is to conduct a thorough safety assessment.
This involves asking direct questions about her suicidal thoughts, plans, and intent.
It is also important to have a safety plan in place and to ensure she has access to mental health support.
This assessment will help determine the level of risk and guide appropriate interventions to ensure her safety and well-being.
</SOLUTION>

Uploaded finetuned model

• Developed by: lastmass
• License: apache-2.0
• Finetuned from model : unsloth/medgemma-4b-it

This gemma3 model was trained 2x faster with Unsloth and Huggingface's TRL library.

<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>