Think Inside the JSON: Reinforcement Strategy for Strict LLM Schema Adherence

2025-02-18

AI Summary

This work addresses the challenge of enforcing strict adherence to structured output schemas—such as JSON—in large language model (LLM) generation. We propose a lightweight reinforcement learning framework integrating Group Relative Policy Optimization (GRPO) with phased structured reasoning fine-tuning. Leveraging synthetically generated reasoning data, a schema-aware reward function, and supervised fine-tuning (SFT), our method achieves zero-schema-deviation performance on a 1.5B-parameter model. Trained for only 20 hours on 8×H100 GPUs, it surpasses state-of-the-art models—including DeepSeek-R1 (671B), Qwen series, and Gemini 2.0 Flash (70B)—in structured generation accuracy. Our key contributions are: (i) the first application of GRPO to structured generation tasks; and (ii) demonstrating that stringent schema compliance—previously attainable only by trillion-parameter models—can be achieved with an order-of-magnitude smaller architecture, enabling efficient, scalable, and deployable structured LLMs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

In this paper, we address the challenge of enforcing strict schema adherence in large language model (LLM) generation by leveraging LLM reasoning capabilities. Building on the DeepSeek R1 reinforcement learning framework, our approach trains structured reasoning skills of a 1.5B parameter model through a novel pipeline that combines synthetic reasoning dataset construction with custom reward functions under Group Relative Policy Optimization (GRPO). Specifically, we first perform R1 reinforcement learning on a 20K sample unstructured-to-structured dataset, mirroring the original DeepSeek R1 methods, to establish core reasoning abilities. Subsequently, we performed supervised fine-tuning on a separate 10K reasoning sample dataset, focusing on refining schema adherence for downstream tasks. Despite the relatively modest training scope, requiring approximately 20 hours on an 8xH100 GPU cluster for GRPO training and 3 hours on 1xA100 for SFT, our model demonstrates robust performance in enforcing schema consistency. We compare our ThinkJSON approach against the original DeepSeek R1 (671B), distilled versions of DeepSeek R1 (Qwen-1.5B and Qwen-7B), and Gemini 2.0 Flash (70B), showcasing its effectiveness in real-world applications. Our results underscore the practical utility of a resource-efficient framework for schema-constrained text generation.

Problem

• Enforcing strict schema adherence in LLM generation.
• Training structured reasoning skills via reinforcement learning.
• Ensuring schema consistency in text generation tasks.

Innovation

• Reinforcement learning for schema adherence
• Synthetic dataset with custom rewards
• Group Relative Policy Optimization training