Qwen3.5-9B · LoRA fine-tune · document structure extraction

Teaching a 9B vision model to turn insurance policies into nested JSON, page by page

A LoRA fine-tune of Qwen3.5-9B that reads one page image of a complex document and writes it out as a nested JSON tree (parts, sections, clauses, items), page after page, so a whole policy reassembles into one tree for retrieval. It runs only under its own extraction prompt; otherwise it behaves like the original model.

Interactive comparison

Original vs fine-tuned, page by page

Neither document was used in training. Both models see the same page image and system prompt (greedy decoding, thinking off), and each page's prompt carries the model's own open path from the previous page, exactly as in real use. The original model is shown in its best setting (JSON-constrained decoding); the fine-tuned one in the final pipeline (free decoding, then a bracket-repairing parser and a tolerant merge).

Loading…

Generalisation

Six more documents it never saw

Loading…

Benchmark

How it compares with Marker, MinerU and Docling

Loading…

Method

How the page-by-page format works

1 · One page per call

A policy runs to 50k+ tokens of JSON, too much for one answer. Pages are the only unit you can see at inference without already knowing the structure (sections would need the answer first), so the model works one page image at a time.

2 · Carry the open path

Each page's prompt lists the nodes still open at the end of the previous page: root → last section → last clause → … with the tail of the last text. The model attaches this page's content to that path.

3 · Continuation stubs

Nodes that are already open are written as {"cont": true, type, label}, only down the left edge. A stub's "text" continues a sentence cut by the page break; its children are new. Everything else is new content.

Reassembly is a pure function

Merging page outputs needs no fuzzy matching: walk down the stub chain, append new nodes. Slicing every teacher tree into pages and merging them back reproduces all 16 documents exactly, the round-trip test for both the training data and the inference code.

def merge(acc, page):              # acc = document so far, page = this page's JSON
    if page.get("text"):           # page starts mid-sentence: continue the open node's text
        acc["text"] = join(acc.get("text", ""), page["text"])
    kids = page.get("children", [])
    if kids and kids[0].get("cont"):           # re-enter the open node one level down
        open_child = acc["children"][-1]
        check_echo(open_child, kids[0])        # stub must name the open node
        merge(open_child, kids[0])
        kids = kids[1:]
    for k in kids:                              # everything else is new content
        acc.setdefault("children", []).append(k)

def open_path(acc):                # the next page's prompt context
    path = [acc]
    while path[-1].get("children"):
        path.append(path[-1]["children"][-1])
    return path

Node schema

type	document · part · section · subsection · clause · definition · subclause · item · paragraph
label	the enumerator as printed: "A.", "1.", "(a)", "(i)"
title	a heading, including a bold run-in heading
text	the node's own prose, verbatim
children	nested nodes, following visual hierarchy

No ids are generated (they are assigned after reassembly). Running headers, footers, form numbers and copyright lines are left out. The first numbered level under a section is a clause, deeper levels are items, and a list nests under the lead-in that introduces it.

Data

How the dataset was built

Source documents

Specimen policy forms (mostly public state insurance regulator filings). Each one was hand-labelled into a whole-document teacher tree from its page images and text layer. Training documents are listed by kind only; the two documents held out entirely for evaluation are named.

Documents	Count	Layout	Role	Pages	Training pages	Nodes	Max depth

From teacher trees to page examples

Find every page break inside each tree. A 12-character shingle alignment against each page's text, corrected with the first and last printed line of every page. breaks: confirmed from both page edges, from one (all checked against the images).
Audit coverage. Every page's target must contain everything printed on it except page furniture. Pages a tree didn't fully cover were excluded (see the table) or fixed.
Fix the trees. One parser-built auto policy jacket was missing six whole Part 1 sections and had seven lists flattened into text, so it was patched (430 → 517 nodes). Node types were normalized across documents (594 nodes). A federal regulation excerpt was dropped (every page also printed unlabelled sections).
Slice into page targets. pages resume at two nesting depths and at three or more; start mid-sentence; stub chains reach levels deep.
Build records. training pages (5 near-duplicate ISO pages dropped at 4-gram Jaccard ≥ 0.7) and held-out GEICO pages for validation.

Token budget (measured with the Qwen3.5 tokenizer)

An early assumption of a 4–5k-token output limit was wrong: the model card recommends 32,768 output tokens. Page units were kept anyway because they are visible at inference, keep training sequences short and make reassembly deterministic.

Contrast and retention examples keep the model normal outside its extraction prompt. Their answers were written by the untrained model itself(self-distillation, vLLM on an L40S) so training reinforces its own behaviour. Images came from 50 rows of HuggingFaceM4/the_cauldron (charts, documents, photos, OCR, diagrams) and text prompts from OpenAssistant. 75% are in thinking mode, so fine-tuning on no-thinking extraction doesn't erode reasoning. Hatched segments are thinking-off.

Fine-tuning

Training

Hyperparameters

