Tiny Reasoner: Engineering a 135M Model for Medical Expertis

1. Introduction

Large Language Models (LLMs) like GPT-4 are generalists. But for specific, high-stakes domains like medical coding (assigning ICD-10/CPT codes to clinical notes), we don't always need a 100B+ parameter giant. We need a specialist.

The Problem: Medical Coding is Complex

Medical coding is the process of translating clinical narratives into standardized billing codes. Consider this example:

"Patient presents with uncontrolled Type 2 diabetes mellitus with diabetic retinopathy."

A human coder must:

This requires domain expertise, reasoning through clinical guidelines, and understanding code hierarchies. A wrong code can lead to claim denials, compliance issues, or lost revenue.

Why Not Just Use GPT-4?

While large models like GPT-4 can perform this task, they have critical limitations:

• Cost: Running inference on 70B+ models at scale (millions of claims/day) is prohibitively expensive.
• Latency: Large models are slow, creating bottlenecks in real-time workflows.
• Hallucination: Without grounding, they can invent plausible-sounding but invalid codes.

My Hypothesis

A tiny, specialized model (135M parameters) can outperform generalist giants on this specific task if we:

This post details my engineering journey building Tiny Reasoner, a project that proves this hypothesis.

The Result (Sneak Peek)

Before diving into the architecture, here is what the final model actually does:

Input: "Patient presents with uncontrolled Type 2 diabetes mellitus
        with mild nonproliferative diabetic retinopathy."

Tiny Reasoner (135M):
+---------------------------------------------------------------+
| <think>                                                       |
| 1. Identify primary condition: Type 2 Diabetes Mellitus.      |
| 2. Identify complications: Retinopathy (Diabetic).            |
| 3. Specify severity: Mild, Nonproliferative.                  |
| 4. Check control status: Uncontrolled (Hyperglycemia).        |
| 5. Construct Code: E11 (Type 2) + .3 (Retinopathy)            |
|    + .2 (Mild Nonproliferative) + .9 (Unspecified eye).       |
| </think>                                                      |
|                                                               |
| Code: E11.329                                                 |
+---------------------------------------------------------------+

It doesn't just guess; it derives the answer.

2. System Architecture

My pipeline is designed to bootstrap reasoning capabilities from zero, verify them rigorously, and then scale them up autonomously.

Design Philosophy

I moved away from the standard "SFT → DPO" pipeline used in most LLM alignment work. Instead, I built a verification-centric architecture where every training decision is grounded in deterministic correctness checks.

Key Insight: In medical coding, "correctness" isn't subjective. A code is either valid (exists in the ICD-10 database) or it isn't. This allows me to replace expensive human preference labeling with a Python script.

High-Level Data Flow

[ Raw Data Sources ]       [ Phase 1: Foundation ]       [ Phase 2: Alignment ]
+------------------+       +---------------------+       +--------------------+
|  CMS Guidelines  |------>|  Extraction Script  |------>|    SFT Trainer     |
|  (Text/PDF)      |       |  (Regex/Parsing)    |       |   (SmolLM2-135M)   |
+------------------+       +---------------------+       +--------------------+
                                    |                              |
+------------------+                v                              v
|   OpenAI API     |------>| Synthetic Generator |       +--------------------+
| (Teacher Model)  |       | (CoT + Validation)  |       |    RLVR Trainer    |
+------------------+       +---------------------+       |    (PPO + Rule)    |
                                                         +--------------------+
                                                                   |
                                                                   v
                                                         [ Phase 4: Scaling ]
                                                         +--------------------+
                                                         |  Self-Amplification|
                                                         |  (Iterative Loop)  |
                                                         +--------------------+

3. Phase 1: The Foundation (Bootstrapping)

Before a model can "reason," it needs to see examples of reasoning. I couldn't just use raw text; I needed Chain-of-Thought (CoT) data.

3.1. Data Extraction Strategy

The Challenge: Most medical coding datasets are either proprietary or lack reasoning traces. I needed examples that showed how to arrive at a code, not just the final answer.

The Solution: I built scripts/extract_cms_guidelines.py to mine "gold standard" examples from official CMS (Centers for Medicare & Medicaid Services) documentation.

CMS guidelines often contain structured examples like:

"Example: Patient admitted for acute bronchitis...

Rationale: The condition is acute (not chronic), so we select from the J20 family...

Code Assignment: J20.9"

I used regex patterns to extract these structured examples:

Raw Text:
"Example: Patient with Type 2 diabetes... Rationale: Code E11.9 is selected because..."

      |
      v  (Regex Parsing)
      |

Structured JSONL:
{
  "narrative": "Patient with Type 2 diabetes...",
  "reasoning": "Code E11.9 is selected because...",
  "code": "E11.9"
}

3.2. Synthetic Teacher Generation

The Challenge: CMS guidelines only yielded ~5,000 examples. I needed 15,000+ to train a robust model.

The Solution: I used GPT-4o as a "teacher" to generate synthetic examples. But not just any examples—I enforced a strict Chain-of-Thought (CoT) schema.

The Prompt Strategy

I designed a prompt that forces the teacher model to:

prompt = """
You are a medical coding expert. Generate a realistic clinical narrative,
step-by-step reasoning process, and correct medical billing code.

Requirements:
- Narrative should be 50-200 words
- Include 4-6 reasoning steps in <think> tags
- Assign ONE code (ICD-10, CPT, or HCPCS)
- Code must be valid and commonly used

Focus on: {specialty}
Code family: {code_family}
Complexity level: {complexity}
"""

The Validation Gate

Critical Problem: LLMs hallucinate. GPT-4 might generate "ICD-10: Z99.999 (Patient abducted by aliens)."

My Solution: Every generated example passes through CodeValidator before being added to the training set.

# scripts/generate_synthetic.py
for attempt in range(max_attempts):
    response = gpt4.generate(prompt)
    parsed = parse_response(response)

    # CRITICAL: Validate the code
    if not validator.validate(parsed['code'], parsed['code_type']):
        continue  # Discard hallucinated codes

    training_data.append(parsed)

This filtering step is crucial. In my experiments, ~8-10% of GPT-4's generated codes were invalid or non-existent.

How to Run Phase 1:

python scripts/extract_cms_guidelines.py
python scripts/generate_synthetic.py --count 5000
python data/prepare_foundation.py

4. Phase 2: RLVR (The "Secret Sauce")

This is where the project diverges from standard LLM training.

Why Not DPO?

Direct Preference Optimization (DPO) is the current state-of-the-art for alignment. But it has a fatal flaw for medical coding:

DPO requires preference pairs: "Response A is better than Response B."

Creating these pairs requires:

This is:

• Expensive: ~$500-700 for 40k preference pairs.
• Noisy: Human annotators might disagree.
• Gameable: The reward model can be fooled by plausible-sounding but incorrect codes.

The RLVR Alternative

For medical coding, I had a ground truth: The Code Book.

I implemented RLVR (Reinforcement Learning with Verifiable Rewards) in post_training/rlvr.py. Instead of a neural reward model, I used a deterministic Python function.

4.1. The Reward Function

Instead of a neural network, my reward function is a Python script:

      Model Output                Ground Truth
     "Code: E11.9"               "Code: E11.9"
           |                           |
           v                           v
    +-------------+             +-------------+
    |  Parser     |             |  Parser     |
    +-------------+             +-------------+
           |                           |
           v                           v
      [E11.9] ==?== [E11.9] ---->  +1.0 Reward
           |
      [E11]   ==?== [E11]   ---->  +0.5 Reward (Partial)

def compute_reward(prediction, ground_truth):
    # 1. Extract code from model output
    pred_code = extract_code(prediction)

    # 2. Check Validity (Is it a real ICD-10 code?)
    if not validator.validate(pred_code):
        return 0.0

    # 3. Check Correctness
    if pred_code == ground_truth:
        return 1.0
    elif pred_code.family == ground_truth.family:
        return 0.3  # Partial credit!

    return 0.1 # Valid but wrong

4.2. PPO Training Loop

I used Proximal Policy Optimization (PPO) to optimize the model against this rule-based reward.

How PPO Works (Simplified):

The Training Loop:

for batch in dataloader:
    # 1. Generate responses
    prompts = batch['prompts']
    responses = model.generate(prompts, temperature=0.8)

    # 2. Compute rewards
    rewards = []
    for response, ground_truth in zip(responses, batch['ground_truths']):
        reward = verifier(response, ground_truth)
        rewards.append(reward)

    # 3. PPO step (update model weights)
    ppo_trainer.step(prompts, responses, rewards)

Why This Works: The model learns: "I don't just need to sound confident; I need to output codes that pass the validator."

   +-------------+        Action (Text)       +-------------+
   |   Policy    |--------------------------->| Environment |
   | (The Model) |                            | (Verifier)  |
   +-------------+                            +-------------+
          ^                                          |
          |             Reward (Float)               |
          +------------------------------------------+

This teaches the model: "I don't know exactly what to write, but I know that if I output a valid code that matches the family, I get a reward."

How to Run Phase 2:

# First, SFT
python post_training/sft.py --config-path post_training/config/foundation_sft.yaml
# Then, RLVR
python post_training/rlvr.py --config post_training/config/rlvr.yaml

5. Phase 4: Self-Amplification (Scaling Up)

Once the model was "good enough" (Phase 2), I used it to teach itself. This is the Self-Amplification loop implemented in scripts/iterative_training.py.

5.1. The Loop

      +------------------+
      |  Current Model   |
      +------------------+
               |
      (1) Generate (Temp=0.8)
               v
      +------------------+      +------------------+
      |  Pseudo-Labels   |----->|  CodeValidator   |
      +------------------+      +------------------+
                                         |
                                    (2) Filter
                                         v
      +------------------+      +------------------+
      |   Next Model     |<-----|  Clean Dataset   |
      +------------------+      +------------------+
             (3) Train

5.2. Why This Works

The Intuition: The model has learned how to reason about medical codes (from Phase 1 & 2). Now, when it sees a new narrative, it can generate a plausible reasoning path. Sometimes, it discovers a correct path that wasn't in the original training data.

The Mechanism:

The Flywheel Effect:

• Better model → Better pseudo-labels → Better model → ...

Scaling Results:

Iteration 1: 20k labeled examples → 85% accuracy
Iteration 2: 20k labeled + 30k pseudo → 89% accuracy
Iteration 3: 20k labeled + 80k pseudo → 92% accuracy

Key Insight: I didn't need to label 100k examples manually. I labeled 20k, and the model generated the rest (with validation).

Iteration 1:  20k Examples  ->  85% Accuracy
      |
      v (Generate & Filter)
      |
Iteration 2:  50k Examples  ->  89% Accuracy
      |
      v (Generate & Filter)
      |
Iteration 3:  100k Examples ->  92% Accuracy

How to Run Phase 4:

python scripts/iterative_training.py \\
    --foundation-model outputs/rlvr/final \\
    --iterations 3

6. Evaluation

I don't just check accuracy; I check reasoning quality.

Evaluation Metrics

Evaluation Script

My evaluation/final_evaluation.py script:

python evaluation/final_evaluation.py \\
    --model-path outputs/iterative/iter_3/final \\
    --test-sets data/processed/foundation_val.jsonl data/test/cms_holdout.jsonl

Key Takeaway: The model is not only accurate but also explainable (shows reasoning) and efficient (runs on CPU).

7. Engineering Challenges & Solutions

Challenge 1: Reward Hacking

Problem: Early in RLVR training, the model learned to output any valid code to get partial credit, ignoring the clinical narrative.

Example:

• Prompt: "Patient has a broken arm."
• Model Output: E11.9 (Type 2 diabetes) — Valid code, but completely wrong!

Solution: I tuned the KL-divergence penalty in PPO. This forces the model to stay close to the original language distribution (the SFT model) while still optimizing for the reward. I also reduced the partial credit reward from 0.5 to 0.3.

Challenge 2: Slow Pseudo-Label Generation

Problem: Generating 80,000 examples for self-training was taking 3+ days on a single GPU.

Solution: I rewrote scripts/generate_pseudo_labels.py to use batch generation:

• Instead of generating one example at a time, I batch 8 prompts together.
• Used masked attention to handle variable-length sequences.
• Result: 8x speedup (3 days → 9 hours).

Challenge 3: "Empty Thinking"

Problem: The model sometimes output <think></think> (empty tags) and then the code, gaming the "reasoning presence" metric.

Solution: I added a heuristic to the reward function:

if len(reasoning_text) < 50:  # Too short
    return 0.0  # Zero reward

This forced the model to generate substantive reasoning.

8. What Didn't Work (Lessons from Failed Experiments)

Not everything worked on the first try. Here are some approaches I tried that failed:

Failed Approach 1: Pure Synthetic Data

What I Tried: Initially, I attempted to train the model entirely on GPT-4 generated synthetic data (no CMS extraction).

