Piensa dentro del JSON: Estrategia de Refuerzo para la Adherencia Estricta al Esquema en Modelos de Lenguaje

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

February 18, 2025

Autores: Bhavik Agarwal, Ishan Joshi, Viktoria Rojkova

cs.AI

Resumen

En este artículo, abordamos el desafío de garantizar la adherencia estricta a esquemas en la generación de modelos de lenguaje de gran escala (LLM) aprovechando las capacidades de razonamiento de estos modelos. Basándonos en el marco de aprendizaje por refuerzo DeepSeek R1, nuestro enfoque entrena habilidades de razonamiento estructurado en un modelo de 1.500 millones de parámetros mediante una novedosa canalización que combina la construcción de un conjunto de datos sintéticos de razonamiento con funciones de recompensa personalizadas bajo la Optimización de Política Relativa de Grupo (GRPO). Específicamente, primero realizamos aprendizaje por refuerzo R1 en un conjunto de datos de 20.000 muestras que van de lo no estructurado a lo estructurado, siguiendo los métodos originales de DeepSeek R1, para establecer habilidades básicas de razonamiento. Posteriormente, llevamos a cabo un ajuste fino supervisado en un conjunto de datos separado de 10.000 muestras de razonamiento, centrándonos en refinar la adherencia al esquema para tareas posteriores. A pesar del alcance relativamente modesto del entrenamiento, que requirió aproximadamente 20 horas en un clúster de 8xH100 GPU para el entrenamiento GRPO y 3 horas en 1xA100 para el ajuste fino supervisado (SFT), nuestro modelo demuestra un rendimiento robusto en la aplicación de consistencia de esquemas. Comparamos nuestro enfoque ThinkJSON con el DeepSeek R1 original (671B), versiones destiladas de DeepSeek R1 (Qwen-1.5B y Qwen-7B) y Gemini 2.0 Flash (70B), destacando su eficacia en aplicaciones del mundo real. Nuestros resultados subrayan la utilidad práctica de un marco eficiente en recursos para la generación de texto restringida por esquemas.

English

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.