Base model	Qwen/Qwen3.5-9B hybrid Gated DeltaNet + attention, vision encoder, 262k context
Method	LoRA (bf16) QLoRA avoided: not recommended for Qwen3.5
Trainable	language model, all linear layers q/k/v/o, in_proj_qkv/z/b/a, out_proj, gate/up/down; vision encoder + projector frozen
Rank / alpha / dropout	16 / 32 / 0.05 43.3M trainable parameters (0.46% of 9.45B)
Learning rate	sweep 5e-5 · 1e-4 · 2e-4 winner 2e-4; LoRA's optimum is ~10× full fine-tuning's
Schedule	cosine, 5% warm-up decays to 0 at step 183
Optimizer	AdamW (fused) β 0.9 / 0.95, ε 1e-8, weight decay 0.1, gradient clip 1.0
Batch	1 × 8 accumulation = 8 61 steps per epoch, 183 in total
Epochs	3 checkpoint + validation after each
Loss	next-token cross-entropy answer tokens only; the empty <think></think> block of no-thinking answers is masked
Sequence / image	8,192 tokens / ≤ 4,096 per image pages enter at full resolution
Memory	gradient checkpointing, SDPA attention peak 68.5 GB of the H100's 80 GB
Stack	ms-swift 4.5.3 · transformers 5.16.1 · peft 0.20.0 flash-linear-attention 0.5.2, PyTorch 2.14, Modal H100

Checked before training

ms-swift trains with its own Qwen3.5 template, so the first encoded sample was inspected:

masked (not trained)trained

4,413 tokens masked: system prompt, the page (exactly 3,657 image tokens), open path and the empty think block.
736 tokens trained: the JSON answer and the end-of-turn token.
The rendered text equals the inference prompt with thinking off, plus the answer.

Runs

Learning rate	Train time	Final val loss

3 runs in parallel on separate H100s. One GPU memory warning at step 66 (an 8k-token sample) was recovered automatically.

Results

Evaluation and choosing a checkpoint

Loading…

Decisions and pivots

The journey

Oct 12–13
Scope and schema
Nested JSON with {type, label, title, text, children}, no ids. Page images instead of the PDF text layer, because two-column layouts interleave when extracted as text. 3 gold documents, then 15 more hand-labelled into whole-document teacher trees.
Oct 13–14
First plan: section-sized units
Training examples were whole sections (≤ 20k, then ≤ 8k characters) with a prompt naming the section; near-duplicate sections dropped at Jaccard ≥ 0.7.
Oct 15
Rethink: pages, not sections
Section units need the structure to find the sections, so inference couldn't use them. Switched to one page per call with an open-path prompt and continuation stubs, so reassembly is deterministic.
Oct 15
Checked the assumptions
The assumed 4–5k output-token limit was wrong (the model card recommends 32k; a forward pass emits one token). Two gotchas surfaced: thinking is on by default, and the recommended presence penalty would corrupt verbatim JSON.
Oct 15
Finding page breaks
Built a shingle-alignment aligner, then fixed what it exposed: text layers that dropped the top inch of every page, single-column pages scrambled by a gutter split, one document's unreadable font layer (switched to pdfium), and repeated form titles mistaken for headings. Three wrong breaks caught by viewing the images.
Oct 15
Auditing the labels
Coverage checks showed one parser-built auto policy was missing six whole sections; it was patched. A federal regulation excerpt was dropped and 14 pages excluded. Node types were normalized across documents.
Oct 15
Records and the mix
Real tokenizer counts: 3,657 image tokens per page, ~1.1k-token answers. The old prompt's false rule ("headings carry no text") was corrected. Research put 55% extraction comfortably within safe ranges; thinking share raised to 75%; retention answers self-distilled from the base model.
Oct 15
Trainer and hyperparameters
ms-swift (Qwen-native, masks the empty think block automatically) over Unsloth and TRL. Settings from LoRA Without Regret: all layers, rank 16, ~10× full-fine-tune learning rate, small batches.
Oct 15
Running on Modal
One volume for data, weights and checkpoints. Fixed a peft/transformers version clash, a missing torchvision, and a vLLM sampler that needed a compiler. A 4-step smoke test confirmed the encoding and masking token by token.
Oct 15–16
The sweep
Three learning rates in parallel on three H100s, 34 minutes each. Validation loss fell from 0.062 to 0.0116 with no overfitting.
Oct 16
Evaluation and fixes
lr 2e-4 at epoch 3 won on both held-out policies and passed every guardrail. Bracket-counting failures were first fixed with JSON-grammar decoding (28 → 33 of 34 pages) and a tolerant merge; the remaining gap traced to premature closes on two-column pages.
Oct 16
Repairing parser
Decoding freely and repairing the brackets afterwards keeps the content the grammar had to cut off: 34 of 34 pages attach, and Encompass node F1 rises from 0.85 to 0.92. This became the final inference pipeline.

Compute

What it cost

Step	Hardware	Cost
Model download and data upload	CPU	≈ $0.02
Self-distilled answers (225 prompts, plus a test run)	L40S	≈ $0.50
Smoke tests and first-run fixes	H100	≈ $1.00
3 training runs in parallel (~37 min each incl. start-up)	H100 × 3	≈ $8.20
Evaluations (test, sweep of 10 models, constrained and repair re-runs)	H100	≈ $2.50
Six more unseen documents, both models (88 pages each, ~13 min)	H100	≈ $0.85
Benchmark: Marker, MinerU, Docling and the original Qwen on 122 pages	L40S × 3, H100	≈ $1.40
Storage (≈ 25 GB on the volume)	Volume	free (1 TiB included)
Hosting this site	Vercel static	free
Total		≈ $14.55

Estimated from Modal's published per-second rates (H100 $3.95/h, L40S $1.95/h, plus CPU and memory). The Modal dashboard has the exact bill.