Why It Failed: The model learned GPT-4's "style" but lacked grounding in official guidelines. When evaluated on real CMS examples, accuracy dropped by ~12%. The model was mimicking a teacher that itself didn't have access to authoritative sources.

Lesson: Domain expertise matters. Even with a powerful teacher model, you need gold-standard examples from authoritative sources.

Failed Approach 2: Standard DPO with Human Preferences

What I Tried: Before implementing RLVR, I experimented with collecting preference pairs (having annotators rank "Response A vs Response B").

Why It Failed:

• Annotators frequently disagreed on which response was "better" when both codes were valid but different levels of specificity.
• The reward model learned to prefer longer, more verbose reasoning, not necessarily correct codes.
• Cost was prohibitive ($700 for just 5k pairs).

Lesson: When you have ground truth, use it. Don't introduce subjectivity where objectivity exists.

Failed Approach 3: Single-Shot Self-Training

What I Tried: I attempted to generate all 80k pseudo-labels in one go (no iterations).

Why It Failed: The model at 85% accuracy generated too many incorrect examples. Even with filtering, ~30% of pseudo-labels were subtly wrong (e.g., correct family, wrong specificity). Training on this degraded performance.

Lesson: Iterate gradually. Self-amplification works best when you start with a strong foundation and scale incrementally.

9. Lessons Learned

After 8 weeks and countless experiments, here are the key insights I gained:

1. Small Models Need Quality, Not Quantity

A 135M model has limited capacity. Feeding it 100k mediocre examples is worse than 20k high-quality examples with clear reasoning traces. Quality > Quantity for specialized domains.

2. Deterministic Verification Beats Learned Rewards

For tasks with objective correctness (math, coding, medical codes), a simple Python function outperforms a neural reward model. It's:

• Cheaper (no training needed)
• More reliable (can't be fooled)
• Easier to debug (just read the code)

3. Self-Amplification Only Works with Strict Filtering

The model will happily learn from its own mistakes if you let it. The validator must be unforgiving—99% valid isn't good enough. I set the threshold at 100% code validity.

4. Reasoning Traces Are Worth the Effort

Adding <think> tags increased training complexity, but the payoff was huge:

• Easier debugging (I could see why the model chose a code)
• Better generalization (the model learned patterns, not memorization)
• Trust from end-users (medical coders want to see the reasoning)

5. Engineering Matters More Than Algorithms

The biggest performance gains came from:

• Better data extraction (CMS guidelines)
• Smarter filtering (validation gates)
• Batch optimization (8x speedup)

Not from tweaking hyperparameters or trying fancier RL algorithms.

10. Future Directions

• Thought-Level Verification: Grading the reasoning steps themselves, not just the final code (Phase 3 implementation ready).
• Model Distillation: Compressing the reasoning patterns into an even smaller model (<50M parameters) for edge deployment.
• Multi-Code Support: Extending the system to handle cases requiring multiple codes (e.g., diabetes + hypertension).
• Real-Time Feedback Loop: Integrating with hospital EHR systems for continuous learning from coder corrections.

Conclusion

By combining symbolic verification (the code validator) with neural generation (the LLM), I built a system that is:

The Key Insight

This project proves that for specialized domains, smart engineering beats raw parameter count.

You don't need a 70B model to be an expert. You need:

• High-quality data (CMS guidelines + validated synthetic).
• Verifiable rewards (deterministic correctness checks).
• Self-amplification (the model teaching itself).

The future of AI isn't just bigger models—it's smarter training pipelines.

Personal Reflection

When I started this project, I thought the hard part would be the RL training. I was wrong.

The hard part was data engineering—extracting structured examples from messy PDFs, designing prompts that forced reasoning, and building filters that caught subtle errors. The RL part was almost trivial once I had high-quality data and a reliable validator.

The biggest surprise? How well self-amplification worked. I expected diminishing returns after iteration 1, but the model kept improving. By iteration 3, it was generating reasoning paths I hadn't seen in the training data—it had genuinely learned how to think about medical codes, not just memorize patterns.

What would I do differently? I'd start with the validator earlier. I initially built it as an afterthought for evaluation, but it became the cornerstone of the entire pipeline. If I were to do this again, I'd design the validator first, then build the training pipeline around it.

This project taught me that constraints breed creativity. Working with a tiny 135M model forced me to be smarter about data, training, and evaluation. And in the end, that smartness beat brute force.

Check out the full code on GitHub