Custom shared-state model¶

This is the "encode the text once, then cross-attention + small heads" architecture from the v1 spec, built on the same Qwen/Qwen3-Reranker-0.6B backbone as the stock backend. It was implemented, trained, tested and benchmarked, and it is a dead end. Code at tag archive/pre-cleanup-2026-09-27: src/personal_jev/custom.py, train_custom.py, tests/test_custom.py.

Bottom line. - As specified, it scores 39.0% question accuracy on the 3,471-question test split. For comparison, the unmodified stock reranker scores 61.0% and the stock reranker + LoRA 73.5%. - With an opt-in similarity (MaxSim) term trained jointly with LoRA, it reaches 58.2%. It matches or beats the stock LoRA on Banking77 and AG News, but still loses on unseen label sets and fails on yes/no questions. - Speed is where it wins. The text is encoded once per request, so with 16 questions × 3 candidates on an 8K-token text it answers in 626 ms versus 24,024 ms on the same A10G GPU. With one short question it is about as fast as the stock model.

Every number below is read from the file cited next to it.

Results (test split, 3,471 questions)¶

Source: reports/<run>/test/report.json; side-by-side table in reports/comparison_custom.md.

Accuracy % by family. "Unseen" families never appear in training, and neither do their label sets.

The seven eval_* families (17–32 test questions each) are in the full comparison table.

Paired exact McNemar tests (questions only one of the two models gets right; computed from the same reports): - stock + LoRA vs custom + similarity + joint LoRA: 749 vs 217 (p ≈ 4.3e-69). - stock base vs custom + similarity + joint LoRA: 577 vs 478 (p ≈ 0.0025). - custom (spec) vs custom + similarity: 217 vs 531 (p ≈ 3.2e-31).

What is computed¶

memory = LN(W_m · z(Qwen(state)))                         [S, L, 256]  each distinct state encoded once per call
q      = W_c · z(Qwen(task + question + proposed answer))  [N, Lc, 256] one short sequence per candidate, same weights
q      = q + Attn(LN(q), [null; memory]);  q = q + FFN(LN(q))     × 2 new blocks, 8 heads, FFN 1,024
score  = head(mean over the candidate's tokens of LN(q))   binary_head: binary + multilabel · choice_head: multiclass
         (+ w_type · MaxSim, only with the opt-in "similarity" flag)

• Backbone. The reranker's transformer body (AutoModel: 595,776,512 parameters, 28 layers, width 1,024) with its stock causal attention. The vocabulary head is not used and there are no yes/no logits. The state and the candidates never share a sequence.
• Text format (custom-v1, sha b96b6174a187).
• State: Document:\n{state}.
• Each candidate: Task: {type}\nQuestion: {instruction}\nProposed answer: {description}, where binary questions use the answer Yes.
• Tokenized with split_special_tokens=True, so <|im_end|> typed inside a state stays plain text.
• Sequences are right-padded, and the checkpoint refuses to load if the format changes.
• New trainable modules (2,172,418 parameters):

The similarity flag adds 2 parameters. LoRA r = 16 on q/k/v/o adds 4,587,520. - Typed outputs. Scores go through the same classify.decide as the stock backend: - binary and multilabel: sigmoid; - multiclass: softmax over the question's own candidates; - held-out temperatures and thresholds apply as before. - Sharing. Identical states in a call are encoded once. Candidate tokens are packed per state, so K/V are projected once per state per block and never copied per candidate. Candidates never attend to each other. Every response's meta counts state_sequences, memory_rows and tokens, and tests/test_custom.py (archived) checked all of this. - Checkpoint. A complete checkpoint holds modules.safetensors (all new modules plus the standardization buffers), adapter/ (LoRA) or backbone.safetensors (full fine-tune), and config.json (base revision, arch, format). Reloading reproduced scores within 9.5e-7 to 5.7e-6 (reload_check in each reports/train_meta/custom_*.json). The checkpoints themselves were not kept.

Additions beyond the spec, each measured¶

Training runs¶

Source: the training logs reports/train_meta/<run>.json; configs configs/<run>.json at the tag. - Training data: 10,112 questions (the same selection as the stock LoRA pilot: seed 13, at most 1,600 per family). - Validation: 983 questions. - Precision: backbone in bf16, new modules in fp32. - Optimizer: AdamW, lr 1e-3 for the new modules (3e-4 in the custom_lora Stage B run) and 2e-4 for LoRA, 20 warm-up steps, then linear decay. - Checkpoint selection: by validation loss.

"Mixed labels" means data/distill.jsonl (scripts/build_distill.py; both at the tag): - Each training state gets a multiclass question whose candidates are drawn from the descriptions of all training families. - Labels come from the local stock reranker + LoRA acting as teacher. 6,559 of 12,260 questions were kept (teacher max p ≥ 0.7). - The pool has 4,545 train-split descriptions; 937 held-out and eval descriptions were excluded (reports/custom_diagnostics/logs/build_distill.log). - Teacher labels are not ground truth.

Why the spec version scored 39%¶

1. It memorizes training labels instead of learning to match. The frozen model reaches 87.7% on AG News (in training) but 20.7% on CLINC and 25.7% on DBpedia, whose label sets it never saw (family table above).
2. The frozen features already match unseen labels; the training throws that away. On 100 test questions per family (maxsim_heldout_frozen.log), a parameter-free MaxSim over the same standardized features compares to the trained model like this:

Speed¶

One CUDA GPU (AWS A10G), bf16, same session, both models with an unmerged LoRA adapter. - End-to-end p50 latency in ms, one request at a time. - Source: reports/bench/gpu_{stock_base,stock_lora,custom_lora,stock_long,custom_long}_bf16/ and reports/bench/gpu_environment.txt (driver 595.91.07, torch 2.14.0+cu130, transformers 5.17.0). - The custom model here is custom_sim_lora.

Apple M5 Pro (MPS), bf16, same afternoon. - Source: reports/bench/lora_bf16/ (stock + LoRA), reports/bench/custom_lora_bf16_nobuckets/ (exact shapes) and reports/bench/custom_lora_bf16/ (MPS shape buckets, the default). - The custom model here is the spec architecture with an unmerged, zero-initialized LoRA adapter, so it has the same compute as a trained one.

• Where the time goes. The stock model's cost is proportional to pairs × state tokens. The custom model reads the state once, then adds a short sequence per candidate, so its cost barely depends on the number of questions. With one question on a long text, both do one pass over the same tokens and take about the same time.
• Shape buckets on MPS cost 11% to 76% extra latency in this table (the smallest request pays the most). Without them, memory grows with every new batch shape, as measured in training above.
• MPS long-context timings drifted during the day on this laptop, so the MPS 16K/32K rows (reports/bench/custom_long_context_bf16/) are not compared with this morning's stock run. Use the CUDA table.

Calibration¶

Source: temperatures from calib/<run>.json (at the tag), fit on the calibration split, with thresholds chosen on validation; ECE and accuracy from reports/<run>/test/ and test_calibrated/.

The uncalibrated similarity model is already well calibrated on multiclass (ECE 0.030). Its binary scores carry almost no signal, so no temperature can fix them: held-out temperature scaling made its binary ECE worse (0.141 → 0.154).

Limitations¶

• Binary is unsolved in this architecture. Test AUROC is 0.45–0.54 in all variants, so route yes/no questions to the stock backend.
• Unseen label sets are still weaker than the stock model. Examples: CLINC 79.3% vs 89.3%, TREC 43.3% vs 68.0%.
• Every run used one seed and one hyperparameter setting. The Stage B run on the spec model was stopped early.
• Eval data caveats are the same as the stock backend's. The eval_* families are LLM-written and small, and the public sets may overlap Qwen's pretraining data.
• Some runs are on a different machine. The distill-LoRA and similarity runs trained on an A10G; the others on the M5 Pro.

Reproduce¶

From a checkout of tag archive/pre-cleanup-2026-09-27 (package personal_jev, CLI pjev); the checkpoints were not kept, so the classify line needs the custom_sim step first:

uv run pytest tests/test_custom.py                 # new-model tests; the 3 real-model ones load the pinned base Qwen
scripts/run_experiments.sh custom                  # spec version, frozen backbone (~24 min on the M5 Pro)
scripts/run_experiments.sh custom_distill          # mixed-label data + frozen and joint-LoRA runs
scripts/run_experiments.sh custom_sim              # similarity variant, frozen and joint LoRA, then compare
scripts/run_experiments.sh custom_bench            # CUDA speed comparison
uv run pjev classify examples/request.json --checkpoint runs/custom_sim_lora/checkpoint
uv run pjev train-custom configs/custom_lora.json  # Stage B on top of custom_frozen (spec schedule)

The diagnostic scripts and their logs are in reports/custom_diagnostics/ (the scripts import personal_jev: run them from the tag): - features.py: anisotropy and MaxSim; - maxsim_heldout.py: MaxSim vs the trained model; - families.py: interference and init probes; - shapes.py: MPS memory; - mixture.py, generalize.py, overfit.py: earlier probes.