Tool-Use DPO: Schema-Constrained Alignment via Identity Preference Optimization

Overview

Tool-Use DPO addresses a critical production failure mode: LLMs generating tool calls that are syntactically valid JSON but violate the API schema — wrong types ("5" instead of 5), hallucinated parameters, invalid enum values, missing required fields. The system implements a two-stage pipeline (SFT cold start → DPO with Identity Preference Optimization) trained on GPT-4o-generated hard negative preference pairs — rejected responses that are valid JSON but fail schema validation in exactly the ways real models fail. Using Qwen2.5-Coder-7B-Instruct with RSLoRA (rank 32, all 7 projection modules), the method achieves a 3.1× improvement in Strict Schema Pass Rate (7.48% → 23.19%) with 67% fewer JSON syntax errors and 22% fewer missing required fields on 802 diverse tool-calling samples.

The Hard Negative Philosophy

The core insight: easy negatives don’t teach discrimination. Malformed JSON or gibberish is trivially distinguishable — models already avoid obvious garbage. Instead, the training data contains rejected responses that are plausible, syntactically correct JSON that fails in precisely the four ways production models fail:

Each training triplet (prompt, chosen, rejected) carries its own JSON Schema, enabling the model to generalize across diverse API contracts rather than memorizing a fixed schema.

Synthetic Data Generation

GPT-4o generates triplets with a structured system prompt enforcing the hard negative constraint. Every generated sample passes through a strict validation gate before inclusion:

chosen_valid, _ = validate_tool_call(chosen_str, schema)    # Must pass
rejected_valid, _ = validate_tool_call(rejected_str, schema)  # Must fail

if chosen_valid and not rejected_valid:
    accept(sample)   # Correct: chosen passes, rejected fails
else:
    discard(sample)   # GPT-4o made an error — self-correcting pipeline

Validation uses the jsonschema library with full Draft 4/6/7 support, including additionalProperties checks. The extract_json() utility handles models wrapping output in markdown code blocks or embedding JSON in surrounding text via brace-depth matching.

Two-Stage Training Pipeline

Stage 1: SFT Cold Start

Format-tuning phase teaching the model to output raw JSON with proper ChatML structure (<|im_start|>/<|im_end|> delimiters). Uses only (prompt, chosen) pairs from the DPO triplets.

Why SFT first? DPO alone struggles without foundational format knowledge. The SFT stage ensures the model can reliably produce JSON before learning which JSON to prefer — without this, the DPO signal is too noisy to distinguish “bad format” from “bad content.”

Stage 2: DPO with IPO Loss

Preference optimization using the SFT checkpoint as initialization.

Why IPO over Standard DPO

IPO (Azar et al., 2023) replaces DPO’s log-sigmoid loss with a squared error objective targeting an explicit margin:

L_IPO = ( log(π(y_w|x)/π_ref(y_w|x)) − log(π(y_l|x)/π_ref(y_l|x)) − 1/(2β) )²

With β = 0.05, the target margin 1/(2β) = 10 creates strong pressure to separate chosen from rejected. IPO’s non-saturating loss is critical with only ~500 synthetic training samples where standard DPO would overfit.

Implicit reference model: Setting ref_model=None in TRL’s DPOTrainer leverages the PEFT adapter architecture — the frozen base model weights serve as the implicit reference while the LoRA delta represents the policy. This halves GPU memory by avoiding a separate reference model copy.

Model Architecture

Rank 32 is unusually high for LoRA (typical: 8–16), giving the adapter significantly more capacity — justified for a task requiring fine-grained structural understanding of JSON schema constraints across diverse API contracts.

Evaluation: Strict Schema Pass Rate (SSPR)

SSPR is a binary, all-or-nothing metric: a response either fully conforms to the schema or it does not. No partial credit.

Failures are classified into 6 categories via jsonschema.validate() error messages: json_error, hallucinated_param ("Additional properties"), type_mismatch ("is not of type"), enum_violation ("is not one of"), missing_required ("is a required property"), and other_schema_error.

Evaluation uses greedy decoding (max_new_tokens=128) with the same ChatML prompt template used during training. Both baseline and DPO models are evaluated on identical data for direct comparison.

Results

Evaluated on 802 diverse tool-calling samples. The 3.1× improvement demonstrates that hard negative preference learning effectively teaches subtle schema constraint satisfaction that SFT alone cannot achieve.

Tech Stack

Python, PyTorch (2.2+), Hugging Face Transformers, TRL (DPOTrainer, DPOConfig), Unsloth (FastLanguageModel), PEFT (RSLoRA), bitsandbytes (NF4), jsonschema, OpenAI GPT-4o, Datasets, Accelerate