Decoupling Reasoning and Reward: A Modular Approach for Stable Alignment of Small Clinical Language Models

This paper introduces a modular, adapter-based framework that decouples reasoning supervision from reward tuning to stabilize the alignment of small clinical language models, achieving robust, auditable, and accurate performance without sacrificing privacy or efficiency.

The Gist

Imagine you are trying to teach a very smart, but small, robot assistant how to be a doctor. You want this robot to be accurate (give the right medical advice), honest (show its work so you can check it), and lightweight (so it can run on a tablet in a clinic without needing a massive supercomputer).

The problem is, when you try to teach these small robots using standard methods, they often get confused. They might start hallucinating facts or stop showing their work, making it impossible to trust them.

This paper introduces a new, smarter way to train these robots called "Decoupling Reasoning and Reward." Here is how it works, using some everyday analogies:

The Old Way: The "Swiss Army Knife" Problem

Imagine you are trying to teach a student two very different skills at the exact same time:

1. How to think (like a detective solving a mystery step-by-step).
2. How to get a gold star (learning what the teacher likes to hear).

In the old method, you force the student to learn both skills using the same notebook. The paper calls this a "monolithic" approach.

• The Result: The student gets confused. They try to solve the mystery while simultaneously worrying about the gold star. For a small student (a small AI model), this causes a mental breakdown. They stop thinking clearly just to chase the reward, or they get the reward but forget how to think.

The New Way: The "Specialized Interns" Approach

The authors propose a modular approach. Instead of one confused student, they hire two specialized interns and give them separate notebooks (called LoRA Adapters).

1. Intern A (The Reasoning Expert): Their only job is to learn how to think logically and show their work step-by-step. They practice on thousands of medical cases until they are perfect at explaining how they got an answer.
2. Intern B (The Reward Tuner): Their only job is to look at the answers and say, "Yes, that's the right answer," or "No, that's wrong." They learn to give "gold stars" for accuracy.

The Magic Trick:
When the robot needs to answer a question, you don't just use one intern. You snap the two notebooks together.

• Intern A says: "Here is my step-by-step thinking."
• Intern B says: "Great thinking, and here is the final correct answer."

Because they are separate, they don't get in each other's way. Intern A doesn't get distracted by the gold stars, and Intern B doesn't get confused by the thinking process.

Why This Matters for Small Robots (Small AI Models)

The paper tested this on robots of different sizes, from tiny (0.5 billion "brain cells") to large (7 billion).

• The Tiny Robots (0.5B - 1.5B): These are like junior interns. If you try to teach them everything at once (the old way), they crash and burn. They become unstable and make up facts. But with the Modular approach, they become surprisingly stable and accurate. They can show their work clearly and get the right answer.
• The Big Robots (7B): These are like senior doctors. They are so smart they can handle learning everything at once without crashing. However, even for them, the modular approach works just as well, if not better.

The "Audit" Analogy

In a hospital, if a doctor makes a mistake, you need to be able to look at their notes to see why they made that decision. This is called auditability.

• Old Way: The robot's notes are a messy scribble because it was trying to do two things at once. You can't read them.
• New Way: Because the "Reasoning" part is separate, the robot always writes its thoughts in a neat, structured box (like a labeled folder). Even if the final answer is wrong, you can look at the folder and see exactly how it got there. This makes it safe to use in real hospitals.

The Big Takeaway

The authors built a massive library of medical questions with "show your work" answers (over 100,000 of them) to train these robots. They found that by separating the "thinking" from the "grading," they could create small, efficient AI doctors that are:

1. Stable: They don't crash or go crazy during training.
2. Trustworthy: They always show their work in a clear format.
3. Accurate: They give the right medical advice.

In short: Don't try to teach a small robot to think and get a reward at the same time. Teach it to think first, then teach it to get the reward, and then combine the two. It's a simple switch that makes small AI models safe enough to use in real life.

Technical Summary

1. Problem Statement

Deploying Large Language Models (LLMs) in clinical settings faces a fundamental tension between three competing demands:

1. Accuracy & Reliability: Providing evidence-based, factually correct answers.
2. Auditability & Verifiability: Ensuring transparent, structurally predictable reasoning (e.g., Chain-of-Thought) for safety-critical review.
3. Efficiency & Privacy: The need for models to run on edge devices or behind institutional firewalls, necessitating smaller model sizes (e.g., <7B parameters).

While smaller models are ideal for privacy and efficiency, they often suffer from training instability and objective conflicts when subjected to standard alignment procedures like Group Relative Policy Optimization (GRPO). Existing approaches typically entangle reasoning supervision (Chain-of-Thought, or CoT) and reward tuning within a single, monolithic model. This monolithic structure limits adaptability, hinders debugging, and often leads to convergence issues, particularly in small models where the model struggles to simultaneously learn complex reasoning patterns and optimize for reward signals.

2. Methodology

The authors propose a modular, adapter-based alignment framework that decouples reasoning supervision and reward tuning into separate, composable stages using LoRA (Low-Rank Adaptation) adapters.

Core Framework

The approach utilizes Parameter-Efficient Fine-Tuning (PEFT) to separate the learning objectives:

• Stage 1 (Reasoning): Supervised Fine-Tuning (SFT) with Chain-of-Thought (CoT) traces.
• Stage 2 (Reward): Group Relative Policy Optimization (GRPO) to align the model with verifiable, factual responses.

Experimental Configurations

The study systematically compares five alignment configurations across Qwen2.5 models (ranging from 0.5B to 7B parameters):

1. Base: Original instruction-tuned model (Zero-shot).
2. SFT Only: LoRA adapter trained on CoT-augmented data (optimized for reasoning traces, no reward alignment).
3. GRPO Only: LoRA adapter trained via GRPO from the base model (no intermediate reasoning supervision).
4. Modular (Decoupled): A two-stage approach where a CoT-SFT adapter is frozen, and a second, independent LoRA adapter is trained using GRPO. These can be composed at inference.
5. Unified (Coupled): A single LoRA adapter trained sequentially with both SFT and GRPO objectives.

Data & Rewards

• Datasets: The study uses a mix of medical QA datasets (MedQA, MedMCQA, AlpaCare) for SFT and a specific verifiable dataset for GRPO. Evaluation is performed on held-out test sets: MedQA, OpenBookQA, and ARC Challenge.
• Reward Mechanism: The GRPO reward function combines a format reward (ensuring output adheres to <thought> and <answer> tags) and an accuracy reward. Notably, the accuracy reward uses fuzzy string matching (Levenshtein ratio) rather than exact matching, allowing for minor phrasing variations while maintaining strict answer correctness.

3. Key Contributions

1. Modular PEFT Pipeline: Introduction of a framework that separates reasoning and reward adapters, enhancing training stability and structural auditability, particularly for small models.
2. Comprehensive Benchmarking: Evaluation of five alignment strategies across model scales (0.5B–7B), demonstrating that decoupling is critical for smaller models but less critical for larger ones.
3. Dataset Release: Creation and release of a structured dataset comprising over 100K medically grounded QA pairs with CoT traces, generated and filtered for correctness, to facilitate reproducible research.
4. Stability Insights: Empirical evidence showing that decoupling objectives prevents training collapse in small models, a common failure mode in monolithic alignment.

4. Key Results

The experiments reveal distinct behaviors based on model size and configuration:

• Training Stability:

  • Small Models (0.5B, 1.5B): The Unified approach (single adapter) often suffers from training collapse between steps 500–900 due to conflicting objectives. The Modular approach remains stable, achieving higher final rewards.
  • Large Models (3B, 7B): Both Unified and Modular approaches converge smoothly, suggesting larger models can internally disentangle objectives even when co-trained.

• Structural Auditability (Format Correctness):

  • The Modular configuration consistently achieves near-perfect format adherence (correct use of <thought> and <answer> tags) across all model sizes.
  • GRPO Only models frequently fail to maintain the required reasoning structure, especially in smaller models.
  • SFT Only models follow the format but lack the reinforcement to maintain it in out-of-domain scenarios.

• Factual Accuracy:

  • Decoupling Advantage: For small and mid-sized models (0.5B, 1.5B), the Modular strategy significantly outperforms Unified and GRPO-only baselines in factual accuracy.
  • Scale Effect: In 7B models, both multi-stage configurations yield similar high accuracy, indicating diminishing returns from modularity at higher capacities.
  • Generalization: Models trained with both SFT and GRPO (especially Modular) show better generalization to out-of-domain science reasoning tasks (OpenBookQA, ARC) compared to SFT-only models.

5. Significance and Implications

• Solving the Small Model Dilemma: The paper demonstrates that small clinical LLMs can achieve high accuracy and auditability if the alignment process is modular. This enables the deployment of privacy-preserving, verifiable AI on edge devices without sacrificing safety.
• Operational Flexibility: The modular approach allows for adapter fusion. Clinical institutions can update the "reward" adapter to align with new medical guidelines or standards of care without retraining the foundational "reasoning" adapter, significantly reducing computational costs and deployment time.
• Robustness: By decoupling reasoning supervision from reward tuning, the framework mitigates the risk of "reward hacking" or training collapse, ensuring that the model's reasoning process remains transparent and auditable—a critical requirement for healthcare AI.

In conclusion, the authors argue that decoupling reasoning and reward offers a flexible, robust foundation for building clinical LLMs that successfully balance the competing demands of accuracy, auditability, and efficiency.