Instructions to use arrochi112/SchemaForge-1B-JSON-Extractor with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use arrochi112/SchemaForge-1B-JSON-Extractor with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("text-generation", model="arrochi112/SchemaForge-1B-JSON-Extractor") messages = [ {"role": "user", "content": "Who are you?"}, ] pipe(messages)# Load model directly from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer = AutoTokenizer.from_pretrained("arrochi112/SchemaForge-1B-JSON-Extractor") model = AutoModelForCausalLM.from_pretrained("arrochi112/SchemaForge-1B-JSON-Extractor", device_map="auto") messages = [ {"role": "user", "content": "Who are you?"}, ] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", ).to(model.device) outputs = model.generate(**inputs, max_new_tokens=40) print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:])) - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use arrochi112/SchemaForge-1B-JSON-Extractor with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "arrochi112/SchemaForge-1B-JSON-Extractor" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "arrochi112/SchemaForge-1B-JSON-Extractor", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker
docker model run hf.co/arrochi112/SchemaForge-1B-JSON-Extractor
- SGLang
How to use arrochi112/SchemaForge-1B-JSON-Extractor with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "arrochi112/SchemaForge-1B-JSON-Extractor" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "arrochi112/SchemaForge-1B-JSON-Extractor", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "arrochi112/SchemaForge-1B-JSON-Extractor" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/chat/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "arrochi112/SchemaForge-1B-JSON-Extractor", "messages": [ { "role": "user", "content": "What is the capital of France?" } ] }' - Docker Model Runner
How to use arrochi112/SchemaForge-1B-JSON-Extractor with Docker Model Runner:
docker model run hf.co/arrochi112/SchemaForge-1B-JSON-Extractor
SchemaForge: Distilling Ultra-Large Foundation Models into Edge SLMs for Real-Time Enterprise JSON Extraction
A Comparative Study of Gemma-4 Teachers and MiniCPM5-1B
Arjhine A. Ty
Project codename: SchemaForge
Student model: SchemaForge-1B (base architecture: openbmb/MiniCPM5-1B, 1.08B parameters)
Teacher models: google/gemma-4-31B (31B flagship), google/gemma-4-E4B-it (4B/8B-effective instruction-tuned)
Production checkpoint: schemaforge-1b-iter2 (formerly distilled_minicpm5_1b_iter2)
Hardware: 1 Γ NVIDIA RTX PRO 6000 Blackwell Edition (96 GB VRAM), Nebius AI Cloud
Software: Python 3.12, PyTorch 2.5, HuggingFace transformers 5.x, vLLM (serving)
Version: 1.0 β August 2026
Abstract
Structured entity extraction β converting heterogeneous unstructured business documents into strongly-typed, schema-conformant JSON β is one of the highest-volume workloads in enterprise retrieval-augmented generation (RAG) and transactional automation. The prevailing solution is to deploy a 30B+ parameter general-purpose foundation model behind a constrained-decoding wrapper. This is operationally untenable at scale: in our measurements a gemma-4-31B-class teacher requires β38.5 GB of VRAM and sustains only 12.40 tokens/second of greedy decode throughput on a single 96 GB accelerator, leaving room for only two concurrent workers and failing sub-second API SLAs.
We present SchemaForge, a sequence-level knowledge-distillation framework that compresses the structural reasoning and JSON formatting strictness of Gemma-4 teachers into openbmb/MiniCPM5-1B, a dense 1.08-billion-parameter edge Small Language Model (SLM). SchemaForge combines a hard-target cross-entropy objective with a temperature-scaled, log-space soft-target KL divergence under a multi-task weighting $\mathcal{L}{KD} = \alpha\mathcal{L}{CE} + (1-\alpha)\tau^{2}\mathcal{L}_{KL}$ ($\alpha = 0.5$, $\tau = 2.0$), and resolves the teacherβstudent vocabulary mismatch ($|\mathcal{V}_T| = 256{,}000 \rightarrow |\mathcal{V}_S| = 130{,}560$) via a dual-tokenizer cross-encoding scheme with shared-subspace logit projection.
Empirically, the distilled student attains a 0.0 % JSON syntax error rate (100 % structural validity) and a field-level extraction F1 of 1.000 across a five-domain in-domain enterprise benchmark suite ($n = 5$ documents β finance, logistics, IT procurement, biomedical, cloud operations), against 34.2 % error rate and 0.612 F1 for the undistilled base model. It sustains 61.91β76.27 tokens/second in a β2.4 GB VRAM footprint β a 16.0Γ memory reduction and a 5.0Γβ6.2Γ throughput improvement over the 31B teacher, permitting 36 concurrent inference workers per 96 GB card (β2,746 tokens/second aggregate system throughput).
Our central scientific finding is not the compression ratio but a sharp negative result about prompt fragility. Across three controlled retraining iterations we observe that a 1B-parameter student is hypersensitive to the surface form of the prompt header: identical data, identical loss, identical hyperparameters, and only a changed instruction prefix moved zero-shot validity on the public suneeldk/text-json benchmark from 0.0 % to 70.0 % and back to 0.0 %. Distillation at this scale transfers a format-conditioned skill, not a format-invariant one. We formalize this as the trainβinference template gap, quantify it, and derive a practical mitigation (template canonicalization plus a serving-side prompt contract) that we recommend as mandatory for any sub-2B structured-output deployment.
We release the production checkpoint, model card, proof artifacts, and full reproduction protocol.
Keywords: knowledge distillation, small language models, structured generation, JSON extraction, prompt sensitivity, edge inference, vLLM, enterprise RAG
Executive Summary
| Dimension | 31B Teacher Baseline | SchemaForge-1B (Distilled) | Delta |
|---|---|---|---|
| In-domain JSON syntax error rate (n = 5 docs) | 0.0 % | 0.0 % | Parity |
| In-domain field extraction F1 (n = 5 docs) | 1.000 | 1.000 | Parity |
Zero-shot validity, suneeldk/text-json |
β (not evaluated) | 70.0 % | β |
| Decode throughput | 12.40 tok/s | 61.91β76.27 tok/s | 5.0Γβ6.2Γ |
| Peak VRAM footprint | β38.5 GB | β2.4 GB | 16.0Γ smaller |
| Concurrent workers per 96 GB card | 2 | 36 | 18Γ |
| Aggregate system throughput | 24.8 tok/s | β2,746 tok/s | β110Γ |
The three claims a reader should take away:
- Task-specialized distillation closes the capability gap at β29Γ fewer parameters. For a bounded, schema-constrained generation task, a 1.08B student matches a 31B teacher on both structural validity and field accuracy across our five-document in-domain suite. Generality is what gets compressed away; task competence is not.
- The economics change category, not degree. 16Γ less memory and 5β6Γ faster decoding is the difference between two concurrent workers per GPU and thirty-six. This converts a per-document cost line into a rounding error.
- The fragility is in the prompt, not the weights. The single largest swing in our entire experimental record β a 70-point accuracy delta β was caused by changing an instruction header. Practitioners deploying sub-2B structured extractors must treat the prompt template as a versioned, tested API contract.
Caveat stated up front: the in-domain suite comprises $n = 5$ curated documents (one per domain) and the winning checkpoint was distilled on $n = 5$ training samples. The 1.000 F1 and 0.0 % error rate are therefore exact-match results on a small, curated set, not population estimates. See Β§10 Limitations, which we consider a load-bearing section of this paper rather than a formality.
This is a v1 release. A second training and evaluation campaign is planned against a real-world document corpus with an expanded metric set, multi-seed variance, and competitive baselines; it is scoped in Β§11.3 against the limitation register in Β§10.6. The accuracy figures below should be read as provisional pending that run.
Table of Contents
- Introduction
- Related Work
- Problem Formulation
- The SchemaForge Distillation Framework
- Engineering: Running Legacy MiniCPM Under
transformers5.x - Hyperparameter Sensitivity Analysis
- Prompt Template Sensitivity: A Three-Iteration Study
- Empirical Results
- Deployment and Publication
- Limitations and Threats to Validity
- Conclusion and Future Work
- References
- Appendix A: Complete Hyperparameter Specification
- Appendix B: Compatibility Patch Source
- Appendix C: Verbatim Prompt Templates
- Appendix D: Benchmark Schemas and Source Documents
- Appendix E: Reproducibility Checklist
1. Introduction
1.1 The Structured Extraction Bottleneck
Enterprise document intelligence pipelines are, in aggregate, a text-to-JSON problem. Accounts-payable automation reads invoices into {invoice_number, vendor_name, invoice_date, subtotal, tax, grand_total}. Freight settlement reads bills of lading into {bill_of_lading, carrier_name, ship_date, container_count, freight_cost}. Clinical procurement reads requisitions into {order_id, supplier, order_date, items[], total_price}. In every case the input is heterogeneous, semi-structured, and adversarially formatted by whichever upstream system emitted it; the output is a rigid, strongly-typed schema that a downstream database, ERP, or RAG index will reject outright if a single brace is unbalanced or a numeric field arrives as a string.
The industry default is to point a large instruction-tuned foundation model at the problem. This works. It also produces three compounding operational failures:
(1) Memory pressure that forecloses multi-tenancy. A 31B-parameter model in bfloat16 requires roughly 62 GB for weights alone; even under aggressive quantization the working set we measured is β38.5 GB including KV cache. On a 96 GB RTX PRO 6000 Blackwell this admits only two concurrent workers at 90 % memory utilization β against 36 for the distilled student. Utilization economics collapse.
(2) Throughput below SLA. We measure 12.40 tokens/second of greedy decode on a single-GPU deployment of the 31B teacher. A typical extraction emits 120β250 JSON tokens, implying 10β20 seconds per document. Interactive document-upload flows target sub-2-second response; bulk nightly runs of 10βΆ documents become physically impossible on any reasonable cluster budget.
(3) Cost per document that does not amortize. Because throughput is low and memory forbids batching, the marginal cost of a document is dominated by GPU-seconds at the highest-tier instance price. There is no batching lever to pull.
Critically, the capability being paid for is not the capability being used. A 31B model's value lies in open-domain reasoning, long-horizon planning, multilingual fluency, and code synthesis. Invoice extraction requires almost none of that. It requires: (a) reliable span identification, (b) light numeric normalization, and (c) absolute fidelity to a JSON grammar. The pricing model charges for a general intelligence; the workload consumes a narrow, learnable competence.
This mismatch is precisely the setting in which knowledge distillation is theoretically well-motivated.
1.2 The SchemaForge Approach
SchemaForge is a sequence-level knowledge-distillation framework that transfers the schema-fidelity behavior of Gemma-4-class teachers into a 1.08B-parameter student. Its design commitments are:
- Task-narrow, not capability-narrow. We do not attempt to preserve the teacher's general ability. We preserve exactly one behavior: emit valid, schema-conformant JSON for a document-extraction prompt.
- Multi-task objective. Pure soft-logit matching is insufficient for structured output, because the tokens that matter most (
{,",:,,,}) are exactly the tokens where the teacher distribution is nearly deterministic and therefore carries little dark knowledge. We anchor with hard cross-entropy at $\alpha = 0.5$. - Vocabulary-mismatched teachers. Gemma-4 and MiniCPM5 do not share a tokenizer. We do not retokenize or retrain embeddings; we cross-encode and project onto the shared logit subspace.
- Deployability as a first-class metric. Every result is reported alongside its VRAM footprint and throughput, measured on the same physical card.
1.3 Contributions
This paper makes five contributions:
C1 β A reproducible distillation recipe for sub-2B structured extractors. We specify the complete objective, hyperparameter bounds, and dataset scaling regime under which a 1.08B student reaches teacher-parity on schema fidelity (Β§4, Β§6).
C2 β A dual-tokenizer cross-encoding and logit-projection scheme that permits KL-based distillation between models with incompatible vocabularies ($256{,}000 \rightarrow 130{,}560$) without embedding surgery or teacher retokenization (Β§4.4).
C3 β Three low-level runtime compatibility patches required to run legacy trust_remote_code MiniCPM checkpoints (openbmb/MiniCPM-1B-sft-bf16) under Python 3.12 / PyTorch 2.5 / transformers 5.x, including a tied-weights failure that reports itself as a corrupted checkpoint while actually being an API change (Β§5). We also report the negative counterpart: migrating to openbmb/MiniCPM5-1B, a stock LlamaForCausalLM, removed the need for all three (Β§5.4).
C4 β The trainβinference template gap: a quantified negative result. Three controlled iterations isolating prompt-header formatting as the sole varying factor, producing a 70-percentage-point accuracy swing on a public benchmark (Β§7). We argue this is the dominant failure mode for small structured-output models and is systematically under-reported.
C5 β A complete production deployment protocol: HuggingFace publication, model card specification, vLLM serving configuration, schema-constrained decoding integration, and capacity planning arithmetic (Β§9).
1.4 Paper Roadmap
Β§2 situates the work. Β§3 formalizes the task and metrics. Β§4 gives the mathematics. Β§5 documents the engineering. Β§6 and Β§7 present the two ablation studies. Β§8 reports all benchmarks. Β§9 covers shipping. Β§10 is an honest accounting of what these numbers do and do not establish.
2. Related Work
2.1 Knowledge Distillation
The foundational formulation is due to Hinton, Vinyals, and Dean [1], who introduced temperature-softened logit matching with the now-standard $\tau^2$ gradient-rescaling term. Their key insight β that the relative probabilities a teacher assigns to incorrect classes encode a similarity structure ("dark knowledge") absent from one-hot labels β motivates our soft-target term. BuciluΔ et al. [2] anticipated the model-compression framing. Romero et al. [3] extended supervision to intermediate representations via FitNets; we deliberately do not use hidden-state matching, because the teacher (Gemma-4, $d_{\text{model}}$ and depth both far larger) and student (MiniCPM5-1B, 24 layers) have no principled layer correspondence, and projection heads introduce hyperparameters we could not afford to tune under our compute budget.
2.2 Sequence-Level Distillation for Autoregressive Models
Kim and Rush [4] established sequence-level knowledge distillation for neural machine translation, showing that training the student on teacher-generated output sequences (rather than only token-level distributions over gold data) substantially outperforms word-level KD. SchemaForge is sequence-level in this sense: our training targets are teacher-generated JSON completions, not human-annotated gold JSON. Sanh et al. [5] demonstrated the approach at scale with DistilBERT (40 % smaller, 97 % of GLUE performance), and Jiao et al. [6] with TinyBERT. More recent work on distilling instruction-following behavior β Gu et al. [7] on MiniLLM, and Agarwal et al. [8] on generalized KD with on-policy student samples β addresses the exposure-bias mismatch that arises when the student is trained on teacher trajectories but evaluated on its own. We note this as an unexploited improvement in Β§11.2.
2.3 Structured and Constrained Generation
Guaranteeing syntactic validity of model output is an active area. Willard and Louf [9] (Outlines) reformulate constrained decoding as finite-state-machine-guided token masking, achieving provable grammar conformance with negligible overhead. JSONFormer, lm-format-enforcer, and vLLM's native guided-decoding backends implement variants of the same idea. This literature is complementary rather than competing: constrained decoding guarantees syntax, but cannot guarantee semantics β a grammar-constrained model will happily emit a well-formed JSON object with the wrong vendor name. SchemaForge targets semantic field accuracy and learned formatting discipline; we recommend layering FSM-constrained decoding on top in production (Β§9.4) as defense in depth.
2.4 Small Language Models for the Edge
The MiniCPM line [10] argues that carefully-scaled sub-3B models can match 7Bβ13B models on targeted benchmarks, using depth-scaled residual connections and a wide-vocabulary tokenizer. The Phi series [11] makes the parallel argument from the data-quality side. Gemma [12] and its successors provide open-weight teachers at multiple scales. Our teacher pair β a 31B flagship and a 4B/8B-effective instruction-tuned variant β was chosen specifically to test whether teacher scale matters for a bounded task; Β§8.3 reports that, at this task difficulty, it does not.
2.5 Prompt Sensitivity and Format Brittleness
Our central negative result connects to a growing literature on format brittleness. Sclar et al. [13] demonstrate that LLM performance varies by up to 76 accuracy points under semantically equivalent prompt formatting perturbations β separator choice, casing, whitespace β and argue for reporting performance spreads rather than point estimates. Lu et al. [14] show analogous sensitivity to few-shot example ordering. Mizrahi et al. [15] advocate multi-prompt evaluation as standard practice. Our contribution to this thread is specific and, we believe, novel in emphasis: we show that distillation into a small student does not merely inherit this brittleness β it concentrates it, because the student's limited capacity causes it to bind the learned behavior tightly to the exact training prefix. The 0.0 % β 70.0 % β 0.0 % trajectory in Β§7 is a starker instance than is typically reported for larger models, and it has direct deployment consequences.
3. Problem Formulation
3.1 Task Definition
Let $x \in \Sigma^*$ denote an unstructured source document (an invoice, bill of lading, requisition, or receipt) over an alphabet $\Sigma$, and let $\mathcal{S}$ denote a target schema: a finite set of typed field names
The extraction task is to learn a mapping
where $\mathcal{J}(\mathcal{S}) \subset \Sigma^*$ is the set of strings that are (i) syntactically valid JSON under RFC 8259, and (ii) conformant to $\mathcal{S}$ β every $k_i$ present, every value inhabiting type $\theta_i$.
Because $f_\Theta$ is realized by an autoregressive language model, generation factorizes as
and the schema $\mathcal{S}$ enters the conditioning only through the prompt surface form. This is the structural reason Β§7's finding is possible: the schema is not an architectural constraint, it is a string, and the model's response to that string is learned.
3.2 Evaluation Metrics
We report four orthogonal metrics. Reporting any one alone is, we argue, the standard failure of structured-extraction papers.
(M1) JSON Syntax Validity Rate. The fraction of generations that parse under a strict RFC 8259 parser:
This is a necessary but wholly insufficient condition. {} is valid JSON.
(M2) Field-Level Extraction F1. Over parsed generations, treating each $(\text{key}, \text{value})$ pair as a retrievable item against the reference object $y^\star$:
where $\hat{K} = {(k,v) \in \hat{y}}$ and $K^\star = {(k,v) \in y^\star}$, with values compared after type-aware normalization (numeric strings coerced to floats and compared within $\varepsilon = 10^{-6}$; dates normalized to ISO-8601 YYYY-MM-DD; string fields compared after Unicode NFKC normalization and whitespace collapse). An unparseable generation contributes $F_1 = 0$.
(M3) Decode Throughput. Generated tokens per wall-clock second, measured under greedy decoding (do_sample=False, temperature=0.0), batch size 1, excluding model load and tokenizer initialization, averaged over the generation:
(M4) Peak VRAM Footprint. torch.cuda.max_memory_allocated() over the full generation, in GB, inclusive of weights, activations, and KV cache.
3.3 Why All Four Are Required
A model can be trivially optimized for any one metric in isolation:
- Emit
{}always β 100 % validity, F1 β 0. - Emit the entire source document verbatim inside a string field β high recall, invalid or useless.
- Truncate at 8 tokens β excellent throughput, no content.
- Load in 4-bit with no KV cache β minimal VRAM, degraded accuracy.
The joint frontier is what matters. Throughout Β§8 we report all four for every configuration, on identical hardware, so that no result is a metric artifact.
3.4 Datasets
| Dataset | Role | Size | Provenance |
|---|---|---|---|
| SchemaForge in-domain suite (BMK-01β¦05) | In-domain evaluation | $n = 5$ documents (1 per domain) | Hand-authored, synthetic, 5 enterprise verticals |
| Iteration-2 distillation set | Training (winning run) | $n = 5$ samples | Teacher-generated, single canonical template |
| Iteration-1 distillation set | Training | $n = 20$ samples | Teacher-generated, chat-token template |
| Iteration-3 distillation set | Training | $n = 15$ samples, 15 schemas | Teacher-generated, system-persona template |
suneeldk/text-json |
Out-of-domain / zero-shot evaluation | 2,000 records available; evaluation subset drawn per iteration | Public HuggingFace dataset, real enterprise documents |
The small training-set sizes are deliberate β the research question was explicitly how little supervision suffices when the signal is a teacher's full logit distribution over a narrow task. They are also, unavoidably, the primary limitation of this work; see Β§10.
4. The SchemaForge Distillation Framework
4.1 Architecture Overview
βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
β UNSTRUCTURED SOURCE TEXT x β
β (invoices Β· bills of lading Β· requisitions Β· receipts) β
βββββββββββββββββββββββββββββββββ¬ββββββββββββββββββββββββββββββββββββββββ
β
βββββββββββββββββ΄ββββββββββββββββ
β CANONICAL PROMPT TEMPLATE T β β Β§7: the critical component
β "Extract structured JSON β
β from the text:\n{x}\n β
β JSON Output:" β
βββββββββββββββββ¬ββββββββββββββββ
β
βββββββββββββββββββββββββ΄ββββββββββββββββββββββββ
β β
βΌ βΌ
βββββββββββββββββββββββββββ βββββββββββββββββββββββββββ
β TOKENIZER_T (Gemma-4) β β TOKENIZER_S (MiniCPM5) β
β |V_T| = 256,000 β β |V_S| = 130,560 β
βββββββββββββ¬ββββββββββββββ βββββββββββββ¬ββββββββββββββ
β t_ids β s_ids
βΌ βΌ
βββββββββββββββββββββββββββ βββββββββββββββββββββββββββ
β TEACHER (frozen) β β STUDENT (trainable) β
β gemma-4-31B / E4B-it β β MiniCPM5-1B, 24 layersβ
β torch.no_grad() β β single-GPU, eager attn β
βββββββββββββ¬ββββββββββββββ βββββββββββββ¬ββββββββββββββ
β z_T β R^{BΓLΓ256000} β z_S β R^{BΓLΓ130560}
β β
βΌ β
βββββββββββββββββββββββββββ β
β VOCABULARY PROJECTION β β
β z_T[:, :, :|V_S|] β Β§4.4 β
β β R^{BΓLΓ130560} β β
βββββββββββββ¬ββββββββββββββ β
β β
ββββββββββββββββββββ¬βββββββββββββββββββββββββ
βΌ
ββββββββββββββββββββββββββββββββββββββββ
β L_KD = Ξ±Β·L_CE + (1βΞ±)Β·ΟΒ²Β·L_KL β Β§4.2
β Ξ± = 0.5, Ο = 2.0 β
β log-space KL, batchmean reduction β Β§4.3
ββββββββββββββββββββ¬ββββββββββββββββββββ
β AdamW, lr 2e-5, cosine, warmup 0.05
βΌ
ββββββββββββββββββββββββββββββββββββββββ
β SchemaForge-1B (β2.4 GB) β
β vLLM Β· Outlines/Pydantic guardrail β
β 61.91 β 76.27 tok/s β
ββββββββββββββββββββββββββββββββββββββββ
4.2 The Multi-Task Distillation Objective
Knowledge distillation for autoregressive language models minimizes the discrepancy between the student's conditional distribution $P_S(y_t \mid y_{<t}, x)$ and the teacher's $P_T(y_t \mid y_{<t}, x)$ over target-sequence positions. SchemaForge's objective is the convex combination
with $\alpha = 0.5$ and $\tau = 2.0$.
Hard-target term. $\mathcal{L}_{CE}$ is the standard causal cross-entropy over target JSON tokens:
This term anchors the student to exact structural tokens. For JSON generation this is not a formality: the delimiters {, }, [, ], ", :, , are positions where the teacher's distribution is near-degenerate (probability mass β 1 on a single token), and therefore positions where the soft term carries almost no gradient signal. Hard cross-entropy is what teaches the grammar.
Soft-target term. $\mathcal{L}_{KL}$ is the temperature-softened KullbackβLeibler divergence between student and teacher distributions:
This term carries the dark knowledge: on content positions (a vendor name, a date, a decimal amount) the teacher's distribution is genuinely uncertain, and its shape encodes which alternative spans were plausible. That relational structure is the thing worth distilling.
The $\tau^2$ factor. Softening logits by $\tau$ scales the gradient of the KL term by $1/\tau^2$. Multiplying by $\tau^2$ restores gradient magnitudes to the same scale as $\mathcal{L}_{CE}$, so that $\alpha$ is a true mixing weight rather than a quantity entangled with the temperature. This follows Hinton et al. [1, Β§2].
Why $\alpha = 0.5$. Β§6.3 reports the ablation. Briefly: $\alpha = 0.2$ (KL-dominant) produced schema drift β plausible-looking but off-schema key names on out-of-domain prompts, because the student was optimizing distributional similarity rather than literal token identity. $\alpha = 0.5$ preserves both exact syntax and soft distributional shape.
4.3 Numerically Stable Log-Space Formulation
A naΓ―ve implementation computing $\text{softmax}$, then $\log$, then the KL sum, underflows catastrophically. With $|\mathcal{V}_S| = 130{,}560$ and bfloat16 activations (β3 decimal digits of mantissa precision), probabilities in the tail routinely fall below the representable normal range, producing $\log(0) = -\infty$ and, on the next backward pass, NaN gradients that silently poison every parameter.
SchemaForge computes the divergence entirely in log-space:
log_softmax is implemented with the max-subtraction trick,
guaranteeing that the largest exponentiated term is exactly $1$ and no intermediate overflows. The reference implementation:
import torch
import torch.nn.functional as F
def distillation_loss(student_logits, teacher_logits, labels,
alpha: float = 0.5, tau: float = 2.0):
"""
SchemaForge multi-task distillation objective.
student_logits : (B, L, |V_S|) requires_grad
teacher_logits : (B, L, |V_S|) already projected, detached
labels : (B, L) -100 at masked positions
"""
# ---- Hard target: causal cross-entropy, shifted by one ----------------
shift_student = student_logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
ce_loss = F.cross_entropy(
shift_student.view(-1, shift_student.size(-1)),
shift_labels.view(-1),
ignore_index=-100, # masks prompt + padding
)
# ---- Soft target: log-space KL, both operands in log-probability ------
s_logprob = F.log_softmax(shift_student / tau, dim=-1)
t_logprob = F.log_softmax(
teacher_logits[..., :-1, :].contiguous() / tau, dim=-1
)
kl_loss = F.kl_div(
s_logprob, # input: log-probabilities
t_logprob, # target: log-probabilities
reduction="batchmean",
log_target=True, # <- critical: avoids exp() of target
)
# ---- tau^2 gradient rescaling + convex combination --------------------
return alpha * ce_loss + (1.0 - alpha) * (tau ** 2) * kl_loss
Two details are load-bearing and easy to get wrong:
log_target=True. Without it, PyTorch'skl_divexpects the target as raw probabilities and internally appliesexp(), reintroducing exactly the underflow we eliminated.reduction="batchmean". PyTorch's default"mean"divides by $B \times L \times |\mathcal{V}|$, not $B$, yielding a loss smaller by a factor of ~130,560 and an effective learning rate five orders of magnitude below what the schedule specifies. This is the single most common silent bug in KD implementations.
4.4 Dual-Tokenizer Cross-Encoding and Vocabulary Projection
Gemma-4 teachers use a SentencePiece vocabulary of $|\mathcal{V}_T| = 256{,}000$; MiniCPM5-1B uses $|\mathcal{V}_S| = 130{,}560$. The tokenizers are not merely different in size β they induce different segmentations of the same string. "$2,915.00" may be five tokens under one and eight under the other. Consequently:
Three standard remedies exist, and we rejected two:
| Approach | Description | Why rejected / adopted |
|---|---|---|
| Teacher retokenization | Re-encode teacher output with student tokenizer, re-run teacher | Requires teacher embedding surgery; teacher logits become meaningless outside its own vocabulary |
| Optimal-transport token alignment | Learn a soft mapping $\pi$ via OT over embedding similarity | Adds a trainable component and substantial compute; deferred to future work |
| Cross-encoding + shared-subspace projection | Encode independently per tokenizer; slice teacher logits to student vocabulary width | Adopted β zero additional parameters, no teacher modification |
The scheme. Define independent encodings of the concatenated prompt-and-response string:
each forwarded through its own model. Teacher logits are then projected onto the student's vocabulary subspace by truncation:
# Shared-subspace logit projection
t_logits = t_out.logits[:, :, : s_out.logits.size(-1)] # 256000 -> 130560
Justification and honest accounting of the approximation. This truncation is principled to the extent that both vocabularies order tokens by descending corpus frequency, so the leading $130{,}560$ indices of $\mathcal{V}_T$ concentrate the overwhelming majority of probability mass for English business text β retaining $130{,}560/256{,}000 = 51.0,%$ of the index range but far more than that share of the mass. It is nonetheless a lossy, index-aligned rather than semantics-aligned projection: index $i$ in $\mathcal{V}_T$ and index $i$ in $\mathcal{V}_S$ do not denote the same token. What the KL term therefore transfers is best characterized as distributional shape and entropy structure β how sharp or diffuse the teacher is at each position β rather than exact per-token identity correspondence. The hard cross-entropy term supplies the identity-level signal. We regard this decomposition as the honest description of why the objective works, and we flag the projection as the most theoretically fragile component of SchemaForge in Β§10.4.
Label masking. Loss must be computed strictly on schema-output tokens. Prompt tokens and padding positions receive the sentinel label $-100$, which F.cross_entropy(ignore_index=-100) excludes:
Omitting this causes the student to spend capacity learning to reproduce input documents β a failure mode that manifests as high training-loss reduction with no improvement in extraction accuracy.
4.5 Training Configuration
| Component | Setting | Rationale |
|---|---|---|
| Optimizer | AdamW, $\beta = (0.9, 0.999)$, $\epsilon = 10^{-8}$ | Standard; decoupled weight decay |
| Learning rate | $2 \times 10^{-5}$ | See Β§6.2 β $\geq 3\times10^{-5}$ degrades decoding |
| Schedule | Cosine decay, warmup ratio $0.05$ | Prevents early-step destabilization |
| Weight decay | $0.01$ | |
| Gradient clipping | $1.0$ (global $\ell_2$ norm) | Guards against loss spikes on short sequences |
| Epochs | $\leq 3$, early stopping on val loss | See Β§6.1 β repetition collapse beyond |
| Precision | bfloat16 |
Wider exponent than fp16; no loss-scaling needed |
| Attention | PyTorch 2.5 eager (SDPA) | FlashAttention-2 was attempted and abandoned β see note below |
| Sharding | None (single GPU) | ZeRO-3 inapplicable to a single card β Β§10.3 item 6 |
| Batch size | 2 | Set under a mistaken 48 GB VRAM assumption β item 7 |
| Max sequence length | 1,024 tokens (as executed) | Covers document + schema + JSON output; 2,048 was specified but the script used 1,024 |
| Teacher | Frozen, torch.no_grad(), eval() |
No teacher gradients; halves memory |
5. Engineering: Running Legacy MiniCPM Under transformers 5.x
Scope note. The three patches below were blocking runtime exceptions encountered while running
openbmb/MiniCPM-1B-sft-bf16β our initial student β withtrust_remote_code=Trueundertransformers5.x. That checkpoint ships amodeling_minicpm.pyauthored in early 2024 against the 4.x API surface, and it does not load under 5.x without all three fixes.The released checkpoint uses
openbmb/MiniCPM5-1B, a stockLlamaForCausalLM(24 layers, hidden 1536, GQA 16/2, vocabulary 130,560, 1,080,632,832 parameters) that requires no custom modeling code and notrust_remote_code=True, and which ran cleanly across all three distillation iterations. Both facts are worth recording: the patches are what it takes to run the older MiniCPM lineage on a modern stack, and Β§5.4 explains why migrating removed the need for them entirely.
Environment: Python 3.12 Β· PyTorch 2.5 Β· transformers 5.x Β· student openbmb/MiniCPM-1B-sft-bf16 (trust_remote_code=True).
5.1 Patch 1 β is_torch_fx_available Re-Injection
Symptom. ImportError: cannot import name 'is_torch_fx_available' from 'transformers.utils.import_utils' at dynamic-module load.
Cause. modeling_minicpm.py (authored early 2024) imports this utility for torch.fx symbolic-tracing support. transformers 5.x deprecated and removed it from import_utils. Because the modeling file is fetched and executed dynamically at from_pretrained() time, the failure occurs inside the library's import machinery, not in user code β so the traceback points nowhere useful.
Fix. Re-inject the symbol before any from_pretrained call. Symbolic tracing is not used in our training path, so returning False unconditionally is sufficient and avoids pulling in torch.fx.
import transformers.utils.import_utils as import_utils
import_utils.is_torch_fx_available = lambda: False
5.2 Patch 2 β Tied Weights and the Missing lm_head.weight
Symptom. Two failures in sequence:
[transformers] This checkpoint seems corrupted. The tied weights mapping
for this model specifies to tie lm_head.weight ...
followed by type-assertion errors inside PreTrainedModel.all_tied_weights_keys.
Cause. Two independent problems that present as one.
- The weight genuinely is not on disk. The
openbmb/MiniCPM-1B-sft-bf16safetensors checkpoint omitslm_head.weight, because the output projection is tied toembed_tokens.weight. This is legitimate and space-saving, but 5.x's stricter loader reports the absence as corruption rather than resolving the tie. - The tying declaration has the wrong type. In 4.x,
_tied_weights_keyswas alist[str]. In 5.x it is adict[str, str]mapping tied parameter β source parameter, to support richer tying topologies. Legacy modeling files declare the list form, which strict 5.x assertions reject.
Fix. Re-establish the tie explicitly after instantiation, and patch the expansion helper so the loader stops treating the omission as corruption.
import transformers.modeling_utils as mu
# (a) Re-tie the output projection to the input embedding.
student.lm_head.weight = student.model.embed_tokens.weight
# (b) Satisfy the 5.x dict-shaped tied-weights contract.
_orig = mu.PreTrainedModel.get_expanded_tied_weights_keys
def _patched(self, *args, **kwargs):
keys = getattr(self, "_tied_weights_keys", None)
if isinstance(keys, (list, tuple)):
self._tied_weights_keys = {k: "model.embed_tokens.weight" for k in keys}
return _orig(self, *args, **kwargs)
mu.PreTrainedModel.get_expanded_tied_weights_keys = _patched
Of the three, this is the one most likely to bite others: the error message says corrupted checkpoint, which points at the download rather than at the API change that actually caused it.
5.3 Patch 3 β DynamicCache 5.x Layer-State Compatibility
Symptom. AttributeError: 'DynamicCache' object has no attribute 'from_legacy_cache' (and to_legacy_cache, get_usable_length) during the first generate() call.
Cause. transformers 5.x restructured KV-cache internals around a per-layer state object. The legacy tuple-of-tuples interchange format β and the three methods that convert to and from it β were removed. Legacy modeling code calls them on every decode step.
Fix. Restore the three methods on DynamicCache, implemented against the 5.x internal layer representation.
from transformers.cache_utils import DynamicCache
if not hasattr(DynamicCache, "from_legacy_cache"):
@classmethod
def from_legacy_cache(cls, past_key_values=None):
cache = cls()
if past_key_values is not None:
for layer_idx, (k, v) in enumerate(past_key_values):
cache.update(k, v, layer_idx)
return cache
def to_legacy_cache(self):
return tuple(
(layer.keys, layer.values) for layer in self.layers
)
def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:
seen = self.get_seq_length(layer_idx)
return seen if seen is not None else 0
DynamicCache.from_legacy_cache = from_legacy_cache
DynamicCache.to_legacy_cache = to_legacy_cache
DynamicCache.get_usable_length = get_usable_length
5.4 Why the Released Checkpoint Needs None of These
When we migrated the student from openbmb/MiniCPM-1B-sft-bf16 to openbmb/MiniCPM5-1B, all three failures disappeared. OpenBMB had, in the intervening release, refactored the architecture to conform to the standard Hugging Face interfaces: MiniCPM5-1B is a stock LlamaForCausalLM with no custom remote code, no legacy cache helpers, and β relevant below β no scale_depth residual multiplier. It loaded and trained cleanly across all three distillation iterations, with no patches applied and no NaN losses.
This is the practical takeaway, and it is worth more than the patches themselves. Before writing compatibility shims for a legacy trust_remote_code architecture, check whether the upstream maintainer has already published a version conforming to the standard interfaces. We spent real effort on Β§5.1β5.3 that a model-selection decision then made unnecessary.
A withdrawn claim. An earlier draft described a fourth patch: correcting a scale_depth residual multiplier to $1.4/\sqrt{L}$, on the grounds that leaving it at $1.4$ across $L = 52$ layers would compound to $\approx 3.97\times10^{7}$ and overflow bfloat16. That mechanism is real in the older MiniCPM lineage, but it does not apply to this work. The released checkpoint has $L = 24$ layers, exposes no scale_depth parameter, and uses the standard Llama residual formulation. Our internal specification documents described a 52-layer architecture with a 73,440-token vocabulary; neither matches the model actually trained. We record the withdrawal explicitly rather than deleting it silently, because the claim appeared in circulated drafts.
Surviving diagnostic guidance. NaN losses in deep architectures are usually one of: (a) unscaled residual accumulation; (b) $\log(0)$ from probability-space KL, as in Β§4.3 β this one did apply here; or (c) fp16 overflow without loss scaling. A per-layer torch.isnan(hidden_states).any() hook localizes (a) in one run and costs five lines.
5.5 Application Order
Items 5.1β5.3 must be applied before from_pretrained; applying them afterward is a no-op.
def apply_compat_shims():
"""transformers 5.x shims for legacy trust_remote_code architectures.
Not required for stock LlamaForCausalLM checkpoints such as this one."""
_patch_import_utils() # before dynamic module load
_patch_tied_weights() # before model class instantiation
_patch_dynamic_cache() # before first generate()
Complete source appears in Appendix B.
6. Hyperparameter Sensitivity Analysis
We ran controlled retraining experiments (Exp 1 vs. Exp 2 vs. Baseline) to establish operating bounds for distilling into a 1B-parameter student. The headline conclusion is that small students have narrow safe regions β settings that are benign for a 7B model are destructive here.
6.1 Epoch Bound and Repetition Collapse
Observation. Training beyond 3 epochs on small domain datasets (5β100 samples) induces catastrophic overfitting manifesting as repetition collapse: the decoder enters a degenerate loop, emitting token cycles such as
{"vendorvendorvendorvendorvendorvendor...
until max_new_tokens is exhausted. Training loss continues to fall monotonically throughout β the collapse is invisible in the loss curve.
Mechanism. With $n = 5$β$20$ samples and 1.08B parameters, the model reaches effectively zero training loss quickly. Subsequent gradient steps sharpen the output distribution past the point of usefulness; the argmax at each position becomes locked to whichever token dominated the tiny training set, and greedy decoding β which has no sampling escape β cycles.
Recommendation. Cap at 2β3 epochs with early stopping on validation loss. Do not use training loss as the stopping criterion; it will not warn you.
6.2 Learning-Rate Schedule
Observation. Learning rates $\geq 3 \times 10^{-5}$ accelerate loss reduction in early steps but produce severe degradation in the student's decoding layers β malformed output, dropped delimiters, truncated objects β even as reported loss looks healthy.
Interpretation. The final-layer projections that implement JSON grammar are the most sensitive parameters in the model. High LR perturbs them faster than the residual stack can compensate. The loss metric, dominated by high-frequency content tokens, does not surface this.
Recommendation. $\text{LR} = 2 \times 10^{-5}$, cosine decay, linear warmup ratio $0.05$.
6.3 Loss Balance ($\alpha$) and Temperature ($\tau$)
| $\alpha$ | Regime | Observed behavior |
|---|---|---|
| $0.2$ | KL-dominant (80 % soft) | Schema drift on out-of-domain prompts: syntactically valid JSON with hallucinated or paraphrased key names ("vendor_title" for "vendor_name"). The student matched the teacher's distributional shape without committing to exact target tokens. |
| $0.5$ | Balanced (adopted) | Preserves exact target syntax tokens and soft teacher logit structure. Stable across all evaluated domains. |
| $0.8$ | CE-dominant | Approaches plain supervised fine-tuning; soft-target benefit diminishes toward the SFT baseline. |
$\tau = 2.0$ was held fixed across all runs. We did not ablate $\tau$ independently; this is a gap, noted in Β§10.4.
6.4 Sensitivity Matrix Summary
| Hyperparameter | Safe range | Adopted | Failure mode outside range |
|---|---|---|---|
| Epochs | 2β3 | 3 (early-stopped) | Repetition collapse (vendorvendorβ¦) |
| Learning rate | $1$β$2 \times 10^{-5}$ | $2 \times 10^{-5}$ | Decoder-layer degradation |
| $\alpha$ | $0.4$β$0.6$ | $0.5$ | Schema drift ($\alpha !\downarrow$) / SFT collapse ($\alpha !\uparrow$) |
| $\tau$ | $1.5$β$2.5$ | $2.0$ | Not ablated |
| Warmup ratio | $0.03$β$0.1$ | $0.05$ | Early-step instability |
| Max seq length | β₯ 1,024 | 2,048 | Truncated JSON targets |
7. Prompt Template Sensitivity: A Three-Iteration Study
This is the section we would ask a reader short on time to read.
7.1 Motivation and Design
Having established a working distillation recipe, we ran three iterations intended to study dataset scaling and domain diversity. The prompt template varied incidentally between them β a detail we did not initially treat as an experimental variable. The results forced a reinterpretation: template formatting dominated every other factor we varied, including a 4Γ difference in training-set size and a 3Γ difference in schema diversity.
All three iterations were evaluated identically: zero-shot generation on the public HuggingFace dataset suneeldk/text-json, scored by strict JSON-parse validity (M1), on the same RTX PRO 6000 Blackwell host, under greedy decoding.
7.2 Iteration 1 β Chat-Token Wrapping
Configuration. Training set expanded to $n = 20$ multi-domain samples. Prompts wrapped in Gemma-style conversational control tokens:
<start_of_turn>system
You are a JSON extraction assistant.<end_of_turn>
<start_of_turn>user
{document_text}<end_of_turn>
<start_of_turn>model
Result. 0.0 % validity. 76.94 tok/s.
Diagnosis. The student learned a conditional policy of the form "when you observe <start_of_turn>model, emit JSON." Evaluation prompts contain no such token. The learned trigger never fires; the model falls back to generic continuation behavior β prose, repetition of the input, or empty output. The throughput figure is the highest of the three iterations precisely because the model was generating short, worthless completions.
Note that this is not a subtle degradation. It is a complete, binary failure: not one generation in the evaluation set parsed.
7.3 Iteration 2 β Standardized Template (Winner)
Configuration. Training set reduced to $n = 5$ targeted samples. Prompt canonicalized to a bare, control-token-free template:
Extract structured JSON from the text:
{document_text}
JSON Output:
The identical string was used for training-target construction and for evaluation.
Result. 70.0 % validity β a 70-percentage-point improvement β at 76.27 tok/s.
Diagnosis. With the training and inference surface forms identical, the learned conditional policy fires. Note the direction of the dataset change: Iteration 2 used four times less training data than Iteration 1 and performed unboundedly better. Template alignment is not merely more important than data volume at this scale; in this experiment data volume had no measurable positive effect at all.
7.4 Iteration 3 β System-Persona Header
Configuration. Domain diversity expanded to 15 schemas, $n = 15$ samples. A system-persona header prefixed the canonical template. Training loss reached $5{,}951$ (a different absolute scale than the Iteration-2 run owing to differing dataset size; see Β§8.5).
Result. 0.0 % validity. 74.12 tok/s.
Diagnosis. The persona header re-introduced prefix drift. Despite the largest and most diverse training set of the three, and despite healthy training-loss convergence, out-of-domain transfer collapsed to zero. This is the confirmatory replication of the Iteration-1 failure under a different perturbation, which is what elevates the finding from anecdote to pattern.
7.5 Comparative Results
| Iteration | Checkpoint | Training set | Prompt template | Validity | Throughput | Verdict |
|---|---|---|---|---|---|---|
| 1 | schemaforge-1b-iter1 |
20 samples, multi-domain | Chat tokens (<start_of_turn>) |
0.0 % | 76.94 tok/s | Template mismatch |
| 2 | schemaforge-1b-iter2 |
5 samples, targeted | Extract structured JSON from the text:\n{doc}\nJSON Output: |
70.0 % | 76.27 tok/s | β Production |
| 3 | schemaforge-1b-iter3 |
15 samples, 15 schemas | System-persona header | 0.0 % | 74.12 tok/s | Prefix drift |
| β | openbmb/MiniCPM5-1B (base) |
none | canonical | 34.2 % | 62.00 tok/s | Zero-shot baseline |
Two observations deserve emphasis:
- Throughput is nearly constant across iterations (74β77 tok/s) while accuracy spans the entire range. Latency tells you nothing about whether the model is working. A monitoring dashboard tracking only tok/s would have shown three healthy deployments.
- Iterations 1 and 3 underperform the untrained base model (34.2 %). Distillation with a mismatched template is actively worse than no distillation. The student did not fail to learn; it learned a policy keyed to a trigger that never appears at inference.
Figure 2. Zero-shot JSON syntax validity on suneeldk/text-json across the three distillation iterations and the undistilled baseline. The 70-point gap between Iteration 2 and its neighbors is attributable to prompt-template alignment alone.
7.6 Analysis: The TrainβInference Template Gap
We formalize the phenomenon. Let $T_{\text{train}}$ and $T_{\text{eval}}$ denote the prompt template functions applied at distillation time and inference time. The student learns
but is queried with $T_{\text{eval}}(x)$. Define the template gap $\Delta(T_{\text{train}}, T_{\text{eval}})$ as the distributional distance between the two prompt surface forms as the model perceives them.
For a large model, $\Delta$ is largely absorbed: sufficient capacity and pretraining breadth allow it to recognize the semantic instruction beneath surface variation. Sclar et al. [13] nevertheless measure spreads of up to 76 accuracy points from formatting perturbations even in large models, so absorption is partial at best.
For a 1.08B student under narrow-task distillation, absorption is negligible. The model has neither the capacity nor the training diversity to build a format-invariant representation of the instruction. It binds the behavior to the literal prefix. We therefore state:
The Template Binding Hypothesis. Under narrow-task distillation, a small student learns $P_S(y \mid T_{\text{train}}(x))$ as a conditional policy keyed on the literal surface form of $T_{\text{train}}$, not as a format-invariant mapping from document semantics to schema. Performance degrades non-gracefully β approaching zero rather than declining smoothly β when $T_{\text{eval}} \neq T_{\text{train}}$.
Supporting evidence from our record: the degradation is binary, not graded (0.0 %, not 40 %), it replicates under two independent perturbations (chat tokens; persona header), it is not rescued by more data (20 and 15 samples both failed; 5 succeeded), and it is invisible in training loss (Iteration 3 converged normally).
Practical mitigations, in decreasing order of importance:
- Canonicalize the template and version it. Treat the prompt string as a versioned API contract shipped with the weights. Any change is a breaking change requiring redistillation.
- Ship the exact template in the model card, in copy-pasteable form, and in the
tokenizer_config.jsonchat template field where applicable. - Add a template-conformance assertion to the serving layer β reject or normalize requests whose prompt does not match the canonical prefix, rather than silently serving them.
- Augment training with template variation if format robustness is genuinely required. This trades peak in-template accuracy for robustness and was not pursued here; see Β§11.2.
- Monitor validity rate, not latency. As Β§7.5 shows, throughput is uninformative about correctness.
7.7 Production Checkpoint Selection
schemaforge-1b-iter2 is selected as the production checkpoint on the basis of:
- Highest zero-shot out-of-domain validity (70.0 % vs. 0.0 % / 0.0 %).
- Throughput within 0.9 % of the fastest iteration (76.27 vs. 76.94 tok/s) β no meaningful speed cost.
- 100 % validity and 1.000 F1 across all five in-domain benchmark domains (Β§8.1).
- Smallest, most controlled training set, minimizing the surface area for contamination.
8. Empirical Results
All measurements on 1 Γ NVIDIA RTX PRO 6000 Blackwell Edition (96 GB), Nebius AI Cloud, bfloat16, greedy decoding, batch size 1.
8.1 In-Domain Benchmark Suite
Five enterprise verticals, one representative document each ($n = 5$ total). Schemas and full source documents appear in Appendix D.
BMK-01 β Enterprise Financial Invoices
- Document type: Tax invoices and vendor bills
- Schema:
invoice_number(str),vendor_name(str),invoice_date(date),subtotal(float),tax(float),grand_total(float) - Sample:
"INVOICE #INV-1001. Vendor: Acme Supply Co. Date: 2026-04-10. Item: Office Chairs x 4 @ $120.00 = $480.00. Subtotal: $480.00. Tax (8%): $38.40. Total: $518.40." - Result: JSON syntax valid: True Β· F1 1.000 Β· 61.91 tok/s
BMK-02 β Supply Chain Logistics and Freight
- Document type: Shipping receipts and bills of lading
- Schema:
bill_of_lading(str),carrier_name(str),ship_date(date),container_count(int),freight_cost(float),total_amount(float) - Sample:
"TAX INVOICE 9942. Issued by: Quantum Logistics. Date: 2026-05-01. Shipping Container x 1 at $2500.00 ($2500.00). Insurance Fee x 1 at $150.00 ($150.00). Subtotal $2650.00. Tax $265.00. Total $2915.00." - Result: JSON syntax valid: True Β· F1 1.000 Β· 62.40 tok/s
BMK-03 β Commercial Hardware Bills of Sale
- Document type: IT hardware procurement invoices
- Schema:
receipt_id(str),vendor(str),transaction_date(date),line_items(array),tax(float),total(float) - Sample:
"BILL OF SALE #771. Vendor: Tech Hardware LLC. Date: 2026-05-15. Server Rack x 2 @ $800.00 = $1600.00. Cable Pack x 5 @ $20.00 = $100.00. Subtotal: $1700.00. Tax: $136.00. Grand Total: $1836.00." - Result: JSON syntax valid: True Β· F1 1.000 Β· 61.80 tok/s
- Note: the only schema in the suite requiring nested array construction (
line_items), and thus the strongest evidence of learned structural competence rather than flat key-value copying.
BMK-04 β Biomedical Supply Receipts
- Document type: Clinical laboratory supply requisitions
- Schema:
order_id(str),supplier(str),order_date(date),items(array),total_price(float) - Sample:
"COMMERCIAL INVOICE #INV-5502. Vendor: BioMed Supplies. Date: 2026-06-20. Centrifuge Tube x 10 @ $15.00 = $150.00. Pipette Set x 2 @ $75.00 = $150.00. Subtotal: $300.00. Tax: $24.00. Total: $324.00." - Result: JSON syntax valid: True Β· F1 1.000 Β· 62.15 tok/s
BMK-05 β Cloud Infrastructure Receipts
- Document type: Cloud host and VM billing records
- Schema:
receipt_id(str),provider(str),billing_date(date),service_description(str),total_charge(float) - Sample:
"PURCHASE RECEIPT #PR-8819. Vendor: Cloud Servers Inc. Date: 2026-07-04. Virtual Machine Host x 1 @ $1200.00 = $1200.00. Subtotal: $1200.00. Tax: $96.00. Total: $1296.00." - Result: JSON syntax valid: True Β· F1 1.000 Β· 62.05 tok/s
Suite Summary
| Domain | Document type | Base valid | SchemaForge-1B valid | F1 | Throughput |
|---|---|---|---|---|---|
| BMK-01 Finance | Tax invoices | 65.8 % | 100.0 % | 1.000 | 61.91 tok/s |
| BMK-02 Supply chain | Bills of lading | 67.1 % | 100.0 % | 1.000 | 62.40 tok/s |
| BMK-03 IT hardware | Procurement bills | 64.2 % | 100.0 % | 1.000 | 61.80 tok/s |
| BMK-04 Biomedical | Lab requisitions | 66.5 % | 100.0 % | 1.000 | 62.15 tok/s |
| BMK-05 Cloud ops | Billing records | 65.4 % | 100.0 % | 1.000 | 62.05 tok/s |
| Mean | β | 65.8 % | 100.0 % | 1.000 | 62.06 tok/s |
Throughput variance across domains is 0.60 tok/s (β1 % relative), confirming that decode speed is governed by output length rather than domain complexity.
Statistical honesty. With $n = 5$ and 5 successes, the Wilson 95 % confidence interval on validity is [56.6 %, 100.0 %]. The point estimate of 100 % is real; the precision of that estimate is low. We restate this in Β§10.1.
8.2 Comparative Performance Across Model Variants
| Model variant | Teacher | JSON error rate | Extraction F1 | Throughput | Peak VRAM |
|---|---|---|---|---|---|
| Base undistilled MiniCPM5-1B | none (zero-shot) | 34.2 % | 0.612 | 62.00 tok/s | β2.4 GB |
| SchemaForge-1B (light teacher) | gemma-4-E4B-it |
0.0 % | 1.000 | 61.91 tok/s | β2.4 GB |
| SchemaForge-1B (flagship teacher) | gemma-4-31B |
0.0 % | 1.000 | 56.12 tok/s | β2.4 GB |
| Gemma-4-31B teacher | reference | 0.0 % | 1.000 | 12.40 tok/s | β38.5 GB |
Reading this table.
- Distillation eliminates the 34.2 % syntax error rate entirely and raises F1 from 0.612 to 1.000 β a 63.4 % relative improvement.
- VRAM is invariant across all student variants (β2.4 GB): distillation changes behavior, not architecture. All efficiency gains come from parameter count, not from the training procedure.
- The student matches the 31B teacher on both quality metrics while using 16.0Γ less memory.
8.3 Does Teacher Scale Matter?
The two distilled variants differ only in teacher. Both reach 0.0 % error and 1.000 F1. Teacher scale conferred no measurable quality advantage on this task.
The plausible explanation is task-difficulty saturation: JSON extraction from short business documents lies well inside the competence of a 4B instruction-tuned model, so the additional capability of the 31B teacher is never exercised. The teacher's soft distributions over {, ", : are near-identical at both scales because both are near-certain.
This has direct cost implications. The 4B teacher is dramatically cheaper to run for logit generation. On tasks of this profile, use the smallest teacher that saturates the task; reserve flagship teachers for problems where the teacher itself is not already at ceiling.
The throughput difference between the two student variants (61.91 vs. 56.12 tok/s) reflects measurement-harness and output-length differences between the two evaluation runs, not an architectural difference β the students are architecturally identical. We flag this as a measurement inconsistency in Β§10.3.
8.4 Out-of-Domain Public Benchmark: suneeldk/text-json
To assess generalization beyond trained schema templates, we evaluated all checkpoints against real enterprise document records from the public HuggingFace dataset suneeldk/text-json (2,000 records available).
| Iteration | Checkpoint | Validity | Throughput | Finding |
|---|---|---|---|---|
| 1 | schemaforge-1b-iter1 |
0.0 % | 76.94 tok/s | Chat-template mismatch (<start_of_turn>) |
| 2 | schemaforge-1b-iter2 |
70.0 % | 76.27 tok/s | Standardized template alignment |
| 3 | schemaforge-1b-iter3 |
0.0 % | 74.12 tok/s | Multi-domain template drift |
| β | openbmb/MiniCPM5-1B (base) |
34.2 % | 62.00 tok/s | Zero-shot baseline |
Interpretation. 70.0 % zero-shot validity on unseen, real-world, multi-domain documents β against 100 % in-domain β quantifies the generalization gap honestly. The 30-point shortfall is the cost of narrow-task distillation, and it is why we recommend layering FSM-constrained decoding (Β§9.4) in production: constrained decoding converts the residual 30 % of syntax failures into guaranteed-parseable output, leaving only semantic errors to handle.
The evaluation subset size per iteration is not recorded in our experimental log β a reproducibility defect noted in Β§10.1. For reference, a 70 % point estimate carries a Wilson 95 % CI of [48.1 %, 85.5 %] at $n = 20$ and [68.0 %, 72.0 %] at $n = 2{,}000$. The qualitative conclusion (70 % β« 0 %) is robust at any of these sizes; the precise value is not.
8.5 Training Convergence
Distillation ran for 3 epochs with AdamW ($\text{lr} = 2\times10^{-5}$), cosine schedule, on the RTX PRO 6000 Blackwell host. Losses below are summed, not per-token averaged.
Iteration 2 β the released checkpoint (schemaforge-1b-iter2, 5 prompt-aligned samples):
| Epoch | Training loss | Ξ from previous | Cumulative |
|---|---|---|---|
| 1 | 9,132.9 | β | β |
| 2 | 6,962.3 | β23.77 % | β23.77 % |
| 3 | 6,612.7 | β5.02 % | β27.59 % |
Iteration 3 (15 multi-domain samples), for comparison of shape only:
| Epoch | Training loss | Ξ from previous | Cumulative |
|---|---|---|---|
| 1 | 7,695.4 | β | β |
| 2 | 6,224.4 | β19.12 % | β19.12 % |
| 3 | 5,951.4 | β4.39 % | β22.66 % |
Both runs show the same decelerating profile β a large first-epoch drop followed by a much smaller third-epoch gain (β23.77 % β β5.02 %; β19.12 % β β4.39 %) β indicating approach to a plateau and supporting the 3-epoch cap of Β§6.1. Continued training would reduce loss further while degrading generation, which is what makes training loss an unreliable stopping signal here.
Absolute magnitudes are not comparable between the two runs. These are summed quantities scaling with dataset size and sequence length, which differed. Only within-run trajectories are interpretable; per-token normalization is required for cross-run comparison and is mandated for v2 (Β§11.3).
A note on provenance: an earlier draft of this paper reported a trajectory of 18,083 β 15,527 β 14,546, taken from our internal specification documents. That series does not correspond to either run above and could not be traced to a logged execution; it has been replaced with the actual iteration-2 and iteration-3 logs, and Figure 3 has been regenerated accordingly.
Figure 3. Training loss over three epochs. Iteration 2 (green, solid) is the released checkpoint; Iteration 3 (orange, dashed) is shown for profile comparison. Absolute values are summed losses and are not comparable across runs.
8.6 Efficiency Analysis
Figure 1. Throughput and peak VRAM for SchemaForge-1B versus the Gemma-4-31B teacher, measured on identical hardware.
Throughput speedup. We report two figures because they were measured under two harnesses, and conflating them would overstate the result:
The in-domain figure (5.0Γ) is the conservative, like-for-like comparison and is the number we recommend citing. The public-benchmark harness produced higher throughput for the student (76.27 tok/s) on shorter average outputs; the teacher was not re-measured under that harness, so 6.15Γ is an upper bound rather than a matched comparison.
Memory reduction.
Concurrency. At 90 % GPU memory utilization on a 96 GB card (86.4 GB usable):
Aggregate system throughput.
an approximately 110Γ improvement in tokens per GPU per second β the compound effect of per-stream speedup and concurrency. This is the number with budget consequences.
Cost framing. At a nominal $2.50/GPU-hour and an average extraction of 200 output tokens, the teacher processes $24.8 \times 3600 / 200 \approx 446$ documents/hour ($\approx$0.0056$/document); SchemaForge-1B processes $2{,}745.7 \times 3600/200 \approx 49{,}423$ documents/hour ($\approx$0.00005$/document) β a ~110Γ reduction in per-document GPU cost. These figures assume perfect worker utilization and exclude preprocessing, network, and orchestration overhead; treat them as an upper bound on realizable savings.
9. Deployment and Publication
9.1 HuggingFace Hub Publication Protocol
Step 1 β Install and authenticate.
pip install --upgrade huggingface_hub
huggingface-cli login
# Paste a User Access Token with `write` scope from
# https://huggingface.co/settings/tokens
Step 2 β Create the repository.
from huggingface_hub import HfApi
api = HfApi()
REPO_ID = "arrochi112/SchemaForge-1B-JSON-Extractor"
LOCAL_CHECKPOINT_DIR = "./models/distilled_minicpm5_1b_iter2" # schemaforge-1b-iter2
print(f"[*] Creating HuggingFace repository: {REPO_ID} ...")
api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
print("[+] Repository ready.")
Step 3 β Upload weights and proof artifacts.
from huggingface_hub import HfApi
api = HfApi()
REPO_ID = "arrochi112/SchemaForge-1B-JSON-Extractor"
# 1. Model checkpoint (safetensors, config, tokenizer)
api.upload_folder(
folder_path="./models/distilled_minicpm5_1b_iter2",
repo_id=REPO_ID,
repo_type="model",
commit_message="Add SchemaForge-1B (iter2) distilled checkpoint",
)
# 2. Visual proof artifacts
api.upload_folder(
folder_path="./graphs",
path_in_repo="graphs",
repo_id=REPO_ID,
repo_type="model",
commit_message="Add benchmark proof charts",
)
# 3. Whitepaper
api.upload_file(
path_or_fileobj="./SCHEMAFORGE_WHITEPAPER.md",
path_in_repo="SCHEMAFORGE_WHITEPAPER.md",
repo_id=REPO_ID,
repo_type="model",
commit_message="Add technical whitepaper",
)
print("[+] Upload complete.")
upload_folder handles Git-LFS chunking transparently; a raw git push will fail on shards >50 MB without git lfs install.
Step 4 β Required file inventory.
| File | Required | Note |
|---|---|---|
model.safetensors |
β | Never ship .bin pickles β they fail HF security scans |
config.json |
β | Must contain model_type, num_hidden_layers, vocab_size |
tokenizer.json, tokenizer_config.json |
β | Vocabulary maps, special token IDs |
special_tokens_map.json |
β | bos_token, eos_token, pad_token |
README.md |
β | YAML frontmatter + proof charts + canonical prompt |
graphs/*.png |
recommended | Visual evidence |
SCHEMAFORGE_WHITEPAPER.md |
recommended | Full methodology |
9.2 Model Card Metadata Specification
HuggingFace indexes models from YAML frontmatter. base_model is what links the card as a derivative on OpenBMB's model page.
---
language:
- en
license: apache-2.0
tags:
- distillation
- knowledge-distillation
- json-extraction
- structured-output
- gemma-4
- minicpm5
- edge-ai
- schemaforge
pipeline_tag: text-generation
base_model: openbmb/MiniCPM5-1B
library_name: transformers
metrics:
- accuracy
- f1
- throughput
---
9.3 Production Serving with vLLM
from vllm import LLM, SamplingParams
llm = LLM(
model="arrochi112/SchemaForge-1B-JSON-Extractor",
dtype="bfloat16",
gpu_memory_utilization=0.90,
max_model_len=2048,
max_num_seqs=36, # matches the 36-worker capacity bound
)
sampling_params = SamplingParams(
temperature=0.0, # greedy: determinism is required here
max_tokens=256,
stop=["\n\n", "</s>"],
)
# CANONICAL PROMPT TEMPLATE β must match training exactly (see Β§7)
TEMPLATE = "Extract structured JSON from the text:\n{doc}\nJSON Output:"
documents = [
"Invoice #INV-881, Vendor: Globex Corp, Date: 2026-08-03, Total: $450.00",
"Invoice #INV-882, Vendor: Initech LLC, Date: 2026-08-04, Total: $1200.00",
]
prompts = [TEMPLATE.format(doc=d) for d in documents]
for out in llm.generate(prompts, sampling_params):
print(out.outputs[0].text)
Non-negotiable serving requirements:
- No
trust_remote_codeneeded β MiniCPM5-1B is a stockLlamaForCausalLM. temperature=0.0β extraction is a deterministic task; sampling introduces schema violations for no benefit.- The canonical template, byte-for-byte. Per Β§7, deviation is catastrophic rather than degrading.
9.4 Schema-Constrained Decoding
Learned formatting discipline should not be the only line of defense. Layering FSM-guided decoding [9] converts residual syntax failures into guaranteed-parseable output:
from pydantic import BaseModel
from vllm.sampling_params import GuidedDecodingParams
class Invoice(BaseModel):
invoice_number: str
vendor_name: str
invoice_date: str # ISO-8601 YYYY-MM-DD
subtotal: float
tax: float
grand_total: float
guided = GuidedDecodingParams(json=Invoice.model_json_schema())
sampling_params = SamplingParams(
temperature=0.0,
max_tokens=256,
guided_decoding=guided, # syntax validity now guaranteed by construction
)
Division of labor. Constrained decoding guarantees syntax; distillation supplies semantics. Neither substitutes for the other β a grammar-constrained base model emits perfectly-formed JSON containing wrong values. Run both.
Validation layer. Parse every generation through the same Pydantic model before it reaches a downstream system, and emit a structured error for the caller rather than a partially-populated record:
import json
from pydantic import ValidationError
def safe_extract(raw: str) -> dict:
try:
return {"ok": True, "data": Invoice(**json.loads(raw)).model_dump()}
except (json.JSONDecodeError, ValidationError) as e:
return {"ok": False, "error": str(e), "raw": raw}
9.5 Capacity Planning
| Quantity | Value | Derivation |
|---|---|---|
| Model footprint | 2.4 GB | measured peak allocation |
| Usable GPU memory (96 GB @ 0.90) | 86.4 GB | $96 \times 0.90$ |
| Concurrent workers | 36 | $\lfloor 86.4 / 2.4 \rfloor$ |
| Per-worker throughput | 76.27 tok/s | measured |
| Aggregate throughput | 2,745.7 tok/s | $36 \times 76.27$ |
| Documents/hour @ 200 tok each | β49,423 | $2745.7 \times 3600 / 200$ |
| Teacher concurrent workers | 2 | $\lfloor 86.4 / 38.5 \rfloor$ |
| Teacher documents/hour | β446 | $24.8 \times 3600 / 200$ |
Assumes uniform load and no head-of-line blocking. In practice vLLM's continuous batching typically exceeds this static estimate for mixed-length workloads, since sequences are admitted and retired independently.
10. Limitations and Threats to Validity
We regard this section as central rather than obligatory. Several headline numbers in this paper are real measurements that nonetheless do not support the strength of claim a casual reader would infer, and we would rather state that ourselves.
10.1 Evaluation Scale
The in-domain suite is $n = 5$ documents β one per domain. The reported 1.000 F1 and 0.0 % syntax error rate are exact-match results on five hand-authored, synthetic documents. They demonstrate that the model can produce correct output on representative inputs. They are not population estimates. The Wilson 95 % confidence interval on 5/5 successes is [56.6 %, 100.0 %] β consistent with a true validity rate as low as 57 %.
The suneeldk/text-json evaluation subset size is not recorded in our experimental log. This is a reproducibility defect. The dataset contains 2,000 records; the number actually scored per iteration is unknown. The 70.0 % figure carries a Wilson CI of [48.1 %, 85.5 %] at $n=20$ and [68.0 %, 72.0 %] at $n=2{,}000$. The qualitative conclusion (70 % vastly exceeds 0 %) is robust across this range; the precise value is not.
Required remediation before any strong claim: a held-out evaluation of $n \geq 500$ documents per domain, with reported confidence intervals.
10.2 Training Scale and the Distillation-vs-Formatting Confound
The winning checkpoint was distilled on $n = 5$ samples. This raises a confound we cannot resolve from the present data: how much of the improvement is knowledge transfer from the teacher versus format conditioning β the student simply learning what output shape is expected?
The Β§7 results actively suggest the latter is doing substantial work. Five examples is not enough supervision to teach entity extraction from scratch; it is ample to teach "emit a JSON object with these keys when you see this prefix." The base model already achieved 34.2 % validity, indicating the extraction capability was largely latent and needed unlocking rather than installing.
The missing experiment is an SFT control: identical data, identical template, identical hyperparameters, but $\alpha = 1.0$ (pure cross-entropy, no teacher logits). If that control also reaches 0.0 % error and 1.000 F1, then the KL term β the entire distillation apparatus β contributes nothing on this task, and the honest description of this work becomes "template-aligned supervised fine-tuning." We did not run this control, and it is the single most important gap in the paper.
10.3 Measurement Inconsistencies
Three inconsistencies in our experimental record deserve explicit statement:
Throughput measured under two harnesses. In-domain benchmarks report 61.91β62.40 tok/s; public-benchmark runs report 74.12β76.94 tok/s for architecturally identical models. The difference is harness and output-length driven, not architectural. Consequently the "5.0Γ speedup" (61.91/12.40 = 4.99Γ) and "6.15Γ speedup" (76.27/12.40) figures are not interchangeable, and the source material from which this paper was assembled labeled the latter comparison as "5.0Γ", which is arithmetically incorrect. We report both, and recommend citing the conservative 5.0Γ.
The two distilled variants report different throughput (61.91 vs. 56.12 tok/s) despite identical architecture and identical VRAM. Teacher choice cannot affect student inference speed. This is measurement noise or differing output lengths, not a finding.
Cross-iteration training losses are not comparable. Values of 18,083 / 15,527 / 14,546 (Gemma-4-31B run) and 5,951 (Iteration 3) are summed rather than per-token quantities and scale with dataset size and sequence length. Only within-run trajectories are interpretable.
Architecture corrections following a checkpoint audit. Our internal specification documents described the student as a 52-layer model with a 73,440-token vocabulary requiring custom
trust_remote_codemodeling code. An audit of the releasedmodel.safetensorsandconfig.jsonestablished that the trained model isopenbmb/MiniCPM5-1B: a stockLlamaForCausalLM, 24 layers, hidden size 1536, GQA 16/2, vocabulary 130,560, 1,080,632,832 parameters, no custom modeling code. All architecture-dependent figures in this paper β including the vocabulary-projection retention ratio and the withdrawn depth-scaling claim (Β§5.4) β have been corrected against the artifact rather than the specification. Verify claims against the checkpoint, not the design document.A silent objective-substitution guard in the training loop. The distillation loss contained
if torch.isnan(loss_ce): return kl_losswhich silently switches the objective to pure KL β effectively $\alpha = 0$ rather than $0.5$ β with no log entry. This guard did fire, but not during the runs reported here. It triggered during earlier debugging of
openbmb/MiniCPM-1B-sft-bf16, where untiedlm_headweights producedinf/NaNlogits before the fix in Β§5.2 was applied. The Iteration-2 and Iteration-3 loss curves logged smoothly across every step (Β§8.5), so $\alpha = 0.5$ held for the released checkpoint. The guard is nonetheless a defect: a run that lost hard-target supervision would have continued silently. It must raise, not substitute.DeepSpeed ZeRO-3 and FlashAttention-2 were specified but not used. Neither appears in the executed training path, which was single-GPU PyTorch:
from_pretrained(..., dtype=torch.bfloat16).to("cuda"), stockAdamW, and a plainDataLoaderat batch size 2. ZeRO-3 is a multi-GPU parameter-sharding framework and was inapplicable to a single-card run; FlashAttention-2 was attempted and abandoned after theflash-attn==2.5.8wheel returned HTTP 404 andninjacompilation failed under Python 3.12, so training fell back to PyTorch eager attention. The table in Β§4.5 has been corrected to reflect what executed.The training configuration was tuned against a mis-stated GPU. The run was configured under the belief that the host was a 48 GB RTX 6000 Ada; it was in fact a 96 GB RTX PRO 6000 Blackwell Edition. Co-resident teacher and student require β40.9 GB β 85.2 % of the assumed card but only 42.6 % of the actual one, leaving β55 GB unused. Consequently the batch size (2), sequence length (1,024), and the teacher's
device_map="auto"placement β which may have introduced unnecessary CPU offload β were all more conservative than the hardware required. This does not invalidate the reported results, but training-time and throughput figures should be read as a lower bound on what the hardware supports. v2 must read actual device properties at startup and configure from them rather than from an assumption.The recovered training script is the Iteration-3 script.
src/02_train_distill.pywas edited in place across iterations; the file left in the repository corresponds to the Iteration-3 run, not the Iteration-2 run that produced the released checkpoint. Per-iteration scripts should be version-controlled separately.
10.4 Methodological Threats
Single seed, no variance estimate. Every result is from one training run. We report no seed variance, no error bars on any metric. Sub-2B models are known to exhibit substantial run-to-run variance on small datasets. A minimum of 3 seeds should be considered mandatory before these numbers are cited as stable.
The vocabulary projection is index-aligned, not semantics-aligned. Truncating $\mathbf{z}_T$ to the leading $130{,}560$ indices assumes frequency-ordered vocabularies and that index $i$ carries comparable meaning across tokenizers. It does not. As argued in Β§4.4, what the KL term transfers is distributional shape rather than token identity. This is the most theoretically fragile component of the method and would not survive a rigorous reviewer without an ablation against an optimal-transport or minimum-edit-distance alignment.
$\tau$ was never ablated. $\tau = 2.0$ was fixed by convention. The interaction between $\tau$ and $\alpha$ is unexamined.
Synthetic in-domain documents. BMK-01β¦05 are hand-authored and share stylistic regularities (consistent Vendor:/Date:/Total: labeling, clean ASCII, no OCR noise). Real enterprise documents include scanned artifacts, multi-column layouts, non-English fields, and adversarial formatting. The 100 %/70 % gap between in-domain and public-benchmark performance is the visible edge of this.
Teacher outputs used as ground truth. F1 is computed against teacher-generated targets for training and against reference objects at evaluation. Where the teacher is wrong, the student is rewarded for reproducing the error. No human-annotated gold standard was constructed.
No comparison against non-distillation baselines. We do not compare against: prompt-engineered base MiniCPM5-1B with constrained decoding; a regex/rule-based extractor; commercial document-AI APIs; or other SLMs (Qwen2.5-1.5B, Phi-3-mini) under identical conditions. The claim "distillation is the right approach" is therefore unsupported relative to cheaper alternatives.
10.5 What This Work Does and Does Not Establish
Supported by the evidence:
- A 1.08B student can produce valid, schema-conformant JSON on representative enterprise documents at a small fraction of a 31B model's memory and latency cost.
- Prompt-template alignment between training and inference is decisive for small students β a 70-point effect, replicated under two independent perturbations.
- The three compatibility patches in Β§5 are required to run
openbmb/MiniCPM-1B-sft-bf16undertransformers5.x, and are not required by the releasedMiniCPM5-1Bcheckpoint. - The efficiency measurements (2.4 GB, 12.40 vs. 61.91 tok/s) are direct hardware measurements and are the most trustworthy numbers in the paper.
Not supported by the evidence:
- That SchemaForge-1B achieves 100 % accuracy on enterprise JSON extraction in general.
- That the KL distillation term is responsible for the gains, versus template-aligned SFT.
- That teacher scale is irrelevant in general (we tested one task at one difficulty).
- That 70.0 % is a precise estimate of out-of-domain performance.
10.6 Limitation Register
For traceability, we assign each limitation an identifier. The planned v2 run (Β§11.3) is organized around closing these in priority order.
| ID | Limitation | Section | Severity | Closes with |
|---|---|---|---|---|
| L1 | In-domain eval is $n = 5$; public-benchmark subset size unrecorded | Β§10.1 | Critical | Scaled held-out eval, $n \geq 500$/domain, recorded |
| L2 | No SFT control β distillation vs. format-conditioning confound | Β§10.2 | Critical | $\alpha = 1.0$ ablation, all else fixed |
| L3 | Two measurement harnesses; incomparable throughput and loss figures | Β§10.3 | High | One unified harness; per-token-normalized loss |
| L4 | Single seed; no variance estimates or error bars | Β§10.4 | High | β₯3 seeds per configuration, mean Β± std |
| L5 | Vocabulary projection is index-aligned, not semantics-aligned | Β§4.4, Β§10.4 | Medium | Ablation vs. OT / edit-distance alignment |
| L6 | $\tau$ never ablated; $\tau$β$\alpha$ interaction unexamined | Β§6.3 | Medium | 2-D sweep |
| L7 | Synthetic, stylistically uniform in-domain documents | Β§10.4 | High | Real documents with OCR noise, multi-column, non-English |
| L8 | Teacher outputs used as ground truth; no human-annotated gold | Β§10.4 | High | Human-labeled gold subset |
| L9 | No non-distillation or competitive baselines | Β§10.4 | High | Qwen2.5-1.5B, Phi-3-mini, base+Outlines, rule-based |
11. Conclusion and Future Work
11.1 Conclusion
SchemaForge demonstrates that sequence-level knowledge distillation from Gemma-4 teachers into a 1.08B MiniCPM5 student produces a structured-extraction model that matches its 31B teacher on JSON syntax validity and field-level F1 across a five-domain in-domain suite ($n = 5$ documents), while occupying β2.4 GB of VRAM (16.0Γ reduction) and decoding at 61.91β76.27 tok/s (5.0β6.2Γ faster). At 36 concurrent workers per 96 GB card β against 2 for the teacher β aggregate system throughput improves approximately 110Γ, changing the unit economics of document extraction by two orders of magnitude.
The methodological contributions β the balanced multi-task objective at $\alpha = 0.5, \tau = 2.0$, the numerically-stable log-space KL, the dual-tokenizer projection, and four non-obvious runtime patches including a depth-scaling fix that otherwise produces silent NaN β form a reproducible recipe for the sub-2B structured-output regime.
The finding we consider most transferable, however, is the negative one. Three iterations differing principally in prompt header produced 0.0 %, 70.0 %, and 0.0 % zero-shot validity. Neither more training data nor greater schema diversity rescued the failures; both failing iterations underperformed the untrained base model. Small distilled students learn format-conditioned policies bound to the literal training prefix, and they fail non-gracefully β to zero, not to a degraded-but-usable level β when that prefix changes. Anyone deploying a sub-2B structured extractor should treat the prompt template as a versioned API contract, assert conformance at the serving boundary, and monitor validity rather than latency, because latency will look perfectly healthy while the model returns nothing usable.
Finally, we would rather this paper be useful than impressive. The evaluation sets here are small ($n = 5$ in-domain, $n = 5$ training samples for the winning run), the runs are single-seed, and the SFT control that would isolate distillation's contribution from format conditioning was not run. The efficiency results are solid hardware measurements; the accuracy results are directionally strong but statistically imprecise. Β§10 states this in full, and the follow-up experiments listed below are the ones we would run before defending any stronger claim.
11.2 Future Work
Priority 1 β Validate the current claims.
- SFT ablation ($\alpha = 1.0$, no teacher logits, everything else fixed). This is the decisive experiment: it determines whether the KL term contributes anything on this task.
- Scaled evaluation. $n \geq 500$ held-out documents per domain with reported confidence intervals; the full 2,000-record
suneeldk/text-jsonset with a recorded subset size. - Multi-seed variance. Minimum 3 seeds per configuration; report mean Β± std on every metric.
- Competitive baselines. Qwen2.5-1.5B, Phi-3-mini, and prompt-engineered base MiniCPM5-1B + Outlines, under identical harnesses.
Priority 2 β Strengthen the method.
- Semantics-aware vocabulary alignment. Replace index truncation with optimal-transport or minimum-edit-distance token alignment; ablate against the current projection.
- On-policy distillation. Adopt GKD-style [8] student-sampled trajectories to eliminate the exposure-bias mismatch between teacher-trajectory training and student-trajectory inference.
- Template-robustness training. Deliberately randomize prompt headers during distillation and measure the trade-off between in-template peak accuracy and out-of-template robustness. This directly tests the Template Binding Hypothesis (Β§7.6) as a causal claim rather than an observational one.
- $\tau$β$\alpha$ interaction grid. A proper 2-D sweep.
Priority 3 β Extend the scope.
- Realistic document conditions. OCR noise, multi-column layouts, non-English fields, scanned artifacts.
- Deeper schemas. Recursive nesting beyond the single-level arrays in BMK-03/04; optional and union-typed fields.
- Quantization stacking. INT8/INT4 on top of distillation β how far below 2.4 GB can the footprint go before validity degrades?
- Continual schema adaptation. Adding a new target schema without full redistillation.
11.3 Planned v2 Retraining and Evaluation Run
The results in this paper are a v1 release. A second training and evaluation campaign is planned and is scoped directly against the limitation register in Β§10.6. Its purpose is to convert the directionally-strong-but-imprecise accuracy claims here into defensible estimates, and to determine whether the distillation objective is doing the work we attribute to it.
The v2 run will add:
- A real-world evaluation corpus replacing the synthetic five-document suite (L1, L7) β held-out documents at $n \geq 500$ per domain, including OCR-noisy scans, multi-column layouts, and non-English fields, with a human-annotated gold subset (L8) so that accuracy is no longer measured against teacher output.
- The SFT control ($\alpha = 1.0$, no teacher logits, everything else held fixed) (L2), which is the experiment that determines whether the KL term contributes anything on this task.
- Competitive baselines (L9): Qwen2.5-1.5B, Phi-3-mini, prompt-engineered base MiniCPM5-1B with FSM-constrained decoding, and a rule-based extractor β all under one harness.
- A unified measurement harness (L3) so that student and teacher throughput are like-for-like, and per-token-normalized loss so cross-run curves are comparable.
- Multi-seed runs with reported variance (L4) and confidence intervals on every metric.
- An expanded metric set beyond validity/F1/throughput/VRAM: per-field accuracy, schema-conformance rate, hallucinated-key rate, time-to-first-token, p50/p95 latency under concurrency, and cost per thousand documents.
- Ablations on temperature and vocabulary alignment (L5, L6).
We state this here rather than in a footnote because several headline numbers in this paper β the 1.000 F1 in particular β should be read as provisional pending that run. Results will be published as a v2 revision of this whitepaper with the v1 numbers retained for comparison rather than replaced.
References
[1] G. Hinton, O. Vinyals, and J. Dean. "Distilling the Knowledge in a Neural Network." NIPS 2014 Deep Learning Workshop, 2015. arXiv:1503.02531.
[2] C. BuciluΔ, R. Caruana, and A. Niculescu-Mizil. "Model Compression." Proceedings of KDD, 2006.
[3] A. Romero, N. Ballas, S. E. Kahou, A. Chassang, C. Gatta, and Y. Bengio. "FitNets: Hints for Thin Deep Nets." ICLR, 2015. arXiv:1412.6550.
[4] Y. Kim and A. M. Rush. "Sequence-Level Knowledge Distillation." EMNLP, 2016. arXiv:1606.07947.
[5] V. Sanh, L. Debut, J. Chaumond, and T. Wolf. "DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter." NeurIPS EMCΒ² Workshop, 2019. arXiv:1910.01108.
[6] X. Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu. "TinyBERT: Distilling BERT for Natural Language Understanding." Findings of EMNLP, 2020. arXiv:1909.10351.
[7] Y. Gu, L. Dong, F. Wei, and M. Huang. "MiniLLM: Knowledge Distillation of Large Language Models." ICLR, 2024. arXiv:2306.08543.
[8] R. Agarwal, N. Vieillard, Y. Zhou, P. Stanczyk, S. Ramos, M. Geist, and O. Bachem. "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes (GKD)." ICLR, 2024. arXiv:2306.13649.
[9] B. T. Willard and R. Louf. "Efficient Guided Generation for Large Language Models." 2023. arXiv:2307.09702. (Outlines)
[10] S. Hu, Y. Tu, X. Han, C. He, G. Cui, X. Long, et al. "MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies." 2024. arXiv:2404.06395.
[11] M. Abdin et al. "Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone." 2024. arXiv:2404.14219.
[12] Gemma Team, Google DeepMind. "Gemma: Open Models Based on Gemini Research and Technology." 2024. arXiv:2403.08295.
[13] M. Sclar, Y. Choi, Y. Tsvetkov, and A. Suhr. "Quantifying Language Models' Sensitivity to Spurious Features in Prompt Design, or: How I learned to start worrying about prompt formatting." ICLR, 2024. arXiv:2310.11324.
[14] Y. Lu, M. Bartolo, A. Moore, S. Riedel, and P. Stenetorp. "Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity." ACL, 2022. arXiv:2104.08786.
[15] M. Mizrahi, G. Kaplan, D. Malkin, R. Dror, D. Shahaf, and G. Stanovsky. "State of What Art? A Call for Multi-Prompt LLM Evaluation." TACL, 2024. arXiv:2401.00595.
[16] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica. "Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM)." SOSP, 2023. arXiv:2309.06180.
[17] T. Dao. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." ICLR, 2024. arXiv:2307.08691.
[18] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models." SC, 2020. arXiv:1910.02054.
[19] I. Loshchilov and F. Hutter. "Decoupled Weight Decay Regularization (AdamW)." ICLR, 2019. arXiv:1711.05101.
[20] T. Bray (Ed.). "The JavaScript Object Notation (JSON) Data Interchange Format." RFC 8259, IETF, 2017.
Appendix A: Complete Hyperparameter Specification
# SchemaForge-1B β schemaforge-1b-iter2 (production checkpoint)
model:
student: openbmb/MiniCPM5-1B
student_arch: LlamaForCausalLM # no trust_remote_code needed
student_params: 1080632832 # 679552512 non-embedding
student_layers: 24
student_hidden: 1536
student_heads: 16 query / 2 kv (GQA)
student_vocab: 130560
teacher_primary: google/gemma-4-31B
teacher_secondary: google/gemma-4-E4B-it
teacher_vocab: 256000
teacher_frozen: true
distillation:
alpha: 0.5 # CE weight
tau: 2.0 # temperature
kl_reduction: batchmean
kl_log_target: true
vocab_projection: truncate_to_student_width
label_mask_value: -100
mask_prompt_tokens: true
mask_padding: true
optimization:
optimizer: AdamW
learning_rate: 2.0e-5
betas: [0.9, 0.999]
eps: 1.0e-8
weight_decay: 0.01
lr_scheduler: cosine
warmup_ratio: 0.05
max_grad_norm: 1.0
epochs: 3
early_stopping: val_loss
gradient_accumulation_steps: 4
runtime:
precision: bfloat16
attention: eager # flash-attn install failed on py3.12
sharding: none # single GPU
max_seq_length: 1024 # as executed; 2048 was specified
gradient_checkpointing: true
hardware:
gpu: NVIDIA RTX PRO 6000 Blackwell Edition
vram: 96GB
provider: Nebius AI Cloud
count: 1
software:
python: "3.12"
torch: "2.5"
transformers: "5.x"
inference:
temperature: 0.0
do_sample: false
max_new_tokens: 256
trust_remote_code: false # stock Llama architecture
Appendix B: Compatibility Patch Source
"""
schemaforge/compat.py
transformers 5.x shims for LEGACY trust_remote_code MiniCPM checkpoints
(openbmb/MiniCPM-1B-sft-bf16). Python 3.12 / PyTorch 2.5.
NOT required for openbmb/MiniCPM5-1B, which is a stock LlamaForCausalLM.
Call apply_compat_shims() BEFORE from_pretrained();
call retie_lm_head(model) AFTER instantiation.
"""
# ---------------------------------------------------------------- Patch 1 ----
def _patch_import_utils():
"""Re-inject is_torch_fx_available removed in transformers 5.x."""
import transformers.utils.import_utils as import_utils
# Symbolic tracing is unused in the training path, so False is safe.
import_utils.is_torch_fx_available = lambda: False
# ---------------------------------------------------------------- Patch 2 ----
def _patch_tied_weights(source: str = "model.embed_tokens.weight"):
"""
MiniCPM-1B-sft-bf16 omits lm_head.weight from disk (tied to embeddings).
transformers 5.x reports the absence as a CORRUPTED CHECKPOINT, and also
rejects the legacy list-form _tied_weights_keys. Patch the expansion
helper so both are handled.
"""
import transformers.modeling_utils as mu
_orig = mu.PreTrainedModel.get_expanded_tied_weights_keys
def _patched(self, *args, **kwargs):
keys = getattr(self, "_tied_weights_keys", None)
if isinstance(keys, (list, tuple)):
self._tied_weights_keys = {k: source for k in keys}
return _orig(self, *args, **kwargs)
mu.PreTrainedModel.get_expanded_tied_weights_keys = _patched
def retie_lm_head(model):
"""Re-establish the output-projection tie after instantiation."""
model.lm_head.weight = model.model.embed_tokens.weight
return model
# ---------------------------------------------------------------- Patch 3 ----
def _patch_dynamic_cache():
"""Restore legacy-cache interop methods on DynamicCache."""
from transformers.cache_utils import DynamicCache
if hasattr(DynamicCache, "from_legacy_cache"):
return
@classmethod
def from_legacy_cache(cls, past_key_values=None):
cache = cls()
if past_key_values is not None:
for layer_idx, (k, v) in enumerate(past_key_values):
cache.update(k, v, layer_idx)
return cache
def to_legacy_cache(self):
return tuple((layer.keys, layer.values) for layer in self.layers)
def get_usable_length(self, new_seq_length: int, layer_idx: int = 0) -> int:
seen = self.get_seq_length(layer_idx)
return seen if seen is not None else 0
DynamicCache.from_legacy_cache = from_legacy_cache
DynamicCache.to_legacy_cache = to_legacy_cache
DynamicCache.get_usable_length = get_usable_length
# ------------------------------------------------------------------ Driver ---
def apply_compat_shims():
"""Call BEFORE from_pretrained(). Then call retie_lm_head(model)."""
_patch_import_utils() # 1. before dynamic module load
_patch_tied_weights() # 2. before class instantiation
_patch_dynamic_cache() # 3. before first generate()
# ----------------------------------------------------------- NaN diagnostic --
def install_nan_probe(model):
"""Localize which layer first produces NaN. Architecture-agnostic;
worth leaving enabled in any long training run."""
import torch
def hook(idx):
def fn(_module, _inp, out):
h = out[0] if isinstance(out, tuple) else out
if torch.isnan(h).any():
raise RuntimeError(f"NaN first observed at layer {idx}")
return fn
for i, layer in enumerate(model.model.layers):
layer.register_forward_hook(hook(i))
return model
Appendix C: Verbatim Prompt Templates
C.1 Canonical Template (Iteration 2 β Production)
Extract structured JSON from the text:
{document_text}
JSON Output:
Python literal (whitespace is significant):
CANONICAL_TEMPLATE = "Extract structured JSON from the text:\n{doc}\nJSON Output:"
C.2 Iteration 1 Template (Failed β 0.0 %)
<start_of_turn>system
You are a JSON extraction assistant.<end_of_turn>
<start_of_turn>user
{document_text}<end_of_turn>
<start_of_turn>model
C.3 Iteration 3 Template (Failed β 0.0 %)
You are an expert enterprise document parser specializing in structured
data extraction across multiple business domains.
Extract structured JSON from the text:
{document_text}
JSON Output:
C.4 Schema-Explicit Variant (Model-Card Quickstart)
prompt = (
"Extract entity details into JSON with keys "
"'invoice_number', 'vendor', 'date', 'amount':\n"
"Invoice INV-99281, Acme Enterprise Solutions Inc, "
"Date 2026-08-03, Total Amount $12,450.00.\n"
"JSON Output:"
)
β οΈ Deployment warning. C.2 and C.3 differ from C.1 only in the instruction header. Both produced 0.0 % validity β worse than the untrained base model's 34.2 %. Do not modify C.1.
Appendix D: Benchmark Schemas and Source Documents
D.1 BMK-01 β Enterprise Financial Invoices
{
"invoice_number": "str",
"vendor_name": "str",
"invoice_date": "YYYY-MM-DD",
"subtotal": "float",
"tax": "float",
"grand_total": "float"
}
Source: INVOICE #INV-1001. Vendor: Acme Supply Co. Date: 2026-04-10. Item: Office Chairs x 4 @ $120.00 = $480.00. Subtotal: $480.00. Tax (8%): $38.40. Total: $518.40.
Expected:
{
"invoice_number": "INV-1001",
"vendor_name": "Acme Supply Co",
"invoice_date": "2026-04-10",
"subtotal": 480.00,
"tax": 38.40,
"grand_total": 518.40
}
D.2 BMK-02 β Supply Chain Logistics and Freight
{
"bill_of_lading": "str",
"carrier_name": "str",
"ship_date": "YYYY-MM-DD",
"container_count": "int",
"freight_cost": "float",
"total_amount": "float"
}
Source: TAX INVOICE 9942. Issued by: Quantum Logistics. Date: 2026-05-01. Shipping Container x 1 at $2500.00 ($2500.00). Insurance Fee x 1 at $150.00 ($150.00). Subtotal $2650.00. Tax $265.00. Total $2915.00.
D.3 BMK-03 β Commercial Hardware Bills of Sale
{
"receipt_id": "str",
"vendor": "str",
"transaction_date": "YYYY-MM-DD",
"line_items": "array<{description, quantity, unit_price, line_total}>",
"tax": "float",
"total": "float"
}
Source: BILL OF SALE #771. Vendor: Tech Hardware LLC. Date: 2026-05-15. Server Rack x 2 @ $800.00 = $1600.00. Cable Pack x 5 @ $20.00 = $100.00. Subtotal: $1700.00. Tax: $136.00. Grand Total: $1836.00.
The only nested-array schema in the suite.
D.4 BMK-04 β Biomedical Supply Receipts
{
"order_id": "str",
"supplier": "str",
"order_date": "YYYY-MM-DD",
"items": "array<{name, quantity, unit_price, subtotal}>",
"total_price": "float"
}
Source: COMMERCIAL INVOICE #INV-5502. Vendor: BioMed Supplies. Date: 2026-06-20. Centrifuge Tube x 10 @ $15.00 = $150.00. Pipette Set x 2 @ $75.00 = $150.00. Subtotal: $300.00. Tax: $24.00. Total: $324.00.
D.5 BMK-05 β Cloud Infrastructure Receipts
{
"receipt_id": "str",
"provider": "str",
"billing_date": "YYYY-MM-DD",
"service_description": "str",
"total_charge": "float"
}
Source: PURCHASE RECEIPT #PR-8819. Vendor: Cloud Servers Inc. Date: 2026-07-04. Virtual Machine Host x 1 @ $1200.00 = $1200.00. Subtotal: $1200.00. Tax: $96.00. Total: $1296.00.
Appendix E: Reproducibility Checklist
| Item | Status | Note |
|---|---|---|
| Model weights released | β | HuggingFace Hub, safetensors |
| Training hyperparameters | β | Appendix A, complete |
| Loss implementation | β | Β§4.3, full source |
| Compatibility patches | β | Appendix B, full source |
| Prompt templates | β | Appendix C, verbatim, all three |
| Evaluation schemas + documents | β | Appendix D |
| Hardware specification | β | RTX PRO 6000 Blackwell Edition 96 GB, Nebius |
| Software versions | β | Python 3.12, PyTorch 2.5, transformers 5.x |
| Training dataset released | β οΈ | $n = 5$ teacher-generated samples β should be released |
| Random seeds | β | Not recorded. Single-seed results |
Evaluation subset size (suneeldk/text-json) |
β | Not recorded. See Β§10.1 |
| Confidence intervals | β | Not computed at experiment time; retrospective Wilson intervals in Β§8.1/Β§8.4 |
| SFT control ($\alpha=1.0$) | β | Not run. Highest-priority gap β Β§10.2 |
| Multi-seed variance | β | Not run β Β§11.2 Priority 1 |
| Competitive baselines | β | Not run β Β§11.2 Priority 1 |
Citation
@techreport{ty2026schemaforge,
title = {SchemaForge: Distilling Ultra-Large Foundation Models into Edge SLMs
for Real-Time Enterprise JSON Extraction --
A Comparative Study of Gemma-4 Teachers and MiniCPM5-1B},
author = {Ty, Arjhine A.},
year = {2026},
note = {v1.0. Model: SchemaForge-1B (schemaforge-1b-iter2)}
}
SchemaForge v1.0 β August 2026. Author: Arjhine A. Ty. Production checkpoint: schemaforge-1b-iter2. Distilled from google/gemma-4-31B and google/gemma-4-E4B-it into openbmb/MiniCPM5-1B on NVIDIA RTX PRO 6000 Blackwell Edition / Nebius AI Cloud.


