Fine-Tuning a 350M Model for Better Structured Outputs in 100 GRPO Steps

Small Language Models (SLMs) with around 350 million parameters have become the holy grail for high-throughput, low-latency, and cost-effective AI deployments. However, off-the-shelf SLMs often struggle with strict structural requirements such as valid JSON syntax, Pydantic schema validation, and tool-calling execution parameters. Traditionally, developers relied on Supervised Fine-Tuning (SFT) over thousands of examples or heavy prompt engineering.

With the advent of Group Relative Policy Optimization (GRPO)—the policy gradient technique popularized by DeepSeek-R1—fine-tuning compact models for format adherence has become dramatically more efficient. By substituting traditional critic models with group-relative advantage estimation, GRPO can align a 350M parameter model to produce deterministic structured outputs in as few as 100 optimization steps.

In this guide, we break down why GRPO works so effectively on small parameter scales, detail the mathematical and technical mechanics, walk through a complete implementation using Hugging Face TRL and PEFT, benchmark performance against standard SFT, and explore production deployment patterns combining edge models with centralized aggregators like n1n.ai.

Why GRPO Outperforms SFT for Structured Outputs

Supervised Fine-Tuning (SFT) forces a model to minimize token-level cross-entropy loss against target outputs. While effective for learning language style, SFT suffers from two major structural weaknesses:

1. Format Rigidity vs. Generalization: SFT penalizes valid architectural variations that deviate slightly from training targets, even when the resulting JSON structure is mathematically valid.
2. Exposure Bias: During token-by-token generation, early syntax mistakes cascade into complete output corruption without recovery mechanisms.

Reinforcement Learning (RL) mitigates these issues by scoring the entire generated completion end-to-end. Traditional RLHF algorithms like Proximal Policy Optimization (PPO) require maintaining a separate Value Model (Critic) alongside the Policy Model (Actor). For small models, running a Critic model doubles memory consumption and slows training velocity.

The Mechanics of GRPO

Group Relative Policy Optimization (GRPO) removes the Critic altogether. Instead of evaluating state values through a secondary network, GRPO samples a group of $G$ outputs $\{q_1, q_2, \dots, q_G\}$ for a given input prompt $P$. It computes rewards $R = \{r_1, r_2, \dots, r_G\}$ across the sample group and determines the relative advantage $A_i$ for each completion:

$A_i = \frac\{r_i - \text\{mean\}(R)\}\{\text\{std\}(R)\}$

The objective function optimizes the policy parameter $\theta$ by scaling token probabilities according to group-relative advantage, regularized by Kullback-Leibler (KL) divergence against the initial reference policy $\theta_\{ref\}$:

$\mathcal\{L\}_\{GRPO\}(\theta) = -\frac\{1\}\{G\} \sum_\{i=1\}^\{G\} \left( \min \left( \frac\{\pi_\theta(q_i|P)\}\{\pi_\{\theta_\{old\}\}(q_i|P)\} A_i, \text\{clip\}\left(\frac\{\pi_\theta(q_i|P)\}\{\pi_\{\theta_\{old\}\}(q_i|P)\}, 1-\epsilon, 1+\epsilon\right) A_i \right) - \beta D_\{KL\}(\pi_\theta || \pi_\{ref\}) \right)$

Because structured validation rewards (such as is_valid_json()) are binary or deterministic step functions, GRPO rapidly suppresses non-conforming tokens and reinforces valid syntactical structures within 100 training iterations.

Designing Multi-Tier Reward Functions for JSON & Schema Adherence

To achieve deterministic JSON outputs, your reward pipeline should evaluate generations across three distinct layers:

1. Syntax Validity Reward: Checks if the string can be parsed by json.loads().
2. Schema Adherence Reward: Evaluates whether parsed objects match required keys, types, and nested constraints (e.g., via Pydantic).
3. Value Accuracy Reward: Penalizes hallucinated keys or extra markdown wrappers (such as backticks) that break production parsers.

import json
import re
from typing import Dict, Any, List

def json_format_reward_func(completions: List[str], **kwargs) -> List[float]: