From Retrieval to Typed Decisions: Calibrated System One Models from Biomedical Sentence Encoders

Motivation

• Because software consumes these decisions, useful outputs must follow a schema and include confidence that can be thresholded to act, defer, or escalate.The overview describes one forward pass producing a calibrated distribution over each typed question’s options; its example values are illustrative.Source: §1 · p. 1, Figure 1 · p. 1

Research questions

• The study asks whether retrieval training improves zero-shot or fine-tuned decisions, how conversion heads trade accuracy and other properties, and which RLCD ingredients matter after temperature fitting.The stated trade-offs include robustness to option order, cost, and retrieval retention.Source: §1 · p. 2

Contributions

• SBERT2S1 converts Sentence-Transformers encoders into typed decision models, including a prior-fused residual head that starts from a calibrated zero-shot retriever.The paper also introduces BIODECIDE, MEDLINE-S1 labels derived from NLM indexing, and a controlled study of conversions and objectives.Source: §1 · p. 2

System One models

• Jev introduced choice, score, and noul primitives, while Laya is an open implementation trained with cross-entropy and a group-baselined policy gradient on noise-perturbed logits [Shao et al., 2024].Other cited work studies System One models in several application settings, but not biomedical conversions, retrieval initialization, or RLCD ablations.Source: §2 · p. 2

Calibration and scoring rules

• A strictly proper scoring rule is uniquely optimized in expectation by the true distribution, and temperature scaling is a strong post-hoc calibration remedy.The passage argues that claimed calibration benefits should be assessed after temperature scaling.Source: §2 · p. 2

Sentence encoders as decision models

• Biomedical sentence encoders trained contrastively on search logs or MS MARCO provide strong retrievers, while comparing text with label descriptions also underlies entailment-based zero-shot classification.This paper uses such encoders as decision-model initializations to measure what contrastive training adds.Source: §2 · p. 2

3.1 Typed decisions

• A typed request pairs a state with a question type, instruction, and allowed options; the model returns logits that become thresholdable confidence after type- and option-count-specific temperature scaling.Choice uses the highest-probability option, score uses the expected rubric level, and noul uses the probability that the statement is true.Source: §3.1 · pp. 2–3

3.2 Converting an encoder

• The conversions include frozen Z, trainable bi-encoder B, cross-head C, and PFR, which adds a cross score to the retrieval prior and initially equals Z.B permits precomputed state embeddings for corpus screening; freezing PFR’s prior preserves the original retriever at twice the memory.Source: §3.2 · p. 3

3.3 Training objectives and the RLCD estimator

• RLCD adds a noisy policy-gradient reward update to cross-entropy, but its self-including reward baseline shrinks the estimate and division by reward spread makes it grow roughly as 1/σ.As σ is annealed, the noisy term can come to dominate fixed-weight cross-entropy; PG-LOO instead uses the other draws’ mean reward and omits that division.Source: §3.3 · p. 3

3.4 Calibration and long states

• The method fits one post-hoc temperature per question type and option-count bucket, and evaluates retrieve-then-decide for long states by selecting three instruction-similar note chunks.The long-state approach requires the converted model to remain a retriever.Source: §3.4 · p. 4

BIODECIDE

• BIODECIDE covers grounded biomedical tasks and a separate knowledge track, while PUBHEALTH, BIOSSES, and MMLU are never trained on and MTSamples is a held-out task family.The knowledge-track exam answers depend on memorized knowledge and do not drive headline numbers.Source: §4 · p. 4

BIODECIDE conversion

• Source examples become records containing a state, one or more typed questions, and gold option distributions, with label descriptions and task-specific conversions.The conversion removes answer-revealing fields, retains official splits where available, and carves a calibration slice from training data before training.Source: §4 · p. 4

MEDLINE-S1

• MEDLINE-S1 derives weak multi-question labels from NLM indexing rather than an LLM teacher, yielding 242,986 decisions across 40,000 training abstracts and 5,000 test abstracts from disjoint files.Indexing may omit tags, absence need not mean false, and indexers may use the full article.Source: §4 · pp. 4–5

Clinical track

• The credentialed clinical track includes MedNLI, human-reviewed trial-eligibility questions, and name-mention questions over MIMIC notes; the released model uses public data only.MIMIC-III discharge notes average 12.9k characters and form the long-state test.Source: §4 · p. 5

Encoders and training

• The study evaluates eleven base-size checkpoints, including six matched parent-retriever pairs, while noting these pairs are not independent replications of one intervention.Models are trained with C and PFR; the full head-by-objective grid uses S-PubMedBERT-MS-MARCO.Source: §5 · p. 5

Training

• Training uses three seeds, with 2% and 10% learning curves for paired checkpoints and a matched grid of two heads and five objectives.The budget is 8,000 steps at batch size 32, with AdamW and gradient clipping.Source: §5 · p. 5

Baselines

• Baselines include zero-shot and fine-tuned Laya-large plus zero-shot Qwen and Gemma; all systems are temperature-scaled on shared calibration splits.Exam-track LLM scores are descriptive because exam questions may occur in pretraining data.Source: §5 · p. 5

Metrics and evaluation

• Chance-normalised accuracy is the headline metric, averaged over question groups within tasks and then over tasks; seen, held-out, and clinical tracks define evaluation scope.The protocol also reports calibration, ordinal error, option-order flips, and retrieval retention, with paired document-cluster bootstrap intervals and Holm correction.Source: §5 · p. 5, Table 2 · p. 6

Data integrity

• Removing audited test inputs lowers seen-task scores by 0.3–0.4 points for every model but changes no ranking.The audit found some duplicated training inputs in calibration and test splits, as well as repeated inputs within test sets.Source: §5 · p. 5

6.1 RQ1: Does retrieval training help?

• At step zero, contrastive checkpoints score 15–24 chance-normalised points on content-differing options, versus 0–7 for MLM parents; label-only options remain near chance.The reported cosine gaps for label-only options are about 10^-3, and no encoder is reliably above chance.Source: §6.1 · p. 6, Table 3 · p. 7
• Across 15 pair-and-data-size comparisons, retrieval training significantly helps PFR in 10 and hurts it in one, while helping C in one and hurting it in five.At 10% and 2% training data, PFR retrievers lead in four of five pairs; Figure 2 plots seen-task accuracy against training data.Source: §6.1 · p. 6, Figure 2 · p. 7

6.2 RQ2: Which conversion works best?

• C exceeds PFR by 1.3–2.4 seen-task points and 3.8–10.2 held-out points across all five shared objectives.The matched comparison uses S-PubMedBERT-MS-MARCO.Source: §6.2 · p. 7, Table 4 · p. 7
• Random option permutation changes 0.3% of S-PubMedBERT PFR answers versus 15.7% for C, with C’s flips concentrated on unseen inputs.PFR also has lower ordinal error on Druglib for every objective.Source: §6.2 · p. 7
• Fine-tuning reduces SciFact nDCG@10 from 67.6 for the pooled embedding to 43.3 for C, 35.6 for PFR, and 34.3 for B; only a frozen prior preserves retrieval.A ten-question request takes 6.6 ms for C and 15.2 ms for PFR, versus 22.7 ms for Laya-large.Source: §6.2 · p. 7

6.3 RQ3: Does RLCD improve accuracy or calibration?

• The released PG recipe trails CE by 3.0 points with PFR and 2.5 with C; PG-LOO recovers 2.5 and 1.5 points over PG, respectively.The analysis attributes the deficit mainly to reward normalisation, which makes the noisy score-function term 3.6–15 times too large.Source: §6.3 · pp. 7–8
• After temperature scaling, objective ECE values lie between 3.6 and 4.4, and PG has the worst NLL; held-out objective rankings are noisier.PG’s raw ECE is lower than CE’s with PFR, but that difference does not persist after scaling.Source: §6.3 · p. 8

6.4 LLMs and confidence-based routing

• On seen tasks, S-PubMedBERT C with CE and Laya-large exceed the zero-shot LLM baselines, while LLMs perform better on held-out, clinical, and exam tasks.Routing the 20% least-confident questions to Gemma raises macro accuracy in the reported retrospective test-set analysis, which is not a validated deployment threshold.Source: §6.4 · p. 8

• Retrieval training helps PFR, especially with scarce labels, but offers no consistent accuracy gain for cross heads.Source: §7 · p. 8
• Figure 3 escalates the least-confident fraction by temperature-scaled confidence; lines show three-seed means and bands show standard deviations.Source: Figure 3 · p. 8
• For per-request decisions, the discussion recommends a cross head; PFR suits low option-order sensitivity or degenerate mask-slot states, while B suits screening many questions.Source: §7 · p. 8
• Every trained conversion degrades the embedding as a retriever unless the prior is frozen.Source: §7 · p. 8

• Cross-head models reach 66–73 seen-task points with top-label ECE below 6 after temperature scaling, except MedCPT-Query, and approach fine-tuned Laya-large at about a third of its size.Source: §8 · p. 8
• Retrieval training helps PFR in 10 of 15 comparisons, especially with little labelled data, but does not consistently help C.Source: §8 · p. 8
• Plain cross-entropy with temperature scaling matches or outperforms the released RLCD recipe, whose deficit is attributed to reward normalisation; leave-one-out is recommended as an independent baseline.Source: §8 · p. 8
• Fast encoders and LLMs are complementary: the calibrated model answers confident cases first and escalates the rest.Source: §8 · p. 8

• Results cover base-size encoders, one fixed training budget and inherited hyperparameters; the head-by-objective grid used one encoder, and parent-retriever comparisons are not independent replications.Source: Limitations · p. 9
• MEDLINE-S1 labels can omit tags and use full articles; exact training inputs recur in calibration and test splits, without de-duplicated retraining or temperature refitting.Source: Limitations · p. 9
• Low seen-task ECE does not imply calibration on new tasks, and cross-head option-order sensitivity grows under shift; thresholds need in-domain validation.Source: Limitations · p. 9
• Clinical tests are small, from one hospital system, and use surrogate names; the released model is not a validated device and was not trained on clinical notes.Source: Limitations · p. 9

Ethics Statement

• Credentialed clinical data stayed on institutional infrastructure, while MEDLINE-S1 labels are released without abstracts and health-related low-confidence or escalated decisions are recommended for human review.The released MEDLINE-S1 materials include PMIDs, questions, and labels; a script retrieves abstracts from PubMed.Source: Ethics Statement · p. 9

A Extended Related Work

• Biomedical QA benchmarks are dominated by generative LLMs, and calibrated System One models are described as a natural first stage in cascades.The cited QA benchmarks include Jin et al. (2019, 2021) and Pal et al. (2022); the cascade discussion cites Chen et al. (2023).Source: A Extended Related Work · p. 12

B Task Conversions

• Each clinical conversion poses typed decisions over notes, including MedNLI, trial-eligibility questions, answerability, and a name-mention gate.The public conversions are listed in Table 5; clinical conversions are described in the text.Source: B Task Conversions · p. 12

C RLCD: Formal Statement and Proof

• The leave-one-out baseline uses the other G−1 rewards without reward normalisation, giving an unbiased score-function estimate on the smoothed-score gradient’s scale.The formal statement contrasts this estimator with the same-sample mean baseline.Source: C RLCD: Formal Statement and Proof · pp. 13–14

D Gradient-Estimator Probe and Optimisation Traces

• 3.6–3.8× at σ = 0.4 and 14–15× at σ = 0.1: released reward normalisation inflates the score estimate as σ anneals, while every objective is clipped at every training step.At matched scale, score-function estimators have per-item SNR around 0.8, versus 85–1,300 for the pathwise estimator; the inflated noisy term dominates the clipped PG direction.Source: D Gradient-Estimator Probe and… · p. 14, Figure 5 · p. 15, Figure 6 · p. 15

E Complete Head × Objective Results; G Full Results by Encoder and Task

• Across every objective, C has higher seen and held-out accuracy and lower NLL, while PFR has lower ordinal error and option-order sensitivity; C’s advantage is concentrated on text-comparison and specialty-routing tasks.On SciFact, C scores 55–61 versus PFR’s 35–40; on held-out MTSamples, S-PubMedBERT scores 69–71 versus 51.Source: E Complete Head × Objective Results · p. 14, Table 7 · p. 15, G Full Results by Encoder and Task · p. 16

F Metrics and Statistics

• Paired bootstrap intervals resample document clusters but average correctness over seeds first, so they are pointwise and omit seed uncertainty; accuracy and ECE use argmax targets.ECE uses 15 equal-width bins; the saved metrics also include equal-mass and class-wise ECE, Brier score, and ranking and ordinal measures.Source: F Metrics and Statistics · pp. 14, 16

Option order

• 10–17% of cross-head answers change under random option permutations, while B and shared-prior PFR are essentially invariant; frozen-prior PFR flips 13%.For S-PubMedBERT C, flips rise from 3% on seen tasks to 17% held-out and 28% on exams.Source: H Robustness and Long Notes · p. 16

Long notes

• For MIMIC-III long-note trial questions, no strategy consistently wins; sliding windows help most models on three-way answers, but long-note decisions remain challenging for base encoders.The three-way sample has 532 answer items and 1,000 answerability items; most answerability strategies remain close to chance.Source: H Robustness and Long Notes · p. 16

Privacy gate

• About 75 to 99.8: training on the name gate’s own windows, with clinical data, raises AUROC; both LLMs are perfect zero-shot.Common-word names have AUROC similar to ordinary names, differing by at most two points.Source: H Robustness and Long Notes · p. 16

I Data Integrity and Sensitivity

• Nine DDI and seven Druglib calibration inputs, plus two SciFact and 13 Druglib test inputs, recur in training; excluding affected test inputs does not repair calibration contamination.The reported sensitivity analysis leaves checkpoints and temperatures unchanged.Source: I Data Integrity and Sensitivity · pp. 16, 18, Table 13 · p. 19

J LLM Prompt

• LLMs receive a chat-template prompt requesting an option letter; option probabilities use next-token log-probabilities, and states longer than 6,000 tokens are truncated.The bare and space-prefixed token probabilities for each letter are combined by log-sum-exp.Source: J LLM Prompt · p. 18

K Hyperparameters and Compute

• Base cross sequences are limited to 512 tokens, with at most 192 question-and-option tokens; reported single-run throughput is 265 items/s for C and 326 items/s for B.These throughput figures use batch size 32.Source: K Hyperparameters and Compute · pp. 18–19

L Released Software and Model

• 70.5 seen-task accuracy for the released best-seed model, versus a three-seed mean of 70.2 ± 0.4; it is S-PubMedBERT-MS-MARCO with a CE-trained cross head on public data.The release includes conversion, temperature-fitting, data-building, and evaluation code alongside the model.Source: L Released Software and Model · pp. 19–20