Pensa all'interno del JSON: Strategia di Rinforzo per la Rigida Adesione allo Schema LLM

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

February 18, 2025

Autori: Bhavik Agarwal, Ishan Joshi, Viktoria Rojkova

cs.AI

Abstract

In questo articolo, affrontiamo la sfida di imporre una stretta aderenza allo schema nella generazione di modelli linguistici di grandi dimensioni (LLM) sfruttando le capacità di ragionamento degli LLM. Basandoci sul framework di apprendimento per rinforzo DeepSeek R1, il nostro approccio addestra le abilità di ragionamento strutturato di un modello da 1,5 miliardi di parametri attraverso una pipeline innovativa che combina la costruzione di un dataset sintetico di ragionamento con funzioni di ricompensa personalizzate nell'ambito dell'ottimizzazione delle politiche relative di gruppo (GRPO). Nello specifico, eseguiamo prima l'apprendimento per rinforzo R1 su un dataset di 20K campioni non strutturati-strutturati, seguendo i metodi originali di DeepSeek R1, per stabilire le capacità di ragionamento di base. Successivamente, abbiamo eseguito un fine-tuning supervisionato su un dataset separato di 10K campioni di ragionamento, concentrandoci sul perfezionamento dell'aderenza allo schema per i task downstream. Nonostante l'ambito di addestramento relativamente modesto, che richiede circa 20 ore su un cluster GPU 8xH100 per l'addestramento GRPO e 3 ore su 1xA100 per il SFT, il nostro modello dimostra una performance robusta nell'imporre la coerenza dello schema. Confrontiamo il nostro approccio ThinkJSON con l'originale DeepSeek R1 (671B), le versioni distillate di DeepSeek R1 (Qwen-1.5B e Qwen-7B) e Gemini 2.0 Flash (70B), evidenziandone l'efficacia nelle applicazioni reali. I nostri risultati sottolineano l'utilità pratica di un framework efficiente in termini di risorse per la generazione di testo vincolato da schema.

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.