Skip to content

Latest commit

 

History

1 Commit

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Structured-Output Reliability — a faithful re-implementation

This repository re-implements and reproduces the experiments of:

When Correct Isn't Usable: Improving Structured Output Reliability in Small Language Models — Galeone, Park, Ettorre, Ligorio (Alomana, 2026), arXiv:2605.02363.

The paper's thesis: a deployed model output must be both mathematically correct and parseable by the downstream system. It defines

output accuracy = task_correct AND json_valid

and shows a systematic format failure in small models — high task accuracy but 0% output accuracy at baseline — which an iterative black-box prompt optimizer (AloLab) closes without any fine-tuning.

This codebase provides all of the machinery to measure that gap and to run AloLab yourself, against live models over the NVIDIA NIM API.


What's implemented

Paper concept Here
task_accuracy, json_valid, output_accuracy (Section 2.2) slo/metrics.py — strict, deterministic, unit-tested
NAIVE / REFERENCE / STRICT prompts (Section 2.3, Appendix A) slo/prompts.py
CONSTRAINED / REF+CONSTRAINED (JSON-grammar decoding) slo/clients.py via NIM guided_json
AloLab 5-node loop: Loader→Solver→Evaluator→Analyzer→Optimizer (Section 2.4, Fig 1) slo/data.py, slo/solver.py, slo/analyzer.py, slo/optimizer.py, slo/alolab.py
Model card given to the meta-agent slo/config.py ModelCard
5 runs, best-validation checkpoint, bootstrap CIs (Section 4) scripts/run_alolab.py, slo/stats.py
Paired McNemar significance (Section 4.6) slo/stats.py
Meta-agent quality ablation (Section 4.8) --meta weak
History-only ablation (Section 2.4) --no-model-card

Models

The paper uses Llama 3.1-8B, Gemma 2-9B and Qwen 2.5-7B served on Vertex AI. The NVIDIA NIM catalogue used here exposes Llama 3.1-8B exactly, plus same-tier small models that reproduce the same pathologies. Substitutions are recorded in each ModelCard.note.

Target models (7–9B deployed tier): llama-8b (exact paper model), gemma-2b (reproduces the Gemma markdown-fence failure), gemma-3n, nemotron-mini.

Meta-agent (the AloLab Analyzer + Optimizer — paper uses Claude Sonnet 4.5): strong = qwen/qwen3-next-80b-a3b-instruct (produced the reported results, but can be intermittently overloaded on NIM), strong-nemotron = nvidia/nemotron-3-super-120b-a12b (recommended for fresh runs — reliable, ~3.5s/call), strong-deepseek = deepseek-v4-flash, weak = llama-3.1-8b (the paper's Haiku ablation role). The backend is pluggable — set --meta-backend anthropic with an ANTHROPIC_API_KEY to use Claude directly.


Setup

python -m venv .venv && . .venv/Scripts/activate      # or your existing .venv
pip install -r requirements.txt

.env (already present) must define an OpenAI-compatible endpoint:

OPENAI_API_KEY=nvapi-...
OPENAI_BASE_URL=https://integrate.api.nvidia.com/v1

Usage

# 1. Baselines on the test split (reproduce the format gap)
python scripts/run_baseline.py --model llama-8b --strategy naive
python scripts/run_baseline.py --model gemma-2b --strategy reference --workers 8

# 2. AloLab optimization (5 runs) + test evaluation
python scripts/run_alolab.py --model gemma-2b --meta strong --runs 5 --workers 8

# 3. Ablations
python scripts/run_alolab.py --model llama-8b --meta weak      # meta-quality
python scripts/run_alolab.py --model llama-8b --no-model-card  # history-only

# 4. Everything, resumably, then a report
python scripts/run_all.py --workers 8 --runs 5 --ablations
python scripts/aggregate.py --dataset gsm8k     # -> results/REPORT_gsm8k.md

Scale toward the paper by overriding split sizes: --optimize 150 --validation 100 --test 1319. Defaults (slo/config.py DEFAULT_SPLITS) are smaller so a run finishes on a free-tier key.

Tests

python -m pytest tests/ -q

The metrics tests encode the paper's definitions directly — e.g. a markdown-fenced-but-correct JSON object scores task_correct=True, json_valid=False, output_correct=False (correct-but-unusable).

Results layout

results/<dataset>/<model_slug>/
  naive|reference|strict|constrained|ref+constrained/
      traces.jsonl   summary.json
  alolab[_meta-weak|_nocard]/
      run<k>.json         # full optimization record incl. every rewritten prompt
      test_run<k>.jsonl   # test traces for that run's best prompt
      summary.json        # mean output accuracy + 95% bootstrap CI + std
  REPORT_<dataset>.md     # generated tables vs. the paper

Each run<k>.json contains the actual prompts the meta-agent wrote at every epoch, so you can read the optimization trajectory (e.g. the moment the optimizer adds the no-markdown-fence directive for Gemma — the paper's Appendix A story).

Differences from the paper (honest notes)

  • Target models are the closest available NIM substitutes for Gemma-2-9B / Qwen-2.5-7B; Llama-3.1-8B is exact. The pathologies reproduce (Gemma fences).
  • Meta-agent is a strong open model, not Claude Sonnet 4.5 (pluggable; see above).
  • Default split sizes are smaller than the paper's for tractability; override to match. GSM8K is the primary reproduction target (unambiguous numeric grading); MATH is supported but secondary.
  • CONSTRAINED decoding uses NIM guided_json; availability depends on the served model.

Reproduction scope & caveats

This repo is a faithful mini-reproduction, not a 1:1 rerun of the paper's compute. Read RESULTS.md for the full findings and honest gap analysis. In short:

  • Every qualitative claim reproduces (format gap → 0% output accuracy; REFERENCE fails for fence-wrapping models; AloLab fixes it with significant McNemar gains; meta-agent quality matters; the model card doesn't). The capable model closest to the paper's tier (gemma-3n) lands in the paper's 84–87% band.
  • Absolute numbers run lower than the paper's, for measured reasons: smaller target models (task-ceiling limited), a weaker/flakier open meta-agent than Claude Sonnet 4.5, and reduced split sizes. All quantified in RESULTS.md §5.
  • Endpoint rate-limiting is real: keep --workers 3–4 for meta/llama-3.1-8b-instruct — higher concurrency triggers 503s that silently depress its scores. Gemma models tolerate --workers 8. The *_hifi result folders are the clean low-concurrency, validation-n=100 runs and are the canonical numbers (marked ‡ in the generated report).
  • Open next steps (not yet done): a Claude Sonnet meta-agent run to chase the last ~16pp, full paper-scale splits (--test 1319), and the MATH track.

This is an independent research re-implementation; see LICENSE. It is not affiliated with the paper's authors, and contains no code from the original.

About

Reproduction & re-implementation of the AloLab paper (arXiv:2605.02363): closing the structured-output reliability gap in small LLMs via iterative black-box prompt optimization. GSM8K, 4 models, ablations, McNemar significance.

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages