{"id":606531,"date":"2025-07-08T20:59:20","date_gmt":"2025-07-09T01:59:20","guid":{"rendered":"https:\/\/towardsdatascience.com\/?p=606531"},"modified":"2025-07-08T20:59:32","modified_gmt":"2025-07-09T01:59:32","slug":"how-to-finetune-small-language-models-to-think-with-reinforcement-learning","status":"publish","type":"post","link":"https:\/\/towardsdatascience.com\/how-to-finetune-small-language-models-to-think-with-reinforcement-learning\/","title":{"rendered":"How to Fine-Tune Small Language Models to Think with Reinforcement Learning"},"content":{"rendered":"\n

Reasoning models<\/strong> are currently<\/mdspan> in fashion. DeepSeek-R1, Gemini-2.5-Pro, OpenAI’s O-series models, Anthropic’s Claude, Magistral, and Qwen3 \u2014 there is a new one every month. When you ask these models a question, they go into a chain of thought<\/em> before generating an answer.<\/p>\n\n\n\n

\"\"
A simple demonstration of what reasoning looks like. When asked a question, the Language Model (LM) generates a chain of thought first, followed by the answer. (Illustration by the Author)<\/figcaption><\/figure>\n\n\n\n

I recently asked myself the question, “Hmm… I wonder if I should write a Reinforcement Learning loop from scratch that teaches this ‘thinking’ behaviour to really<\/em> small models \u2014 like only 135 million<\/em> parameters<\/em>“. It should be easy, right?<\/p>\n\n\n\n

Well, it wasn’t.<\/p>\n\n\n\n

\n

Small models simply do not have the world knowledge that large models do. This makes < 1B parameter model lack the “common sense” to easily reason through complex logical tasks. Therefore, you cannot just rely on compute to train them to reason.

You need additional tricks up your sleeve.<\/strong><\/p>\n<\/blockquote>\n\n\n\n

In this article, I won’t just cover tricks though. I will cover the major ideas behind training reasoning behaviours into language models, share some simple code snippets, and some practical tips to fine-tune Small Language Models (SLMs) with RL.<\/p>\n\n\n\n

This article is divided into 5 sections:<\/strong><\/p>\n\n\n\n

    \n
  1. Intro to RLVR (Reinforcement Learning with Verifiable Rewards) and why it is uber cool<\/li>\n\n\n\n
  2. A visual overview of the GRPO algorithm and the clipped surrogate PPO loss.<\/li>\n\n\n\n
  3. A code walkthrough!<\/li>\n\n\n\n
  4. Supervised fine-tuning and practical tips to train reasoning models<\/li>\n\n\n\n
  5. Results!<\/li>\n<\/ol>\n\n\n\n

    Unless otherwise mentioned, all images used in this article are illustrations produced by the author.<\/em><\/p>\n\n\n\n

    At the end of this article, I will link to the 50-minute companion YouTube video of this article. If you have any queries, that video likely has the answers\/clarification you need. You can also reach out to me on X (@neural_avb<\/a>).<\/em><\/p>\n\n\n\n

    1. Reinforcement Learning with Verifiable Rewards (RLVR)<\/h2>\n\n\n\n

    Before diving into specific challenges with Small models, let’s first introduce some terms.<\/p>\n\n\n\n

    Group Relative Policy Optimization, or GRPO, is a (rather new) Reinforcement Learning (RL) technique that researchers are using to fine-tune Large Language Models (LLMs) on logical and analytical tasks. Since its inception, a new term has been circulating in the LLM research space: <\/span>RLVR<\/strong>, or R<\/strong>einforcement L<\/strong>earning with V<\/strong>erifiable R<\/strong>ewards.<\/p>\n\n\n\n

    To understand what makes RLVR unique, it’s helpful to contrast it with the most common application of RL in language models: RLHF (R<\/strong>einforcement L<\/strong>earning with H<\/strong>uman F<\/strong>eedback). In RLHF, an RL module is trained to maximize scores from a separate reward model, which acts as a proxy for human preferences. This reward model is trained on a dataset where humans have ranked or rated different model responses. <\/p>\n\n\n\n

    \n

    In other words, RLHF is trained so LLMs can output responses that are more aligned with human preferences. It tries to make models follow instructions more closely.<\/strong><\/p>\n\n\n\n

    RLVR tries to solve a different problem. RLVR teaches a model to be verifiably correct, often by learning to generate it’s own chain of thought.<\/strong><\/p>\n<\/blockquote>\n\n\n\n

    Where RLHF had a subjective<\/em> reward model, RLVR uses an objective<\/em> verifier. The core idea is to provide rewards based on whether an answer is demonstrably correct, not on a prediction of what a human might prefer.<\/p>\n\n\n\n

    \"\"
    An illustration of how RLVR works (Illustrated by the Author)<\/figcaption><\/figure>\n\n\n\n

    This is exactly why this system is called ‘RL with verifiable rewards<\/em>‘. Not every question’s answer can be verified easily. Especially open-ended questions like “What iPhone should I buy?<\/em>” or “Where should I go to college?”<\/em>. Some use cases, however, do fit easily in the “verifiable rewards” paradigm, like math, logical tasks, and code-writing, to name a few. In the reasoning-gym<\/code> section below, we will look into how exactly these tasks can be simulated and how the rewards can be generated.<\/p>\n\n\n\n

    But before that, you might ask: well where does “reasoning” fit into all of this?<\/em><\/p>\n\n\n\n

    We will train the LLM to generate arbitrarily long chain of thought reasoning texts before generating the final answer. We instruct the model to wrap its thinking process in <think><\/code> tags and its final conclusion in <answer><\/code> tags. <\/p>\n\n\n\n

    The full language model response will look something like this:<\/p>\n\n\n\n

    <think>\nUser has asked me to count the number of r's in strawberry.\nLet's do a cumulative count.\ns=0, t=0, r=1, a=0, w=0, b=0, e=0, r=2, r=3, y=4\n\nIt seems there are 3 r's in strawberry. \nI notice that there is an r in straw and 2 r's in berry.\nSince 1+2=3 I am more confident there are 3 r's\n<\/think>\n<answer>\n3\n<\/answer><\/code><\/pre>\n\n\n\n
    This structure allows us to easily extract just the final answer and check if it’s correct. The verifier is a single source of truth, and can be a simple piece of code that (literally) counts alphabets.<\/p>\n\n\n\n
    def count_alphabets(word, letter):\n    return sum([1 for l in word if l == letter])\n\nreward = 1 if (lm_answer == count_alphabets("strawberry", "r") else -1<\/code><\/pre>\n\n\n\n
    We will keep a record of the model’s experiences \u2014 its responses and the corresponding rewards received from the verifier. The RL algorithm will then train to promote behaviours that increase the likelihood of correct final answers.<\/p>\n\n\n\n
    \n
    By consistently rewarding correct answers and good formatting, we would increase the likelihood of reasoning tokens that lead to correct answers. <\/strong><\/p>\n\n\n\n
    Get this: we don’t need<\/em> to directly evaluate the intermediate reasoning tokens. By simply rewarding the final answer, we will indirectly elicit reasoning steps into the LLM’s chain of thought that lead to correct answers!<\/strong><\/p>\n<\/blockquote>\n\n\n\n
    \"\"
    Source: Some exercepts from the DeepSeek-R1<\/a> paper (License: Free) <\/figcaption><\/figure>\n\n\n\n
    2. GRPO (Group Relative Policy Optimization)<\/h2>\n\n\n\n
    I am going to skip the usual Reinforcement Learning 101<\/em> intro here, I expect most of you who read this far to understand the basics of RL. There is an agent who observes states from the environment and takes an action \u2014 the environment rewards the agent depending on how good the action was \u2014 the agent stores these experiences and trains to take better actions in the future that lead to higher rewards. RL 101<\/em> class dismissed.<\/p>\n\n\n\n
    But how do we transfer the RL paradigm to language?<\/em><\/p>\n\n\n\n
    Let’s talk about our algorithm of choice \u2014 G<\/strong>roup R<\/strong>elative P<\/strong>olicy O<\/strong>ptimization to understand how. GRPO works in two iteratively self-repeating phases \u2014 an experience collection phase where the Language Model (LM) accumulates experiences in the environment with its current weights. And a training phase where it uses the collected memories to update its weights to improve. After training, it once again goes into an experience collection step with the updated weights.<\/p>\n\n\n\n
    Experience Collection<\/h3>\n\n\n\n
    Let’s dissect each step in the experience collection phase now.<\/p>\n\n\n\n