Думай внутри JSON: Стратегия подкрепления для строгого соблюдения схемы в языковых моделях

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

February 18, 2025

Авторы: Bhavik Agarwal, Ishan Joshi, Viktoria Rojkova

cs.AI

Аннотация

В данной работе мы решаем задачу обеспечения строгого соответствия схеме при генерации текста крупными языковыми моделями (LLM), используя их способности к рассуждению. Основываясь на фреймворке обучения с подкреплением DeepSeek R1, наш подход обучает навыки структурированного рассуждения модели с 1,5 миллиардами параметров через новый конвейер, который сочетает создание синтетического набора данных для рассуждений с пользовательскими функциями вознаграждения в рамках оптимизации групповой относительной политики (GRPO). В частности, мы сначала проводим обучение с подкреплением R1 на наборе данных из 20 тысяч примеров, преобразующих неструктурированные данные в структурированные, следуя оригинальным методам DeepSeek R1, чтобы заложить базовые способности к рассуждению. Затем мы выполняем контролируемую тонкую настройку на отдельном наборе данных из 10 тысяч примеров для рассуждений, сосредоточившись на улучшении соответствия схеме для последующих задач. Несмотря на относительно скромный объем обучения, требующий примерно 20 часов на кластере из 8 GPU H100 для GRPO и 3 часов на одном GPU A100 для SFT, наша модель демонстрирует устойчивую производительность в обеспечении согласованности схемы. Мы сравниваем наш подход ThinkJSON с оригинальной моделью DeepSeek R1 (671B), дистиллированными версиями DeepSeek R1 (Qwen-1.5B и Qwen-7B) и Gemini 2.0 Flash (70B), демонстрируя его эффективность в реальных приложениях. Наши результаты подчеркивают практическую полезность ресурсоэффективного фреймворка для генерации текста с ограничениями по схеме.

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.